fix(grpc): patch loadPackageDefinition

packages/opentelemetry-plugin-grpc/package.json

@@ -38,6 +38,7 @@
     "access": "public"
   },
   "devDependencies": {
+    "@grpc/proto-loader": "^0.4.0",
     "@types/mocha": "^5.2.7",
     "@types/node": "^12.6.9",
     "@types/shimmer": "^1.0.1",

packages/opentelemetry-plugin-grpc/src/grpc.ts

@@ -26,7 +26,6 @@ import {
 import { AttributeNames } from './enums/AttributeNames';
 import {
   grpc,
-  ModuleExportsMapping,
   GrpcPluginOptions,
   ServerCallWithMeta,
   SendUnaryDataCallback,
@@ -60,10 +59,12 @@ export class GrpcPlugin extends BasePlugin<grpc> {
     this._config = {};
   }
 
-  protected readonly _internalFilesList: ModuleExportsMapping = {
-    '0.13 - 1.6': { client: 'src/node/src/client.js' },
-    '^1.7': { client: 'src/client.js' },
-  };
+  // TODO: uncomment once makeClientConstructor patch is uncommented
+  // protected readonly _internalFilesList: ModuleExportsMapping = {
+  //   '0.13 - 1.6': { client: 'src/node/src/client.js' },
+  //   '^1.7': { client: 'src/client.js' },
+  // };
+
   protected readonly _basedir = basedir;
 
   protected patch(): typeof grpcTypes {
@@ -73,6 +74,20 @@ export class GrpcPlugin extends BasePlugin<grpc> {
       this.version
     );
 
+    // TODO: uncomment on completion: https://github.com/open-telemetry/opentelemetry-js/issues/285
+    // if (this._internalFilesExports['client']) {
+    //   grpcClientModule = this._internalFilesExports[
+    //     'client'
+    //   ] as GrpcInternalClientTypes;
+
+    //   // Wrap the internally used client constructor
+    //   shimmer.wrap(
+    //     grpcClientModule,
+    //     'makeClientConstructor',
+    //     this._patchClient()
+    //   );
+    // }
+
     if (this._moduleExports.Server) {
       shimmer.wrap(
         this._moduleExports.Server.prototype,
@@ -91,16 +106,11 @@ export class GrpcPlugin extends BasePlugin<grpc> {
       );
     }
 
-    if (this._internalFilesExports['client']) {
-      grpcClientModule = this._internalFilesExports[
-        'client'
-      ] as GrpcInternalClientTypes;
-
-      // Wrap the internally used client constructor
+    if (this._moduleExports.loadPackageDefinition) {
       shimmer.wrap(
-        grpcClientModule,
-        'makeClientConstructor',
-        this._patchClient()
+        this._moduleExports,
+        'loadPackageDefinition',
+        this._patchLoadPackageDefinition()
       );
     }
 
@@ -121,6 +131,10 @@ export class GrpcPlugin extends BasePlugin<grpc> {
       shimmer.unwrap(this._moduleExports, 'makeGenericClientConstructor');
     }
 
+    if (this._moduleExports.loadPackageDefinition) {
+      shimmer.unwrap(this._moduleExports, 'loadPackageDefinition');
+    }
+
     if (grpcClientModule) {
       shimmer.unwrap(grpcClientModule, 'makeClientConstructor');
     }
@@ -321,6 +335,59 @@ export class GrpcPlugin extends BasePlugin<grpc> {
     return (original as any).call(self, call);
   }
 
+  private _patchLoadPackageDefinition() {
+    const plugin = this;
+    return (original: typeof grpcTypes.loadPackageDefinition) => {
+      plugin._logger.debug('patching loadPackageDefinition');
+      return function loadPackageDefinition(
+        this: grpc,
+        packageDef: grpcTypes.PackageDefinition
+      ) {
+        const result = original.apply(this, arguments as never);
+        // Copied from exports.loadPackageDefintion(...)
+        for (const serviceFqn in packageDef) {
+          const nameComponents = serviceFqn.split('.');
+          const serviceName = nameComponents[nameComponents.length - 1];
+          let current = result;
+          for (const packageName of nameComponents.slice(0, -1)) {
+            if (!current[packageName]) {
+              current[packageName] = {};
+            }
+            current = current[packageName] as grpcTypes.GrpcObject;
+          }
+
+          // if makeClientConstructor was used, patch the client
+          if (
+            current[serviceName].prototype instanceof
+            plugin._moduleExports.Client
+          ) {
+            const serviceClient = current[
+              serviceName
+            ] as grpcTypes.ProtobufMessage;
+            const methodList: string[] = [];
+            (Object.values(
+              serviceClient.prototype.$method_names
+            ) as string[]).forEach(method => {
+              const originalName =
+                serviceClient.service[method as string].originalName;
+
+              // e.g. push both "Capitalize" and "capitalize"
+              originalName && methodList.push(originalName);
+              methodList.push(method);
+            });
+
+            shimmer.massWrap(
+              serviceClient.prototype,
+              methodList as never[],
+              plugin._getPatchedClientMethods() as never
+            );
+          }
+        }
+        return result;
+      };
+    };
+  }
+
   private _patchClient() {
     const plugin = this;
     return (original: typeof grpcTypes.makeGenericClientConstructor): never => {
@@ -471,12 +538,15 @@ export class GrpcPlugin extends BasePlugin<grpc> {
         ((call as unknown) as events.EventEmitter).on(
           'status',
           (status: Status) => {
-            span.setStatus({ code: CanonicalCode.OK });
-            span.setAttribute(
-              AttributeNames.GRPC_STATUS_CODE,
-              status.code.toString()
-            );
-            endSpan();
+            // if an error was emitted, the span will be ended there
+            if (status.code === 0) {
+              span.setStatus({ code: CanonicalCode.OK });
+              span.setAttribute(
+                AttributeNames.GRPC_STATUS_CODE,
+                status.code.toString()
+              );
+              endSpan();
+            }
           }
         );
       }

packages/opentelemetry-plugin-grpc/test/grpc.test.ts

@@ -32,6 +32,12 @@ import * as grpc from 'grpc';
 import * as sinon from 'sinon';
 
 const PROTO_PATH = __dirname + '/fixtures/grpc-test.proto';
+const PROTO_OPTIONS = {
+  keepCae: true,
+  enums: String,
+  defaults: true,
+  oneofs: true,
+};
 const memoryExporter = new InMemorySpanExporter();
 
 type GrpcModule = typeof grpc;
@@ -299,7 +305,7 @@ describe('GrpcPlugin', () => {
 
     it('should patch client constructor makeClientConstructor() and makeGenericClientConstructor()', () => {
       plugin.enable(grpc, new NoopTracer(), new NoopLogger());
-      (plugin['_moduleExports'] as any).makeGenericClientConstructor({});
+      (grpc as any).makeGenericClientConstructor({});
       assert.strictEqual(clientPatchStub.callCount, 1);
     });
   });
@@ -536,8 +542,9 @@ describe('GrpcPlugin', () => {
         // TODO: add plugin options here once supported
       };
       plugin.enable(grpc, tracer, logger, config);
-
-      const proto = grpc.load(PROTO_PATH).pkg_test;
+      const protoLoader = require('@grpc/proto-loader');
+      const definition = protoLoader.loadSync(PROTO_PATH, PROTO_OPTIONS);
+      const proto = grpc.loadPackageDefinition(definition).pkg_test;
       server = startServer(grpc, proto);
       client = createClient(grpc, proto);
     });