feat: add baggage support to the opentracing shim

Signed-off-by: Ruben Vargas <ruben.vp8510@gmail.com>
open-telemetry · Mar 31, 2020 · dc54078 · dc54078
1 parent bc9658d
commit dc54078
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 11 deletions.
diff --git a/packages/opentelemetry-shim-opentracing/src/shim.ts b/packages/opentelemetry-shim-opentracing/src/shim.ts
@@ -22,7 +22,7 @@ import {
   setActiveSpan,
 } from '@opentelemetry/core';
 import * as opentracing from 'opentracing';
-import { defaultSetter } from '@opentelemetry/api';
+import { CorrelationContext, defaultSetter } from '@opentelemetry/api';
 
 function translateReferences(references: opentracing.Reference[]): api.Link[] {
   const links: api.Link[] = [];
@@ -72,10 +72,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: CorrelationContext
+  ) {
     super();
     this._spanContext = spanContext;
+    this._correlationContext = correlationContext;
   }
 
   /**
@@ -85,6 +90,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 +110,18 @@ export class SpanContextShim extends opentracing.SpanContext {
   toSpanId(): string {
     return this._spanContext.spanId;
   }
+
+  getBaggageItem(key: string): string | undefined {
+    const entry: api.EntryValue = this._correlationContext[key];
+    return entry === undefined ? undefined : entry.value;
+  }
+
+  setBaggageItem(key: string, value: string) {
+    this._correlationContext = Object.assign(
+      { [key]: { value } },
+      this._correlationContext
+    );
+  }
 }
 
 /**
@@ -125,11 +149,19 @@ export class TracerShim extends opentracing.Tracer {
       getContextWithParent(options)
     );
 
+    let correlationContext: CorrelationContext = {};
+    if (options.childOf) {
+      if (options.childOf instanceof SpanShim) {
+        const shimContext = options.childOf.context() as SpanContextShim;
+        correlationContext = shimContext.getCorrelationContext();
+      }
+    }
+
     if (options.tags) {
       span.setAttributes(options.tags);
     }
 
-    return new SpanShim(this, span);
+    return new SpanShim(this, span, correlationContext);
   }
 
   _inject(
@@ -156,7 +188,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:
     }
@@ -173,7 +205,8 @@ export class TracerShim extends opentracing.Tracer {
         if (!context) {
           return null;
         }
-        return new SpanContextShim(context);
+        // @todo Extract correlation context
+        return new SpanContextShim(context, {});
       case opentracing.FORMAT_BINARY:
         // @todo: Implement binary format
         this._logger.warn(
@@ -198,10 +231,14 @@ export class SpanShim extends opentracing.Span {
   private readonly _contextShim: SpanContextShim;
   private readonly _tracerShim: TracerShim;
 
-  constructor(tracerShim: TracerShim, span: api.Span) {
+  constructor(
+    tracerShim: TracerShim,
+    span: api.Span,
+    correlationalCtx: CorrelationContext
+  ) {
     super();
     this._span = span;
-    this._contextShim = new SpanContextShim(span.context());
+    this._contextShim = new SpanContextShim(span.context(), correlationalCtx);
     this._tracerShim = tracerShim;
   }
 
@@ -293,12 +330,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;
   }
 

diff --git a/packages/opentelemetry-shim-opentracing/test/Shim.test.ts b/packages/opentelemetry-shim-opentracing/test/Shim.test.ts
@@ -190,7 +190,20 @@ 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 childr baggage item.
+      assert.equal(span.getBaggageItem('key2'), undefined);
     });
   });
 });