Backport OTEL autoinstrumentation from series/2.x

opentelemetry/src/main/scala/zio/telemetry/opentelemetry/ContextStorage.scala

@@ -0,0 +1,62 @@
+package zio.telemetry.opentelemetry
+
+import io.opentelemetry.context.Context
+import zio.{ FiberRef, UIO, ZIO }
+
+private[opentelemetry] trait ContextStorage {
+
+  def get: UIO[Context]
+
+  def set(context: Context): UIO[Unit]
+
+  def getAndSet(context: Context): UIO[Context]
+
+  def updateAndGet(f: Context => Context): UIO[Context]
+
+  def locally[R, E, A](context: Context)(effect: ZIO[R, E, A]): ZIO[R, E, A]
+
+}
+
+private[opentelemetry] object ContextStorage {
+
+  /**
+   * Use provided [[FiberRef]] as a [[ContextStorage]].
+   */
+  def fiberRef(ref: FiberRef[Context]): ContextStorage = new ContextStorage {
+    override def get: UIO[Context]                                                      = ref.get
+    override def set(context: Context): UIO[Unit]                                       = ref.set(context)
+    override def getAndSet(context: Context): UIO[Context]                              = ref.getAndSet(context)
+    override def updateAndGet(f: Context => Context): UIO[Context]                      = ref.updateAndGet(f)
+    override def locally[R, E, A](context: Context)(effect: ZIO[R, E, A]): ZIO[R, E, A] = ref.locally(context)(effect)
+  }
+
+  /**
+   * Use OpenTelemetry's default context storage which is backed by a [[ThreadLocal]]. This makes sense only if
+   * [[PropagatingSupervisor]] is used.
+   */
+  def threadLocal: ContextStorage = new ContextStorage {
+
+    override def get: UIO[Context] = ZIO.succeed(Context.current())
+
+    override def set(context: Context): UIO[Unit] = ZIO.succeed(context.makeCurrent()).unit
+
+    override def getAndSet(context: Context): UIO[Context] =
+      ZIO.succeed {
+        val old = Context.current()
+        val _   = context.makeCurrent()
+        old
+      }.uninterruptible
+
+    override def updateAndGet(f: Context => Context): UIO[Context] =
+      ZIO.succeed {
+        val updated = f(Context.current())
+        val _       = updated.makeCurrent()
+        updated
+      }.uninterruptible
+
+    override def locally[R, E, A](context: Context)(effect: ZIO[R, E, A]): ZIO[R, E, A] =
+      ZIO.bracket(get <* set(context))(set)(_ => effect)
+
+  }
+
+}

opentelemetry/src/main/scala/zio/telemetry/opentelemetry/PropagatingSupervisor.scala

@@ -0,0 +1,50 @@
+package zio.telemetry.opentelemetry
+
+import io.opentelemetry.api.trace.Span
+import io.opentelemetry.context.Context
+import zio.Supervisor.Propagation
+import zio._
+
+import java.util.concurrent.ConcurrentHashMap
+import scala.annotation.nowarn
+
+@nowarn("msg=discarded non-Unit value")
+final class PropagatingSupervisor extends Supervisor[Unit] {
+
+  private val storage = new ConcurrentHashMap[Fiber.Id, Span]()
+
+  override def value: UIO[Unit] = ZIO.unit
+
+  override def unsafeOnStart[R, E, A](
+    environment: R,
+    effect: ZIO[R, E, A],
+    parent: Option[Fiber.Runtime[Any, Any]],
+    fiber: Fiber.Runtime[E, A]
+  ): Propagation = {
+    val span = Span.current()
+    if (span != null) storage.put(fiber.id, span)
+    else storage.put(fiber.id, Span.fromContext(Context.root()))
+
+    Propagation.Continue
+  }
+
+  override def unsafeOnEnd[R, E, A](value: Exit[E, A], fiber: Fiber.Runtime[E, A]): Propagation = {
+    storage.remove(fiber.id)
+    Context.root().makeCurrent()
+
+    Propagation.Continue
+  }
+
+  override def unsafeOnSuspend[E, A](fiber: Fiber.Runtime[E, A]): Unit = {
+    val span = Span.current()
+    if (span != null) storage.put(fiber.id, span)
+    else storage.put(fiber.id, Span.fromContext(Context.root()))
+    Context.root().makeCurrent()
+  }
+
+  override def unsafeOnResume[E, A](fiber: Fiber.Runtime[E, A]): Unit = {
+    val span = storage.get(fiber.id)
+    if (span != null) span.makeCurrent()
+  }
+
+}

