add generator 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,58 @@ class TracingChannel {
       }
     });
   }
+
+  *traceGenerator(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'),
+      });
+    }
+
+    const gen = this.traceSync(fn, context, thisArg, ...args);
+
+    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);
+
+    yield* gen;
+  }
+
+  async *traceAsyncGenerator(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'),
+      });
+    }
+
+    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);
+
+    yield* gen;
+  }
 }
 
 function tracingChannel(nameOrChannels) {