diff --git a/lib/diagnostics_channel.js b/lib/diagnostics_channel.js index 3deb301e7f3cd2..81aeb211f58b92 100644 --- a/lib/diagnostics_channel.js +++ b/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) {