From 127a8452e01b7e24b6d829bca66ddc3f4e53846d Mon Sep 17 00:00:00 2001 From: Michael Nedokushev Date: Wed, 8 Mar 2023 15:35:39 +0000 Subject: [PATCH] Opentelemetry zio integrate log annotations to baggage (#662) * Integrate log annotations into baggage API * Fix README * Fix examples compilation --- .../opentelemetry/example/BackendApp.scala | 2 +- .../opentelemetry/example/ProxyApp.scala | 2 +- .../opentelemetry/baggage/Baggage.scala | 59 +++++++++++----- .../opentelemetry/baggage/BaggageTest.scala | 70 ++++++++++++++++++- 4 files changed, 113 insertions(+), 20 deletions(-) diff --git a/opentelemetry-example/src/main/scala/zio/telemetry/opentelemetry/example/BackendApp.scala b/opentelemetry-example/src/main/scala/zio/telemetry/opentelemetry/example/BackendApp.scala index 209d7c33..92504b60 100644 --- a/opentelemetry-example/src/main/scala/zio/telemetry/opentelemetry/example/BackendApp.scala +++ b/opentelemetry-example/src/main/scala/zio/telemetry/opentelemetry/example/BackendApp.scala @@ -21,7 +21,7 @@ object BackendApp extends ZIOAppDefault { BackendHttpServer.live, BackendHttpApp.live, Tracing.live, - Baggage.live, + Baggage.live(), ContextStorage.fiberRef, JaegerTracer.live ) diff --git a/opentelemetry-example/src/main/scala/zio/telemetry/opentelemetry/example/ProxyApp.scala b/opentelemetry-example/src/main/scala/zio/telemetry/opentelemetry/example/ProxyApp.scala index a65aaa0e..3842fc8d 100644 --- a/opentelemetry-example/src/main/scala/zio/telemetry/opentelemetry/example/ProxyApp.scala +++ b/opentelemetry-example/src/main/scala/zio/telemetry/opentelemetry/example/ProxyApp.scala @@ -29,7 +29,7 @@ object ProxyApp extends ZIOAppDefault { ProxyHttpServer.live, ProxyHttpApp.live, Tracing.live, - Baggage.live, + Baggage.live(), ContextStorage.fiberRef, JaegerTracer.live ) diff --git a/opentelemetry/src/main/scala/zio/telemetry/opentelemetry/baggage/Baggage.scala b/opentelemetry/src/main/scala/zio/telemetry/opentelemetry/baggage/Baggage.scala index 8557bf80..08f76bea 100644 --- a/opentelemetry/src/main/scala/zio/telemetry/opentelemetry/baggage/Baggage.scala +++ b/opentelemetry/src/main/scala/zio/telemetry/opentelemetry/baggage/Baggage.scala @@ -114,39 +114,44 @@ trait Baggage { self => object Baggage { - def live: URLayer[ContextStorage, Baggage] = + def live(logAnnotated: Boolean = false): URLayer[ContextStorage, Baggage] = ZLayer { for { ctxStorage <- ZIO.service[ContextStorage] } yield new Baggage { self => override def getCurrentBaggage(implicit trace: Trace): UIO[Baggaje] = - getCurrentContext.map(Baggaje.fromContext) + injectLogAnnotations *> + getCurrentContext.map(Baggaje.fromContext) override def get(name: String)(implicit trace: Trace): UIO[Option[String]] = - getCurrentBaggage.map(baggage => Option(baggage.getEntryValue(name))) + injectLogAnnotations *> + getCurrentBaggage.map(baggage => Option(baggage.getEntryValue(name))) override def getAll(implicit trace: Trace): UIO[Map[String, String]] = - getCurrentBaggage.map(_.asMap().asScala.toMap.map { case (k, v) => k -> v.getValue }) + injectLogAnnotations *> + getCurrentBaggage.map(_.asMap().asScala.toMap.map { case (k, v) => k -> v.getValue }) override def getAllWithMetadata(implicit trace: Trace): UIO[Map[String, (String, String)]] = - getCurrentBaggage.map( - _.asMap().asScala.toMap.map { case (k, v) => (k, (v.getValue, v.getMetadata.getValue)) } - ) + injectLogAnnotations *> + getCurrentBaggage.map( + _.asMap().asScala.toMap.map { case (k, v) => (k, (v.getValue, v.getMetadata.getValue)) } + ) override def set(name: String, value: String)(implicit trace: Trace): UIO[Unit] = - modifyBuilder(_.put(name, value)).unit + injectLogAnnotations *> modifyBuilder(_.put(name, value)).unit override def setWithMetadata(name: String, value: String, metadata: String)(implicit trace: Trace): UIO[Unit] = - modifyBuilder(_.put(name, value, BaggageEntryMetadata.create(metadata))).unit + injectLogAnnotations *> modifyBuilder(_.put(name, value, BaggageEntryMetadata.create(metadata))).unit override def remove(name: String)(implicit trace: Trace): UIO[Unit] = - modifyBuilder(_.remove(name)).unit + injectLogAnnotations *> modifyBuilder(_.remove(name)).unit override def inject[C]( propagator: BaggagePropagator, carrier: OutgoingContextCarrier[C] )(implicit trace: Trace): UIO[Unit] = for { + _ <- injectLogAnnotations ctx <- getCurrentContext _ <- ZIO.succeed(propagator.instance.inject(ctx, carrier.kernel, carrier)) } yield () @@ -155,16 +160,14 @@ object Baggage { propagator: BaggagePropagator, carrier: IncomingContextCarrier[C] )(implicit trace: Trace): UIO[Unit] = - ZIO.uninterruptible { - modifyContext(ctx => propagator.instance.extract(ctx, carrier.kernel, carrier)).unit - } + injectLogAnnotations *> + ZIO.uninterruptible { + modifyContext(ctx => propagator.instance.extract(ctx, carrier.kernel, carrier)).unit + } private def getCurrentContext(implicit trace: Trace): UIO[Context] = ctxStorage.get - private def modifyContext(body: Context => Context)(implicit trace: Trace): UIO[Context] = - ctxStorage.updateAndGet(body) - private def modifyBuilder(body: BaggageBuilder => BaggageBuilder)(implicit trace: Trace): UIO[Context] = modifyContext { ctx => body(Baggaje.fromContext(ctx).toBuilder) @@ -172,7 +175,31 @@ object Baggage { .storeInContext(ctx) } + private def modifyContext(body: Context => Context)(implicit trace: Trace): UIO[Context] = + ctxStorage.updateAndGet(body) + + private def injectLogAnnotations(implicit trace: Trace): UIO[Unit] = + ZIO + .when(logAnnotated) { + for { + annotations <- ZIO.logAnnotations + annotationsWithMetadata = annotations.map { case (k, v) => + (k, (v, BaggageEntryMetadata.create("zio log annotation"))) + } + current <- getCurrentContext + .map(Baggaje.fromContext) + .map(_.asMap().asScala.toMap.map { case (k, v) => (k, (v.getValue, v.getMetadata)) }) + _ <- modifyBuilder { builder => + (annotationsWithMetadata ++ current).foreach { case (k, (v, m)) => builder.put(k, v, m) } + builder + } + } yield () + } + .unit } } + def logAnnotated: URLayer[ContextStorage, Baggage] = + live(logAnnotated = true) + } diff --git a/opentelemetry/src/test/scala/zio/telemetry/opentelemetry/baggage/BaggageTest.scala b/opentelemetry/src/test/scala/zio/telemetry/opentelemetry/baggage/BaggageTest.scala index c790e929..d05ec159 100644 --- a/opentelemetry/src/test/scala/zio/telemetry/opentelemetry/baggage/BaggageTest.scala +++ b/opentelemetry/src/test/scala/zio/telemetry/opentelemetry/baggage/BaggageTest.scala @@ -11,13 +11,17 @@ import scala.collection.mutable object BaggageTest extends ZIOSpecDefault { def baggageLayer: ULayer[Baggage] = - ContextStorage.fiberRef >>> Baggage.live + ContextStorage.fiberRef >>> Baggage.live() + + def logAnnotatedBaggageLayer: ULayer[Baggage] = + (ContextStorage.fiberRef >>> Baggage.logAnnotated) def spec = suite("zio opentelemetry")( suite("Baggage")( operationsSpec, - propagationSpec + propagationSpec, + logAnnotatedSpec ) ) @@ -82,4 +86,66 @@ object BaggageTest extends ZIOSpecDefault { } ) + private def logAnnotatedSpec = + suite("log annotated")( + test("get") { + ZIO.serviceWithZIO[Baggage] { baggage => + ZIO.logAnnotate("zio", "annotation") { + for { + result <- baggage.get("zio") + } yield assert(result)(isSome(equalTo("annotation"))) + } + } + }, + test("getAll") { + ZIO.serviceWithZIO[Baggage] { baggage => + ZIO.logAnnotate(LogAnnotation("foo", "bar"), LogAnnotation("dog", "fox")) { + for { + result <- baggage.getAll + } yield assert(result)(equalTo(Map("foo" -> "bar", "dog" -> "fox"))) + } + } + }, + test("set overrides a value of a key taken from log annotations") { + ZIO.serviceWithZIO[Baggage] { baggage => + ZIO.logAnnotate(LogAnnotation("foo", "bar"), LogAnnotation("dog", "fox")) { + for { + _ <- baggage.set("some", "thing") + _ <- baggage.set("foo", "bark") + result <- baggage.getAll + } yield assert(result)(equalTo(Map("foo" -> "bark", "dog" -> "fox", "some" -> "thing"))) + } + } + }, + test("remove doesn't work for keys provided by log annotations") { + ZIO.serviceWithZIO[Baggage] { baggage => + ZIO.logAnnotate(LogAnnotation("foo", "bar")) { + for { + _ <- baggage.remove("foo") + result <- baggage.getAll + } yield assert(result)(equalTo(Map("foo" -> "bar"))) + } + } + }, + test("getAllWithMetadata returns a metadata provided by log annotations") { + ZIO.serviceWithZIO[Baggage] { baggage => + ZIO.logAnnotate(LogAnnotation("foo", "bar")) { + for { + result <- baggage.getAllWithMetadata + } yield assert(result)(equalTo(Map("foo" -> ("bar" -> "zio log annotation")))) + } + } + }, + test("setWithMetadata overrides a value with metadata taken from log annotations") { + ZIO.serviceWithZIO[Baggage] { baggage => + ZIO.logAnnotate(LogAnnotation("foo", "bar")) { + for { + _ <- baggage.setWithMetadata("foo", "bar", "baz") + result <- baggage.getAllWithMetadata + } yield assert(result)(equalTo(Map("foo" -> ("bar" -> "baz")))) + } + } + } + ).provideLayer(logAnnotatedBaggageLayer) + }