Merge pull request #149 from iRevive/joinOrContinue

Add `Tracer#joinOrRoot` utility method
typelevel · Apr 18, 2023 · c3b88c6 · c3b88c6
2 parents af59af6 + 2040c9c
commit c3b88c6
Show file tree

Hide file tree

Showing 9 changed files with 175 additions and 28 deletions.
diff --git a/core/trace/src/main/scala/org/typelevel/otel4s/trace/Tracer.scala b/core/trace/src/main/scala/org/typelevel/otel4s/trace/Tracer.scala
@@ -76,6 +76,47 @@ trait Tracer[F[_]] extends TracerMacro[F] {
         fa
     }
 
+  /** Creates a new tracing scope if a parent can be extracted from the given
+    * `carrier`. A newly created non-root span will be a child of the extracted
+    * parent.
+    *
+    * If the context cannot be extracted from the `carrier`, the given effect
+    * `fa` will be executed within the '''root''' span.
+    *
+    * To make the propagation and extraction work, you need to configure the
+    * OpenTelemetry SDK. For example, you can use `OTEL_PROPAGATORS` environment
+    * variable. See the official
+    * [[https://opentelemetry.io/docs/reference/specification/sdk-environment-variables/#general-sdk-configuration SDK configuration guide]].
+    *
+    * ==Examples==
+    *
+    * ===Propagation via [[https://www.w3.org/TR/trace-context W3C headers]]:===
+    * {{{
+    * val w3cHeaders: Map[String, String] =
+    *   Map("traceparent" -> "00-80f198ee56343ba864fe8b2a57d3eff7-e457b5a2e4d86bd1-01")
+    *
+    * Tracer[F].joinOrRoot(w3cHeaders) {
+    *   Tracer[F].span("child").use { span => ??? } // a child of the external span
+    * }
+    * }}}
+    *
+    * ===Start a root span as a fallback:===
+    * {{{
+    * Tracer[F].span("process").surround {
+    *   Tracer[F].joinOrRoot(Map.empty) { // cannot extract the context from the empty map
+    *     Tracer[F].span("child").use { span => ??? } // a child of the new root span
+    *   }
+    * }
+    * }}}
+    *
+    * @param carrier
+    *   the carrier to extract the context from
+    *
+    * @tparam C
+    *   the type of the carrier
+    */
+  def joinOrRoot[A, C: TextMapGetter](carrier: C)(fa: F[A]): F[A]
+
   /** Creates a new root tracing scope. The parent span will not be available
     * inside. Thus, a span created inside of the scope will be a root one.
     *
@@ -161,5 +202,6 @@ object Tracer {
       def noopScope[A](fa: F[A]): F[A] = fa
       def childScope[A](parent: SpanContext)(fa: F[A]): F[A] = fa
       def spanBuilder(name: String): SpanBuilder.Aux[F, Span[F]] = builder
+      def joinOrRoot[A, C: TextMapGetter](carrier: C)(fa: F[A]): F[A] = fa
     }
 }
diff --git a/examples/src/main/scala/TracingExample.scala b/examples/src/main/scala/TracingExample.scala
@@ -22,11 +22,8 @@ import cats.effect.std.Console
 import cats.syntax.all._
 import io.opentelemetry.api.GlobalOpenTelemetry
 import org.typelevel.otel4s.Otel4s
-import org.typelevel.otel4s.TextMapPropagator
 import org.typelevel.otel4s.java.OtelJava
-import org.typelevel.otel4s.trace.SpanContext
 import org.typelevel.otel4s.trace.Tracer
-import org.typelevel.vault.Vault
 
 import scala.concurrent.duration._
 
@@ -35,13 +32,11 @@ trait Work[F[_]] {
 }
 
 object Work {
-  def apply[F[_]: Monad: Tracer: TextMapPropagator: Console]: Work[F] =
+  def apply[F[_]: Monad: Tracer: Console]: Work[F] =
     new Work[F] {
       def request(headers: Map[String, String]): F[Unit] = {
-        val vault =
-          implicitly[TextMapPropagator[F]].extract(Vault.empty, headers)
         Tracer[F].currentSpanContext.flatMap { current =>
-          Tracer[F].childOrContinue(SpanContext.fromContext(vault)) {
+          Tracer[F].joinOrRoot(headers) {
             val builder = Tracer[F].spanBuilder("Work.DoWork")
             current.fold(builder)(builder.addLink(_)).build.use { span =>
               Tracer[F].currentSpanContext
@@ -71,8 +66,6 @@ object TracingExample extends IOApp.Simple {
 
   def run: IO[Unit] = {
     globalOtel4s.use { (otel4s: Otel4s[IO]) =>
-      implicit val textMapProp: TextMapPropagator[IO] =
-        otel4s.propagators.textMapPropagator
       otel4s.tracerProvider.tracer("example").get.flatMap {
         implicit tracer: Tracer[IO] =>
           val resource: Resource[IO, Unit] =

diff --git a/java/all/src/main/scala/org/typelevel/otel4s/java/OtelJava.scala b/java/all/src/main/scala/org/typelevel/otel4s/java/OtelJava.scala
@@ -56,15 +56,16 @@ object OtelJava {
   def local[F[_]](
       jOtel: JOpenTelemetry
   )(implicit F: Async[F], L: Local[F, Vault]): Otel4s[F] = {
+    val contentPropagators = new ContextPropagatorsImpl[F](
+      jOtel.getPropagators,
+      ContextConversions.toJContext,
+      ContextConversions.fromJContext
+    )
+
     val metrics = Metrics.forAsync(jOtel)
-    val traces = Traces.local(jOtel)
+    val traces = Traces.local(jOtel, contentPropagators)
     new Otel4s[F] {
-      def propagators: ContextPropagators[F] =
-        new ContextPropagatorsImpl[F](
-          jOtel.getPropagators,
-          ContextConversions.toJContext,
-          ContextConversions.fromJContext
-        )
+      def propagators: ContextPropagators[F] = contentPropagators
       def meterProvider: MeterProvider[F] = metrics.meterProvider
       def tracerProvider: TracerProvider[F] = traces.tracerProvider
     }

diff --git a/...evel/otel4s/java/ContextConversions.scala → ...evel/otel4s/java/ContextConversions.scala b/...evel/otel4s/java/ContextConversions.scala → ...evel/otel4s/java/ContextConversions.scala
diff --git a/java/trace/src/main/scala/org/typelevel/otel4s/java/trace/TracerBuilderImpl.scala b/java/trace/src/main/scala/org/typelevel/otel4s/java/trace/TracerBuilderImpl.scala
@@ -18,10 +18,12 @@ package org.typelevel.otel4s.java.trace
 
 import cats.effect.Sync
 import io.opentelemetry.api.trace.{TracerProvider => JTracerProvider}
+import org.typelevel.otel4s.ContextPropagators
 import org.typelevel.otel4s.trace._
 
 private[java] final case class TracerBuilderImpl[F[_]: Sync](
     jTracerProvider: JTracerProvider,
+    propagators: ContextPropagators[F],
     scope: TraceScope[F],
     name: String,
     version: Option[String] = None,
@@ -38,7 +40,7 @@ private[java] final case class TracerBuilderImpl[F[_]: Sync](
     val b = jTracerProvider.tracerBuilder(name)
     version.foreach(b.setInstrumentationVersion)
     schemaUrl.foreach(b.setSchemaUrl)
-    new TracerImpl(b.build(), scope)
+    new TracerImpl(b.build(), scope, propagators)
   }
 
 }
diff --git a/java/trace/src/main/scala/org/typelevel/otel4s/java/trace/TracerImpl.scala b/java/trace/src/main/scala/org/typelevel/otel4s/java/trace/TracerImpl.scala
@@ -21,15 +21,18 @@ import cats.syntax.flatMap._
 import cats.syntax.functor._
 import io.opentelemetry.api.trace.{Span => JSpan}
 import io.opentelemetry.api.trace.{Tracer => JTracer}
-import org.typelevel.otel4s.java.trace.WrappedSpanContext
+import org.typelevel.otel4s.ContextPropagators
+import org.typelevel.otel4s.TextMapGetter
 import org.typelevel.otel4s.trace.Span
 import org.typelevel.otel4s.trace.SpanBuilder
 import org.typelevel.otel4s.trace.SpanContext
 import org.typelevel.otel4s.trace.Tracer
+import org.typelevel.vault.Vault
 
 private[java] class TracerImpl[F[_]: Sync](
     jTracer: JTracer,
-    scope: TraceScope[F]
+    scope: TraceScope[F],
+    propagators: ContextPropagators[F]
 ) extends Tracer[F] {
 
   private val runner: SpanRunner[F, Span[F]] = SpanRunner.span(scope)
@@ -59,4 +62,15 @@ private[java] class TracerImpl[F[_]: Sync](
 
   def noopScope[A](fa: F[A]): F[A] =
     scope.noopScope(fa)
+
+  def joinOrRoot[A, C: TextMapGetter](carrier: C)(fa: F[A]): F[A] = {
+    val context = propagators.textMapPropagator.extract(Vault.empty, carrier)
+
+    SpanContext.fromContext(context) match {
+      case Some(parent) =>
+        childScope(parent)(fa)
+      case None =>
+        rootScope(fa)
+    }
+  }
 }
diff --git a/java/trace/src/main/scala/org/typelevel/otel4s/java/trace/TracerProviderImpl.scala b/java/trace/src/main/scala/org/typelevel/otel4s/java/trace/TracerProviderImpl.scala
@@ -19,25 +19,28 @@ package org.typelevel.otel4s.java.trace
 import cats.effect.Sync
 import cats.mtl.Local
 import io.opentelemetry.api.trace.{TracerProvider => JTracerProvider}
+import org.typelevel.otel4s.ContextPropagators
 import org.typelevel.otel4s.trace.TracerBuilder
 import org.typelevel.otel4s.trace.TracerProvider
 import org.typelevel.vault.Vault
 
 private[java] class TracerProviderImpl[F[_]: Sync](
     jTracerProvider: JTracerProvider,
+    propagators: ContextPropagators[F],
     scope: TraceScope[F]
 ) extends TracerProvider[F] {
   def tracer(name: String): TracerBuilder[F] =
-    TracerBuilderImpl(jTracerProvider, scope, name)
+    TracerBuilderImpl(jTracerProvider, propagators, scope, name)
 }
 
 private[java] object TracerProviderImpl {
 
   def local[F[_]](
-      jTracerProvider: JTracerProvider
+      jTracerProvider: JTracerProvider,
+      propagators: ContextPropagators[F]
   )(implicit F: Sync[F], L: Local[F, Vault]): TracerProvider[F] = {
     val traceScope = TraceScope.fromLocal[F]
-    new TracerProviderImpl(jTracerProvider, traceScope)
+    new TracerProviderImpl(jTracerProvider, propagators, traceScope)
   }
 
 }
diff --git a/java/trace/src/main/scala/org/typelevel/otel4s/java/trace/Traces.scala b/java/trace/src/main/scala/org/typelevel/otel4s/java/trace/Traces.scala
@@ -14,7 +14,8 @@
  * limitations under the License.
  */
 
