From 73fc0e7be68992071c21d840c4e5e43bad4b6419 Mon Sep 17 00:00:00 2001
From: Wojciech <duxetlg@gmail.com>
Date: Tue, 21 Dec 2021 15:09:41 +0100
Subject: [PATCH] Add baggage API #338 (#462)

---
 .../zio/telemetry/opentelemetry/Tracing.scala | 19 +++++++++++++++++++
 .../opentelemetry/TracingSyntax.scala         |  3 +++
 .../telemetry/opentelemetry/TracingTest.scala |  7 +++++++
 3 files changed, 29 insertions(+)

diff --git a/opentelemetry/src/main/scala/zio/telemetry/opentelemetry/Tracing.scala b/opentelemetry/src/main/scala/zio/telemetry/opentelemetry/Tracing.scala
index 6054f98c..1bd6c9ec 100644
--- a/opentelemetry/src/main/scala/zio/telemetry/opentelemetry/Tracing.scala
+++ b/opentelemetry/src/main/scala/zio/telemetry/opentelemetry/Tracing.scala
@@ -1,5 +1,7 @@
 package zio.telemetry.opentelemetry
 
+import io.opentelemetry.api.baggage.Baggage
+
 import java.util.concurrent.TimeUnit
 import scala.concurrent.ExecutionContext
 import io.opentelemetry.api.common.{ AttributeKey, Attributes }
@@ -305,6 +307,23 @@ object Tracing {
     getCurrentSpan.map(_.setAttribute(AttributeKey.doubleArrayKey(name), v))
   }
 
+  /**
+   * Sets a baggage entry in the current context
+   */
+  def setBaggage(name: String, value: String): URIO[Tracing, Context] =
+    for {
+      contextRef <- currentContext
+      context    <- contextRef.updateAndGet(context =>
+                      Baggage.fromContext(context).toBuilder.put(name, value).build().storeInContext(context)
+                    )
+    } yield context
+
+  /**
+   * Gets the baggage from current context
+   */
+  def getCurrentBaggage: URIO[Tracing, Baggage] =
+    getCurrentContext.map(Baggage.fromContext)
+
   /**
    * Gets the current SpanContext
    */
diff --git a/opentelemetry/src/main/scala/zio/telemetry/opentelemetry/TracingSyntax.scala b/opentelemetry/src/main/scala/zio/telemetry/opentelemetry/TracingSyntax.scala
index 9875ecc0..353be2ee 100644
--- a/opentelemetry/src/main/scala/zio/telemetry/opentelemetry/TracingSyntax.scala
+++ b/opentelemetry/src/main/scala/zio/telemetry/opentelemetry/TracingSyntax.scala
@@ -91,5 +91,8 @@ object TracingSyntax {
       i3: DummyImplicit
     ): ZIO[Tracing with R, E, A] =
       effect <* Tracing.setAttribute(name, values)(i1, i2, i3)
+
+    def setBaggage(name: String, value: String): ZIO[Tracing with R, E, A] =
+      effect <* Tracing.setBaggage(name, value)
   }
 }
diff --git a/opentelemetry/src/test/scala/zio/telemetry/opentelemetry/TracingTest.scala b/opentelemetry/src/test/scala/zio/telemetry/opentelemetry/TracingTest.scala
index a22334b6..a56b7f97 100644
--- a/opentelemetry/src/test/scala/zio/telemetry/opentelemetry/TracingTest.scala
+++ b/opentelemetry/src/test/scala/zio/telemetry/opentelemetry/TracingTest.scala
@@ -297,6 +297,13 @@ object TracingTest extends DefaultRunnableSpec {
             )
             assert(tags)(equalTo(expected))
           }
+        },
+        testM("baggaging") {
+          for {
+            _         <- UIO.unit.setBaggage("some", "thing")
+            baggage   <- Tracing.getCurrentBaggage
+            entryValue = Option(baggage.getEntryValue("some"))
+          } yield assert(entryValue)(equalTo(Some("thing")))
         }
       ).provideCustomLayer(tracingMockLayer)
     )