Introduce a TraceResource type class

modules/core/shared/src/main/scala/Trace.scala

@@ -45,30 +45,7 @@ object Trace {
 
   /** A `Trace` instance that uses `IOLocal` internally. */
   def ioTrace(rootSpan: Span[IO]): IO[Trace[IO]] =
-    IOLocal(rootSpan).map { local =>
-      new Trace[IO] {
-
-        def put(fields: (String, TraceValue)*): IO[Unit] =
-          local.get.flatMap(_.put(fields: _*))
-
-        def kernel: IO[Kernel] =
-          local.get.flatMap(_.kernel)
-
-        def span[A](name: String)(k: IO[A]): IO[A] =
-          local.get.flatMap { parent =>
-            parent.span(name).flatMap { child =>
-              Resource.make(local.set(child))(_ => local.set(parent))
-            } .use { _ => k }
-          }
-
-        def traceId: IO[Option[String]] =
-          local.get.flatMap(_.traceId)
-
-        def traceUri: IO[Option[URI]] =
-          local.get.flatMap(_.traceUri)
-
-      }
-    }
+    IOLocal(rootSpan).map(new IOTrace(_) {})
 
   object Implicits {
 
@@ -78,15 +55,7 @@ object Trace {
      * effect type.
      */
     implicit def noop[F[_]: Applicative]: Trace[F] =
-      new Trace[F] {
-        final val void = ().pure[F]
-        val kernel: F[Kernel] = Kernel(Map.empty).pure[F]
-        def put(fields: (String, TraceValue)*): F[Unit] = void
-        def span[A](name: String)(k: F[A]): F[A] = k
-        def traceId: F[Option[String]] = none.pure[F]
-        def traceUri: F[Option[URI]] = none.pure[F]
-      }
-
+      new NoopTrace[F] {}
   }
 
   /**
@@ -137,101 +106,141 @@ object Trace {
 
     def traceUri: Kleisli[F,Span[F],Option[URI]] =
       Kleisli(_.traceUri)
+      new NoopTrace[F] {}
 
   }
 
   implicit def liftKleisli[F[_], E](implicit trace: Trace[F]): Trace[Kleisli[F, E, *]] =
-    new Trace[Kleisli[F, E, *]] {
+    new KleisliETrace[F, E] {}
 
-      def put(fields: (String, TraceValue)*): Kleisli[F, E, Unit] =
-        Kleisli.liftF(trace.put(fields: _*))
+  implicit def liftStateT[F[_], S](implicit F: Monad[F], trace: Trace[F]): Trace[StateT[F, S, *]] =
+    new StateTTrace[F, S] {}
 
-      def kernel: Kleisli[F, E, Kernel] =
-        Kleisli.liftF(trace.kernel)
+  implicit def liftEitherT[F[_], E](implicit F: Functor[F], trace: Trace[F]): Trace[EitherT[F, E, *]] =
+    new EitherTTrace[F, E] {}
 
-      def span[A](name: String)(k: Kleisli[F, E, A]): Kleisli[F, E, A] =
-        Kleisli(e => trace.span[A](name)(k.run(e)))
+  implicit def liftOptionT[F[_]](implicit F: Functor[F], trace: Trace[F]): Trace[OptionT[F, *]] =
+    new OptionTTrace[F] {}
+
+  implicit def liftNested[F[_], G[_]](implicit F: Functor[F], G: Applicative[G], trace: Trace[F]): Trace[Nested[F, G, *]] =
+    new NestedTrace[F, G] {}
+}
 
-      def traceId: Kleisli[F, E, Option[String]] =
-        Kleisli.liftF(trace.traceId)
+private[natchez] abstract class IOTrace(local: IOLocal[Span[IO]]) extends Trace[IO] {
 
-      def traceUri: Kleisli[F, E, Option[URI]] =
-        Kleisli.liftF(trace.traceUri)
+  def put(fields: (String, TraceValue)*): IO[Unit] =
+    local.get.flatMap(_.put(fields: _*))
+
+  def kernel: IO[Kernel] =
+    local.get.flatMap(_.kernel)
+
+  def span[A](name: String)(k: IO[A]): IO[A] =
+    local.get.flatMap { parent =>
+      parent.span(name).flatMap { child =>
+        Resource.make(local.set(child))(_ => local.set(parent))
+      } .use { _ => k }
     }
 
-  implicit def liftStateT[F[_]: Monad, S](implicit trace: Trace[F]): Trace[StateT[F, S, *]] =
-    new Trace[StateT[F, S, *]] {
+  def traceId: IO[Option[String]] =
+    local.get.flatMap(_.traceId)
+
+  def traceUri: IO[Option[URI]] =
+    local.get.flatMap(_.traceUri)
 
-      def put(fields: (String, TraceValue)*): StateT[F, S, Unit] =
-        StateT.liftF(trace.put(fields: _*))
+}
 
-      def kernel: StateT[F, S, Kernel] =
-        StateT.liftF(trace.kernel)
+private[natchez] abstract class NoopTrace[F[_]: Applicative] extends Trace[F] {
+  final val void = ().pure[F]
+  val kernel: F[Kernel] = Kernel(Map.empty).pure[F]
+  def put(fields: (String, TraceValue)*): F[Unit] = void
+  def span[A](name: String)(k: F[A]): F[A] = k
+  def traceId: F[Option[String]] = none.pure[F]
+  def traceUri: F[Option[URI]] = none.pure[F]
+}
 
-      def span[A](name: String)(k: StateT[F, S, A]): StateT[F, S, A] =
-        StateT(s => trace.span[(S, A)](name)(k.run(s)))
+private[natchez] abstract class KleisliETrace[F[_], E](implicit trace: Trace[F])
+  extends Trace[Kleisli[F, E, *]]
+{
+  def put(fields: (String, TraceValue)*): Kleisli[F, E, Unit] =
+    Kleisli.liftF(trace.put(fields: _*))
 
-      def traceId: StateT[F, S, Option[String]] =
-        StateT.liftF(trace.traceId)
+  def kernel: Kleisli[F, E, Kernel] =
+    Kleisli.liftF(trace.kernel)
 
-      def traceUri: StateT[F, S, Option[URI]] =
-        StateT.liftF(trace.traceUri)
-    }
+  def span[A](name: String)(k: Kleisli[F, E, A]): Kleisli[F, E, A] =
+    Kleisli(e => trace.span[A](name)(k.run(e)))
 
-  implicit def liftEitherT[F[_]: Functor, E](implicit trace: Trace[F]): Trace[EitherT[F, E, *]] =
-    new Trace[EitherT[F, E, *]] {
+  def traceId: Kleisli[F, E, Option[String]] =
+    Kleisli.liftF(trace.traceId)
 
-      def put(fields: (String, TraceValue)*): EitherT[F, E, Unit] =
-        EitherT.liftF(trace.put(fields: _*))
+  def traceUri: Kleisli[F, E, Option[URI]] =
+    Kleisli.liftF(trace.traceUri)
+}
 
-      def kernel: EitherT[F, E, Kernel] =
-        EitherT.liftF(trace.kernel)
+private[natchez] abstract class StateTTrace[F[_], S](implicit F: Monad[F], trace: Trace[F]) extends Trace[StateT[F, S, *]] {
+  def put(fields: (String, TraceValue)*): StateT[F, S, Unit] =
+    StateT.liftF(trace.put(fields: _*))
 
-      def span[A](name: String)(k: EitherT[F, E, A]): EitherT[F, E, A] =
-        EitherT(trace.span(name)(k.value))
+  def kernel: StateT[F, S, Kernel] =
+    StateT.liftF(trace.kernel)
 
-      def traceId: EitherT[F, E, Option[String]] =
-        EitherT.liftF(trace.traceId)
+  def span[A](name: String)(k: StateT[F, S, A]): StateT[F, S, A] =
+    StateT(s => trace.span[(S, A)](name)(k.run(s)))
 
-      def traceUri: EitherT[F, E, Option[URI]] =
-        EitherT.liftF(trace.traceUri)
-    }
+  def traceId: StateT[F, S, Option[String]] =
+    StateT.liftF(trace.traceId)
 
-  implicit def liftOptionT[F[_]: Functor](implicit trace: Trace[F]): Trace[OptionT[F, *]] =
-    new Trace[OptionT[F, *]] {
+  def traceUri: StateT[F, S, Option[URI]] =
+    StateT.liftF(trace.traceUri)
+}
 
-      def put(fields: (String, TraceValue)*): OptionT[F, Unit] =
-        OptionT.liftF(trace.put(fields: _*))
+private[natchez] abstract class EitherTTrace[F[_], E](implicit F: Functor[F], trace: Trace[F]) extends Trace[EitherT[F, E, *]] {
+  def put(fields: (String, TraceValue)*): EitherT[F, E, Unit] =
+    EitherT.liftF(trace.put(fields: _*))
 
-      def kernel: OptionT[F, Kernel] =
-        OptionT.liftF(trace.kernel)
+  def kernel: EitherT[F, E, Kernel] =
+    EitherT.liftF(trace.kernel)
 
-      def span[A](name: String)(k: OptionT[F, A]): OptionT[F, A] =
-        OptionT(trace.span(name)(k.value))
+  def span[A](name: String)(k: EitherT[F, E, A]): EitherT[F, E, A] =
+    EitherT(trace.span(name)(k.value))
 
-      def traceId: OptionT[F, Option[String]] =
-        OptionT.liftF(trace.traceId)
+  def traceId: EitherT[F, E, Option[String]] =
+    EitherT.liftF(trace.traceId)
 
-      def traceUri: OptionT[F, Option[URI]] =
-        OptionT.liftF(trace.traceUri)
-    }
+  def traceUri: EitherT[F, E, Option[URI]] =
+    EitherT.liftF(trace.traceUri)
+}
 
-  implicit def liftNested[F[_]: Functor, G[_]: Applicative](implicit trace: Trace[F]): Trace[Nested[F, G, *]] =
-    new Trace[Nested[F, G, *]] {
+private[natchez] abstract class OptionTTrace[F[_]](implicit F: Functor[F], trace: Trace[F]) extends Trace[OptionT[F, *]] {
+  def put(fields: (String, TraceValue)*): OptionT[F, Unit] =
+    OptionT.liftF(trace.put(fields: _*))
 
-      def put(fields: (String, TraceValue)*): Nested[F, G, Unit] =
-        trace.put(fields: _*).map(_.pure[G]).nested
+  def kernel: OptionT[F, Kernel] =
+    OptionT.liftF(trace.kernel)
 
-      def kernel: Nested[F, G, Kernel] =
-        trace.kernel.map(_.pure[G]).nested
+  def span[A](name: String)(k: OptionT[F, A]): OptionT[F, A] =
+    OptionT(trace.span(name)(k.value))
 
-      def span[A](name: String)(k: Nested[F, G, A]): Nested[F, G, A] =
-        trace.span(name)(k.value).nested
+  def traceId: OptionT[F, Option[String]] =
+    OptionT.liftF(trace.traceId)
 
-      def traceId: Nested[F, G, Option[String]] =
-        trace.traceId.map(_.pure[G]).nested
+  def traceUri: OptionT[F, Option[URI]] =
+    OptionT.liftF(trace.traceUri)
+}
 
-      def traceUri: Nested[F, G, Option[URI]] =
-        trace.traceUri.map(_.pure[G]).nested
-    }
-}
+private[natchez] abstract class NestedTrace[F[_], G[_]](implicit F: Functor[F], G: Applicative[G], trace: Trace[F]) extends Trace[Nested[F, G, *]] {
+  def put(fields: (String, TraceValue)*): Nested[F, G, Unit] =
+    trace.put(fields: _*).map(_.pure[G]).nested
+
+  def kernel: Nested[F, G, Kernel] =
+    trace.kernel.map(_.pure[G]).nested
+
+  def span[A](name: String)(k: Nested[F, G, A]): Nested[F, G, A] =
+    trace.span(name)(k.value).nested
+
+  def traceId: Nested[F, G, Option[String]] =
+    trace.traceId.map(_.pure[G]).nested
+
+  def traceUri: Nested[F, G, Option[URI]] =
+    trace.traceUri.map(_.pure[G]).nested
+}

modules/core/shared/src/main/scala/TraceResource.scala

@@ -0,0 +1,78 @@
+// Copyright (c) 2019-2020 by Rob Norris and Contributors
+// This software is licensed under the MIT License (MIT).
+// For more information see LICENSE or https://opensource.org/licenses/MIT
+
+package natchez
+
+import cats._
+import cats.data._
+import cats.effect.{Trace => _, _}
+
+/** Extends `Trace` with the capability to span a resource. */
+trait TraceResource[F[_]] extends Trace[F] {
+  /** Create a new span, and within it run the returned resource.  The
+    * new span becomes the ambient span for the life of the resource.
+    */
+  def spanResource(name: String): Resource[F, Unit]
+}
+
+object TraceResource {
+
+  def apply[F[_]](implicit ev: TraceResource[F]): ev.type = ev
+
+  /** A `TraceResource` instance that uses `IOLocal` internally. */
+  def ioTrace(rootSpan: Span[IO]): IO[TraceResource[IO]] =
+    IOLocal(rootSpan).map { local =>
+      new IOTrace(local) with TraceResource[IO] {
+        def spanResource(name: String): Resource[IO, Unit] =
+          Resource.eval(local.get).flatMap { parent =>
+            parent.span(name).flatMap { child =>
+              Resource.make(local.set(child))(_ => local.set(parent))
+            }
+          }
+      }
+    }
+
+  object Implicits {
+
+    /**
+     * A no-op `TraceResource` implementation is freely available for
+     * any applicative effect. This lets us add a `TraceResource`
+     * constraint to most existing code without demanding anything new
+     * from the concrete effect type.
+     */
+    implicit def noop[F[_]: Applicative]: TraceResource[F] =
+      new NoopTrace[F] with TraceResource[F] {
+        def spanResource(name: String): Resource[F, Unit] = Resource.unit
+      }
+
+  }
+
+  implicit def liftKleisli[F[_], E](implicit F: MonadCancelThrow[F], trace: TraceResource[F]): TraceResource[Kleisli[F, E, *]] =
+    new KleisliETrace[F, E] with TraceResource[Kleisli[F, E, *]] {
+
+      def spanResource(name: String): Resource[Kleisli[F, E, *], Unit] =
+        trace.spanResource(name).mapK(Kleisli.liftK[F, E])
+    }
+
+  implicit def liftStateT[F[_], S](implicit F: MonadCancelThrow[F], trace: TraceResource[F]): TraceResource[StateT[F, S, *]] =
+    new StateTTrace[F, S] with TraceResource[StateT[F, S, *]] {
+
+      def spanResource(name: String): Resource[StateT[F, S, *], Unit] =
+        trace.spanResource(name).mapK(StateT.liftK[F, S])
+    }
+
+  implicit def liftEitherT[F[_], E](implicit F: MonadCancelThrow[F], trace: TraceResource[F]): TraceResource[EitherT[F, E, *]] =
+    new EitherTTrace[F, E] with TraceResource[EitherT[F, E, *]] {
+
+      def spanResource(name: String): Resource[EitherT[F, E, *], Unit] =
+        trace.spanResource(name).mapK(EitherT.liftK[F, E])
+    }
+
+  implicit def liftOptionT[F[_]](implicit F: MonadCancelThrow[F], trace: TraceResource[F]): TraceResource[OptionT[F, *]] =
+    new OptionTTrace[F] with TraceResource[OptionT[F, *]] {
+
+      def spanResource(name: String): Resource[OptionT[F, *], Unit] =
+        trace.spanResource(name).mapK(OptionT.liftK[F])
+    }
+}