Opentelemetry zio integrate log annotations to baggage (#662)

* Integrate log annotations into baggage API * Fix README * Fix examples compilation
zio · Mar 8, 2023 · 127a845 · 127a845
1 parent 7b42d02
commit 127a845
Show file tree

Hide file tree

Showing 4 changed files with 113 additions and 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
@@ -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
@@ -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
@@ -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,24 +160,46 @@ 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)
               .build()
               .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
@@ -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)
+
 }