otel drop in support for baggage

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

@@ -1,7 +1,7 @@
 'use strict'
 
 const { storage } = require('../../../datadog-core')
-const { trace, ROOT_CONTEXT } = require('@opentelemetry/api')
+const { trace, ROOT_CONTEXT, propagation } = require('@opentelemetry/api')
 const DataDogSpanContext = require('../opentracing/span_context')
 
 const SpanContext = require('./span_context')
@@ -18,17 +18,40 @@ class ContextManager {
     const context = (activeSpan && activeSpan.context()) || store || ROOT_CONTEXT
 
     if (!(context instanceof DataDogSpanContext)) {
+      const span = trace.getSpan(context)
+      // span instanceof NonRecordingSpan
+      if (span && span._spanContext && span._spanContext._ddContext && span._spanContext._ddContext._baggageItems) {
+        const baggages = span._spanContext._ddContext._baggageItems
+        const entries = {}
+        for (const [key, value] of Object.entries(baggages)) {
+          entries[key] = { value }
+        }
+        const otelBaggages = propagation.createBaggage(entries)
+        return propagation.setBaggage(context, otelBaggages)
+      }
       return context
     }
 
+    const baggages = JSON.parse(activeSpan.getAllBaggageItems())
+    const entries = {}
+    for (const [key, value] of Object.entries(baggages)) {
+      entries[key] = { value }
+    }
+    const otelBaggages = propagation.createBaggage(entries)
+
     if (!context._otelSpanContext) {
       const newSpanContext = new SpanContext(context)
       context._otelSpanContext = newSpanContext
     }
     if (store && trace.getSpanContext(store) === context._otelSpanContext) {
-      return store
+      return otelBaggages
+        ? propagation.setBaggage(store, otelBaggages)
+        : store
     }
-    return trace.setSpanContext(store || ROOT_CONTEXT, context._otelSpanContext)
+    const wrappedContext = trace.setSpanContext(store || ROOT_CONTEXT, context._otelSpanContext)
+    return otelBaggages
+      ? propagation.setBaggage(wrappedContext, otelBaggages)
+      : wrappedContext
   }
 
   with (context, fn, thisArg, ...args) {
@@ -38,9 +61,26 @@ class ContextManager {
       const cb = thisArg == null ? fn : fn.bind(thisArg)
       return this._store.run(context, cb, ...args)
     }
+    const baggages = propagation.getBaggage(context)
+    let baggageItems = []
+    if (baggages) {
+      baggageItems = baggages.getAllEntries()
+    }
     if (span && span._ddSpan) {
+      // does otel always override datadog?
+      span._ddSpan.removeAllBaggageItems()
+      for (const baggage of baggageItems) {
+        span._ddSpan.setBaggageItem(baggage[0], baggage[1].value)
+      }
       return ddScope.activate(span._ddSpan, run)
     }
+    // span instanceof NonRecordingSpan
+    if (span && span._spanContext && span._spanContext._ddContext && span._spanContext._ddContext._baggageItems) {
+      span._spanContext._ddContext._baggageItems = {}
+      for (const baggage of baggageItems) {
+        span._spanContext._ddContext._baggageItems[baggage[0]] = baggage[1].value
+      }
+    }
     return run()
   }
 

packages/dd-trace/test/opentelemetry/context_manager.spec.js

@@ -4,8 +4,17 @@ require('../setup/tap')
 
 const { expect } = require('chai')
 const ContextManager = require('../../src/opentelemetry/context_manager')
-const { ROOT_CONTEXT } = require('@opentelemetry/api')
+const TracerProvider = require('../../src/opentelemetry/tracer_provider')
+const { context, propagation, trace, ROOT_CONTEXT } = require('@opentelemetry/api')
 const api = require('@opentelemetry/api')