opentelemetry/src/main/scala/zio/telemetry/opentelemetry/Tracing.scala

@@ -18,7 +18,7 @@ import scala.jdk.CollectionConverters._
 object Tracing {
   trait Service {
     private[opentelemetry] def currentNanos: UIO[Long]
-    private[opentelemetry] val currentContext: FiberRef[Context]
+    private[opentelemetry] val currentContext: ContextStorage
     private[opentelemetry] def createRoot(spanName: String, spanKind: SpanKind): UManaged[Context]
     private[opentelemetry] def createChildOf(parent: Context, spanName: String, spanKind: SpanKind): UManaged[Context]
     private[opentelemetry] def createChildOfUnsafe(parent: Context, spanName: String, spanKind: SpanKind): UIO[Context]
@@ -29,7 +29,7 @@ object Tracing {
   private def currentNanos: URIO[Tracing, Long] =
     ZIO.serviceWith[Tracing.Service](_.currentNanos)
 
-  private def currentContext: URIO[Tracing, FiberRef[Context]] =
+  private def currentContext: URIO[Tracing, ContextStorage] =
     ZIO.access[Tracing](_.get.currentContext)
 
   private def createRoot(spanName: String, spanKind: SpanKind): URManaged[Tracing, Context] =
@@ -331,78 +331,87 @@ object Tracing {
     getCurrentSpan.map(_.getSpanContext())
 
   def managed(tracer: Tracer): URManaged[Clock, Service] = {
-    class Live(defaultContext: FiberRef[Context], clock: Clock.Service) extends Service {
-      private def endSpan(span: Span): UIO[Unit] = currentNanos.map(span.end(_, TimeUnit.NANOSECONDS))
-
-      def currentNanos: UIO[Long] = clock.currentTime(TimeUnit.NANOSECONDS)
-
-      def createRoot(spanName: String, spanKind: SpanKind): UManaged[Context] =
-        for {
-          nanoSeconds <- currentNanos.toManaged_
-          span        <- ZManaged.make(
-                           UIO(
-                             tracer
-                               .spanBuilder(spanName)
-                               .setNoParent()
-                               .setSpanKind(spanKind)
-                               .setStartTimestamp(nanoSeconds, TimeUnit.NANOSECONDS)
-                               .startSpan()
-                           )
-                         )(endSpan)
-        } yield span.storeInContext(Context.root())
-
-      def createChildOf(parent: Context, spanName: String, spanKind: SpanKind): UManaged[Context] =
-        for {
-          nanoSeconds <- currentNanos.toManaged_
-          span        <- ZManaged.make(
-                           UIO(
-                             tracer
-                               .spanBuilder(spanName)
-                               .setParent(parent)
-                               .setSpanKind(spanKind)
-                               .setStartTimestamp(nanoSeconds, TimeUnit.NANOSECONDS)
-                               .startSpan()
-                           )
-                         )(endSpan)
-        } yield span.storeInContext(parent)
-
-      def createChildOfUnsafe(parent: Context, spanName: String, spanKind: SpanKind): UIO[Context] =
-        for {
-          nanoSeconds <- currentNanos
-          span        <-
-            UIO(
-              tracer
-                .spanBuilder(spanName)
-                .setParent(parent)
-                .setSpanKind(spanKind)
-                .setStartTimestamp(nanoSeconds, TimeUnit.NANOSECONDS)
-                .startSpan()
-            )
-        } yield span.storeInContext(parent)
-
-      override private[opentelemetry] def end: UIO[Any] =
-        for {
-          nanos   <- currentNanos
-          context <- currentContext.get
-          span     = Span.fromContext(context)
-        } yield span.end(nanos, TimeUnit.NANOSECONDS)
-
-      override private[opentelemetry] def getTracer: UIO[Tracer] =
-        UIO.succeed(tracer)
-
-      val currentContext: FiberRef[Context] = defaultContext
-    }
-
     val tracing: URIO[Clock, Service] =
       for {
-        clock          <- ZIO.access[Clock](_.get)
-        defaultContext <- FiberRef.make[Context](Context.root())
-      } yield new Live(defaultContext, clock)
+        clock      <- ZIO.access[Clock](_.get)
+        contextRef <- FiberRef.make[Context](Context.root())
+      } yield new Live(tracer, ContextStorage.fiberRef(contextRef), clock)
 
     ZManaged.make(tracing)(_.end)
   }
 
-  def live: URLayer[Clock with Has[Tracer], Tracing] = ZLayer.fromManaged(
-    ZIO.access[Has[Tracer]](_.get).toManaged_.flatMap(managed)
-  )
+  def live: URLayer[Clock with Has[Tracer], Tracing] =
+    ZLayer.fromManaged(ZIO.service[Tracer].toManaged_.flatMap(managed))
+
+  def propagating: URLayer[Clock with Has[Tracer], Tracing] = {
+    val tracing = for {
+      clock  <- ZIO.access[Clock](_.get)
+      tracer <- ZIO.service[Tracer]
+    } yield new Live(tracer, ContextStorage.threadLocal, clock)
+
+    ZLayer.fromManaged(ZManaged.make(tracing)(_.end))
+  }
+
+  private class Live(tracer: Tracer, defaultContext: ContextStorage, clock: Clock.Service) extends Service {
+    private def endSpan(span: Span): UIO[Unit] = currentNanos.map(span.end(_, TimeUnit.NANOSECONDS))
+
+    def currentNanos: UIO[Long] = clock.currentTime(TimeUnit.NANOSECONDS)
+
+    def createRoot(spanName: String, spanKind: SpanKind): UManaged[Context] =
+      for {
+        nanoSeconds <- currentNanos.toManaged_
+        span        <- ZManaged.make(
+                         UIO(
+                           tracer
+                             .spanBuilder(spanName)
+                             .setNoParent()
+                             .setSpanKind(spanKind)
+                             .setStartTimestamp(nanoSeconds, TimeUnit.NANOSECONDS)
+                             .startSpan()
+                         )
+                       )(endSpan)
+      } yield span.storeInContext(Context.root())
+
+    def createChildOf(parent: Context, spanName: String, spanKind: SpanKind): UManaged[Context] =
+      for {
+        nanoSeconds <- currentNanos.toManaged_
+        span        <- ZManaged.make(
+                         UIO(
+                           tracer
+                             .spanBuilder(spanName)
+                             .setParent(parent)
+                             .setSpanKind(spanKind)
+                             .setStartTimestamp(nanoSeconds, TimeUnit.NANOSECONDS)
+                             .startSpan()
+                         )
+                       )(endSpan)
+      } yield span.storeInContext(parent)
+
+    def createChildOfUnsafe(parent: Context, spanName: String, spanKind: SpanKind): UIO[Context] =
+      for {
+        nanoSeconds <- currentNanos
+        span        <-
+          UIO(
+            tracer
+              .spanBuilder(spanName)
+              .setParent(parent)
+              .setSpanKind(spanKind)
+              .setStartTimestamp(nanoSeconds, TimeUnit.NANOSECONDS)
+              .startSpan()
+          )
+      } yield span.storeInContext(parent)
+
+    override private[opentelemetry] def end: UIO[Any] =
+      for {
+        nanos   <- currentNanos
+        context <- currentContext.get
+        span     = Span.fromContext(context)
+      } yield span.end(nanos, TimeUnit.NANOSECONDS)
+
+    override private[opentelemetry] def getTracer: UIO[Tracer] =
+      UIO.succeed(tracer)
+
+    val currentContext: ContextStorage = defaultContext
+  }
+
 }