OpenTelemetry: add Tracing.spanScoped method (#709)

zio · Jun 5, 2023 · 4c10b0c · 4c10b0c
1 parent 1cd8a7b
commit 4c10b0c
Show file tree

Hide file tree

Showing 3 changed files with 153 additions and 3 deletions.
diff --git a/opentelemetry/src/main/scala/zio/telemetry/opentelemetry/context/ContextStorage.scala b/opentelemetry/src/main/scala/zio/telemetry/opentelemetry/context/ContextStorage.scala
@@ -15,6 +15,7 @@ trait ContextStorage {
 
   def locally[R, E, A](context: Context)(zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, A]
 
+  def locallyScoped(context: Context)(implicit trace: Trace): ZIO[Scope, Nothing, Unit]
 }
 
 object ContextStorage {
@@ -45,6 +46,9 @@ object ContextStorage {
                 trace: Trace
               ): ZIO[R, E, A] =
                 ref.locally(context)(zio)
+
+              override def locallyScoped(context: Context)(implicit trace: Trace): ZIO[Scope, Nothing, Unit] =
+                ref.locallyScoped(context)
             }
           }
         }
@@ -79,6 +83,10 @@ object ContextStorage {
 
         override def locally[R, E, A](context: Context)(zio: ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, A] =
           ZIO.acquireReleaseWith(get <* set(context))(set)(_ => zio)
+
+        override def locallyScoped(context: Context)(implicit trace: Trace): ZIO[Scope, Nothing, Unit] =
+          ZIO.acquireRelease(get <* set(context))(set).unit
+
       }
     }
 

diff --git a/opentelemetry/src/main/scala/zio/telemetry/opentelemetry/tracing/Tracing.scala b/opentelemetry/src/main/scala/zio/telemetry/opentelemetry/tracing/Tracing.scala
@@ -161,6 +161,28 @@ trait Tracing { self =>
     links: Seq[SpanContext] = Seq.empty
   )(zio: => ZIO[R, E, A])(implicit trace: Trace): ZIO[R, E, A]
 
