feat: use well-known format for propagating trace context thru grpc

src/constants.ts

@@ -19,6 +19,9 @@
 /** Constant values. */
 // tslint:disable-next-line:variable-name
 export const Constants = {
+  /** The metadata key under which trace context  */
+  TRACE_CONTEXT_GRPC_METADATA_NAME: 'grpc-trace-bin',
+
   /** Header that carries trace context across Google infrastructure. */
   TRACE_CONTEXT_HEADER_NAME: 'x-cloud-trace-context',
 

src/plugins/plugin-grpc.ts

@@ -108,6 +108,25 @@ function unpatchMetadata() {
 }
 
 function patchClient(client: ClientModule, api: TraceAgent) {
+  /**
+   * Set trace context on a Metadata object if it exists.
+   * @param metadata The Metadata object to which a trace context should be
+   * added.
+   * @param stringifiedTraceContext The stringified trace context. If this is
+   * a falsey value, metadata will not be modified.
+   */
+  function setTraceContextFromString(
+      metadata: Metadata, stringifiedTraceContext: string): void {
+    const traceContext =
+        api.traceContextUtils.decodeFromString(stringifiedTraceContext);
+    if (traceContext) {
+      const metadataValue =
+          api.traceContextUtils.encodeAsByteArray(traceContext);
+      metadata.set(
+          api.constants.TRACE_CONTEXT_GRPC_METADATA_NAME, metadataValue);
+    }
+  }
+
   /**
    * Wraps a callback so that the current span for this trace is also ended when
    * the callback is invoked.
@@ -194,8 +213,7 @@ function patchClient(client: ClientModule, api: TraceAgent) {
       // TS: Safe cast as we either found the index of the Metadata argument
       //     or spliced it in at metaIndex.
       const metadata = args[metaIndex] as Metadata;
-      metadata.set(
-          api.constants.TRACE_CONTEXT_HEADER_NAME, span.getTraceContext());
+      setTraceContextFromString(metadata, span.getTraceContext());
       const call: EventEmitter = method.apply(this, args);
       // Add extra data only when call successfully goes through. At this point
       // we know that the arguments are correct.
@@ -264,7 +282,29 @@ function unpatchClient(client: ClientModule) {
 }
 
 function patchServer(server: ServerModule, api: TraceAgent) {
-  const traceContextHeaderName = api.constants.TRACE_CONTEXT_HEADER_NAME;
+  /**
+   * 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.
+   * @param metadata The Metadata object from which trace context should be
+   * retrieved.
+   */
+  function getStringifiedTraceContext(metadata: grpcModule.Metadata): string|
+      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);
+  }
 
   /**
    * A helper function to record metadata in a trace span. The return value of
@@ -301,15 +341,12 @@ function patchServer(server: ServerModule, api: TraceAgent) {
       return function serverMethodTrace(
           this: Server, call: ServerUnaryCall<S>,
           callback: ServerUnaryCallback<T>) {
-        // TODO(kjin): Is it possible for a metadata value to be a buffer?
-        // This needs to be investigated in order to avoid the cast here and
-        // in other server wrapper functions.
         const rootSpanOptions = {
           name: requestName,
           url: requestName,
-          traceContext: call.metadata.getMap()[traceContextHeaderName],
+          traceContext: getStringifiedTraceContext(call.metadata),
           skipFrames: SKIP_FRAMES
-        } as RootSpanOptions;
+        };
         return api.runInRootSpan(rootSpanOptions, (rootSpan) => {
           if (!api.isRealSpan(rootSpan)) {
             return serverMethod.call(this, call, callback);
@@ -362,7 +399,7 @@ function patchServer(server: ServerModule, api: TraceAgent) {
         const rootSpanOptions = {
           name: requestName,
           url: requestName,
-          traceContext: stream.metadata.getMap()[traceContextHeaderName],
+          traceContext: getStringifiedTraceContext(stream.metadata),
           skipFrames: SKIP_FRAMES
         } as RootSpanOptions;
         return api.runInRootSpan(rootSpanOptions, (rootSpan) => {
@@ -425,7 +462,7 @@ function patchServer(server: ServerModule, api: TraceAgent) {
         const rootSpanOptions = {
           name: requestName,
           url: requestName,
-          traceContext: stream.metadata.getMap()[traceContextHeaderName],
+          traceContext: getStringifiedTraceContext(stream.metadata),
           skipFrames: SKIP_FRAMES
         } as RootSpanOptions;
         return api.runInRootSpan(rootSpanOptions, (rootSpan) => {
@@ -486,7 +523,7 @@ function patchServer(server: ServerModule, api: TraceAgent) {
         const rootSpanOptions = {
           name: requestName,
           url: requestName,
-          traceContext: stream.metadata.getMap()[traceContextHeaderName],
+          traceContext: getStringifiedTraceContext(stream.metadata),
           skipFrames: SKIP_FRAMES
         } as RootSpanOptions;
         return api.runInRootSpan(rootSpanOptions, (rootSpan) => {

test/plugins/test-trace-grpc.ts

@@ -45,20 +45,28 @@ var SEND_METADATA = 131;
 var EMIT_ERROR = 13412;
 
 // Regular expression matching client-side metadata labels
-var metadataRegExp =
-  /^{"a":"b","x-cloud-trace-context":"[a-f0-9]{32}\/[0-9]+;o=1"}$/;
+var metadataRegExp = /"a":"b"/;
 
 // Whether asserts in checkServerMetadata should be run
-// Turned on only for the test that checks propagated tract context
+// Turned on only for the test that checks propagated trace context
 var checkMetadata;
 
+// When trace IDs are checked in checkServerMetadata, they should have this
+// exact value. This only applies in the test "should support distributed
+// context".
+const COMMON_TRACE_ID = 'ffeeddccbbaa99887766554433221100';
+
 function checkServerMetadata(metadata) {
   if (checkMetadata) {
-    var traceContext = metadata.getMap()[Constants.TRACE_CONTEXT_HEADER_NAME];
-    assert.ok(/[a-f0-9]{32}\/[0-9]+;o=1/.test(traceContext));
-    var parsedContext = util.parseContextFromHeader(traceContext);
+    var traceContext = metadata.getMap()[Constants.TRACE_CONTEXT_GRPC_METADATA_NAME];
+    var parsedContext = util.deserializeTraceContext(traceContext);
     assert.ok(parsedContext);
     var root = asRootSpanData(cls.get().getContext() as Span);
+    // Check that we were able to propagate trace context.
+    assert.strictEqual(parsedContext!.traceId, COMMON_TRACE_ID);
+    assert.strictEqual(root.trace.traceId, COMMON_TRACE_ID);
+    // Check that we correctly assigned the parent ID of the current span to
+    // that of the incoming span ID.
     assert.strictEqual(root.span.parentSpanId, parsedContext!.spanId);
   }
 }
@@ -205,7 +213,7 @@ function callClientStream(client, grpc, metadata, cb) {
   if (Object.keys(metadata).length > 0) {
     var m = new grpc.Metadata();
     for (var key in metadata) {
-      m.set(key, metadata[key]);
+      m.add(key, metadata[key]);
     }
     args.unshift(m);
   }
@@ -454,9 +462,10 @@ Object.keys(versions).forEach(function(version) {
     it('should support distributed trace context', function(done) {
       function makeLink(fn, meta, next) {
         return function() {
-          common.runInTransaction(function(terminate) {
+          agent.runInRootSpan({ name: '', traceContext: `${COMMON_TRACE_ID}/0;o=1` }, function(span) {
+            assert.strictEqual(span.type, agent.spanTypes.ROOT);
             fn(client, grpc, meta, function() {
-              terminate();
+              span.endSpan();
               next();
             });
           });
@@ -465,28 +474,25 @@ Object.keys(versions).forEach(function(version) {
       // Enable asserting properties of the metdata on the grpc server.
       checkMetadata = true;
       var next;
-      common.runInTransaction(function (endTransaction) {
-        var metadata = { a: 'b' };
-        next = function() {
-          endTransaction();
-          checkMetadata = false;
-          done();
-        };
-        // Try without supplying metadata (call* will not supply metadata to
-        // the grpc client methods at all if no fields are present).
-        // The plugin should automatically create a new Metadata object and
-        // populate it with trace context data accordingly.
-        next = makeLink(callUnary, {}, next);
-        next = makeLink(callClientStream, {}, next);
-        next = makeLink(callServerStream, {}, next);
-        next = makeLink(callBidi, {}, next);
-        // Try with metadata. The plugin should simply add trace context data
-        // to it.
-        next = makeLink(callUnary, metadata, next);
-        next = makeLink(callClientStream, metadata, next);
-        next = makeLink(callServerStream, metadata, next);
-        next = makeLink(callBidi, metadata, next);
-      });
+      var metadata = { a: 'b' };
+      next = function() {
+        checkMetadata = false;
+        done();
+      };
+      // Try without supplying metadata (call* will not supply metadata to
+      // the grpc client methods at all if no fields are present).
+      // The plugin should automatically create a new Metadata object and
+      // populate it with trace context data accordingly.
+      next = makeLink(callUnary, {}, next);
+      next = makeLink(callClientStream, {}, next);
+      next = makeLink(callServerStream, {}, next);
+      next = makeLink(callBidi, {}, next);
+      // Try with metadata. The plugin should simply add trace context data
+      // to it.
+      next = makeLink(callUnary, metadata, next);
+      next = makeLink(callClientStream, metadata, next);
+      next = makeLink(callServerStream, metadata, next);
+      next = makeLink(callBidi, metadata, next);
       next();
     });