feat: Async hooks based context tracking (#538)

Introduces a lower overhead, more correct context tracking
mechanism based on async hooks that is available when using the
trace agent on newer versions of node (>=8). To enable async
hooks based context, set the GCLOUD_TRACE_NEW_CONTEXT environment
variable. This change is not expected to have observable behavior
difference but due to the nature of context propagation it is
difficult to say for sure without real world testing.

PR-URL: #538

.travis.yml

@@ -19,7 +19,7 @@ before_install:
   - psql -c "alter user postgres with password 'Password12!';" -U postgres
 
 env:
-  - CXX=g++-4.8
+  - CXX=g++-4.8 GCLOUD_TRACE_NEW_CONTEXT=1
 
 script:
   - ./bin/run-test.sh -c

appveyor.yml

@@ -23,8 +23,8 @@ install:
   - npm install
   - ps: ForEach ($fix in Get-ChildItem -Directory test/plugins/fixtures/*) { cd $fix; npm install --silent }
   - ps: cd ..\..\..\..
-  # set GCLOUD_PROJECT
   - SET GCLOUD_PROJECT=0
+  - SET GCLOUD_TRACE_NEW_CONTEXT=1
 
 before_test:
   - SET PATH=C:\Program Files\MySql\MySQL Server 5.7\bin;%PATH%

circle.yml

@@ -5,3 +5,4 @@ machine:
     - redis
   environment:
     GCLOUD_PROJECT: 0
+    GCLOUD_TRACE_NEW_CONTEXT: 1

index.js

@@ -20,7 +20,10 @@ var filesLoadedBeforeTrace = Object.keys(require.cache);
 
 // Load continuation-local-storage first to ensure the core async APIs get
 // patched before any user-land modules get loaded.
-require('continuation-local-storage');
+if (require('semver').satisfies(process.version, '<8') &&
+    process.env.GCLOUD_TRACE_NEW_CONTEXT) {
+  require('continuation-local-storage');
+}
 
 var path = require('path');
 var cls = require('./src/cls.js');

src/cls-ah.js

@@ -0,0 +1,113 @@
+/**
+ * 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';
+
+const asyncHook = require('async_hooks');
+
+const wrappedSymbol = Symbol('context_wrapped');
+let contexts = {};
+let current = {};
+
+asyncHook.createHook({init, before, destroy}).enable();
+
+function Namespace() {}
+
+Namespace.prototype.get = function(k) {
+  return current[k];
+};
+
+Namespace.prototype.set = function(k, v) {
+  current[k] = v;
+};
+
+Namespace.prototype.run = function(fn) {
+  const oldContext = current;
+  current = {};
+  const res = fn();
+  current = oldContext;
+  return res;
+};
+
+Namespace.prototype.runAndReturn = Namespace.prototype.run;
+
+Namespace.prototype.bind = function(cb) {
+  if (cb[wrappedSymbol] || !current) {
+    return cb;
+  }
+  const boundContext = current;
+  const contextWrapper = function() {
+    const oldContext = current;
+    current = boundContext;
+    const res = cb.apply(this, arguments);
+    current = oldContext;
+    return res;
+  };
+  contextWrapper[wrappedSymbol] = true;
+  Object.defineProperty(contextWrapper, 'length', {
+    enumerable: false,
+    configurable: true,
+    writable: false,
+    value: cb.length
+  });
+  return contextWrapper;
+};
+
+const eventEmitterMethods =
+  [ 'addListener', 'on', 'once', 'prependListener', 'prependOncelistener' ];
+
+// This function is not technically needed and all tests currently pass without it
+// (after removing call sites). While it is not a complete solution, restoring
+// correct context before running every request/response event handler reduces
+// the number of situations in which userspace queuing will cause us to lose context.
+Namespace.prototype.bindEmitter = function(ee) {
+  var ns = this;
+  eventEmitterMethods.forEach(function(method) {
+    const oldMethod = ee[method];
+    ee[method] = function(e, f) { return oldMethod.call(this, e, ns.bind(f)); };
+  });
+};
+
+var namespace = new Namespace();
+
+// AsyncWrap Hooks
+
+function init(uid, provider, parentUid, parentHandle) {
+  contexts[uid] = current;
+}
+
+function before(uid) {
+  if (contexts[uid]) {
+    current = contexts[uid];
+  }
+}
+
+function destroy(uid) {
+  delete contexts[uid];
+}
+
+module.exports = {
+  createNamespace: function() {
+    return namespace;
+  },
+
+  destroyNamespace: function() {
+    current = {};
+    contexts = {};
+  },
+
+  getNamespace: function() { return namespace; }
+};

src/cls.js

@@ -16,8 +16,10 @@
 
 'use strict';
 
-
-var cls = require('continuation-local-storage');
+var semver = require('semver');
+var cls = semver.satisfies(process.version, '>=8') &&
+          process.env.GCLOUD_TRACE_NEW_CONTEXT ?
+            require('./cls-ah.js') : require('continuation-local-storage');
 
 /** @const {string} */
 var TRACE_NAMESPACE = 'com.google.cloud.trace';

src/trace-api.js

@@ -24,6 +24,9 @@ var Trace = require('./trace.js');
 var SpanData = require('./span-data.js');
 var uuid = require('uuid');
 var TracingPolicy = require('./tracing-policy.js');
+var semver = require('semver');
+
+var ROOT_SPAN_STACK_OFFSET = semver.satisfies(process.version, '>=8') ? 0 : 2;
 
 /**
  * Phantom implementation of the trace api. When disabled, a TraceAgent instance
@@ -170,7 +173,7 @@ TraceAgent.prototype.runInRootSpan = function(options, fn) {
       options.name, /* Span name */
       parentId, /* Parent's span ID */
       true, /* Is root span */
-      options.skipFrames ? options.skipFrames + 2 : 2);
+      ROOT_SPAN_STACK_OFFSET + (options.skipFrames || 0));
     rootContext.span.kind = 'RPC_SERVER';
     cls.setRootContext(rootContext);
     return fn(rootContext);

