lib: add tracing channel to diagnostics_channel

Author: Stephen Belanger

doc/api/diagnostics_channel.md

@@ -4,9 +4,14 @@ const {
   ArrayPrototypeIndexOf,
   ArrayPrototypePush,
   ArrayPrototypeSplice,
+  FunctionPrototypeBind,
   ObjectCreate,
   ObjectGetPrototypeOf,
   ObjectSetPrototypeOf,
+  PromisePrototypeThen,
+  PromiseReject,
+  ReflectApply,
+  SafeMap,
   SymbolHasInstance,
 } = primordials;
 
@@ -23,11 +28,40 @@ const { triggerUncaughtException } = internalBinding('errors');
 
 const { WeakReference } = internalBinding('util');
 
+function decRef(channel) {
+  channel._weak.decRef();
+  if (channel._weak.getRef() === 0) {
+    delete channels[channel.name];
+  }
+}
+
+function markActive(channel) {
+  // eslint-disable-next-line no-use-before-define
+  ObjectSetPrototypeOf(channel, ActiveChannel.prototype);
+  channel._subscribers = [];
+  channel._stores = new SafeMap();
+}
+
+function maybeMarkInactive(channel) {
+  // When there are no more active subscribers, restore to fast prototype.
+  if (!channel._subscribers.length && !channel._stores.size) {
+    // eslint-disable-next-line no-use-before-define
+    ObjectSetPrototypeOf(channel, Channel.prototype);
+    channel._subscribers = undefined;
+    channel._stores = undefined;
+  }
+}
+
+function wrapStoreRun(store, data, next, transform = (v) => v) {
+  return () => store.run(transform(data), next);
+}
+
 // TODO(qard): should there be a C++ channel interface?
 class ActiveChannel {
   subscribe(subscription) {
     validateFunction(subscription, 'subscription');
     ArrayPrototypePush(this._subscribers, subscription);
+    this._weak.incRef();
   }
 
   unsubscribe(subscription) {
@@ -36,12 +70,28 @@ class ActiveChannel {
 
     ArrayPrototypeSplice(this._subscribers, index, 1);
 
-    // When there are no more active subscribers, restore to fast prototype.
-    if (!this._subscribers.length) {
-      // eslint-disable-next-line no-use-before-define
-      ObjectSetPrototypeOf(this, Channel.prototype);
+    decRef(this);
+    maybeMarkInactive(this);
+
+    return true;
+  }
+
+  bindStore(store, transform) {
+    const replacing = this._stores.has(store);
+    if (!replacing) this._weak.incRef();
+    this._stores.set(store, transform);
+  }
+
+  unbindStore(store) {
+    if (!this._stores.has(store)) {
+      return false;
     }
 
+    this._stores.delete(store);
+
+    decRef(this);
+    maybeMarkInactive(this);
+
     return true;
   }
 
@@ -61,11 +111,28 @@ class ActiveChannel {
       }
     }
   }
+
+  runStores(data, fn, thisArg, ...args) {
+    this.publish(data);
+
+    // Bind base fn first due to AsyncLocalStorage.run not having thisArg
+    fn = FunctionPrototypeBind(fn, thisArg, ...args);
+
+    for (const entry of this._stores.entries()) {
+      const store = entry[0];
+      const transform = entry[1];
+      fn = wrapStoreRun(store, data, fn, transform);
+    }
+
+    return fn();
+  }
 }
 
 class Channel {
   constructor(name) {
     this._subscribers = undefined;
+    this._stores = undefined;
+    this._weak = undefined;
     this.name = name;
   }
 
@@ -76,20 +143,32 @@ class Channel {
   }
 
   subscribe(subscription) {
-    ObjectSetPrototypeOf(this, ActiveChannel.prototype);
-    this._subscribers = [];
+    markActive(this);
     this.subscribe(subscription);
   }
 
   unsubscribe() {
     return false;
   }
 
+  bindStore(store, transform = (v) => v) {
+    markActive(this);
+    this.bindStore(store, transform);
+  }
+
+  unbindStore() {
+    return false;
+  }
+
   get hasSubscribers() {
     return false;
   }
 
   publish() {}
+
+  runStores(data, fn, thisArg, ...args) {
+    return ReflectApply(fn, thisArg, args);
+  }
 }
 
 const channels = ObjectCreate(null);
