Add baggage API #338 (#462)

zio · Dec 21, 2021 · 73fc0e7 · 73fc0e7
1 parent 130d199
commit 73fc0e7
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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)
     )