+const tracer = require('../../').init()
+
+function makeSpan (...args) {
+  const tracerProvider = new TracerProvider()
+  tracerProvider.register()
+  const tracer = tracerProvider.getTracer()
+  return tracer.startSpan(...args)
+}
 
 describe('OTel Context Manager', () => {
   let contextManager
@@ -114,4 +123,75 @@ describe('OTel Context Manager', () => {
     })
     expect(ret).to.equal('return value')
   })
+
+  it('should propagate baggage from an otel span to a datadog span', () => {
+    const entries = {
+      foo: { value: 'bar' }
+    }
+    const baggage = propagation.createBaggage(entries)
+    const contextWithBaggage = propagation.setBaggage(context.active(), baggage)
+    const span = makeSpan('otel-to-dd')
+    const contextWithSpan = trace.setSpan(contextWithBaggage, span)
+    api.context.with(contextWithSpan, () => {
+      expect(tracer.scope().active().getBaggageItem('foo')).to.be.equal('bar')
+    })
+  })
+
+  it('should propagate baggage from a datadog span to an otel span', () => {
+    const baggageKey = 'raccoon'
+    const baggageVal = 'chunky'
+    const ddSpan = tracer.startSpan('dd-to-otel')
+    ddSpan.setBaggageItem(baggageKey, baggageVal)
+    tracer.scope().activate(ddSpan, () => {
+      const baggages = propagation.getActiveBaggage().getAllEntries()
+      expect(baggages.length).to.equal(1)
+      const baggage = baggages[0]
+      expect(baggage[0]).to.equal(baggageKey)
+      expect(baggage[1].value).to.equal(baggageVal)
+    })
+  })
+
+  it('should handle dd-otel baggage conflict', () => {
+    const ddSpan = tracer.startSpan('dd')
+    ddSpan.setBaggageItem('key1', 'dd1')
+    let contextWithUpdatedBaggages
+    tracer.scope().activate(ddSpan, () => {
+      let baggages = propagation.getBaggage(api.context.active())
+      baggages = baggages.setEntry('key1', { value: 'otel1' })
+      baggages = baggages.setEntry('key2', { value: 'otel2' })
+      contextWithUpdatedBaggages = propagation.setBaggage(api.context.active(), baggages)
+    })
+    expect(JSON.parse(ddSpan.getAllBaggageItems())).to.deep.equal({ key1: 'dd1' })
+    api.context.with(contextWithUpdatedBaggages, () => {
+      expect(JSON.parse(ddSpan.getAllBaggageItems())).to.deep.equal(
+        { key1: 'otel1', key2: 'otel2' }
+      )
+      ddSpan.setBaggageItem('key2', 'dd2')
+      expect(propagation.getActiveBaggage().getAllEntries()).to.deep.equal(
+        [['key1', { value: 'otel1' }], ['key2', { value: 'dd2' }]]
+      )
+    })
+  })
+
+  it('should handle dd-otel baggage removal', () => {
+    const ddSpan = tracer.startSpan('dd')
+    ddSpan.setBaggageItem('key1', 'dd1')
+    ddSpan.setBaggageItem('key2', 'dd2')
+    let contextWithUpdatedBaggages
+    tracer.scope().activate(ddSpan, () => {
+      let baggages = propagation.getBaggage(api.context.active())
+      baggages = baggages.removeEntry('key1')
+      contextWithUpdatedBaggages = propagation.setBaggage(api.context.active(), baggages)
+    })
+    expect(JSON.parse(ddSpan.getAllBaggageItems())).to.deep.equal(
+      { key1: 'dd1', key2: 'dd2' }
+    )
+    api.context.with(contextWithUpdatedBaggages, () => {
+      expect(JSON.parse(ddSpan.getAllBaggageItems())).to.deep.equal(
+        { key2: 'dd2' }
+      )
+      ddSpan.removeBaggageItem('key2')
+      expect(propagation.getActiveBaggage().getAllEntries()).to.deep.equal([])
+    })
+  })
 })