Merge pull request #1160 from iRevive/core-trace/with-current-span

core-trace: add `withCurrentSpanOrNoop` utility to `Tracer`

Author: iRevive

core/trace/src/main/scala/org/typelevel/otel4s/trace/Tracer.scala

@@ -20,6 +20,7 @@ package trace
 import cats.Applicative
 import cats.ApplicativeThrow
 import cats.effect.kernel.MonadCancelThrow
+import cats.syntax.flatMap._
 import cats.syntax.functor._
 import cats.~>
 import org.typelevel.otel4s.context.propagation.TextMapGetter
@@ -64,6 +65,15 @@ sealed trait Tracer[F[_]] extends TracerMacro[F] {
     */
   def currentSpanOrThrow: F[Span[F]]
 
+  /** Applies `f` to the current span if one exists in the local scope, or a no-op span otherwise.
+    *
+    * This is a convenience for `currentSpanOrNoop.flatMap(f)`.
+    *
+    * @see
+    *   [[currentSpanOrNoop]]
+    */
+  def withCurrentSpanOrNoop[A](f: Span[F] => F[A]): F[A]
+
   /** Creates a new [[SpanBuilder]]. The builder can be used to make a fully customized [[Span]].
     *
     * @param name
@@ -232,12 +242,13 @@ object Tracer {
   def noop[F[_]: Applicative]: Tracer[F] =
     new Tracer[F] {
       private val noopBackend = Span.Backend.noop
+      private val noopSpan = Span.fromBackend(noopBackend)
       private val builder = SpanBuilder.noop(noopBackend)
       val meta: InstrumentMeta[F] = InstrumentMeta.disabled
       val currentSpanContext: F[Option[SpanContext]] = Applicative[F].pure(None)
-      val currentSpanOrNoop: F[Span[F]] =
-        Applicative[F].pure(Span.fromBackend(noopBackend))
+      val currentSpanOrNoop: F[Span[F]] = Applicative[F].pure(noopSpan)
       def currentSpanOrThrow: F[Span[F]] = currentSpanOrNoop
+      def withCurrentSpanOrNoop[A](f: Span[F] => F[A]): F[A] = f(noopSpan)
       def rootScope[A](fa: F[A]): F[A] = fa
       def noopScope[A](fa: F[A]): F[A] = fa
       def childScope[A](parent: SpanContext)(fa: F[A]): F[A] = fa
@@ -259,6 +270,8 @@ object Tracer {
       kt.liftK(tracer.currentSpanOrNoop.map(_.liftTo[G]))
     def currentSpanOrThrow: G[Span[G]] =
       kt.liftK(tracer.currentSpanOrThrow.map(_.liftTo[G]))
+    def withCurrentSpanOrNoop[A](f: Span[G] => G[A]): G[A] =
+      currentSpanOrNoop.flatMap(f)
     def spanBuilder(name: String): SpanBuilder[G] =
       tracer.spanBuilder(name).liftTo[G]
     def childScope[A](parent: SpanContext)(ga: G[A]): G[A] =

core/trace/src/test/scala/org/typelevel/otel4s/trace/BaseTracerSuite.scala

@@ -311,6 +311,55 @@ abstract class BaseTracerSuite[Ctx, K[X] <: Key[X]](implicit
     }
   }
 
+  sdkTest("`withCurrentSpan` outside of a span (root scope)") { sdk =>
+    for {
+      tracer <- sdk.provider.get("tracer")
+      isValid <- tracer.withCurrentSpanOrNoop(span => IO.pure(span.context.isValid))
+      spans <- sdk.finishedSpans
+    } yield {
+      assert(!isValid)
+      assertEquals(spans.length, 0)
+    }
+  }
+
+  sdkTest("`withCurrentSpan` in noop scope") { sdk =>
+    for {
+      tracer <- sdk.provider.get("tracer")
+      isValid <- tracer.noopScope {
+        tracer.withCurrentSpanOrNoop(span => IO.pure(span.context.isValid))
+      }
+      spans <- sdk.finishedSpans
+    } yield {
+      assert(!isValid)
+      assertEquals(spans.length, 0)
+    }
+  }
+
+  sdkTest("`withCurrentSpan` inside a span") { sdk =>
+    def expected(now: FiniteDuration) =
+      List(SpanTree(SpanInfo("span", now, now)))
+
+    val attribute =
+      Attribute("string-attribute", "value")
+
+    TestControl.executeEmbed {
+      for {
+        now <- IO.monotonic.delayBy(1.second) // otherwise returns 0
+        tracer <- sdk.provider.get("tracer")
+        _ <- tracer.span("span").surround {
+          tracer.withCurrentSpanOrNoop { span =>
+            span.addAttribute(attribute) >> IO(assert(span.context.isValid))
+          }
+        }
+        spans <- sdk.finishedSpans
+        tree <- IO.pure(treeOf(spans))
+      } yield {
+        assertEquals(tree, expected(now))
+        assertEquals(spans.map(_.attributes), List(Attributes(attribute)))
+      }
+    }
+  }
+
   sdkTest("`currentSpanOrThrow` outside of a span (root scope)") { sdk =>
     {
       for {

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

@@ -132,6 +132,11 @@ class TracerSuite extends CatsEffectSuite {
       yield assert(!span.context.isValid)
   }
 
+  test("`withCurrentSpan` uses a no-op span when instrument is noop") {
+    val tracer = Tracer.noop[IO]
+    tracer.withCurrentSpanOrNoop(span => IO(assert(!span.context.isValid)))
+  }
+
   // utility
 
   private sealed trait BuilderOp
@@ -173,6 +178,7 @@ class TracerSuite extends CatsEffectSuite {
     def currentSpanContext: F[Option[SpanContext]] = underlying.currentSpanContext
     def currentSpanOrNoop: F[Span[F]] = underlying.currentSpanOrNoop
     def currentSpanOrThrow: F[Span[F]] = underlying.currentSpanOrThrow
+    def withCurrentSpanOrNoop[A](f: Span[F] => F[A]): F[A] = underlying.withCurrentSpanOrNoop(f)
     def childScope[A](parent: SpanContext)(fa: F[A]): F[A] = underlying.childScope(parent)(fa)
     def joinOrRoot[A, C: TextMapGetter](carrier: C)(fa: F[A]): F[A] = underlying.joinOrRoot(carrier)(fa)
     def rootScope[A](fa: F[A]): F[A] = underlying.rootScope(fa)

oteljava/trace/src/main/scala/org/typelevel/otel4s/oteljava/trace/TracerImpl.scala

@@ -65,6 +65,9 @@ private[oteljava] class TracerImpl[F[_]](
         .getOrElse(Tracer.raiseNoCurrentSpan)
     }.flatten
 
+  def withCurrentSpanOrNoop[A](f: Span[F] => F[A]): F[A] =
+    currentSpanOrNoop.flatMap(f)
+
   def spanBuilder(name: String): SpanBuilder[F] =
     SpanBuilderImpl[F](jTracer, name, meta, runner, traceScope)