diagnostics_channel: add iterator support to tracing channel

Author: rochdev

lib/diagnostics_channel.js

@@ -281,6 +281,8 @@ function tracingChannelFrom(nameOrChannels, name) {
 }
 
 class TracingChannel {
+  #generatorChannel
+
   constructor(nameOrChannels) {
     for (let i = 0; i < traceEvents.length; ++i) {
       const eventName = traceEvents[i];
@@ -428,6 +430,84 @@ class TracingChannel {
       }
     });
   }
+
+  traceIterator(fn, context = {}, thisArg, ...args) {
+    if (!this.hasSubscribers) {
+      return ReflectApply(fn, thisArg, args);
+    }
+
+    const { channel, start, end, asyncStart, asyncEnd, error } = this;
+
+    if (!this.#nextChannel) {
+      this.#nextChannel = this.tracingChannel({
+        start: channel(start.name.slice(0, -6) + ':next:start'),
+        end: channel(end.name.slice(0, -4) + ':next:end'),
+        asyncStart: channel(asyncStart.name.slice(0, -11) + ':next:asyncStart'),
+        asyncEnd: channel(asyncEnd.name.slice(0, -9) + ':next:asyncEnd'),
+        error: channel(error.name.slice(0, -6) + ':next:error'),
+      });
+    }
+
+    const iter = this.#traceMaybePromise(fn, context, thisArg, ...args);
+
+    return iter instanceof Promise
+      ? iter.then((iter, method) => {
+        const { next: iterNext, return: iterReturn, throw: iterThrow } = iter;
+
+        iter.next = (...args) =>
+          this.#nextChannel.#traceMaybePromise(iterNext, ctx, iter, ...args);
+        iter.return = (...args) =>
+          this.#nextChannel.#traceMaybePromise(iterReturn, ctx, iter, ...args);
+        iter.throw = (...args) =>
+          this.#nextChannel.#traceMaybePromise(iterThrow, ctx, iter, ...args);
+
+        return iter;
+      })
+      : iter;
+  }
+
+  #traceMaybePromise(fn, context = {}, thisArg, ...args) {
+    if (!this.hasSubscribers) {
+      return ReflectApply(fn, thisArg, args);
+    }
+
+    const { start, end, asyncStart, asyncEnd, error } = this;
+
+    function reject(err) {
+      context.error = err;
+      error.publish(context);
+      asyncStart.publish(context);
+      // TODO: Is there a way to have asyncEnd _after_ the continuation?
+      asyncEnd.publish(context);
+      return PromiseReject(err);
+    }
+
+    function resolve(result) {
+      context.result = result;
+      asyncStart.publish(context);
+      // TODO: Is there a way to have asyncEnd _after_ the continuation?
+      asyncEnd.publish(context);
+      return result;
+    }
+
+    return start.runStores(context, () => {
+      try {
+        const result = ReflectApply(fn, thisArg, args);
+        // TODO: Should tracePromise just always do this?
+        if (!(result instanceof Promise)) {
+          context.result = result
+          return result
+        }
+        return PromisePrototypeThen(result, resolve, reject);
+      } catch (err) {
+        context.error = err;
+        error.publish(context);
+        throw err;
+      } finally {
+        end.publish(context);
+      }
+    });
+  }
 }
 
 function tracingChannel(nameOrChannels) {