Merge pull request #526 from rossabaker/spanR

Add Trace.spanR

build.sbt

@@ -1,6 +1,6 @@
 import com.typesafe.tools.mima.core._
 
-ThisBuild / tlBaseVersion := "0.1"
+ThisBuild / tlBaseVersion := "0.2"
 
 val scala212Version        = "2.12.17"
 val scala213Version        = "2.13.10"
@@ -10,6 +10,7 @@ val collectionCompatVersion = "2.8.1"
 
 val catsVersion = "2.9.0"
 val catsEffectVersion = "3.4.1"
+val fs2Version = "3.3.0"
 
 // Publishing
 
@@ -96,6 +97,7 @@ lazy val core = crossProject(JSPlatform, JVMPlatform, NativePlatform)
       "org.typelevel" %%% "cats-core"   % catsVersion,
       "org.typelevel" %%% "cats-effect-kernel" % catsEffectVersion,
       "org.typelevel" %%% "cats-effect" % catsEffectVersion,
+      "co.fs2"        %%% "fs2-io"      % fs2Version,
     )
   )
   .nativeSettings(commonNativeSettings)
@@ -302,7 +304,7 @@ lazy val xray = crossProject(JSPlatform, JVMPlatform)
     description := "AWS X-Ray bindings implementation",
     libraryDependencies ++= Seq(
       "io.circe"          %%% "circe-core"      % "0.14.3",
-      "co.fs2"            %%% "fs2-io"          % "3.3.0",
+      "co.fs2"            %%% "fs2-io"          % fs2Version,
       "com.comcast"       %%% "ip4s-core"       % "3.2.0",
       "org.scodec"        %%% "scodec-bits"     % "1.1.34"
     )

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

@@ -5,9 +5,11 @@
 package natchez
 
 import cats._
+import cats.arrow.FunctionK
 import cats.data._
 import cats.effect._
 import cats.syntax.all._
+import fs2.Stream
 import java.net.URI
 
 /** A tracing effect, which always has a current span. */