@@ -105,27 +184,17 @@ function channel(name) {
   }
 
   channel = new Channel(name);
-  channels[name] = new WeakReference(channel);
+  channel._weak = new WeakReference(channel);
+  channels[name] = channel._weak;
   return channel;
 }
 
 function subscribe(name, subscription) {
-  const chan = channel(name);
-  channels[name].incRef();
-  chan.subscribe(subscription);
+  return channel(name).subscribe(subscription);
 }
 
 function unsubscribe(name, subscription) {
-  const chan = channel(name);
-  if (!chan.unsubscribe(subscription)) {
-    return false;
-  }
-
-  channels[name].decRef();
-  if (channels[name].getRef() === 0) {
-    delete channels[name];
-  }
-  return true;
+  return channel(name).unsubscribe(subscription);
 }
 
 function hasSubscribers(name) {
@@ -139,10 +208,168 @@ function hasSubscribers(name) {
   return channel.hasSubscribers;
 }
 
+const traceEvents = [
+  'start',
+  'end',
+  'asyncStart',
+  'asyncEnd',
+  'error',
+];
+
+function assertChannel(value, name) {
+  if (!(value instanceof Channel)) {
+    throw new ERR_INVALID_ARG_TYPE(name, ['Channel'], value);
+  }
+}
+
+class TracingChannel {
+  constructor(nameOrChannels) {
+    if (typeof nameOrChannels === 'string') {
+      this.start = channel(`tracing:${nameOrChannels}:start`);
+      this.end = channel(`tracing:${nameOrChannels}:end`);
+      this.asyncStart = channel(`tracing:${nameOrChannels}:asyncStart`);
+      this.asyncEnd = channel(`tracing:${nameOrChannels}:asyncEnd`);
+      this.error = channel(`tracing:${nameOrChannels}:error`);
+    } else if (typeof nameOrChannels === 'object') {
+      const { start, end, asyncStart, asyncEnd, error } = nameOrChannels;
+
+      assertChannel(start, 'nameOrChannels.start');
+      assertChannel(end, 'nameOrChannels.end');
+      assertChannel(asyncStart, 'nameOrChannels.asyncStart');
+      assertChannel(asyncEnd, 'nameOrChannels.asyncEnd');
+      assertChannel(error, 'nameOrChannels.error');
+
+      this.start = start;
+      this.end = end;
+      this.asyncStart = asyncStart;
+      this.asyncEnd = asyncEnd;
+      this.error = error;
+    } else {
+      throw new ERR_INVALID_ARG_TYPE('nameOrChannels',
+                                     ['string', 'object', 'Channel'],
+                                     nameOrChannels);
+    }
+  }
+
+  subscribe(handlers) {
+    for (const name of traceEvents) {
+      if (!handlers[name]) continue;
+
+      this[name]?.subscribe(handlers[name]);
+    }
+  }
+
+  unsubscribe(handlers) {
+    let done = true;
+
+    for (const name of traceEvents) {
+      if (!handlers[name]) continue;
+
+      if (!this[name]?.unsubscribe(handlers[name])) {
+        done = false;
+      }
+    }
+
+    return done;
+  }
+
+  traceSync(fn, ctx = {}, thisArg, ...args) {
+    const { start, end, error } = this;
+
+    try {
+      const result = start.runStores(ctx, fn, thisArg, ...args);
+      ctx.result = result;
+      return result;
+    } catch (err) {
+      ctx.error = err;
+      error.publish(ctx);
+      throw err;
+    } finally {
+      end.publish(ctx);
+    }
+  }
+
+  tracePromise(fn, ctx = {}, thisArg, ...args) {
+    const { start, end, asyncStart, asyncEnd, error } = this;
+
+    function reject(err) {
+      ctx.error = err;
+      error.publish(ctx);
+      asyncStart.publish(ctx);
+      // TODO: Is there a way to have asyncEnd _after_ the continuation?
+      asyncEnd.publish(ctx);
+      return PromiseReject(err);
+    }
+
+    function resolve(result) {
+      ctx.result = result;
+      asyncStart.publish(ctx);
+      // TODO: Is there a way to have asyncEnd _after_ the continuation?
+      asyncEnd.publish(ctx);
+      return result;
+    }
+
+    try {
+      const promise = start.runStores(ctx, fn, thisArg, ...args);
+      return PromisePrototypeThen(promise, resolve, reject);
+    } catch (err) {
+      ctx.error = err;
+      error.publish(ctx);
+      throw err;
+    } finally {
+      end.publish(ctx);
+    }
+  }
+
+  traceCallback(fn, position = 0, ctx = {}, thisArg, ...args) {
+    const { start, end, asyncStart, asyncEnd, error } = this;
+
+    function wrap(fn) {
+      return function wrappedCallback(err, res) {
+        if (err) {
+          ctx.error = err;
+          error.publish(ctx);
+        } else {
+          ctx.result = res;
+        }
+
+        asyncStart.publish(ctx);
+        try {
+          if (fn) {
+            return ReflectApply(fn, this, arguments);
+          }
+        } finally {
+          asyncEnd.publish(ctx);
+        }
+      };
+    }
+
+    if (position >= 0) {
+      args.splice(position, 1, wrap(args.at(position)));
+    }
+
+    try {
+      return start.runStores(ctx, fn, thisArg, ...args);
+    } catch (err) {
+      ctx.error = err;
+      error.publish(ctx);
+      throw err;
+    } finally {
+      end.publish(ctx);
+    }
+  }
+}
+
+function tracingChannel(nameOrChannels) {
+  return new TracingChannel(nameOrChannels);
+}
+
 module.exports = {
   channel,
   hasSubscribers,
   subscribe,
+  tracingChannel,
   unsubscribe,
-  Channel
+  Channel,
+  TracingChannel
 };

lib/diagnostics_channel.js

@@ -0,0 +1,87 @@
+'use strict';
+
+const common = require('../common');
+const assert = require('assert');
+const dc = require('diagnostics_channel');
+const { AsyncLocalStorage } = require('async_hooks');
+
+let n = 0;
+const thisArg = new Date();
+const inputs = [
+  { foo: 'bar' },
+  { baz: 'buz' },
+];
+
+const channel = dc.channel('test');
+
+// Bind a storage directly to published data
+const store1 = new AsyncLocalStorage();
+channel.bindStore(store1);
+
+// Bind a store with transformation of published data
+const store2 = new AsyncLocalStorage();
+channel.bindStore(store2, common.mustCall((data) => {
+  assert.deepStrictEqual(data, inputs[n]);
+  return { data };
+}, 3));
+
+// Regular subscribers should see publishes from runStores calls
+channel.subscribe(common.mustCall((data) => {
+  assert.deepStrictEqual(data, inputs[n]);
+}, 3));
+
+// Verify stores are empty before run
+assert.strictEqual(store1.getStore(), undefined);
+assert.strictEqual(store2.getStore(), undefined);
+
+channel.runStores(inputs[n], common.mustCall(function(a, b) {
+  // Verify this and argument forwarding
+  assert.deepStrictEqual(this, thisArg);
+  assert.strictEqual(a, 1);
+  assert.strictEqual(b, 2);
+
+  // Verify store 1 state matches input
+  assert.deepStrictEqual(store1.getStore(), inputs[n]);
+
+  // Verify store 2 state has expected transformation
+  assert.deepStrictEqual(store2.getStore(), { data: inputs[n] });
+
+  // Should support nested contexts
+  n++;
+  channel.runStores(inputs[n], common.mustCall(function() {
+    // Verify this and argument forwarding
+    assert.strictEqual(this, undefined);
+
+    // Verify store 1 state matches input
+    assert.deepStrictEqual(store1.getStore(), inputs[n]);
+
+    // Verify store 2 state has expected transformation
+    assert.deepStrictEqual(store2.getStore(), { data: inputs[n] });
+  }));
+  n--;
+
+  // Verify store 1 state matches input
+  assert.deepStrictEqual(store1.getStore(), inputs[n]);
+
+  // Verify store 2 state has expected transformation
+  assert.deepStrictEqual(store2.getStore(), { data: inputs[n] });
+}), thisArg, 1, 2);
+
+// Verify stores are empty after run
+assert.strictEqual(store1.getStore(), undefined);
+assert.strictEqual(store2.getStore(), undefined);
+
+// Verify unbinding works
+assert.ok(channel.unbindStore(store1));
+
+// Verify unbinding a store that is not bound returns false
+assert.ok(!channel.unbindStore(store1));
+
+n++;
+channel.runStores(inputs[n], common.mustCall(() => {
+  // Verify after unbinding store 1 will remain undefined
+  assert.strictEqual(store1.getStore(), undefined);
+
+  // Verify still bound store 2 receives expected data
+  assert.deepStrictEqual(store2.getStore(), { data: inputs[n] });
+}));

test/parallel/test-diagnostics-channel-bind-store.js

@@ -0,0 +1,46 @@
+'use strict';
+
+const common = require('../common');
+const dc = require('diagnostics_channel');
+const assert = require('assert');
+
+const channel = dc.tracingChannel('test');
+
+const expectedError = new Error('test');
+const input = { foo: 'bar' };
+const thisArg = { baz: 'buz' };
+
+function check(found) {
+  assert.deepStrictEqual(found, input);
+}
+
+const handlers = {
+  start: common.mustCall(check, 2),
+  end: common.mustCall(check, 2),
+  asyncStart: common.mustCall(check, 2),
+  asyncEnd: common.mustCall(check, 2),
+  error: common.mustCall((found) => {
+    check(found);
+    assert.deepStrictEqual(found.error, expectedError);
+  }, 2)
+};
+
+channel.subscribe(handlers);
+
+channel.traceCallback(function(cb, err) {
+  assert.deepStrictEqual(this, thisArg);
+  setImmediate(cb, err);
+}, 0, input, thisArg, common.mustCall((err, res) => {
+  assert.strictEqual(err, expectedError);
+  assert.strictEqual(res, undefined);
+}), expectedError);
+
+channel.tracePromise(function(value) {
+  assert.deepStrictEqual(this, thisArg);
+  return Promise.reject(value);
+}, input, thisArg, expectedError).then(
+  common.mustNotCall(),
+  common.mustCall((value) => {
+    assert.deepStrictEqual(value, expectedError);
+  })
+);

test/parallel/test-diagnostics-channel-tracing-channel-async-error.js

@@ -0,0 +1,51 @@
+'use strict';
+
+const common = require('../common');
+const dc = require('diagnostics_channel');
+const assert = require('assert');
+
+const channel = dc.tracingChannel('test');
+
+const expectedResult = { foo: 'bar' };
+const input = { foo: 'bar' };
+const thisArg = { baz: 'buz' };
+
+function check(found) {
+  assert.deepStrictEqual(found, input);
+}
+
+const handlers = {
+  start: common.mustCall(check, 2),
+  end: common.mustCall(check, 2),
+  asyncStart: common.mustCall((found) => {
+    check(found);
+    assert.strictEqual(found.error, undefined);
+    assert.deepStrictEqual(found.result, expectedResult);
+  }, 2),
+  asyncEnd: common.mustCall((found) => {
+    check(found);
+    assert.strictEqual(found.error, undefined);
+    assert.deepStrictEqual(found.result, expectedResult);
+  }, 2),
+  error: common.mustNotCall()
+};
+
+channel.subscribe(handlers);
+
+channel.traceCallback(function(cb, err, res) {
+  assert.deepStrictEqual(this, thisArg);
+  setImmediate(cb, err, res);
+}, 0, input, thisArg, common.mustCall((err, res) => {
+  assert.strictEqual(err, null);
+  assert.deepStrictEqual(res, expectedResult);
+}), null, expectedResult);
+
+channel.tracePromise(function(value) {
+  assert.deepStrictEqual(this, thisArg);
+  return Promise.resolve(value);
+}, input, thisArg, expectedResult).then(
+  common.mustCall((value) => {
+    assert.deepStrictEqual(value, expectedResult);
+  }),
+  common.mustNotCall()
+);

test/parallel/test-diagnostics-channel-tracing-channel-async.js

@@ -0,0 +1,39 @@
+'use strict';
+
+const common = require('../common');
+const dc = require('diagnostics_channel');
+const assert = require('assert');
+
+const channel = dc.tracingChannel('test');
+
+const expectedError = new Error('test');
+const input = { foo: 'bar' };
+const thisArg = { baz: 'buz' };
+
+function check(found) {
+  assert.deepStrictEqual(found, input);
+}
+
+const handlers = {
+  start: common.mustCall(check),
+  end: common.mustCall(check),
+  asyncStart: common.mustNotCall(),
+  asyncEnd: common.mustNotCall(),
+  error: common.mustCall((found) => {
+    check(found);
+    assert.deepStrictEqual(found.error, expectedError);
+  })
+};
+
+channel.subscribe(handlers);
+try {
+  channel.traceSync(function(err) {
+    assert.deepStrictEqual(this, thisArg);
+    assert.strictEqual(err, expectedError);
+    throw err;
+  }, input, thisArg, expectedError);
+
+  throw new Error('It should not reach this error');
+} catch (error) {
+  assert.deepStrictEqual(error, expectedError);
+}

test/parallel/test-diagnostics-channel-tracing-channel-sync-error.js

@@ -0,0 +1,38 @@
+'use strict';
+
+const common = require('../common');
+const dc = require('diagnostics_channel');
+const assert = require('assert');
+
+const channel = dc.tracingChannel('test');
+
+const expectedResult = { foo: 'bar' };
+const input = { foo: 'bar' };
+
+function check(found) {
+  assert.deepStrictEqual(found, input);
+}
+
+const handlers = {
+  start: common.mustCall(check),
+  end: common.mustCall((found) => {
+    check(found);
+    assert.deepStrictEqual(found.result, expectedResult);
+  }),
+  asyncStart: common.mustNotCall(),
+  asyncEnd: common.mustNotCall(),
+  error: common.mustNotCall()
+};
+
+assert.strictEqual(channel.start.hasSubscribers, false);
+channel.subscribe(handlers);
+assert.strictEqual(channel.start.hasSubscribers, true);
+channel.traceSync(() => {
+  return expectedResult;
+}, input);
+
+channel.unsubscribe(handlers);
+assert.strictEqual(channel.start.hasSubscribers, false);
+channel.traceSync(() => {
+  return expectedResult;
+}, input);

test/parallel/test-diagnostics-channel-tracing-channel-sync.js

@@ -57,6 +57,8 @@ const customTypesMap = {
   'Module Namespace Object':
     'https://tc39.github.io/ecma262/#sec-module-namespace-exotic-objects',
 
+  'AsyncLocalStorage': 'async_context.html#class-asynclocalstorage',
+
   'AsyncHook': 'async_hooks.html#async_hookscreatehookcallbacks',
   'AsyncResource': 'async_hooks.html#class-asyncresource',
 
@@ -108,6 +110,7 @@ const customTypesMap = {
   'dgram.Socket': 'dgram.html#class-dgramsocket',
 
   'Channel': 'diagnostics_channel.html#class-channel',
+  'TracingChannel': 'diagnostics_channel.html#class-tracingchannel',
 
   'Domain': 'domain.html#class-domain',