Change callback lookup to not depend on gRPC version

googleapis · Jul 1, 2016 · 68038cb · 68038cb
1 parent b15a705
commit 68038cb
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 25 deletions.
diff --git a/lib/hooks/userspace/hook-grpc.js b/lib/hooks/userspace/hook-grpc.js
@@ -25,7 +25,16 @@ var agent;
 
 var SUPPORTED_VERSIONS = '0.13 - 0.15';
 
-function makeClientMethod(useDeprecatedArgumentOrder, method, name) {
+function findFirstIndex(arr, predicate) {
+  for (var i = 0; i < arr.length; ++i) {
+    if (predicate(arr[i])) {
+      return i;
+    }
+  }
+  return -1;
+}
+
+function makeClientMethod(method, name) {
   return function clientMethodTrace() {
     var root = cls.getRootContext();
     if (!root) {
@@ -37,20 +46,13 @@ function makeClientMethod(useDeprecatedArgumentOrder, method, name) {
     var span = agent.startSpan('grpc-call-' + name);
     // Check if the response is through a stream or a callback.
     if (!method.responseStream) {
-      // Grab the callback which is always required.
-      // Depending on the version of grpc, the position of the callback
-      // function differs.
       // We need to wrap the callback with the context, to propagate it.
-      var cbIndex;
-      if (useDeprecatedArgumentOrder) {
-        cbIndex = method.requestStream ? 0 : 1;
-      } else {
-        cbIndex = arguments.length - 1;
-      }
-      // If the arguments are incorrect, we want gRPC to throw the Error
-      // so we do not wrap the callback unnecessarily.
-      if (cbIndex >= 0 && cbIndex < arguments.length &&
-          typeof arguments[cbIndex] === 'function') {
+      // The callback is always required. It should be the only function in the
+      // arguments, since we cannot send a function as an argument through gRPC.
+      var cbIndex = findFirstIndex(arguments, function(arg) {
+        return typeof arg === 'function';
+      });
+      if (cbIndex !== -1) {
         arguments[cbIndex] = wrapCallback(span, arguments[cbIndex]);
       }
     }
@@ -80,12 +82,10 @@ function wrapCallback(span, done) {
   return cls.getNamespace().bind(fn);
 }
 
-function makeClientConstructorWrap(useDeprecatedArgumentOrder,
-                                   makeClientConstructor) {
+function makeClientConstructorWrap(makeClientConstructor) {
   return function makeClientConstructorTrace(methods) {
     var Client = makeClientConstructor.apply(this, arguments);
-    shimmer.massWrap(Client.prototype, Object.keys(methods),
-        makeClientMethod.bind(null, useDeprecatedArgumentOrder));
+    shimmer.massWrap(Client.prototype, Object.keys(methods), makeClientMethod);
     return Client;
   };
 }
@@ -99,14 +99,12 @@ module.exports = function(version_, agent_) {
     'src/node/src/client.js': {
       patch: function(client) {
         agent = agent_;
-        // If version < 0.14, use the old argument order for client methods.
-        var useDeprecatedArgumentOrder = semver.satisfies(version_, '<0.14');
         shimmer.wrap(client, 'makeClientConstructor',
-            makeClientConstructorWrap.bind(null, useDeprecatedArgumentOrder));
+            makeClientConstructorWrap);
       },
       unpatch: function(client) {
-        // Only the client constructor is unwrapped, so that future grpc.load's
-        // will not wrap client methods with tracing. However, existing Client
+        // Only the Client constructor is unwrapped, so that future grpc.load's
+        // will not wrap Client methods with tracing. However, existing Client
         // objects with wrapped prototype methods will continue tracing.
         shimmer.unwrap(client, 'makeClientConstructor');
         agent_.logger.info('gRPC makeClientConstructor: unpatched');

diff --git a/test/standalone/test-grpc-context.js b/test/standalone/test-grpc-context.js
@@ -94,7 +94,7 @@ Object.keys(versions).forEach(function(version) {
       });
 
       client = new proto.Tester('localhost:' + grpcPort,
-        grpc.credentials.createInsecure());
+          grpc.credentials.createInsecure());
 
       server = app.listen(common.serverPort, function() {
         grpcServer = new grpc.Server();
@@ -124,7 +124,7 @@ Object.keys(versions).forEach(function(version) {
           }
         });
         grpcServer.bind('localhost:' + grpcPort,
-          grpc.ServerCredentials.createInsecure());
+            grpc.ServerCredentials.createInsecure());
         grpcServer.start();
         done();
       });