-package org.typelevel.otel4s.java.trace
+package org.typelevel.otel4s
+package java.trace
 
 import cats.effect.IOLocal
 import cats.effect.LiftIO
@@ -32,18 +33,23 @@ trait Traces[F[_]] {
 object Traces {
 
   def local[F[_]](
-      jOtel: JOpenTelemetry
+      jOtel: JOpenTelemetry,
+      propagators: ContextPropagators[F]
   )(implicit F: Sync[F], L: Local[F, Vault]): Traces[F] = {
-    val provider = TracerProviderImpl.local(jOtel.getTracerProvider)
+    val provider =
+      TracerProviderImpl.local(jOtel.getTracerProvider, propagators)
     new Traces[F] {
       def tracerProvider: TracerProvider[F] = provider
     }
   }
 
-  def ioLocal[F[_]: LiftIO: Sync](jOtel: JOpenTelemetry): F[Traces[F]] =
+  def ioLocal[F[_]: LiftIO: Sync](
+      jOtel: JOpenTelemetry,
+      propagators: ContextPropagators[F]
+  ): F[Traces[F]] =
     IOLocal(Vault.empty)
       .map { implicit ioLocal: IOLocal[Vault] =>
-        local(jOtel)
+        local(jOtel, propagators)
       }
       .to[F]
 

diff --git a/java/trace/src/test/scala/org/typelevel/otel4s/java/trace/TracerSuite.scala b/java/trace/src/test/scala/org/typelevel/otel4s/java/trace/TracerSuite.scala
@@ -23,6 +23,10 @@ import cats.effect.testkit.TestControl
 import io.opentelemetry.api.common.{AttributeKey => JAttributeKey}
 import io.opentelemetry.api.common.Attributes
 import io.opentelemetry.api.trace.StatusCode
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator
+import io.opentelemetry.context.propagation.{
+  ContextPropagators => JContextPropagators
+}
 import io.opentelemetry.sdk.common.InstrumentationScopeInfo
 import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter
 import io.opentelemetry.sdk.testing.time.TestClock
@@ -35,6 +39,8 @@ import io.opentelemetry.sdk.trace.`export`.SimpleSpanProcessor
 import io.opentelemetry.sdk.trace.internal.data.ExceptionEventData
 import munit.CatsEffectSuite
 import org.typelevel.otel4s.Attribute
+import org.typelevel.otel4s.java.ContextConversions
+import org.typelevel.otel4s.java.ContextPropagatorsImpl
 import org.typelevel.otel4s.java.instances._
 import org.typelevel.otel4s.trace.Span
 import org.typelevel.otel4s.trace.Tracer
@@ -641,6 +647,80 @@ class TracerSuite extends CatsEffectSuite {
     }
   }
 
+  // external span does not appear in the recorded spans of the in-memory sdk
+  test("joinOrRoot: join an external span when can be extracted") {
+    val traceId = "84b54e9330faae5350f0dd8673c98146"
+    val spanId = "279fa73bc935cc05"
+
+    val headers = Map(
+      "traceparent" -> s"00-$traceId-$spanId-01"
+    )
+
+    TestControl.executeEmbed {
+      for {
+        sdk <- makeSdk()
+        tracer <- sdk.provider.get("tracer")
+        pair <- tracer.span("local").surround {
+          tracer.joinOrRoot(headers) {
+            tracer.currentSpanContext.product(
+              tracer.span("inner").use(r => IO.pure(r.context))
+            )
+          }
+        }
+        (external, inner) = pair
+      } yield {
+        assertEquals(external.map(_.traceIdHex), Some(traceId))
+        assertEquals(external.map(_.spanIdHex), Some(spanId))
+        assertEquals(inner.traceIdHex, traceId)
+      }
+    }
+  }
+
+  test(
+    "joinOrRoot: ignore an external span when cannot be extracted and start a root span"
+  ) {
+    val traceId = "84b54e9330faae5350f0dd8673c98146"
+    val spanId = "279fa73bc935cc05"
+
+    val headers = Map(
+      "some_random_header" -> s"00-$traceId-$spanId-01"
+    )
+
+    // we must always start a root span
+    def expected(now: FiniteDuration) = List(
+      SpanNode(
+        name = "inner",
+        start = now.plus(500.millis),
+        end = now.plus(500.millis).plus(200.millis),
+        children = Nil
+      ),
+      SpanNode(
+        name = "local",
+        start = now,
+        end = now.plus(500.millis).plus(200.millis),
+        children = Nil
+      )
+    )
+
+    TestControl.executeEmbed {
+      for {
+        now <- IO.monotonic.delayBy(1.second) // otherwise returns 0
+        sdk <- makeSdk()
+        tracer <- sdk.provider.get("tracer")
+        _ <- tracer.span("local").surround {
+          tracer
+            .joinOrRoot(headers) {
+              tracer.span("inner").surround(IO.sleep(200.millis))
+            }
+            .delayBy(500.millis)
+        }
+        spans <- sdk.finishedSpans
+        tree <- IO.pure(SpanNode.fromSpans(spans))
+        // _ <- IO.println(tree.map(SpanNode.render).mkString("\n"))
+      } yield assertEquals(tree, expected(now))
+    }
+  }
+
   /*
   test("propagate trace info over stream scopes") {
     def expected(now: FiniteDuration) =
@@ -706,7 +786,13 @@ class TracerSuite extends CatsEffectSuite {
       customize(builder).build()
 
     IOLocal(Vault.empty).map { implicit ioLocal: IOLocal[Vault] =>
-      val provider = TracerProviderImpl.local[IO](tracerProvider)
+      val propagators = new ContextPropagatorsImpl[IO](
+        JContextPropagators.create(W3CTraceContextPropagator.getInstance()),
+        ContextConversions.toJContext,
+        ContextConversions.fromJContext
+      )
+
+      val provider = TracerProviderImpl.local[IO](tracerProvider, propagators)
       new TracerSuite.Sdk(provider, exporter)
     }
   }