feat: add baggage support to the opentracing shim

packages/opentelemetry-shim-opentracing/src/shim.ts

@@ -19,10 +19,11 @@ import {
   getExtractedSpanContext,
   NoopLogger,
   setExtractedSpanContext,
+  setCorrelationContext,
   setActiveSpan,
+  getCorrelationContext,
 } from '@opentelemetry/core';
 import * as opentracing from 'opentracing';
-import { defaultSetter } from '@opentelemetry/api';
 
 function translateReferences(references: opentracing.Reference[]): api.Link[] {
   const links: api.Link[] = [];
@@ -72,10 +73,15 @@ function getContextWithParent(options: opentracing.SpanOptions) {
  */
 export class SpanContextShim extends opentracing.SpanContext {
   private readonly _spanContext: api.SpanContext;
+  private _correlationContext: api.CorrelationContext;
 
-  constructor(spanContext: api.SpanContext) {
+  constructor(
+    spanContext: api.SpanContext,
+    correlationContext: api.CorrelationContext
+  ) {
     super();
     this._spanContext = spanContext;
+    this._correlationContext = correlationContext;
   }
 
   /**
@@ -85,6 +91,13 @@ export class SpanContextShim extends opentracing.SpanContext {
     return this._spanContext;
   }
 
+  /**
+   * Returns the underlying {@link api.CorrelationContext}
+   */
+  getCorrelationContext(): api.CorrelationContext {
+    return this._correlationContext;
+  }
+
   /**
    * Returns the trace ID as a string.
    */
@@ -98,6 +111,16 @@ export class SpanContextShim extends opentracing.SpanContext {
   toSpanId(): string {
     return this._spanContext.spanId;
   }
+
+  getBaggageItem(key: string): string | undefined {
+    return this._correlationContext[key]?.value;
+  }
+
+  setBaggageItem(key: string, value: string) {
+    this._correlationContext = Object.assign({}, this._correlationContext, {
+      [key]: { value },
+    });
+  }
 }
 
 /**
@@ -125,29 +148,40 @@ export class TracerShim extends opentracing.Tracer {
       getContextWithParent(options)
     );
 
+    let correlationContext: api.CorrelationContext = {};
+    if (options.childOf instanceof SpanShim) {
+      const shimContext = options.childOf.context() as SpanContextShim;
+      correlationContext = shimContext.getCorrelationContext();
+    } else if (options.childOf instanceof SpanContextShim) {
+      correlationContext = options.childOf.getCorrelationContext();
+    }
+
     if (options.tags) {
       span.setAttributes(options.tags);
     }
 
-    return new SpanShim(this, span);
+    return new SpanShim(this, span, correlationContext);
   }
 
   _inject(
     spanContext: opentracing.SpanContext,
     format: string,
     carrier: unknown
   ): void {
-    const opentelemSpanContext: api.SpanContext = (spanContext as SpanContextShim).getSpanContext();
+    const spanContextShim: SpanContextShim = spanContext as SpanContextShim;
+    const oTelSpanContext: api.SpanContext = spanContextShim.getSpanContext();
+    const oTelSpanCorrelationContext: api.CorrelationContext = spanContextShim.getCorrelationContext();
+
     if (!carrier || typeof carrier !== 'object') return;
     switch (format) {
       case opentracing.FORMAT_HTTP_HEADERS:
       case opentracing.FORMAT_TEXT_MAP: {
         api.propagation.inject(
           carrier,
-          defaultSetter,
-          setExtractedSpanContext(
-            api.Context.ROOT_CONTEXT,
-            opentelemSpanContext
+          api.defaultSetter,
+          setCorrelationContext(
+            setExtractedSpanContext(api.Context.ROOT_CONTEXT, oTelSpanContext),
+            oTelSpanCorrelationContext
           )
         );
         return;
@@ -156,7 +190,7 @@ export class TracerShim extends opentracing.Tracer {
         this._logger.warn(
           'OpentracingShim.inject() does not support FORMAT_BINARY'
         );
-        // @todo: Implement binary format
+        // @todo: Implement binary formats
         return;
       }
       default:
@@ -167,13 +201,14 @@ export class TracerShim extends opentracing.Tracer {
     switch (format) {
       case opentracing.FORMAT_HTTP_HEADERS:
       case opentracing.FORMAT_TEXT_MAP: {
-        const context = getExtractedSpanContext(
-          api.propagation.extract(carrier)
-        );
-        if (!context) {
+        const context: api.Context = api.propagation.extract(carrier);
+        const spanContext = getExtractedSpanContext(context);
+        const correlationContext = getCorrelationContext(context);
+
+        if (!spanContext) {
           return null;
         }
-        return new SpanContextShim(context);
+        return new SpanContextShim(spanContext, correlationContext || {});
       }
       case opentracing.FORMAT_BINARY: {
         // @todo: Implement binary format
@@ -191,19 +226,23 @@ export class TracerShim extends opentracing.Tracer {
 /**
  * SpanShim wraps an {@link types.Span} and implements the OpenTracing Span API
  * around it.
- * @todo: Out of band baggage propagation is not currently supported.
- */
+ *
+ **/
 export class SpanShim extends opentracing.Span {
   // _span is the original OpenTelemetry span that we are wrapping with
   // an opentracing interface.
   private readonly _span: api.Span;
   private readonly _contextShim: SpanContextShim;
   private readonly _tracerShim: TracerShim;
 
-  constructor(tracerShim: TracerShim, span: api.Span) {
+  constructor(
+    tracerShim: TracerShim,
+    span: api.Span,
+    correlationContext: api.CorrelationContext
+  ) {
     super();
     this._span = span;
-    this._contextShim = new SpanContextShim(span.context());
+    this._contextShim = new SpanContextShim(span.context(), correlationContext);
     this._tracerShim = tracerShim;
   }
 
@@ -295,12 +334,11 @@ export class SpanShim extends opentracing.Span {
   }
 
   getBaggageItem(key: string): string | undefined {
-    // TODO: should this go into the context?
-    return undefined;
+    return this._contextShim.getBaggageItem(key);
   }
 
   setBaggageItem(key: string, value: string): this {
-    // TODO: should this go into the context?
+    this._contextShim.setBaggageItem(key, value);
     return this;
   }
 

packages/opentelemetry-shim-opentracing/test/Shim.test.ts

@@ -22,6 +22,8 @@ import {
   INVALID_SPAN_CONTEXT,
   timeInputToHrTime,
   HttpTraceContext,
+  CompositePropagator,
+  HttpCorrelationContext,
 } from '@opentelemetry/core';
 import { propagation } from '@opentelemetry/api';
 import { performance } from 'perf_hooks';
@@ -32,7 +34,11 @@ describe('OpenTracing Shim', () => {
     provider.getTracer('default')
   );
   opentracing.initGlobalTracer(shimTracer);
-  propagation.setGlobalPropagator(new HttpTraceContext());
+  const compositePropagator = new CompositePropagator({
+    propagators: [new HttpTraceContext(), new HttpCorrelationContext()],
+  });
+
+  propagation.setGlobalPropagator(compositePropagator);
 
   describe('TracerShim', () => {
     let span: opentracing.Span;
@@ -86,6 +92,25 @@ describe('OpenTracing Shim', () => {
         /* const extractedContext = shimTracer.extract(opentracing.FORMAT_BINARY, { buffer: new Uint8Array(carrier)}); */
         /* assert.strictEqual(context.toSpanId(), extractedContext.toSpanId()) */
       });
+
+      it('injects/extracts a span with baggage', () => {
+        const carrier: { [key: string]: unknown } = {};
+        span.setBaggageItem('baggage1', 'value1');
+        span.setBaggageItem('baggage2', 'value2');
+        shimTracer.inject(span, opentracing.FORMAT_HTTP_HEADERS, carrier);
+        const extractedContext = shimTracer.extract(
+          opentracing.FORMAT_HTTP_HEADERS,
+          carrier
+        ) as SpanContextShim;
+        const childSpan = shimTracer.startSpan('other-span', {
+          childOf: extractedContext,
+        }) as SpanShim;
+        assert.ok(extractedContext !== null);
+        assert.strictEqual(context.toTraceId(), extractedContext!.toTraceId());
+        assert.strictEqual(context.toSpanId(), extractedContext!.toSpanId());
+        assert.strictEqual(childSpan.getBaggageItem('baggage1'), 'value1');
+        assert.strictEqual(childSpan.getBaggageItem('baggage2'), 'value2');
+      });
     });
 
     it('creates parent/child relationship using a span object', () => {
@@ -135,7 +160,7 @@ describe('OpenTracing Shim', () => {
 
   describe('SpanContextShim', () => {
     it('returns the correct context', () => {
-      const shim = new SpanContextShim(INVALID_SPAN_CONTEXT);
+      const shim = new SpanContextShim(INVALID_SPAN_CONTEXT, {});
       assert.strictEqual(shim.getSpanContext(), INVALID_SPAN_CONTEXT);
       assert.strictEqual(shim.toTraceId(), INVALID_SPAN_CONTEXT.traceId);
       assert.strictEqual(shim.toSpanId(), INVALID_SPAN_CONTEXT.spanId);
@@ -190,7 +215,36 @@ describe('OpenTracing Shim', () => {
 
     it('can set and retrieve baggage', () => {
       span.setBaggageItem('baggage', 'item');
-      // TODO: baggage
+      const value = span.getBaggageItem('baggage');
+      assert.equal('item', value);
+
+      const childSpan = shimTracer.startSpan('child-span1', {
+        childOf: span,
+      });
+      childSpan.setBaggageItem('key2', 'item2');
+
+      // child should have parent baggage items.
+      assert.equal('item', childSpan.getBaggageItem('baggage'));
+      assert.equal('item2', childSpan.getBaggageItem('key2'));
+
+      // Parent shouldn't have the child baggage item.
+      assert.equal(span.getBaggageItem('key2'), undefined);
+    });
+
+    it('can set and retrieve baggage with same key', () => {
+      span.setBaggageItem('key1', 'value1');
+      const value = span.getBaggageItem('key1');
+      assert.equal('value1', value);
+
+      const childSpan = shimTracer.startSpan('child-span1', {
+        childOf: span,
+      });
+      childSpan.setBaggageItem('key2', 'value2');
+      childSpan.setBaggageItem('key1', 'value3');
+
+      // child should have parent baggage items.
+      assert.equal('value3', childSpan.getBaggageItem('key1'));
+      assert.equal('value2', childSpan.getBaggageItem('key2'));
     });
   });
 });