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)
  )