From 83e481b347c6f84cb2490d1e0cf0a06fb211dd87 Mon Sep 17 00:00:00 2001
From: Kelvin Jin <kelvinjin@google.com>
Date: Wed, 18 Jul 2018 10:55:06 -0700
Subject: [PATCH] feat: use well-known format for propagating trace context
 thru grpc

---
 src/constants.ts           |  3 ++
 src/plugins/plugin-grpc.ts | 58 ++++++++++++++++++++++++++++++--------
 2 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/src/constants.ts b/src/constants.ts
index 14d98df59..11dcfd7f9 100644
--- 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
index 2fd6d35d0..0afceb3a2 100644
--- a/src/plugins/plugin-grpc.ts
+++ b/src/plugins/plugin-grpc.ts
@@ -108,6 +108,24 @@ 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 span The span that contains the trace context to add as a metadata
+   * entry.
+   */
+  function setTraceContextFromString(metadata: Metadata, span: Span): void {
+    const traceContext =
+        api.traceContextUtils.decodeFromString(span.getTraceContext());
+    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 +212,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);
       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 +281,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 +340,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 +398,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 +461,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 +522,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) => {