test/plugins/common.js

@@ -21,9 +21,18 @@ proxyquire('gcp-metadata', {
   'retry-request': require('request')
 });
 
+var semver = require('semver');
+
 var logger = require('@google-cloud/common').logger;
 var trace = require('../..');
 var cls = require('../../src/cls.js');
+if (semver.satisfies(process.version, '>=8')) {
+  // Monkeypatch the monkeypatcher
+  var oldIt = global.it;
+  global.it = function it(title, fn) {
+    return oldIt.call(this, title, cls.createNamespace().bind(fn));
+  };
+}
 var tracePolicy = require('../../src/tracing-policy.js');
 var TraceWriter = require('../../src/trace-writer.js');
 

test/plugins/test-trace-grpc.js

@@ -284,6 +284,19 @@ Object.keys(versions).forEach(function(version) {
 
       grpc = require(versions[version]);
 
+      var oldRegister = grpc.Server.prototype.register;
+      grpc.Server.prototype.register = function register(n, h, s, d, m) {
+        var result = oldRegister.call(this, n, h, s, d, m);
+        var oldFunc = this.handlers[n].func;
+        this.handlers[n].func = function() {
+          if (cls.getRootContext()) {
+            cls.setRootContext(null);
+          }
+          return oldFunc.apply(this, arguments);
+        };
+        return result;
+      };
+
       // This metadata can be used by all test methods.
       metadata = new grpc.Metadata();
       metadata.set('a', 'b');
@@ -432,6 +445,17 @@ Object.keys(versions).forEach(function(version) {
     });
 
     it('should support distributed trace context', function(done) {
+      function makeLink(fn, meta, next) {
+        return function() {
+          cls.setRootContext(null);
+          common.runInTransaction(function(terminate) {
+            fn(client, grpc, meta, function() {
+              terminate();
+              next();
+            });
+          });
+        };
+      }
       // Enable asserting properties of the metdata on the grpc server.
       checkMetadata = true;
       common.runInTransaction(function (endTransaction) {
@@ -445,16 +469,16 @@ Object.keys(versions).forEach(function(version) {
         // the grpc client methods at all if no fields are present).
         // The plugin should automatically create a new Metadata object and
         // populate it with trace context data accordingly.
-        next = callUnary.bind(null, client, grpc, {}, next);
-        next = callClientStream.bind(null, client, grpc, {}, next);
-        next = callServerStream.bind(null, client, grpc, {}, next);
-        next = callBidi.bind(null, client, grpc, {}, next);
+        next = makeLink(callUnary, {}, next);
+        next = makeLink(callClientStream, {}, next);
+        next = makeLink(callServerStream, {}, next);
+        next = makeLink(callBidi, {}, next);
         // Try with metadata. The plugin should simply add trace context data
         // to it.
-        next = callUnary.bind(null, client, grpc, metadata, next);
-        next = callClientStream.bind(null, client, grpc, metadata, next);
-        next = callServerStream.bind(null, client, grpc, metadata, next);
-        next = callBidi.bind(null, client, grpc, metadata, next);
+        next = makeLink(callUnary, metadata, next);
+        next = makeLink(callClientStream, metadata, next);
+        next = makeLink(callServerStream, metadata, next);
+        next = makeLink(callBidi, metadata, next);
         next();
       });
     });

test/test-cls-ah.js

@@ -0,0 +1,102 @@
+/**
+ * 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 assert = require('assert');
+var http = require('http');
+var semver = require('semver');
+
+if (semver.satisfies(process.version, '<8')) {
+  console.log('Skipping cls-ah tests on node version without async hooks');
+  return;
+} else {
+  assert.ok(process.env.GCLOUD_TRACE_NEW_CONTEXT);
+}
+
+var cls = require('../src/cls-ah.js');
+
+describe('test-cls-ah', function() {
+  it('should preserve request context', function(done) {
+    var namespace = cls.createNamespace();
+    var id = 0;
+    var ended = 0;
+    var server = http.createServer(function(req, res) {
+      var reqId = id++;
+      namespace.run(function() {
+        namespace.set('id', reqId);
+        assert.equal(namespace.get('id'), reqId);
+        var count = 0;
+        var i = setInterval(function () {
+          assert.equal(namespace.get('id'), reqId);
+          if (count++ > reqId) {
+            clearInterval(i);
+            res.end('yay');
+            if (++ended === 10) {
+              done();
+            }
+          }
+        }, Math.random() * 50);
+      });
+    });
+    server.listen(8080, function() {
+      for (var i = 0; i < 10; i++) {
+        http.get('http://localhost:8080');
+      }
+    });
+  });
+
+  it('should correctly run context in series', function(done) {
+    var namespace = cls.createNamespace();
+    namespace.run(function() {
+      assert.equal(namespace.get('id'), null);
+      namespace.set('id', 'first');
+      setTimeout(function() {
+        assert.equal(namespace.get('id'), 'first');
+        done();
+      }, 30);
+      assert.equal(namespace.get('id'), 'first');
+    });
+    namespace.run(function() {
+      assert.equal(namespace.get('id'), null);
+      namespace.set('id', 'second');
+      setTimeout(function() {
+        assert.equal(namespace.get('id'), 'second');
+      }, 10);
+      assert.equal(namespace.get('id'), 'second');
+    });
+  });
+
+  it('should correctly run context nested', function(done) {
+    var namespace = cls.createNamespace();
+    namespace.run(function() {
+      assert.equal(namespace.get('id'), null);
+      namespace.set('id', 'first');
+      namespace.run(function() {
+        assert.equal(namespace.get('id'), null);
+        namespace.set('id', 'second');
+        setTimeout(function() {
+          assert.equal(namespace.get('id'), 'second');
+        }, 10);
+        assert.equal(namespace.get('id'), 'second');
+      });
+      setTimeout(function() {
+        assert.equal(namespace.get('id'), 'first');
+        done();
+      }, 30);
+      assert.equal(namespace.get('id'), 'first');
+    });
+  });
+});

test/test-span-data.js

@@ -124,7 +124,7 @@ describe('SpanData', function() {
 
   it('captures stack traces', function() {
     TraceWriter.get().config_.stackTraceLimit = 25;
-    cls.getNamespace().run(function() {
+    cls.getNamespace().run(function awesome() {
       var spanData = createRootSpanData('name', 1, 2, 1);
       assert.ok(!spanData.span.isClosed());
       spanData.endSpan();
@@ -134,19 +134,19 @@ describe('SpanData', function() {
       var frames = JSON.parse(stack);
       assert.ok(frames && frames.stack_frame);
       assert.ok(Array.isArray(frames.stack_frame));
-      assert.strictEqual(frames.stack_frame[1].method_name, 'Namespace.run [as run]');
+      assert.strictEqual(frames.stack_frame[0].method_name, 'awesome');
     });
   });
 
   it('does not limit stack trace', function() {
     TraceWriter.get().config_.maximumLabelValueSize = 10;
-    cls.getNamespace().run(function() {
+    cls.getNamespace().run(function awesome() {
       var spanData = createRootSpanData('name', 1, 2, 1);
       spanData.endSpan();
       var stack = spanData.span.labels[TraceLabels.STACK_TRACE_DETAILS_KEY];
       assert.ok(stack.length > 10);
       var frames = JSON.parse(stack);
-      assert.strictEqual(frames.stack_frame[1].method_name, 'Namespace.run [as run]');
+      assert.strictEqual(frames.stack_frame[0].method_name, 'awesome');
     });
   });