feat: use well-known format for propagating trace context thru grpc (#814)

kjin · web-flow · commit 63b13cac84ab · 2018-07-18T15:07:16.000-07:00
BREAKING CHANGE: The change in distributed trace context propagation across gRPC is not backwards-compatible. In other words, distributed tracing will not work between two Node instances communicating using gRPC with v2 and v3 of the Trace Agent, respectively.

Trace context can now be propagated thru gRPC with the `'grpc-trace-bin'` key, and a binary-encoded value.
diff --git a/src/constants.ts b/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',
 
diff --git a/src/plugins/plugin-grpc.ts b/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) => {
diff --git a/test/plugins/test-trace-grpc.ts b/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();
     });
 
