feat!: support user-specified context header propagation

package.json

@@ -100,6 +100,7 @@
   },
   "dependencies": {
     "@google-cloud/common": "^1.0.0",
+    "@opencensus/propagation-stackdriver": "0.0.11",
     "builtin-modules": "^3.0.0",
     "console-log-level": "^1.4.0",
     "continuation-local-storage": "^3.2.1",

src/config.ts

@@ -52,6 +52,24 @@ export interface TracePolicy {
   shouldTrace: (requestDetails: RequestDetails) => boolean;
 }
 
+export type GetHeaderFunction = {
+  getHeader: (key: string) => string[]|string|undefined;
+};
+export type SetHeaderFunction = {
+  setHeader: (key: string, value: string) => void;
+};
+export interface OpenCensusPropagation {
+  extract: (getHeader: GetHeaderFunction) => {
+    traceId: string;
+    spanId: string;
+    options?: number
+  } | null;
+  inject: (setHeader: SetHeaderFunction, traceContext: {
+    traceId: string; spanId: string;
+    options?: number
+  }) => void;
+}
+
 /**
  * Available configuration options. All fields are optional. See the
  * defaultConfig object defined in this file for default assigned values.
@@ -195,6 +213,13 @@ export interface Config {
    */
   tracePolicy?: TracePolicy;
 
+  /**
+   * If specified, the Trace Agent will use this context header propagation
+   * implementation instead of @opencensus/propagation-stackdriver, the default
+   * trace context header format.
+   */
+  propagation?: OpenCensusPropagation;
+
   /**
    * Buffer the captured traces for `flushDelaySeconds` seconds before
    * publishing to the Stackdriver Trace API, unless the buffer fills up first.

src/index.ts

@@ -146,7 +146,10 @@ function initConfig(userConfig: Forceable<Config>): Forceable<TopLevelConfig> {
       contextHeaderBehavior: mergedConfig.contextHeaderBehavior as
           TraceContextHeaderBehavior
     },
-    overrides: {tracePolicy: mergedConfig.tracePolicy}
+    overrides: {
+      tracePolicy: mergedConfig.tracePolicy,
+      propagation: mergedConfig.propagation
+    }
   };
 }
 

src/plugin-types.ts

@@ -38,11 +38,10 @@ export interface TraceAgentExtension {
  */
 export interface Span {
   /**
-   * Gets the current trace context serialized as a string, or an empty string
-   * if it can't be generated.
-   * @return The stringified trace context.
+   * Gets the current trace context, or null if it can't be retrieved.
+   * @return The trace context.
    */
-  getTraceContext(): string;
+  getTraceContext(): TraceContext|null;
 
   /**
    * Adds a key-value pair as a label to the trace span. The value will be
@@ -106,10 +105,9 @@ export interface RootSpanOptions extends SpanOptions {
   /* A Method associated with the root span, if applicable. */
   method?: string;
   /**
-   * The serialized form of an object that contains information about an
-   * existing trace context, if it exists.
+   * An existing trace context, if it exists.
    */
-  traceContext?: string|null;
+  traceContext?: TraceContext|null;
 }
 
 export interface Tracer {
@@ -191,21 +189,22 @@ export interface Tracer {
   isRealSpan(span: Span): boolean;
 
   /**
-   * Generates a stringified trace context that should be set as the trace
+   * Generates a trace context object that should be set as the trace
    * context header in a response to an incoming web request. This value is
    * based on the trace context header value in the corresponding incoming
    * request, as well as the result from the local trace policy on whether this
    * request will be traced or not.
    * @param incomingTraceContext The trace context that was attached to
    * the incoming web request, or null if the incoming request didn't have one.
-   * @param isTraced Whether the incoming was traced. This is determined
+   * @param isTraced Whether the incoming request was traced. This is determined
    * by the local tracing policy.
    * @returns If the response should contain the trace context within its
-   * header, the string to be set as this header's value. Otherwise, an empty
-   * string.
+   * header, the context object to be serialized as this header's value.
+   * Otherwise, null.
    */
-  getResponseTraceContext(incomingTraceContext: string|null, isTraced: boolean):
-      string;
+  getResponseTraceContext(
+      incomingTraceContext: TraceContext|null, isTraced: boolean): TraceContext
+      |null;
 
   /**
    * Binds the trace context to the given function.
@@ -233,11 +232,21 @@ export interface Tracer {
   readonly spanTypes: typeof SpanType;
   /** A collection of functions for encoding and decoding trace context. */
   readonly traceContextUtils: {
-    encodeAsString: (ctx: TraceContext) => string;
-    decodeFromString: (str: string) => TraceContext | null;
     encodeAsByteArray: (ctx: TraceContext) => Buffer;
     decodeFromByteArray: (buf: Buffer) => TraceContext | null;
   };
+  /**
+   * A collection of functions for dealing with trace context in HTTP headers.
+   */
+  readonly propagation: Propagation;
+}
+
+export type GetHeaderFunction = (key: string) => string[]|string|null|undefined;
+export type SetHeaderFunction = (key: string, value: string) => void;
+export interface Propagation {
+  extract: (getHeader: GetHeaderFunction) => TraceContext | null;
+  inject:
+      (setHeader: SetHeaderFunction, traceContext: TraceContext|null) => void;
 }
 
 export interface Monkeypatch<T> {

src/plugins/plugin-connect.ts

@@ -28,32 +28,23 @@ type Request = IncomingMessage&{originalUrl?: string};
 
 const SUPPORTED_VERSIONS = '3.x';
 
-function getFirstHeader(req: IncomingMessage, key: string): string|null {
-  let headerValue = req.headers[key] || null;
-  if (headerValue && typeof headerValue !== 'string') {
-    headerValue = headerValue[0];
-  }
-  return headerValue;
-}
-
 function createMiddleware(api: PluginTypes.Tracer):
     connect_3.NextHandleFunction {
   return function middleware(req: Request, res, next) {
     const options = {
       name: req.originalUrl ? (urlParse(req.originalUrl).pathname || '') : '',
       url: req.originalUrl,
       method: req.method,
-      traceContext:
-          getFirstHeader(req, api.constants.TRACE_CONTEXT_HEADER_NAME),
+      traceContext: api.propagation.extract((key) => req.headers[key]),
       skipFrames: 1
     };
     api.runInRootSpan(options, (root) => {
       // Set response trace context.
       const responseTraceContext = api.getResponseTraceContext(
-          options.traceContext || null, api.isRealSpan(root));
+          options.traceContext, api.isRealSpan(root));
       if (responseTraceContext) {
-        res.setHeader(
-            api.constants.TRACE_CONTEXT_HEADER_NAME, responseTraceContext);
+        api.propagation.inject(
+            (k, v) => res.setHeader(k, v), responseTraceContext);
       }
 
       if (!api.isRealSpan(root)) {

src/plugins/plugin-express.ts

@@ -35,19 +35,20 @@ function patchModuleRoot(express: Express4Module, api: PluginTypes.Tracer) {
   function middleware(
       req: express_4.Request, res: express_4.Response,
       next: express_4.NextFunction) {
-    const options: PluginTypes.RootSpanOptions = {
+    const options = {
       name: req.path,
-      traceContext: req.get(api.constants.TRACE_CONTEXT_HEADER_NAME),
+      traceContext: api.propagation.extract((key) => req.get(key)),
       url: req.originalUrl,
       method: req.method,
       skipFrames: 1
     };
     api.runInRootSpan(options, (rootSpan) => {
       // Set response trace context.
       const responseTraceContext = api.getResponseTraceContext(
-          options.traceContext || null, api.isRealSpan(rootSpan));
+          options.traceContext, api.isRealSpan(rootSpan));
       if (responseTraceContext) {
-        res.set(api.constants.TRACE_CONTEXT_HEADER_NAME, responseTraceContext);
+        api.propagation.inject(
+            (k, v) => res.setHeader(k, v), responseTraceContext);
       }
 
       if (!api.isRealSpan(rootSpan)) {

src/plugins/plugin-grpc.ts

@@ -19,7 +19,7 @@ import * as grpcModule from 'grpc';  // for types only.
 import {Client, MethodDefinition, ServerReadableStream, ServerUnaryCall, StatusObject} from 'grpc';
 import * as shimmer from 'shimmer';
 
-import {Plugin, RootSpan, RootSpanOptions, Span, Tracer} from '../plugin-types';
+import {Plugin, RootSpan, RootSpanOptions, Span, TraceContext, Tracer} from '../plugin-types';
 
 // Re-definition of Metadata with private fields
 type Metadata = grpcModule.Metadata&{
@@ -116,9 +116,7 @@ function patchClient(client: ClientModule, api: Tracer) {
    * a falsey value, metadata will not be modified.
    */
   function setTraceContextFromString(
-      metadata: Metadata, stringifiedTraceContext: string): void {
-    const traceContext =
-        api.traceContextUtils.decodeFromString(stringifiedTraceContext);
+      metadata: Metadata, traceContext: TraceContext|null): void {
     if (traceContext) {
       const metadataValue =
           api.traceContextUtils.encodeAsByteArray(traceContext);
@@ -292,26 +290,19 @@ function unpatchClient(client: ClientModule) {
 function patchServer(server: ServerModule, api: Tracer) {
   /**
    * Returns a trace context on a Metadata object if it exists and is
-   * well-formed, or null otherwise. The result will be encoded as a string.
+   * well-formed, or null otherwise.
    * @param metadata The Metadata object from which trace context should be
    * retrieved.
    */
-  function getStringifiedTraceContext(metadata: grpcModule.Metadata): string|
-      null {
+  function getTraceContext(metadata: grpcModule.Metadata): TraceContext|null {
     const metadataValue =
         metadata.getMap()[api.constants.TRACE_CONTEXT_GRPC_METADATA_NAME] as
         Buffer;
     // Entry doesn't exist.
     if (!metadataValue) {
       return null;
     }
-    const traceContext =
-        api.traceContextUtils.decodeFromByteArray(metadataValue);
-    // Value is malformed.
-    if (!traceContext) {
-      return null;
-    }
-    return api.traceContextUtils.encodeAsString(traceContext);
+    return api.traceContextUtils.decodeFromByteArray(metadataValue);
   }
 
   /**
@@ -356,7 +347,7 @@ function patchServer(server: ServerModule, api: Tracer) {
       const rootSpanOptions = {
         name: requestName,
         url: requestName,
-        traceContext: getStringifiedTraceContext(call.metadata),
+        traceContext: getTraceContext(call.metadata),
         skipFrames: SKIP_FRAMES
       };
       return api.runInRootSpan(rootSpanOptions, (rootSpan) => {
@@ -410,7 +401,7 @@ function patchServer(server: ServerModule, api: Tracer) {
       const rootSpanOptions = {
         name: requestName,
         url: requestName,
-        traceContext: getStringifiedTraceContext(stream.metadata),
+        traceContext: getTraceContext(stream.metadata),
         skipFrames: SKIP_FRAMES
       } as RootSpanOptions;
       return api.runInRootSpan(rootSpanOptions, (rootSpan) => {
@@ -472,7 +463,7 @@ function patchServer(server: ServerModule, api: Tracer) {
       const rootSpanOptions = {
         name: requestName,
         url: requestName,
-        traceContext: getStringifiedTraceContext(stream.metadata),
+        traceContext: getTraceContext(stream.metadata),
         skipFrames: SKIP_FRAMES
       } as RootSpanOptions;
       return api.runInRootSpan(rootSpanOptions, (rootSpan) => {
@@ -532,7 +523,7 @@ function patchServer(server: ServerModule, api: Tracer) {
       const rootSpanOptions = {
         name: requestName,
         url: requestName,
-        traceContext: getStringifiedTraceContext(stream.metadata),
+        traceContext: getTraceContext(stream.metadata),
         skipFrames: SKIP_FRAMES
       } as RootSpanOptions;
       return api.runInRootSpan(rootSpanOptions, (rootSpan) => {

src/plugins/plugin-hapi.ts

@@ -34,34 +34,26 @@ type Hapi17Request = hapi_17.Request&{
   _execute: Hapi17RequestExecutePrivate;
 };
 
-function getFirstHeader(req: IncomingMessage, key: string): string|null {
-  let headerValue = req.headers[key] || null;
-  if (headerValue && typeof headerValue !== 'string') {
-    headerValue = headerValue[0];
-  }
-  return headerValue;
-}
-
 function instrument<T>(
     api: PluginTypes.Tracer, request: hapi_16.Request|hapi_17.Request,
     continueCb: () => T): T {
   const req = request.raw.req;
   const res = request.raw.res;
   const originalEnd = res.end;
-  const options: PluginTypes.RootSpanOptions = {
+  const options = {
     name: req.url ? (urlParse(req.url).pathname || '') : '',
     url: req.url,
     method: req.method,
-    traceContext: getFirstHeader(req, api.constants.TRACE_CONTEXT_HEADER_NAME),
+    traceContext: api.propagation.extract(key => req.headers[key]),
     skipFrames: 2
   };
   return api.runInRootSpan(options, (root) => {
     // Set response trace context.
-    const responseTraceContext = api.getResponseTraceContext(
-        options.traceContext || null, api.isRealSpan(root));
+    const responseTraceContext =
+        api.getResponseTraceContext(options.traceContext, api.isRealSpan(root));
     if (responseTraceContext) {
-      res.setHeader(
-          api.constants.TRACE_CONTEXT_HEADER_NAME, responseTraceContext);
+      api.propagation.inject(
+          (k, v) => res.setHeader(k, v), responseTraceContext);
     }
 
     if (!api.isRealSpan(root)) {

src/plugins/plugin-http.ts

@@ -142,9 +142,11 @@ function makeRequestTrace(
       // headers.
       options = Object.assign({}, options) as ClientRequestArgs;
       options.headers = Object.assign({}, options.headers);
+      const headers = options.headers;
       // Inject the trace context header.
-      options.headers[api.constants.TRACE_CONTEXT_HEADER_NAME] =
-          span.getTraceContext();
+      api.propagation.inject((key, value) => {
+        headers[key] = value;
+      }, span.getTraceContext());
     }
 
     const req = request(options, (res) => {
@@ -188,19 +190,20 @@ function makeRequestTrace(
     // Inject the trace context header, but only if it wasn't already injected
     // earlier.
     if (!traceHeaderPreinjected) {
-      try {
-        req.setHeader(
-            api.constants.TRACE_CONTEXT_HEADER_NAME, span.getTraceContext());
-      } catch (e) {
-        if (e.code === ERR_HTTP_HEADERS_SENT ||
-            e.message === ERR_HTTP_HEADERS_SENT_MSG) {
-          // Swallow the error.
-          // This would happen in the pathological case where the Expect header
-          // exists but is not detected by hasExpectHeader.
-        } else {
-          throw e;
+      api.propagation.inject((key, value) => {
+        try {
+          req.setHeader(key, value);
+        } catch (e) {
+          if (e.code === ERR_HTTP_HEADERS_SENT ||
+              e.message === ERR_HTTP_HEADERS_SENT_MSG) {
+            // Swallow the error.
+            // This would happen in the pathological case where the Expect
+            // header exists but is not detected by hasExpectHeader.
+          } else {
+            throw e;
+          }
         }
-      }
+      }, span.getTraceContext());
     }
     return req;
   };

src/plugins/plugin-http2.ts

@@ -97,8 +97,8 @@ function makeRequestTrace(
         api.labels.HTTP_METHOD_LABEL_KEY, extractMethodName(newHeaders));
     requestLifecycleSpan.addLabel(
         api.labels.HTTP_URL_LABEL_KEY, extractUrl(authority, newHeaders));
-    newHeaders[api.constants.TRACE_CONTEXT_HEADER_NAME] =
-        requestLifecycleSpan.getTraceContext();
+    api.propagation.inject(
+        (k, v) => newHeaders[k] = v, requestLifecycleSpan.getTraceContext());
     const stream: http2.ClientHttp2Stream = request.call(
         this, newHeaders, ...Array.prototype.slice.call(arguments, 1));
     api.wrapEmitter(stream);