Add new baggage APIs which act on currently active baggage (#5365)

* new baggage api

* added unit tests

* add TS documentation

* inject and extract using storage when propagating baggage without spans

Author: ida613

index.d.ts

@@ -138,6 +138,22 @@ interface Tracer extends opentracing.Tracer {
    * LLM Observability SDK
    */
   llmobs: tracer.llmobs.LLMObs;
+
+  /**
+   * @experimental
+   * Provide same functionality as OpenTelemetry Baggage:
+   * https://opentelemetry.io/docs/concepts/signals/baggage/
+   *
+   * Since the equivalent of OTel Context is implicit in dd-trace-js,
+   * these APIs act on the currently active baggage
+   *
+   * Work with storage('baggage'), therefore do not follow the same continuity as other APIs
+   */
+  setBaggageItem (key: string, value: string): Record<string, string>;
+  getBaggageItem (key: string): string | undefined;
+  getAllBaggageItems (): Record<string, string>;
+  removeBaggageItem (key: string): Record<string, string>;
+  removeAllBaggageItems (): Record<string, string>;
 }
 
 // left out of the namespace, so it

packages/dd-trace/src/baggage.js

@@ -0,0 +1,36 @@
+'use strict'
+
+const { storage } = require('../../datadog-core')
+const baggageStorage = storage('baggage')
+
+function setBaggageItem (key, value) {
+  storage('baggage').enterWith({ ...baggageStorage.getStore(), [key]: value })
+  return storage('baggage').getStore()
+}
+
+function getBaggageItem (key) {
+  return storage('baggage').getStore()?.[key]
+}
+
+function getAllBaggageItems () {
+  return storage('baggage').getStore()
+}
+
+function removeBaggageItem (keyToRemove) {
+  const { [keyToRemove]: _, ...newBaggage } = storage('baggage').getStore()
+  storage('baggage').enterWith(newBaggage)
+  return newBaggage
+}
+
+function removeAllBaggageItems () {
+  storage('baggage').enterWith({})
+  return storage('baggage').getStore()
+}
+
+module.exports = {
+  setBaggageItem,
+  getBaggageItem,
+  getAllBaggageItems,
+  removeBaggageItem,
+  removeAllBaggageItems
+}

packages/dd-trace/src/noop/proxy.js

@@ -16,6 +16,11 @@ class NoopProxy {
     this.appsec = noopAppsec
     this.dogstatsd = noopDogStatsDClient
     this.llmobs = noopLLMObs
+    this.setBaggageItem = () => {}
+    this.getBaggageItem = () => {}
+    this.getAllBaggageItems = () => {}
+    this.removeBaggageItem = () => {}
+    this.removeAllBaggageItems = () => {}
   }
 
   init () {

packages/dd-trace/src/opentelemetry/context_manager.js

@@ -12,6 +12,7 @@ class ContextManager {
     this._store = storage('opentelemetry')
   }
 
+  // converts dd to otel
   active () {
     const activeSpan = tracer.scope().active()
     const store = this._store.getStore()
@@ -54,6 +55,7 @@ class ContextManager {
       : wrappedContext
   }
 
+  // converts otel to dd
   with (context, fn, thisArg, ...args) {
     const span = trace.getSpan(context)
     const ddScope = tracer.scope()

packages/dd-trace/src/opentracing/propagation/text_map.js

@@ -7,6 +7,7 @@ const log = require('../../log')
 const TraceState = require('./tracestate')
 const tags = require('../../../../../ext/tags')
 const { channel } = require('dc-polyfill')
+const { setBaggageItem, getAllBaggageItems, removeAllBaggageItems } = require('../../baggage')
 
 const { AUTO_KEEP, AUTO_REJECT, USER_KEEP } = require('../../../../../ext/priority')
 
@@ -55,9 +56,10 @@ class TextMapPropagator {
   }
 
   inject (spanContext, carrier) {
-    if (!spanContext || !carrier) return
-
+    if (!carrier) return
     this._injectBaggageItems(spanContext, carrier)
+    if (!spanContext) return
+
     this._injectDatadog(spanContext, carrier)
     this._injectB3MultipleHeaders(spanContext, carrier)
     this._injectB3SingleHeader(spanContext, carrier)
@@ -125,7 +127,7 @@ class TextMapPropagator {
 
   _injectBaggageItems (spanContext, carrier) {
     if (this._config.legacyBaggageEnabled) {
-      spanContext._baggageItems && Object.keys(spanContext._baggageItems).forEach(key => {
+      spanContext?._baggageItems && Object.keys(spanContext._baggageItems).forEach(key => {
         carrier[baggagePrefix + key] = String(spanContext._baggageItems[key])
       })
     }
@@ -134,7 +136,9 @@ class TextMapPropagator {
       let itemCounter = 0
       let byteCounter = 0
 
-      for (const [key, value] of Object.entries(spanContext._baggageItems)) {
+      const baggageItems = spanContext ? spanContext._baggageItems : getAllBaggageItems()
+      if (!baggageItems) return
+      for (const [key, value] of Object.entries(baggageItems)) {
         const item = `${this._encodeOtelBaggageKey(String(key).trim())}=${encodeURIComponent(String(value).trim())},`
         itemCounter += 1
         byteCounter += item.length
@@ -627,23 +631,27 @@ class TextMapPropagator {
   _extractBaggageItems (carrier, spanContext) {
     if (!this._hasPropagationStyle('extract', 'baggage')) return
     if (!carrier || !carrier.baggage) return
-    if (!spanContext) return
+    if (!spanContext) removeAllBaggageItems()
     const baggages = carrier.baggage.split(',')
     for (const keyValue of baggages) {
       if (!keyValue.includes('=')) {
-        spanContext._baggageItems = {}
+        if (spanContext) spanContext._baggageItems = {}
         return
       }
       let [key, value] = keyValue.split('=')
       key = this._decodeOtelBaggageKey(key.trim())
       value = decodeURIComponent(value.trim())
       if (!key || !value) {
-        spanContext._baggageItems = {}
+        if (spanContext) spanContext._baggageItems = {}
         return
       }
       // the current code assumes precedence of ot-baggage- (legacy opentracing baggage) over baggage
-      if (key in spanContext._baggageItems) return
-      spanContext._baggageItems[key] = value
+      if (spanContext) {
+        if (key in spanContext._baggageItems) return
+        spanContext._baggageItems[key] = value
+      } else {
+        setBaggageItem(key, value)
+      }
     }
   }
 

packages/dd-trace/src/proxy.js

@@ -10,6 +10,13 @@ const telemetry = require('./telemetry')
 const nomenclature = require('./service-naming')
 const PluginManager = require('./plugin_manager')
 const NoopDogStatsDClient = require('./noop/dogstatsd')
+const {
+  setBaggageItem,
+  getBaggageItem,
+  getAllBaggageItems,
+  removeBaggageItem,
+  removeAllBaggageItems
+} = require('./baggage')
 
 class LazyModule {
   constructor (provider) {
@@ -65,6 +72,11 @@ class Tracer extends NoopProxy {
     this.dogstatsd = new NoopDogStatsDClient()
     this._tracingInitialized = false
     this._flare = new LazyModule(() => require('./flare'))
+    this.setBaggageItem = setBaggageItem
+    this.getBaggageItem = getBaggageItem
+    this.getAllBaggageItems = getAllBaggageItems
+    this.removeBaggageItem = removeBaggageItem
+    this.removeAllBaggageItems = removeAllBaggageItems
 
     // these requires must work with esm bundler
     this._modules = {

packages/dd-trace/test/opentracing/propagation/text_map.spec.js

@@ -7,6 +7,7 @@ const id = require('../../../src/id')
 const SpanContext = require('../../../src/opentracing/span_context')
 const TraceState = require('../../../src/opentracing/propagation/tracestate')
 const { channel } = require('dc-polyfill')
+const { setBaggageItem, getBaggageItem, getAllBaggageItems, removeAllBaggageItems } = require('../../../src/baggage')
 
 const { AUTO_KEEP, AUTO_REJECT, USER_KEEP } = require('../../../../../ext/priority')
 const { SAMPLING_MECHANISM_MANUAL } = require('../../../src/constants')
@@ -142,6 +143,16 @@ describe('TextMapPropagator', () => {
       expect(carrier.baggage).to.equal('raccoon=chunky')
     })
 
+    it('should inject baggage items when spanContext is null', () => {
+      const carrier = {}
+      const spanContext = null
+      removeAllBaggageItems()
+      setBaggageItem('spring', 'blossom')
+
+      propagator.inject(spanContext, carrier)
+      expect(carrier.baggage).to.equal('spring=blossom')
+    })
+
     it('should inject an existing sampling priority', () => {
       const carrier = {}
       const spanContext = createContext({
@@ -492,8 +503,8 @@ describe('TextMapPropagator', () => {
       expect(spanContext._trace.tags).to.not.have.property('_dd.p.tid')
     })
 
-    // temporary test. On the contrary, it SHOULD extract baggage
-    it('should not extract baggage when it is the only propagation style', () => {
+    it('should extract baggage when it is the only propagation style', () => {
+      removeAllBaggageItems()
       config = new Config({
         tracePropagationStyle: {
           extract: ['baggage']
@@ -505,6 +516,8 @@ describe('TextMapPropagator', () => {
       }
       const spanContext = propagator.extract(carrier)
       expect(spanContext).to.be.null
+      expect(getBaggageItem('foo')).to.equal('bar')
+      expect(getAllBaggageItems()).to.deep.equal({ foo: 'bar' })
     })
 
     it('should convert signed IDs to unsigned', () => {

packages/dd-trace/test/proxy.spec.js

@@ -624,6 +624,72 @@ describe('TracerProxy', () => {
       })
     })
 
+    describe('baggage', () => {
+      afterEach(() => {
+        proxy.removeAllBaggageItems()
+      })
+
+      describe('setBaggageItem', () => {
+        it('should set a baggage item', () => {
+          const baggage = proxy.setBaggageItem('key', 'value')
+          expect(baggage).to.deep.equal({ key: 'value' })
+        })
+
+        it('should merge with existing baggage items', () => {
+          proxy.setBaggageItem('key1', 'value1')
+          const baggage = proxy.setBaggageItem('key2', 'value2')
+          expect(baggage).to.deep.equal({ key1: 'value1', key2: 'value2' })
+        })
+      })
+
+      describe('getBaggageItem', () => {
+        it('should get a baggage item', () => {
+          proxy.setBaggageItem('key', 'value')
+          expect(proxy.getBaggageItem('key')).to.equal('value')
+        })
+
+        it('should return undefined for non-existent items', () => {
+          expect(proxy.getBaggageItem('missing')).to.be.undefined
+        })
+      })
+
+      describe('getAllBaggageItems', () => {
+        it('should get all baggage items', () => {
+          proxy.setBaggageItem('key1', 'value1')
+          proxy.setBaggageItem('key2', 'value2')
+          expect(proxy.getAllBaggageItems()).to.deep.equal({ key1: 'value1', key2: 'value2' })
+        })
+
+        it('should return empty object when no items exist', () => {
+          expect(proxy.getAllBaggageItems()).to.be.undefined
+        })
+      })
+
+      describe('removeBaggageItem', () => {
+        it('should remove a specific baggage item', () => {
+          proxy.setBaggageItem('key1', 'value1')
+          proxy.setBaggageItem('key2', 'value2')
+          const baggage = proxy.removeBaggageItem('key1')
+          expect(baggage).to.deep.equal({ key2: 'value2' })
+        })
+
+        it('should handle removing non-existent items', () => {
+          proxy.setBaggageItem('key', 'value')
+          const baggage = proxy.removeBaggageItem('missing')
+          expect(baggage).to.deep.equal({ key: 'value' })
+        })
+      })
+
+      describe('removeAllBaggageItems', () => {
+        it('should remove all baggage items', () => {
+          proxy.setBaggageItem('key1', 'value1')
+          proxy.setBaggageItem('key2', 'value2')
+          const baggage = proxy.removeAllBaggageItems()
+          expect(baggage).to.deep.equal({})
+        })
+      })
+    })
+
     describe('appsec', () => {
       describe('trackUserLoginSuccessEvent', () => {
         it('should call the underlying NoopAppsecSdk method', () => {