+  /**
+   * Sets the current span to be the child of the current span with name 'spanName'.
+   *
+   * Ends the span when the scope closes.
+   *
+   * @param spanName
+   *   name of the child span
+   * @param spanKind
+   *   kind of the child span
+   * @param errorMapper
+   *   error mapper
+   * @param links
+   *   spanContexts of the linked Spans.
+   */
+  def spanScoped(
+    spanName: String,
+    spanKind: SpanKind = SpanKind.INTERNAL,
+    attributes: Attributes = Attributes.empty(),
+    errorMapper: ErrorMapper[Any] = ErrorMapper.default[Any],
+    links: Seq[SpanContext] = Seq.empty
+  )(implicit trace: Trace): ZIO[Scope, Nothing, Unit]
+
   /**
    * Unsafely sets the current span to be the child of the current span with name 'spanName'.
    *
@@ -576,6 +598,27 @@ object Tracing {
               }
             }
 
+          override def spanScoped(
+            spanName: String,
+            spanKind: SpanKind,
+            attributes: Attributes,
+            errorMapper: ErrorMapper[Any] = ErrorMapper.default[Any],
+            links: Seq[SpanContext]
+          )(implicit trace: Trace): ZIO[Scope, Nothing, Unit] =
+            getCurrentContext.flatMap { old =>
+              ZIO.acquireReleaseExit {
+                for {
+                  res <- createChild(old, spanName, spanKind, attributes, links)
+                  _   <- ctxStorage.locallyScoped(res._2)
+                } yield res
+              } { case ((endSpan, ctx), exit) =>
+                (exit match {
+                  case Exit.Success(_)     => ZIO.unit
+                  case Exit.Failure(cause) => setErrorStatus(Span.fromContext(ctx), cause, errorMapper)
+                }) *> endSpan
+              }.unit
+            }
+
           override def spanUnsafe(
             spanName: String,
             spanKind: SpanKind = SpanKind.INTERNAL,

diff --git a/opentelemetry/src/test/scala/zio/telemetry/opentelemetry/tracing/TracingTest.scala b/opentelemetry/src/test/scala/zio/telemetry/opentelemetry/tracing/TracingTest.scala
@@ -43,11 +43,12 @@ object TracingTest extends ZIOSpecDefault {
     suite("zio opentelemetry")(
       suite("Tracing")(
         creationSpec,
-        spansSpec
+        spansSpec,
+        spanScopedSpec
       )
     )
 
-  private def creationSpec =
+  private val creationSpec =
     suite("creation")(
       test("live") {
         for {
@@ -57,7 +58,7 @@ object TracingTest extends ZIOSpecDefault {
       }.provideLayer(inMemoryTracerLayer)
     )
 
-  private def spansSpec =
+  private val spansSpec =
     suite("spans")(
       test("span") {
         ZIO.serviceWithZIO[Tracing] { tracing =>
@@ -435,4 +436,102 @@ object TracingTest extends ZIOSpecDefault {
       }
     ).provideLayer(tracingMockLayer)
 
+  private val spanScopedSpec =
+    suite("scoped spans")(
+      test("span") {
+        ZIO.serviceWithZIO[Tracing] { tracing =>
+          for {
+            _     <- ZIO.scoped[Any](
+                       tracing.spanScoped("Root") *> ZIO.scoped[Any](
+                         tracing.spanScoped("Child")
+                       )
+                     )
+            spans <- getFinishedSpans
+            root   = spans.find(_.getName == "Root")
+            child  = spans.find(_.getName == "Child")
+          } yield assert(root)(isSome(anything)) &&
+            assert(child)(
+              isSome(
+                hasField[SpanData, String](
+                  "parentSpanId",
+                  _.getParentSpanId,
+                  equalTo(root.get.getSpanId)
+                )
+              )
+            )
+        }
+      },
+      test("span single scope") {
+        ZIO.serviceWithZIO[Tracing] { tracing =>
+          for {
+            _     <- ZIO.scoped[Any](
+                       for {
+                         _ <- tracing.spanScoped("Root")
+                         _ <- tracing.spanScoped("Child")
+                       } yield ()
+                     )
+            spans <- getFinishedSpans
+            root   = spans.find(_.getName == "Root")
+            child  = spans.find(_.getName == "Child")
+          } yield assert(root)(isSome(anything)) &&
+            assert(child)(
+              isSome(
+                hasField[SpanData, String](
+                  "parentSpanId",
+                  _.getParentSpanId,
+                  equalTo(root.get.getSpanId)
+                )
+              )
+            )
+        }
+      },
+      test("setError") {
+        ZIO.serviceWithZIO[Tracing] { tracing =>
+          val assertStatusCodeError             =
+            hasField[SpanData, StatusCode]("statusCode", _.getStatus.getStatusCode, equalTo(StatusCode.ERROR))
+          val assertStatusDescriptionError      =
+            hasField[SpanData, String](
+              "statusDescription",
+              _.getStatus.getDescription,
+              containsString("java.lang.RuntimeException: some_error")
+            )
+          val assertRecordedExceptionAttributes = hasField[SpanData, List[(String, String)]](
+            "exceptionAttributes",
+            _.getEvents.asScala.toList
+              .flatMap(_.getAttributes.asMap().asScala.toList.map(x => x._1.getKey -> x._2.toString)),
+            hasSubset(List("exception.message" -> "some_error", "exception.type" -> "java.lang.RuntimeException"))
+          )
+          val assertion                         = assertStatusCodeError && assertRecordedExceptionAttributes && assertStatusDescriptionError
+          val errorMapper                       = ErrorMapper[Any]({ case _ => StatusCode.ERROR }, Some(_.asInstanceOf[Throwable]))
+
+          val failedEffect: ZIO[Any, Throwable, Unit] =
+            ZIO.fail(new RuntimeException("some_error")).unit
+
+          for {
+            _     <- ZIO
+                       .scoped[Any](
+                         tracing.spanScoped("Root", errorMapper = errorMapper) *> ZIO.scoped[Any](
+                           tracing.spanScoped("Child", errorMapper = errorMapper) *> failedEffect
+                         )
+                       )
+                       .ignore
+            spans <- getFinishedSpans
+            root   = spans.find(_.getName == "Root")
+            child  = spans.find(_.getName == "Child")
+          } yield assert(root)(isSome(assertion)) && assert(child)(isSome(assertion))
+        }
+      },
+      test("setAttribute") {
+        ZIO.serviceWithZIO[Tracing] { tracing =>
+          for {
+            _     <- ZIO.scoped[Any](for {
+                       _ <- tracing.spanScoped("foo")
+                       _ <- tracing.setAttribute("string", "bar")
+                     } yield ())
+            spans <- getFinishedSpans
+            tags   = spans.head.getAttributes
+          } yield assert(tags.get(AttributeKey.stringKey("string")))(equalTo("bar"))
+        }
+      }
+    ).provideLayer(tracingMockLayer)
 }