async_hooks: refactor diagnostics_channel integration

Author: Stephen Belanger

lib/async_hooks.js

@@ -28,6 +28,7 @@ const {
 } = require('internal/validators');
 const internal_async_hooks = require('internal/async_hooks');
 const { tracingChannel } = require('diagnostics_channel');
+const console = require('console');
 
 // Get functions
 // For userland AsyncResources, make sure to emit a destroy event when the
@@ -361,31 +362,13 @@ class AsyncLocalStorage {
 
     const store = this;
 
-    function start(ctx) {
-      if (store.enabled) {
-        ctx[store.kResourceStore] = store.getStore();
-      }
-      store.enterWith(transform(ctx.data));
-    }
-
-    function end(ctx) {
-      if (ObjectPrototypeHasOwnProperty(ctx, store.kResourceStore)) {
-        store.enterWith(ctx[store.kResourceStore]);
-      } else {
-        store.disable();
-      }
+    function start(register) {
+      register(store, transform);
     }
 
-    channel.subscribe({
-      start,
-      end
-    });
-
+    channel.subscribe({ start });
     return () => {
-      channel.unsubscribe({
-        start,
-        end
-      });
+      channel.unsubscribe({ start });
     };
   }
 
@@ -394,10 +377,26 @@ class AsyncLocalStorage {
       channel = tracingChannel(channel);
     }
 
-    channel.traceSync(runner, { data });
+    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
 // otherwise.
 module.exports = {