feat(jaeger-propagator): propagate/extract baggage #2137

open-telemetry · Apr 25, 2021 · 09c9302 · 09c9302
1 parent d504f32
commit 09c9302
Show file tree

Hide file tree

Showing 2 changed files with 154 additions and 18 deletions.
diff --git a/packages/opentelemetry-propagator-jaeger/src/JaegerHttpTracePropagator.ts b/packages/opentelemetry-propagator-jaeger/src/JaegerHttpTracePropagator.ts
@@ -16,17 +16,21 @@
 
 import {
   Context,
+  getBaggage,
   getSpanContext,
+  setBaggage,
   setSpanContext,
   SpanContext,
   TraceFlags,
   TextMapGetter,
   TextMapPropagator,
   TextMapSetter,
   isInstrumentationSuppressed,
+  createBaggage,
 } from '@opentelemetry/api';
 
 export const UBER_TRACE_ID_HEADER = 'uber-trace-id';
+export const UBER_BAGGAGE_HEADER_PREFIX = 'uberctx';
 
 /**
  * Propagates {@link SpanContext} through Trace Context format propagation.
@@ -55,31 +59,67 @@ export class JaegerHttpTracePropagator implements TextMapPropagator {
 
   inject(context: Context, carrier: unknown, setter: TextMapSetter) {
     const spanContext = getSpanContext(context);
-    if (!spanContext || isInstrumentationSuppressed(context)) return;
-
-    const traceFlags = `0${(spanContext.traceFlags || TraceFlags.NONE).toString(
-      16
-    )}`;
-
-    setter.set(
-      carrier,
-      this._jaegerTraceHeader,
-      `${spanContext.traceId}:${spanContext.spanId}:0:${traceFlags}`
-    );
+    const baggage = getBaggage(context);
+    if (spanContext && isInstrumentationSuppressed(context) === false) {
+      const traceFlags = `0${(
+        spanContext.traceFlags || TraceFlags.NONE
+      ).toString(16)}`;
+
+      setter.set(
+        carrier,
+        this._jaegerTraceHeader,
+        `${spanContext.traceId}:${spanContext.spanId}:0:${traceFlags}`
+      );
+    }
+
+    if (baggage) {
+      for (const [key, entry] of baggage.getAllEntries()) {
+        setter.set(
+          carrier,
+          `${UBER_BAGGAGE_HEADER_PREFIX}-${key}`,
+          encodeURIComponent(entry.value)
+        );
+      }
+    }
   }
 
   extract(context: Context, carrier: unknown, getter: TextMapGetter): Context {
     const uberTraceIdHeader = getter.get(carrier, this._jaegerTraceHeader);
     const uberTraceId = Array.isArray(uberTraceIdHeader)
       ? uberTraceIdHeader[0]
       : uberTraceIdHeader;
-
-    if (typeof uberTraceId !== 'string') return context;
-
-    const spanContext = deserializeSpanContext(uberTraceId);
-    if (!spanContext) return context;
-
-    return setSpanContext(context, spanContext);
+    const baggageValues = getter
+      .keys(carrier)
+      .filter(key => key.includes(UBER_BAGGAGE_HEADER_PREFIX))
+      .map(key => {
+        const value = getter.get(carrier, key);
+        return {
+          key: key.replace(`${UBER_BAGGAGE_HEADER_PREFIX}-`, ''),
+          value: Array.isArray(value) ? value[0] : value,
+        };
+      });
+
+    let newContext = context;
+    // if the trace id header is present and valid, inject it into the context
+    if (typeof uberTraceId === 'string') {
+      const spanContext = deserializeSpanContext(uberTraceId);
+      if (spanContext) {
+        newContext = setSpanContext(newContext, spanContext);
+      }
+    }
+    if (baggageValues.length === 0) return newContext;
+
+    // if baggage values are present, inject it into the current baggage
+    let currentBaggage = getBaggage(context) ?? createBaggage();
+    for (const baggageEntry of baggageValues) {
+      if (baggageEntry.value === undefined) continue;
+      currentBaggage = currentBaggage.setEntry(baggageEntry.key, {
+        value: decodeURIComponent(baggageEntry.value),
+      });
+    }
+    newContext = setBaggage(newContext, currentBaggage);
+
+    return newContext;
   }
 
   fields(): string[] {

diff --git a/packages/opentelemetry-propagator-jaeger/test/JaegerHttpTracePropagator.test.ts b/packages/opentelemetry-propagator-jaeger/test/JaegerHttpTracePropagator.test.ts
@@ -15,10 +15,13 @@
  */
 
 import {
+  createBaggage,
   defaultTextMapGetter,
   defaultTextMapSetter,
+  getBaggage,
   getSpanContext,
   ROOT_CONTEXT,
+  setBaggage,
   setSpanContext,
   SpanContext,
   suppressInstrumentation,
@@ -29,6 +32,7 @@ import * as assert from 'assert';
 import {
   JaegerHttpTracePropagator,
   UBER_TRACE_ID_HEADER,
+  UBER_BAGGAGE_HEADER_PREFIX,
 } from '../src/JaegerHttpTracePropagator';
 
 describe('JaegerHttpTracePropagator', () => {
@@ -94,6 +98,28 @@ describe('JaegerHttpTracePropagator', () => {
       );
       assert.strictEqual(carrier[UBER_TRACE_ID_HEADER], undefined);
     });
+
+    it('should propagate baggage with url encoded values', () => {
+      const baggage = createBaggage({
+        test: {
+          value: '1',
+        },
+        myuser: {
+          value: '%id%',
+        },
+      });
+
+      jaegerHttpTracePropagator.inject(
+        setBaggage(ROOT_CONTEXT, baggage),
+        carrier,
+        defaultTextMapSetter
+      );
+      assert.strictEqual(carrier[`${UBER_BAGGAGE_HEADER_PREFIX}-test`], '1');
+      assert.strictEqual(
+        carrier[`${UBER_BAGGAGE_HEADER_PREFIX}-myuser`],
+        encodeURIComponent('%id%')
+      );
+    });
   });
 
   describe('.extract()', () => {
@@ -199,6 +225,76 @@ describe('JaegerHttpTracePropagator', () => {
         undefined
       );
     });
+
+    it('should extract baggage from carrier', () => {
+      carrier[`${UBER_BAGGAGE_HEADER_PREFIX}-test`] = 'value';
+      carrier[`${UBER_BAGGAGE_HEADER_PREFIX}-myuser`] = '%25id%25';
+      const extractedBaggage = getBaggage(
+        jaegerHttpTracePropagator.extract(
+          ROOT_CONTEXT,
+          carrier,
+          defaultTextMapGetter
+        )
+      );
+
+      const firstEntry = extractedBaggage?.getEntry('test');
+      assert(typeof firstEntry !== 'undefined');
+      assert(firstEntry.value === 'value');
+      const secondEntry = extractedBaggage?.getEntry('myuser');
+      assert(typeof secondEntry !== 'undefined');
+      assert(secondEntry.value === '%id%');
+    });
+
+    it('should extract baggage from carrier and not override current one', () => {
+      carrier[`${UBER_BAGGAGE_HEADER_PREFIX}-test`] = 'value';
+      carrier[`${UBER_BAGGAGE_HEADER_PREFIX}-myuser`] = '%25id%25';
+      const extractedBaggage = getBaggage(
+        jaegerHttpTracePropagator.extract(
+          setBaggage(ROOT_CONTEXT, createBaggage({ one: { value: 'two' } })),
+          carrier,
+          defaultTextMapGetter
+        )
+      );
+
+      const firstEntry = extractedBaggage?.getEntry('test');
+      assert(typeof firstEntry !== 'undefined');
+      assert(firstEntry.value === 'value');
+      const secondEntry = extractedBaggage?.getEntry('myuser');
+      assert(typeof secondEntry !== 'undefined');
+      assert(secondEntry.value === '%id%');
+      const alreadyExistingEntry = extractedBaggage?.getEntry('one');
+      assert(typeof alreadyExistingEntry !== 'undefined');
+      assert(alreadyExistingEntry.value === 'two');
+    });
+
+    it('should handle invalid baggage from carrier (undefined)', () => {
+      carrier[`${UBER_BAGGAGE_HEADER_PREFIX}-test`] = undefined;
+      const extractedBaggage = getBaggage(
+        jaegerHttpTracePropagator.extract(
+          ROOT_CONTEXT,
+          carrier,
+          defaultTextMapGetter
+        )
+      );
+
+      const firstEntry = extractedBaggage?.getEntry('test');
+      assert(typeof firstEntry === 'undefined');
+    });
+
+    it('should handle invalid baggage from carrier (array)', () => {
+      carrier[`${UBER_BAGGAGE_HEADER_PREFIX}-test`] = ['one', 'two'];
+      const extractedBaggage = getBaggage(
+        jaegerHttpTracePropagator.extract(
+          ROOT_CONTEXT,
+          carrier,
+          defaultTextMapGetter
+        )
+      );
+
+      const firstEntry = extractedBaggage?.getEntry('test');
+      assert(typeof firstEntry !== 'undefined');
+      assert(firstEntry.value === 'one');
+    });
   });
 
   describe('.fields()', () => {