From 4c10b0c59b0a13799b2333075146ad9a125994fb Mon Sep 17 00:00:00 2001 From: Ori Dagan Date: Mon, 5 Jun 2023 23:07:10 +0300 Subject: [PATCH] OpenTelemetry: add Tracing.spanScoped method (#709) --- .../context/ContextStorage.scala | 8 ++ .../opentelemetry/tracing/Tracing.scala | 43 +++++++ .../opentelemetry/tracing/TracingTest.scala | 105 +++++++++++++++++- 3 files changed, 153 insertions(+), 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 index 079489d7..9cbb13ae 100644 --- 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 index c793f8d5..42ee69dd 100644 --- 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 index fab2f85c..6fe2d2f7 100644 --- 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) }