expand support to any iterator

Author: rochdev

lib/diagnostics_channel.js

@@ -431,56 +431,82 @@ class TracingChannel {
     });
   }
 
-  *traceGenerator(fn, context = {}, thisArg, ...args) {
+  traceIterator(fn, context = {}, thisArg, ...args) {
     if (!this.hasSubscribers) {
       return ReflectApply(fn, thisArg, args);
     }
 
-    const { start, end, asyncStart, asyncEnd, error } = this;
+    const { channel, start, end, asyncStart, asyncEnd, error } = this;
 
-    if (!this.#generatorChannel) {
-      this.#generatorChannel = this.tracingChannel({
-        start: this.channel(start.name.slice(0, -6) + ':next:start'),
-        end: this.channel(end.name.slice(0, -4) + ':next:end'),
-        asyncStart: this.channel(asyncStart.name.slice(0, -11) + ':next:asyncStart'),
-        asyncEnd: this.channel(asyncEnd.name.slice(0, -9) + ':next:asyncEnd'),
-        error: this.channel(error.name.slice(0, -6) + ':next:error'),
+    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 gen = this.traceSync(fn, context, thisArg, ...args);
+    const iter = this.#traceMaybePromise(fn, context, thisArg, ...args);
+
+    return iter instanceof Promise
+      ? iter.then((iter, method) => {
+        const { next: iterNext, return: iterReturn, throw: iterThrow } = iter;
 
-    gen.next = (...args) => this.traceSync(gen.next, context, gen, ...args);
-    gen.return = (...args) => this.traceSync(gen.return, context, gen, ...args);
-    gen.throw = (...args) => this.traceSync(gen.throw, context, gen, ...args);
+        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);
 
-    yield* gen;
+        return iter;
+      })
+      : iter;
   }
 
-  async *traceAsyncGenerator(fn, context = {}, thisArg, ...args) {
+  #traceMaybePromise(fn, context = {}, thisArg, ...args) {
     if (!this.hasSubscribers) {
       return ReflectApply(fn, thisArg, args);
     }
 
     const { start, end, asyncStart, asyncEnd, error } = this;
 
-    if (!this.#generatorChannel) {
-      this.#generatorChannel = this.tracingChannel({
-        start: this.channel(start.name.slice(0, -6) + ':next:start'),
-        end: this.channel(end.name.slice(0, -4) + ':next:end'),
-        asyncStart: this.channel(asyncStart.name.slice(0, -11) + ':next:asyncStart'),
-        asyncEnd: this.channel(asyncEnd.name.slice(0, -9) + ':next:asyncEnd'),
-        error: this.channel(error.name.slice(0, -6) + ':next:error'),
-      });
+    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);
     }
 
-    const gen = this.traceSync(fn, context, thisArg, ...args);
-
-    gen.next = (...args) => this.tracePromise(gen.next, context, gen, ...args);
-    gen.return = (...args) => this.tracePromise(gen.return, context, gen, ...args);
-    gen.throw = (...args) => this.tracePromise(gen.throw, context, gen, ...args);
+    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;
+    }
 
-    yield* gen;
+    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);
+      }
+    });
   }
 }