async_hooks: AsyncLocalStorage to diagnostics_channel integration

Author: Stephen Belanger

lib/_http_server.js

@@ -92,6 +92,7 @@ let debug = require('internal/util/debuglog').debuglog('http', (fn) => {
 
 const dc = require('diagnostics_channel');
 const onRequestStartChannel = dc.channel('http.server.request.start');
+const onRequestEndChannel = dc.channel('http.server.request.end');
 const onResponseFinishChannel = dc.channel('http.server.response.finish');
 
 const kServerResponse = Symbol('ServerResponse');
@@ -996,13 +997,18 @@ function parserOnIncoming(server, socket, state, req, keepAlive) {
   res.shouldKeepAlive = keepAlive;
   res[kUniqueHeaders] = server[kUniqueHeaders];
 
-  if (onRequestStartChannel.hasSubscribers) {
-    onRequestStartChannel.publish({
+  const dcMessage = (
+    onRequestStartChannel.hasSubscribers ||
+    onRequestEndChannel.hasSubscribers
+  ) ? {
       request: req,
       response: res,
       socket,
       server
-    });
+    } : undefined;
+
+  if (onRequestStartChannel.hasSubscribers) {
+    onRequestStartChannel.publish(dcMessage);
   }
 
   if (socket._httpMessage) {
@@ -1063,6 +1069,10 @@ function parserOnIncoming(server, socket, state, req, keepAlive) {
     server.emit('request', req, res);
   }
 
+  if (onRequestEndChannel.hasSubscribers) {
+    onRequestEndChannel.publish(dcMessage);
+  }
+
   return 0;  // No special treatment.
 }
 

lib/async_hooks.js

@@ -9,6 +9,7 @@ const {
   FunctionPrototypeBind,
   NumberIsSafeInteger,
   ObjectDefineProperties,
+  ObjectPrototypeHasOwnProperty,
   ObjectIs,
   ReflectApply,
   Symbol,
@@ -26,6 +27,7 @@ const {
   validateString,
 } = require('internal/validators');
 const internal_async_hooks = require('internal/async_hooks');
+const { tracingChannel } = require('diagnostics_channel');
 
 // Get functions
 // For userland AsyncResources, make sure to emit a destroy event when the
@@ -351,6 +353,47 @@ class AsyncLocalStorage {
       return resource[this.kResourceStore];
     }
   }
+
+  bindToTracingChannel(channel, transform = (v) => v) {
+    if (typeof channel === 'string') {
+      channel = tracingChannel(channel);
+    }
+
+    const store = this;
+
+    function start(register) {
+      register(store, transform);
+    }
+
+    channel.subscribe({ start });
+    return () => {
+      channel.unsubscribe({ start });
+    };
+  }
+
+  static runInTracingChannel(channel, data, runner) {
+    if (typeof channel === 'string') {
+      channel = tracingChannel(channel);
+    }
+
+    const bindings = [];
+
+    function register(store, transform) {
+      bindings.push({ store, data: transform(data) });
+    }
+
+    return channel.traceSync(() => {
+      for (const { store, data } of bindings) {
+        runner = wrapRunInStoreBinding(store, data, runner);
+      }
+
+      return runner();
+    }, register);
+  }
+}
+
+function wrapRunInStoreBinding(store, data, next) {
+  return () => store.run(data, next);
 }
 
 // Placing all exports down here because the exported classes won't export

lib/diagnostics_channel.js

@@ -289,5 +289,6 @@ module.exports = {
   subscribe,
   tracingChannel,
   unsubscribe,
-  Channel
+  Channel,
+  TracingChannel
 };

test/parallel/test-async-local-storage-tracing-channel-integration.js

@@ -0,0 +1,60 @@
+'use strict';
+
+const common = require('../common');
+const assert = require('assert');
+const { AsyncLocalStorage } = require('async_hooks');
+
+const contextA = {
+  foo: 'bar'
+};
+const contextB = {
+  baz: 'buz'
+};
+
+const expected = [
+  contextA,
+  contextB
+];
+
+let builds = 0;
+let runs = 0;
+
+const storageA = new AsyncLocalStorage();
+const storageB = new AsyncLocalStorage();
+
+// Check both passthrough bindings and transformed bindings
+storageA.bindToTracingChannel('test');
+storageB.bindToTracingChannel('test', common.mustCall((received) => {
+  assert.deepStrictEqual(received, expected[builds++]);
+  return { received };
+}, 2));
+
+function check(context) {
+  assert.deepStrictEqual(storageA.getStore(), context);
+  assert.deepStrictEqual(storageB.getStore(), {
+    received: context
+  });
+}
+
+// Should have no context before run
+assert.strictEqual(storageA.getStore(), undefined);
+assert.strictEqual(storageB.getStore(), undefined);
+
+AsyncLocalStorage.runInTracingChannel('test', expected[runs], common.mustCall(() => {
+  // Should have set expected context
+  check(expected[runs]);
+
+  // Should support nested contexts
+  runs++;
+  AsyncLocalStorage.runInTracingChannel('test', expected[runs], common.mustCall(() => {
+    check(expected[runs]);
+  }));
+  runs--;
+
+  // Should have restored outer context
+  check(expected[runs]);
+}));
+
+// Should have no context after run
+assert.strictEqual(storageA.getStore(), undefined);
+assert.strictEqual(storageB.getStore(), undefined);