Stateful tracing, revisted

build.sbt

@@ -196,6 +196,7 @@ lazy val examples = project
   .settings(
     name := "otel4s-examples",
     libraryDependencies ++= Seq(
+      "co.fs2" %% "fs2-core" % FS2Version,
       "io.opentelemetry" % "opentelemetry-exporter-otlp" % OpenTelemetryVersion,
       "io.opentelemetry" % "opentelemetry-sdk" % OpenTelemetryVersion,
       "io.opentelemetry" % "opentelemetry-sdk-extension-autoconfigure" % s"${OpenTelemetryVersion}-alpha"

examples/src/main/scala/TracingExample.scala

@@ -21,19 +21,23 @@ import cats.effect.Resource
 import cats.effect.std.Console
 import cats.syntax.all._
 import io.opentelemetry.api.GlobalOpenTelemetry
+import org.typelevel.otel4s.Attribute
 import org.typelevel.otel4s.java.OtelJava
 import org.typelevel.otel4s.trace.Tracer
 
+import scala.concurrent.duration._
+
 trait Work[F[_]] {
-  def doWork: F[Unit]
+  def doWork(i: Int): F[Unit]
 }
 
 object Work {
   def apply[F[_]: MonadCancelThrow: Tracer: Console]: Work[F] =
     new Work[F] {
-      def doWork: F[Unit] =
+      def doWork(i: Int): F[Unit] =
         Tracer[F].span("Work.DoWork").use { span =>
-          span.addEvent("Starting the work.") *>
+          span.addAttribute(Attribute("number", i.toLong)) *>
+            span.addEvent("Starting the work.") *>
             doWorkInternal *>
             span.addEvent("Finished working.")
         }
@@ -50,9 +54,23 @@ object TracingExample extends IOApp.Simple {
       .evalMap(OtelJava.forSync[IO])
       .evalMap(_.tracerProvider.tracer("Example").get)
 
-  def run: IO[Unit] = {
+  def run: IO[Unit] =
     tracerResource.use { implicit tracer: Tracer[IO] =>
-      Work[IO].doWork
-    }
-  }
+      tracer
+        .span("outer")
+        .surround(
+          IO.both(
+            tracer
+              .span("left")
+              .surround(
+                tracer.span("left-worker").use_.replicateA(10)
+              ),
+            tracer
+              .span("right")
+              .surround(
+                tracer.span("right-worker").use_.replicateA(10)
+              )
+          )
+        )
+    }.void
 }

java/trace/src/main/scala/org/typelevel/otel4s/java/trace/TraceScope.scala

@@ -18,9 +18,10 @@ package org.typelevel.otel4s.java.trace
 
 import cats.effect.IOLocal
 import cats.effect.LiftIO
+import cats.effect.Ref
 import cats.effect.Resource
 import cats.effect.Sync
-import cats.syntax.functor._
+import cats.syntax.all._
 import io.opentelemetry.api.trace.{Span => JSpan}
 import io.opentelemetry.context.{Context => JContext}
 import org.typelevel.otel4s.trace.SpanContext
@@ -52,60 +53,63 @@ private[java] object TraceScope {
   ): F[TraceScope[F]] = {
     val scopeRoot = Scope.Root(default)
 
-    IOLocal[Scope](scopeRoot).to[F].map { local =>
-      new TraceScope[F] {
-        val root: F[Scope.Root] =
-          Sync[F].pure(scopeRoot)
-
-        def current: F[Scope] =
-          local.get.to[F]
-
-        def makeScope(span: JSpan): Resource[F, Unit] =
-          for {
-            current <- Resource.eval(current)
-            _ <- createScope(nextScope(current, span))
-          } yield ()
-
-        def rootScope: Resource[F, Unit] =
-          Resource.eval(current).flatMap {
-            case Scope.Root(_) =>
-              createScope(scopeRoot)
-
-            case Scope.Span(_, _, _) =>
-              createScope(scopeRoot)
-
-            case Scope.Noop =>
-              createScope(Scope.Noop)
-          }
-
-        def noopScope: Resource[F, Unit] =
-          createScope(Scope.Noop)
-
-        private def createScope(scope: Scope): Resource[F, Unit] =
-          Resource
-            .make(local.getAndSet(scope).to[F])(p => local.set(p).to[F])
-            .void
-
-        private def nextScope(scope: Scope, span: JSpan): Scope =
-          scope match {
-            case Scope.Root(ctx) =>
-              Scope.Span(
-                ctx.`with`(span),
-                span,
-                WrappedSpanContext(span.getSpanContext)
-              )
-
-            case Scope.Span(ctx, _, _) =>
-              Scope.Span(
-                ctx.`with`(span),
-                span,
-                WrappedSpanContext(span.getSpanContext)
-              )
-
-            case Scope.Noop =>
-              Scope.Noop
-          }
-
+    Ref.of[F, Scope](scopeRoot).flatMap { ref =>
+      IOLocal[Ref[F, Scope]](ref).to[F].map { local =>
+        new TraceScope[F] {
+          val root: F[Scope.Root] =
+            Sync[F].pure(scopeRoot)
+
+          def current: F[Scope] =
+            local.get.to[F].flatMap(_.get)
+
+          def makeScope(span: JSpan): Resource[F, Unit] =
+            for {
+              current <- Resource.eval(current)
+              _ <- createScope(nextScope(current, span))
+            } yield ()
+
+          def rootScope: Resource[F, Unit] =
+            Resource.eval(current).flatMap {
+              case Scope.Root(_) =>
+                createScope(scopeRoot)
+
+              case Scope.Span(_, _, _) =>
+                createScope(scopeRoot)
+
+              case Scope.Noop =>
+                createScope(Scope.Noop)
+            }
+
+          def noopScope: Resource[F, Unit] =
+            createScope(Scope.Noop)
+
+          private def createScope(scope: Scope): Resource[F, Unit] =
+            Resource(
+              local.modify(oldRef =>
+                (Ref.unsafe[F, Scope](scope), () -> local.set(oldRef).to[F])).to[F]
+            )
+
+          private def nextScope(scope: Scope, span: JSpan): Scope =
+            scope match {
+              case Scope.Root(ctx) =>
+                Scope.Span(
+                  ctx.`with`(span),
+                  span,
+                  WrappedSpanContext(span.getSpanContext)
+                )
+
+              case Scope.Span(ctx, _, _) =>
+                Scope.Span(
+                  ctx.`with`(span),
+                  span,
+                  WrappedSpanContext(span.getSpanContext)
+                )
+
+              case Scope.Noop =>
+                Scope.Noop
+            }
+
+        }
       }
     }
   }

java/trace/src/test/scala/org/typelevel/otel4s/java/trace/TracerSuite.scala

@@ -517,6 +517,51 @@ class TracerSuite extends CatsEffectSuite {
     }
   }
 
+  test("propagate trace info over interrupted stream scopes") {
+    def expected(now: FiniteDuration) =
+      SpanNode(
+        "span",
+        now,
+        now,
+        List(
+          SpanNode("span-3", now, now, Nil),
+          SpanNode("span-2", now, now, Nil)
+        )
+      )
+
+    def flow(tracer: Tracer[IO]): Stream[IO, Unit] =
+      for {
+        span <- Stream.resource(tracer.span("span"))
+        _ <- Stream.eval(
+          tracer.currentSpanContext.assertEquals(Some(span.context))
+        )
+        span2 <- Stream.resource(tracer.span("span-2"))
+        _ <- Stream.eval(
+          tracer.currentSpanContext.assertEquals(Some(span2.context))
+        )
+        span3 <- Stream.resource(
+          tracer.spanBuilder("span-3").withParent(span.context).start
+        )
+        _ <- Stream.eval(
+          tracer.currentSpanContext.assertEquals(Some(span3.context))
+        )
+      } yield ()
+
+    TestControl.executeEmbed {
+      for {
+        now <- IO.monotonic.delayBy(1.second) // otherwise returns 0
+        sdk <- makeSdk()
+        tracer <- sdk.provider.tracer("tracer").get
+        _ <- tracer.currentSpanContext.assertEquals(None)
+        _ <- flow(tracer).interruptScope.compile.drain
+        _ <- tracer.currentSpanContext.assertEquals(None)
+        spans <- sdk.finishedSpans
+        tree <- IO.pure(SpanNode.fromSpans(spans))
+        // _ <- IO.println(tree.map(SpanNode.render).mkString("\n"))
+      } yield assertEquals(tree, List(expected(now)))
+    }
+  }
+
   private def assertIdsNotEqual(s1: Span[IO], s2: Span[IO]): Unit = {
     assertNotEquals(s1.context.traceIdHex, s2.context.traceIdHex)
     assertNotEquals(s1.context.spanIdHex, s2.context.spanIdHex)