feat(execution): publish on graphql:subscribe tracing channel

Wraps the public subscribe() entry point with maybeTraceMixed using the
same ExecutionArgs-based context factory as graphql:execute. The
subscription setup path may return sync (when the subscribe resolver is
synchronous) or async (when the resolver returns a promise resolving to
an AsyncIterable), so maybeTraceMixed's sync / promise branching fits
exactly.

Per-subscription-event executions continue to publish on graphql:execute
via executeSubscriptionEvent, which was wired in the previous commit.
The graphql:subscribe context therefore owns the document reference and
covers the setup span; subscribers that need to correlate per-event
execute spans to their parent subscription use AsyncLocalStorage context
propagation provided by Channel.runStores.

Author: logaretm

integrationTests/diagnostics/test.js

@@ -10,68 +10,70 @@ import {
   enableDiagnosticsChannel,
   execute,
   parse,
+  subscribe,
   validate,
 } from 'graphql';
 
 enableDiagnosticsChannel(dc);
 
-// graphql:parse - synchronous
-{
-  const events = [];
-  const handler = {
-    start: (msg) => events.push({ kind: 'start', source: msg.source }),
-    end: (msg) => events.push({ kind: 'end', source: msg.source }),
-    asyncStart: (msg) =>
-      events.push({ kind: 'asyncStart', source: msg.source }),
-    asyncEnd: (msg) => events.push({ kind: 'asyncEnd', source: msg.source }),
-    error: (msg) =>
-      events.push({ kind: 'error', source: msg.source, error: msg.error }),
-  };
-
-  const channel = dc.tracingChannel('graphql:parse');
-  channel.subscribe(handler);
-
-  try {
-    const doc = parse('{ field }');
-    assert.equal(doc.kind, 'Document');
-    assert.deepEqual(
-      events.map((e) => e.kind),
-      ['start', 'end'],
-    );
-    assert.equal(events[0].source, '{ field }');
-    assert.equal(events[1].source, '{ field }');
-  } finally {
-    channel.unsubscribe(handler);
+function runParseCases() {
+  // graphql:parse - synchronous.
+  {
+    const events = [];
+    const handler = {
+      start: (msg) => events.push({ kind: 'start', source: msg.source }),
+      end: (msg) => events.push({ kind: 'end', source: msg.source }),
+      asyncStart: (msg) =>
+        events.push({ kind: 'asyncStart', source: msg.source }),
+      asyncEnd: (msg) => events.push({ kind: 'asyncEnd', source: msg.source }),
+      error: (msg) =>
+        events.push({ kind: 'error', source: msg.source, error: msg.error }),
+    };
+
+    const channel = dc.tracingChannel('graphql:parse');
+    channel.subscribe(handler);
+
+    try {
+      const doc = parse('{ field }');
+      assert.equal(doc.kind, 'Document');
+      assert.deepEqual(
+        events.map((e) => e.kind),
+        ['start', 'end'],
+      );
+      assert.equal(events[0].source, '{ field }');
+      assert.equal(events[1].source, '{ field }');
+    } finally {
+      channel.unsubscribe(handler);
+    }
   }
-}
 
-// graphql:parse - error path fires start, error, end (traceSync finally-emits end)
-{
-  const events = [];
-  const handler = {
-    start: (msg) => events.push({ kind: 'start', source: msg.source }),
-    end: (msg) => events.push({ kind: 'end', source: msg.source }),
-    error: (msg) =>
-      events.push({ kind: 'error', source: msg.source, error: msg.error }),
-  };
-
-  const channel = dc.tracingChannel('graphql:parse');
-  channel.subscribe(handler);
-
-  try {
-    assert.throws(() => parse('{ '));
-    assert.deepEqual(
-      events.map((e) => e.kind),
-      ['start', 'error', 'end'],
-    );
-    assert.ok(events[1].error instanceof Error);
-  } finally {
-    channel.unsubscribe(handler);
+  // graphql:parse - error path fires start, error, end.
+  {
+    const events = [];
+    const handler = {
+      start: (msg) => events.push({ kind: 'start', source: msg.source }),
+      end: (msg) => events.push({ kind: 'end', source: msg.source }),
+      error: (msg) =>
+        events.push({ kind: 'error', source: msg.source, error: msg.error }),
+    };
+
+    const channel = dc.tracingChannel('graphql:parse');
+    channel.subscribe(handler);
+
+    try {
+      assert.throws(() => parse('{ '));
+      assert.deepEqual(
+        events.map((e) => e.kind),
+        ['start', 'error', 'end'],
+      );
+      assert.ok(events[1].error instanceof Error);
+    } finally {
+      channel.unsubscribe(handler);
+    }
   }
 }
 
-// graphql:validate - synchronous, with schema/document context
-{
+function runValidateCase() {
   const schema = buildSchema(`type Query { field: String }`);
   const doc = parse('{ field }');
 
@@ -104,9 +106,7 @@ enableDiagnosticsChannel(dc);
   }
 }
 
-// graphql:execute - sync path, ctx carries operationType, operationName,
-// document, schema.
-{
+function runExecuteCase() {
   const schema = buildSchema(`type Query { hello: String }`);
   const document = parse('query Greeting { hello }');
 
@@ -149,10 +149,67 @@ enableDiagnosticsChannel(dc);
   }
 }
 
-// No-op when nothing is subscribed - parse still succeeds.
-{
+async function runSubscribeCase() {
+  async function* ticks() {
+    yield { tick: 'one' };
+  }
+
+  const schema = buildSchema(`
+    type Query { dummy: String }
+    type Subscription { tick: String }
+  `);
+  // buildSchema doesn't attach a subscribe resolver to fields; inject one.
+  schema.getSubscriptionType().getFields().tick.subscribe = () => ticks();
+
+  const document = parse('subscription Tick { tick }');
+
+  const events = [];
+  const handler = {
+    start: (msg) =>
+      events.push({
+        kind: 'start',
+        operationType: msg.operationType,
+        operationName: msg.operationName,
+      }),
+    end: () => events.push({ kind: 'end' }),
+    asyncStart: () => events.push({ kind: 'asyncStart' }),
+    asyncEnd: () => events.push({ kind: 'asyncEnd' }),
+    error: (msg) => events.push({ kind: 'error', error: msg.error }),
+  };
+
+  const channel = dc.tracingChannel('graphql:subscribe');
+  channel.subscribe(handler);
+
+  try {
+    const result = subscribe({ schema, document });
+    const stream = typeof result.then === 'function' ? await result : result;
+    if (stream[Symbol.asyncIterator]) {
+      await stream.return?.();
+    }
+    // Subscription setup is synchronous here; start/end fire, no async tail.
+    assert.deepEqual(
+      events.map((e) => e.kind),
+      ['start', 'end'],
+    );
+    assert.equal(events[0].operationType, 'subscription');
+    assert.equal(events[0].operationName, 'Tick');
+  } finally {
+    channel.unsubscribe(handler);
+  }
+}
+
+function runNoSubscriberCase() {
   const doc = parse('{ field }');
   assert.equal(doc.kind, 'Document');
 }
 
-console.log('diagnostics integration test passed');
+async function main() {
+  runParseCases();
+  runValidateCase();
+  runExecuteCase();
+  await runSubscribeCase();
+  runNoSubscriberCase();
+  console.log('diagnostics integration test passed');
+}
+
+main();

src/execution/__tests__/subscribe-diagnostics-test.ts

@@ -0,0 +1,130 @@
+import { expect } from 'chai';
+import { afterEach, beforeEach, describe, it } from 'mocha';
+
+import {
+  collectEvents,
+  FakeDc,
+} from '../../__testUtils__/fakeDiagnosticsChannel.js';
+
+import { isPromise } from '../../jsutils/isPromise.js';
+
+import { parse } from '../../language/parser.js';
+
+import { GraphQLObjectType } from '../../type/definition.js';
+import { GraphQLString } from '../../type/scalars.js';
+import { GraphQLSchema } from '../../type/schema.js';
+
+import { enableDiagnosticsChannel } from '../../diagnostics.js';
+
+import { subscribe } from '../execute.js';
+
+function buildSubscriptionSchema(
+  subscribeFn: () => AsyncIterable<{ tick: string }>,
+): GraphQLSchema {
+  return new GraphQLSchema({
+    query: new GraphQLObjectType({
+      name: 'Query',
+      fields: { dummy: { type: GraphQLString } },
+    }),
+    subscription: new GraphQLObjectType({
+      name: 'Subscription',
+      fields: {
+        tick: {
+          type: GraphQLString,
+          subscribe: subscribeFn,
+        },
+      },
+    }),
+  });
+}
+
+async function* twoTicks(): AsyncIterable<{ tick: string }> {
+  await Promise.resolve();
+  yield { tick: 'one' };
+  yield { tick: 'two' };
+}
+
+const fakeDc = new FakeDc();
+const subscribeChannel = fakeDc.tracingChannel('graphql:subscribe');
+
+describe('subscribe diagnostics channel', () => {
+  let active: ReturnType<typeof collectEvents> | undefined;
+
+  beforeEach(() => {
+    enableDiagnosticsChannel(fakeDc);
+  });
+
+  afterEach(() => {
+    active?.unsubscribe();
+    active = undefined;
+  });
+
+  it('emits start and end for a synchronous subscription setup', async () => {
+    active = collectEvents(subscribeChannel);
+
+    const schema = buildSubscriptionSchema(twoTicks);
+    const document = parse('subscription S { tick }');
+
+    const result = subscribe({ schema, document });
+    const resolved = isPromise(result) ? await result : result;
+    if (!(Symbol.asyncIterator in resolved)) {
+      throw new Error('Expected an async iterator');
+    }
+    await resolved.return?.();
+
+    expect(active.events.map((e) => e.kind)).to.deep.equal(['start', 'end']);
+    expect(active.events[0].ctx.operationType).to.equal('subscription');
+    expect(active.events[0].ctx.operationName).to.equal('S');
+    expect(active.events[0].ctx.document).to.equal(document);
+    expect(active.events[0].ctx.schema).to.equal(schema);
+  });
+
+  it('emits the full async lifecycle when subscribe resolver returns a promise', async () => {
+    active = collectEvents(subscribeChannel);
+
+    const asyncResolver = (): Promise<AsyncIterable<{ tick: string }>> =>
+      Promise.resolve(twoTicks());
+    const schema = buildSubscriptionSchema(
+      asyncResolver as unknown as () => AsyncIterable<{ tick: string }>,
+    );
+    const document = parse('subscription { tick }');
+
+    const result = subscribe({ schema, document });
+    const resolved = isPromise(result) ? await result : result;
+    if (!(Symbol.asyncIterator in resolved)) {
+      throw new Error('Expected an async iterator');
+    }
+    await resolved.return?.();
+
+    expect(active.events.map((e) => e.kind)).to.deep.equal([
+      'start',
+      'end',
+      'asyncStart',
+      'asyncEnd',
+    ]);
+  });
+
+  it('emits only start and end for a synchronous validation failure', () => {
+    active = collectEvents(subscribeChannel);
+
+    const schema = buildSubscriptionSchema(twoTicks);
+    // Invalid: no operation.
+    const document = parse('fragment F on Subscription { tick }');
+
+    // eslint-disable-next-line @typescript-eslint/no-floating-promises
+    subscribe({ schema, document });
+
+    expect(active.events.map((e) => e.kind)).to.deep.equal(['start', 'end']);
+  });
+
+  it('does nothing when no subscribers are attached', async () => {
+    const schema = buildSubscriptionSchema(twoTicks);
+    const document = parse('subscription { tick }');
+
+    const result = subscribe({ schema, document });
+    const resolved = isPromise(result) ? await result : result;
+    if (Symbol.asyncIterator in resolved) {
+      await resolved.return?.();
+    }
+  });
+});

src/execution/execute.ts

@@ -282,24 +282,26 @@ export function subscribe(
 ): PromiseOrValue<
   AsyncGenerator<ExecutionResult, void, void> | ExecutionResult
 > {
-  // If a valid execution context cannot be created due to incorrect arguments,
-  // a "Response" with only errors is returned.
-  const validatedExecutionArgs = validateExecutionArgs(args);
+  return maybeTraceMixed('subscribe', buildExecuteCtxFromArgs(args), () => {
+    // If a valid execution context cannot be created due to incorrect
+    // arguments, a "Response" with only errors is returned.
+    const validatedExecutionArgs = validateExecutionArgs(args);
 
-  // Return early errors if execution context failed.
-  if (!('schema' in validatedExecutionArgs)) {
-    return { errors: validatedExecutionArgs };
-  }
+    // Return early errors if execution context failed.
+    if (!('schema' in validatedExecutionArgs)) {
+      return { errors: validatedExecutionArgs };
+    }
 
-  const resultOrStream = createSourceEventStreamImpl(validatedExecutionArgs);
+    const resultOrStream = createSourceEventStreamImpl(validatedExecutionArgs);
 
-  if (isPromise(resultOrStream)) {
-    return resultOrStream.then((resolvedResultOrStream) =>
-      mapSourceToResponse(validatedExecutionArgs, resolvedResultOrStream),
-    );
-  }
+    if (isPromise(resultOrStream)) {
+      return resultOrStream.then((resolvedResultOrStream) =>
+        mapSourceToResponse(validatedExecutionArgs, resolvedResultOrStream),
+      );
+    }
 
-  return mapSourceToResponse(validatedExecutionArgs, resultOrStream);
+    return mapSourceToResponse(validatedExecutionArgs, resultOrStream);
+  });
 }
 
 /**