Context propagation for google-gax (#404)

This should allow tracing of pubsub, vision, monitoring, spanner,
speech, and nlp.

PR-URL: #404

config.js

@@ -45,6 +45,7 @@ module.exports = {
     plugins: {
       'connect': path.join(__dirname, 'src/plugins/plugin-connect.js'),
       'express': path.join(__dirname, 'src/plugins/plugin-express.js'),
+      'google-gax': path.join(__dirname, 'src/plugins/plugin-google-gax.js'),
       'grpc': path.join(__dirname, 'src/plugins/plugin-grpc.js'),
       'hapi': path.join(__dirname, 'src/plugins/plugin-hapi.js'),
       'http': path.join(__dirname, 'src/plugins/plugin-http.js'),

src/plugins/plugin-google-gax.js

@@ -0,0 +1,50 @@
+/**
+ * Copyright 2017 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+'use strict';
+
+var shimmer = require('shimmer');
+
+module.exports = [
+  {
+    file: '',
+    versions: '^0.10.x',
+    patch: function(gax, api) {
+      shimmer.wrap(gax, 'createApiCall', function(createApiCall) {
+        return function createApiCall_trace(funcWithAuth, settings, optDescriptor) {
+          var funcWithAuthThen = funcWithAuth.then;
+          funcWithAuth.then = function(cb) {
+            var result = funcWithAuthThen.call(this, cb);
+            var resultThen = result.then;
+            result.then = function(cb) {
+              return resultThen.call(this, api.wrap(cb));
+            };
+            return result;
+          };
+          var apiCallInner = createApiCall.call(this, funcWithAuth, settings, optDescriptor);
+          return function apiCallInner_trace(request, callOptions, callback) {
+            // This api.wrap is only applied to ensure context is restored when the user callback
+            // is invoked. It is not required to trace google api libraries.
+            return apiCallInner.call(this, request, callOptions, api.wrap(callback));
+          };
+        };
+      });
+    },
+    unpatch: function(gax) {
+      shimmer.unwrap(gax, 'createApiCall');
+    }
+  }
+];

test/plugins/fixtures/google-cloud-speech0.6/index.js

@@ -0,0 +1 @@
+module.exports = require('@google-cloud/speech');

test/plugins/fixtures/google-cloud-speech0.6/package.json

@@ -0,0 +1,8 @@
+{
+  "name": "google-cloud-speech0.6",
+  "version": "1.0.0",
+  "main": "index.js",
+  "dependencies": {
+    "@google-cloud/speech": "^0.6.0"
+  }
+}

test/plugins/test-trace-google-gax.js

@@ -0,0 +1,61 @@
+/**
+ * Copyright 2017 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+'use strict';
+
+var common = require('./common.js');
+var assert = require('assert');
+var path = require('path');
+
+describe('google-gax', function() {
+  var agent;
+  var speech;
+
+  before(function() {
+    agent = require('../..').start({
+      enhancedDatabaseReporting: true,
+      samplingRate: 0
+    });
+    process.env.GOOGLE_APPLICATION_CREDENTIALS =
+        path.join(__dirname, '..', 'fixtures', 'gcloud-credentials.json');
+    speech = require('./fixtures/google-cloud-speech0.6')();
+  });
+
+  after(function() {
+    delete process.env.GOOGLE_APPLICATION_CREDENTIALS;
+  });
+
+  it('should not interfere with google-cloud api tracing', function(done) {
+    common.runInTransaction(agent, function(endRootSpan) {
+      speech.recognize('./index.js', {
+        encoding: 'LINEAR16',
+        sampleRate: 16000
+      }, function(err, res) {
+        endRootSpan();
+        // Authentication will fail due to invalid credentials but a span will still be
+        // generated.
+        assert.equal(err.message, 'invalid_client');
+        assert.equal(err.code, 16);
+        var span = common.getMatchingSpan(agent, function(span) {
+          return span.kind === 'RPC_CLIENT' && span.name.indexOf('grpc:') === 0;
+        });
+        assert.ok(span);
+        assert.equal(span.name, 'grpc:/google.cloud.speech.v1beta1.Speech/SyncRecognize');
+        done();
+      });
+    });
+  });
+
+});

test/test-config-plugins.js

@@ -21,7 +21,7 @@ var trace = require('..');
 
 var common = require('./plugins/common.js');
 
-var instrumentedModules = ['connect', 'express', 'grpc', 'hapi', 'http', 'koa',
+var instrumentedModules = ['connect', 'express', 'google-gax', 'grpc', 'hapi', 'http', 'koa',
   'mongodb-core', 'mysql', 'redis', 'restify'];
 
 describe('plugin configuration', function() {