@@ -22,6 +24,9 @@ trait Trace[F[_]] {
    */
   def kernel: F[Kernel]
 
+  /** Creates a new span, and within it acquires and releases the spanR `r`. */
+  def spanR(name: String): Resource[F, F ~> F]
+
   /** Create a new span, and within it run the continuation `k`. */
   def span[A](name: String)(k: F[A]): F[A]
 
@@ -54,13 +59,18 @@ object Trace {
         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 spanR(name: String): Resource[IO, IO ~> IO] =
+          for {
+            parent <- Resource.eval(local.get)
+            child <- parent.span(name)
+          } yield new (IO ~> IO) {
+            def apply[A](fa: IO[A]): IO[A] =
+              local.set(child).bracket(_ => fa)(_ => local.set(parent))
           }
 
+        def span[A](name: String)(k: IO[A]): IO[A] =
+          spanR(name).surround(k)
+
         def traceId: IO[Option[String]] =
           local.get.flatMap(_.traceId)
 
@@ -82,6 +92,8 @@ object Trace {
         final val void = ().pure[F]
         val kernel: F[Kernel] = Kernel(Map.empty).pure[F]
         def put(fields: (String, TraceValue)*): F[Unit] = void
+        def spanR(name: String): Resource[F, F ~> F] =
+          Resource.pure(FunctionK.id)
         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]
@@ -109,6 +121,15 @@ object Trace {
     def put(fields: (String, TraceValue)*): Kleisli[F, Span[F], Unit] =
       Kleisli(_.put(fields: _*))
 
+    def spanR(name: String): Resource[Kleisli[F, Span[F], *], Kleisli[F, Span[F], *] ~> Kleisli[F, Span[F], *]] =
+      Resource(Kleisli((span: Span[F]) => span.span(name).allocated.map {
+        case (child, release) =>
+          new (Kleisli[F, Span[F], *] ~> Kleisli[F, Span[F], *]) {
+            def apply[A](fa: Kleisli[F, Span[F], A]): Kleisli[F, Span[F], A] =
+              fa.local(_ => child)
+          } -> Kleisli.liftF[F, Span[F], Unit](release)
+      }))
+
     def span[A](name: String)(k: Kleisli[F, Span[F], A]): Kleisli[F,Span[F],A] =
       Kleisli(_.span(name).use(k.run))
 
@@ -121,6 +142,15 @@ object Trace {
         def put(fields: (String, TraceValue)*): Kleisli[F,E,Unit] =
           Kleisli(e => f(e).put(fields: _*))
 
+        def spanR(name: String): Resource[Kleisli[F, E, *], Kleisli[F, E, *] ~> Kleisli[F, E, *]] =
+          Resource(Kleisli((e: E) => f(e).span(name).allocated.map {
+            case (child, release) =>
+              new (Kleisli[F, E, *] ~> Kleisli[F, E, *]) {
+                def apply[A](fa: Kleisli[F, E, A]): Kleisli[F, E, A] =
+                  fa.local(_ => g(e, child))
+              } -> Kleisli.liftF[F, E, Unit](release)
+          }))
+
         def span[A](name: String)(k: Kleisli[F, E, A]): Kleisli[F, E, A] =
           Kleisli(e => f(e).span(name).use(s => k.run(g(e, s))))
 
@@ -140,7 +170,7 @@ object Trace {
 
   }
 
-  implicit def liftKleisli[F[_], E](implicit trace: Trace[F]): Trace[Kleisli[F, E, *]] =
+  implicit def liftKleisli[F[_]: MonadCancelThrow, E](implicit trace: Trace[F]): Trace[Kleisli[F, E, *]] =
     new Trace[Kleisli[F, E, *]] {
 
       def put(fields: (String, TraceValue)*): Kleisli[F, E, Unit] =
@@ -149,6 +179,16 @@ object Trace {
       def kernel: Kleisli[F, E, Kernel] =
         Kleisli.liftF(trace.kernel)
 
+      def spanR(name: String): Resource[Kleisli[F, E, *], Kleisli[F, E, *] ~> Kleisli[F, E, *]] =
+        Resource(
+          Kleisli((e: E) =>
+            trace.spanR(name).allocated.map { case (f, release) =>
+              f.compose(Kleisli.applyK(e)).andThen(Kleisli.liftK[F, E]) ->
+              Kleisli.liftF[F, E, Unit](f(release))
+            }
+          )
+        )
+
       def span[A](name: String)(k: Kleisli[F, E, A]): Kleisli[F, E, A] =
         Kleisli(e => trace.span[A](name)(k.run(e)))
 
@@ -159,15 +199,27 @@ object Trace {
         Kleisli.liftF(trace.traceUri)
     }
 
-  implicit def liftStateT[F[_]: Monad, S](implicit trace: Trace[F]): Trace[StateT[F, S, *]] =
+  implicit def liftStateT[F[_]: MonadCancelThrow, S](implicit trace: Trace[F]): Trace[StateT[F, S, *]] =
     new Trace[StateT[F, S, *]] {
-
       def put(fields: (String, TraceValue)*): StateT[F, S, Unit] =
         StateT.liftF(trace.put(fields: _*))
 
       def kernel: StateT[F, S, Kernel] =
         StateT.liftF(trace.kernel)
 
+      def spanR(name: String): Resource[StateT[F, S, *], StateT[F, S, *] ~> StateT[F, S, *]] =
+        Resource(
+          StateT.liftF(
+            trace.spanR(name).allocated.map { case (f, release) =>
+              new (StateT[F, S, *] ~> StateT[F, S, *]) {
+                def apply[A](fa: StateT[F, S, A]): StateT[F, S, A] =
+                  StateT.applyF(f(fa.runF))
+              } ->
+              StateT.liftF[F, S, Unit](f(release))
+            }
+          )
+        )
+
       def span[A](name: String)(k: StateT[F, S, A]): StateT[F, S, A] =
         StateT(s => trace.span[(S, A)](name)(k.run(s)))
 
@@ -178,7 +230,7 @@ object Trace {
         StateT.liftF(trace.traceUri)
     }
 
-  implicit def liftEitherT[F[_]: Functor, E](implicit trace: Trace[F]): Trace[EitherT[F, E, *]] =
+  implicit def liftEitherT[F[_]: MonadCancelThrow, E](implicit trace: Trace[F]): Trace[EitherT[F, E, *]] =
     new Trace[EitherT[F, E, *]] {
 
       def put(fields: (String, TraceValue)*): EitherT[F, E, Unit] =
@@ -187,6 +239,19 @@ object Trace {
       def kernel: EitherT[F, E, Kernel] =
         EitherT.liftF(trace.kernel)
 
+      def spanR(name: String): Resource[EitherT[F, E, *], EitherT[F, E, *] ~> EitherT[F, E, *]] =
+        Resource(
+          EitherT.liftF(
+            trace.spanR(name).allocated.map { case (f, release) =>
+              new (EitherT[F, E, *] ~> EitherT[F, E, *]) {
+                def apply[A](fa: EitherT[F, E, A]): EitherT[F, E, A] =
+                  EitherT(f(fa.value))
+              } ->
+              EitherT.liftF[F, E, Unit](f(release))
+            }
+          )
+        )
+
       def span[A](name: String)(k: EitherT[F, E, A]): EitherT[F, E, A] =
         EitherT(trace.span(name)(k.value))
 
@@ -197,7 +262,7 @@ object Trace {
         EitherT.liftF(trace.traceUri)
     }
 
-  implicit def liftOptionT[F[_]: Functor](implicit trace: Trace[F]): Trace[OptionT[F, *]] =
+  implicit def liftOptionT[F[_]: MonadCancelThrow](implicit trace: Trace[F]): Trace[OptionT[F, *]] =
     new Trace[OptionT[F, *]] {
 
       def put(fields: (String, TraceValue)*): OptionT[F, Unit] =
@@ -206,6 +271,19 @@ object Trace {
       def kernel: OptionT[F, Kernel] =
         OptionT.liftF(trace.kernel)
 
+      def spanR(name: String): Resource[OptionT[F, *], OptionT[F, *] ~> OptionT[F, *]] =
+        Resource(
+          OptionT.liftF(
+            trace.spanR(name).allocated.map { case (f, release) =>
+              new (OptionT[F, *] ~> OptionT[F, *]) {
+                def apply[A](fa: OptionT[F, A]): OptionT[F, A] =
+                  OptionT(f(fa.value))
+              } ->
+              OptionT.liftF[F, Unit](f(release))
+            }
+          )
+        )
+
       def span[A](name: String)(k: OptionT[F, A]): OptionT[F, A] =
         OptionT(trace.span(name)(k.value))
 
@@ -216,7 +294,7 @@ object Trace {
         OptionT.liftF(trace.traceUri)
     }
 
-  implicit def liftNested[F[_]: Functor, G[_]: Applicative](implicit trace: Trace[F]): Trace[Nested[F, G, *]] =
+  implicit def liftNested[F[_]: MonadCancelThrow, G[_]: Applicative](implicit trace: Trace[F], FG: MonadCancelThrow[Nested[F, G, *]]): Trace[Nested[F, G, *]] =
     new Trace[Nested[F, G, *]] {
 
       def put(fields: (String, TraceValue)*): Nested[F, G, Unit] =
@@ -225,6 +303,19 @@ object Trace {
       def kernel: Nested[F, G, Kernel] =
         trace.kernel.map(_.pure[G]).nested
 
+      def spanR(name: String): Resource[Nested[F, G, *], Nested[F, G, *] ~> Nested[F, G, *]] =
+        Resource(
+          Nested(
+            trace.spanR(name).allocated.map { case (f, release) => (
+              new (Nested[F, G, *] ~> Nested[F, G, *]) {
+                def apply[A](fa: Nested[F, G, A]): Nested[F, G, A] =
+                   Nested(f(fa.value))
+              } ->
+              Nested(f(release).map(_.pure[G]))
+            ).pure[G] }
+          )
+        )
+
       def span[A](name: String)(k: Nested[F, G, A]): Nested[F, G, A] =
         trace.span(name)(k.value).nested
 
@@ -234,4 +325,70 @@ object Trace {
       def traceUri: Nested[F, G, Option[URI]] =
         trace.traceUri.map(_.pure[G]).nested
     }
-}
+
+  implicit def liftResource[F[_]: MonadCancelThrow](implicit trace: Trace[F]): Trace[Resource[F, *]] =
+    new Trace[Resource[F, *]] {
+      def put(fields: (String, TraceValue)*): Resource[F, Unit] =
+        Resource.eval(trace.put(fields: _*))
+
+      def kernel: Resource[F, Kernel] =
+        Resource.eval(trace.kernel)
+
+      def spanR(name: String): Resource[Resource[F, *], Resource[F, *] ~> Resource[F, *]] =
+        Resource(
+          Resource.eval(
+            trace.spanR(name).allocated.map { case (f, release) =>
+              new (Resource[F, *] ~> Resource[F, *]) {
+                def apply[A](fa: Resource[F, A]): Resource[F, A] =
+                  fa.mapK(f)
+              } ->
+              Resource.eval[F, Unit](f(release))
+            }
+          )
+        )
+
+      def span[A](name: String)(k: Resource[F, A]): Resource[F, A] =
+        trace.spanR(name).flatMap { f =>
+          Resource(f(k.allocated).map { case (a, release) =>
+            a -> f(release)
+          })
+        }
+
+      def traceId: Resource[F, Option[String]] =
+        Resource.eval(trace.traceId)
+
+      def traceUri: Resource[F, Option[URI]] =
+        Resource.eval(trace.traceUri)
+    }
+
+  implicit def liftStream[F[_]: MonadCancelThrow](implicit trace: Trace[F]): Trace[Stream[F, *]] =
+    new Trace[Stream[F, *]] {
+      def put(fields: (String, TraceValue)*): Stream[F, Unit] =
+        Stream.eval(trace.put(fields: _*))
+
+      def kernel: Stream[F, Kernel] =
+        Stream.eval(trace.kernel)
+
+      def spanR(name: String): Resource[Stream[F, *], Stream[F, *] ~> Stream[F, *]] =
+        Resource(
+          Stream.eval(
+            trace.spanR(name).allocated.map { case (f, release) =>
+              new (Stream[F, *] ~> Stream[F, *]) {
+                def apply[A](fa: Stream[F, A]): Stream[F, A] =
+                  fa.translate(f)
+              } ->
+              Stream.eval[F, Unit](f(release))
+            }
+          )
+        )
+
+      def span[A](name: String)(k: Stream[F, A]): Stream[F, A] =
+        Stream.resource(trace.spanR(name)).flatMap(k.translate)
+
+      def traceId: Stream[F, Option[String]] =
+        Stream.eval(trace.traceId)
+
+      def traceUri: Stream[F, Option[URI]] =
+        Stream.eval(trace.traceUri)
+    }
+}

modules/core/shared/src/test/scala/ImplicitResolutionTest.scala

@@ -13,12 +13,12 @@ object ImplicitResolutionTest {
   def kleisliTrace[F[_]](implicit ev: MonadCancel[F, Throwable]) =
     Trace[Kleisli[F, Span[F], *]]
 
-  def kleisliLiftedTrace[F[_]: Trace] =
+  def kleisliLiftedTrace[F[_]](implicit ev1: Trace[F], ev2: MonadCancel[F, Throwable]) =
     Trace[Kleisli[F, Int, *]]
 
   def kleisliKleisliTrace[F[_]](implicit ev: MonadCancel[F, Throwable]) =
     Trace[Kleisli[Kleisli[F, Span[F], *], Int, *]]
 
   def kleisliKleisliTrace2[F[_]](implicit ev: MonadCancel[F, Throwable]) =
     Trace[Kleisli[Kleisli[F, Int, *], Span[Kleisli[F, Int, *]], *]]
-}
+}

modules/mtl/shared/src/main/scala/LocalTrace.scala

@@ -6,8 +6,10 @@ package natchez
 
 package mtl
 
+import cats.~>
 import cats.mtl.Local
 import cats.effect.MonadCancel
+import cats.effect.Resource
 import cats.syntax.all._
 import java.net.URI
 
@@ -21,6 +23,15 @@ private[mtl] class LocalTrace[F[_]](local: Local[F, Span[F]])(
     def put(fields: (String, TraceValue)*): F[Unit] =
       local.ask.flatMap(_.put(fields: _*))
 
+    def spanR(name: String): Resource[F, F ~> F] =
+      Resource(local.ask.flatMap(_.span(name).allocated.map {
+        case (child, release) =>
+          new (F ~> F) {
+            def apply[A](fa: F[A]): F[A] =
+              local.scope(fa)(child)
+          } -> release
+      }))
+
     def span[A](name: String)(k: F[A]): F[A] =
       local.ask.flatMap { span =>
         span.span(name).use(local.scope(k))