fix: enable tracing on original client method names (#874)

kjin · web-flow · commit 497c760b732a · 2018-09-28T17:30:40.000-07:00
diff --git a/src/plugins/plugin-grpc.ts b/src/plugins/plugin-grpc.ts
@@ -29,7 +29,7 @@ type Metadata = grpcModule.Metadata&{
 type MetadataModule = typeof grpcModule.Metadata;
 // Type of makeClientConstructor as exported from client.js
 type MakeClientConstructorFunction =
-    (methods: {[key: string]: never;}, serviceName: string,
+    (methods: {[key: string]: {originalName?: string;};}, serviceName: string,
      classOptions: never) => typeof Client;
 // Meta-type of client-side handlers
 type ClientMethod<S, T> = ((typeof Client.prototype.makeUnaryRequest)|
@@ -265,8 +265,16 @@ function patchClient(client: ClientModule, api: Tracer) {
       // Client is a class.
       // tslint:disable-next-line:variable-name
       const Client = makeClientConstructor.apply(this, arguments);
-      shimmer.massWrap(
-          [Client.prototype], Object.keys(methods), makeClientMethod);
+      const methodsToWrap = [
+        ...Object.keys(methods),
+        ...Object.keys(methods)
+                .map(methodName => methods[methodName].originalName)
+                .filter(
+                    originalName => !!originalName &&
+                        Client.prototype.hasOwnProperty(originalName)) as
+            string[]
+      ];
+      shimmer.massWrap([Client.prototype], methodsToWrap, makeClientMethod);
       return Client;
     };
   }
diff --git a/test/plugins/common.ts b/test/plugins/common.ts
@@ -30,6 +30,10 @@ if (semver.satisfies(process.version, '>=8')) {
   // Monkeypatch Mocha's it() to create a fresh context with each test case.
   var oldIt = global.it;
   global.it = Object.assign(function it(title, fn) {
+    // it.skip calls it without a function argument
+    if (!fn) {
+      return oldIt.call(this, title);
+    }
     function wrappedFn() {
       if (cls.exists()) {
         return cls.get().runWithContext(() => fn.apply(this, arguments), TraceCLS.UNCORRELATED);
diff --git a/test/plugins/test-trace-grpc.ts b/test/plugins/test-trace-grpc.ts
@@ -278,6 +278,8 @@ Object.keys(versions).forEach(function(version) {
   var client;
   var shouldTraceArgs: any[] = [];
   describe(version, function() {
+    const skipFor1_6 = version === 'grpc1_6' ? it.skip : it;
+
     before(function() {
       // Set up to record invocations of shouldTrace
       shimmer.wrap(TracingPolicy, 'createTracePolicy', function(original) {
@@ -356,16 +358,6 @@ Object.keys(versions).forEach(function(version) {
       });
     });
 
-    it('should propagate context', function(done) {
-      common.runInTransaction(function(endTransaction) {
-        callUnary(client, grpc, {}, function() {
-          assert.ok(common.hasContext());
-          endTransaction();
-          done();
-        });
-      });
-    });
-
     it('should accurately measure time for client streaming requests', function(done) {
       var start = Date.now();
       common.runInTransaction(function(endTransaction) {
@@ -431,6 +423,39 @@ Object.keys(versions).forEach(function(version) {
       });
     });
 
+    // Older versions of gRPC (<1.7) do not add original names.
+    skipFor1_6('should trace client requests using the original method name', (done) => {
+      common.runInTransaction((endTransaction) => {
+        // The original method name is TestUnary.
+        client.TestUnary({n: 10}, (err, result) => {
+          assert.ifError(err);
+          assert.strictEqual(result.n, 10);
+          endTransaction();
+          var assertTraceProperties = function(predicate) {
+            var trace = common.getMatchingSpan(predicate);
+            assert.ok(trace);
+            assert.strictEqual(trace.labels.argument, '{"n":10}');
+            assert.strictEqual(trace.labels.result, '{"n":10}');
+          };
+          assertTraceProperties(grpcClientPredicate);
+          assertTraceProperties(grpcServerOuterPredicate);
+          // Check that a child span was created in gRPC root span 
+          assert.ok(common.getMatchingSpan(grpcServerInnerPredicate));
+          done();
+        });
+      });
+    });
+
+    it('should propagate context', function(done) {
+      common.runInTransaction(function(endTransaction) {
+        callUnary(client, grpc, {}, function() {
+          assert.ok(common.hasContext());
+          endTransaction();
+          done();
+        });
+      });
+    });
+
     it('should not break if no parent transaction', function(done) {
       callUnary(client, grpc, {}, function() {
         assert.strictEqual(common.getMatchingSpans(grpcClientPredicate).length, 0);
