make Publisher generic

in terms of sources, incremental results, and final payload

Author: yaacovCR

src/execution/execute.ts

@@ -122,7 +122,11 @@ export interface ExecutionContext {
   typeResolver: GraphQLTypeResolver<any, any>;
   subscribeFieldResolver: GraphQLFieldResolver<any, any>;
   errors: Array<GraphQLError>;
-  publisher: Publisher;
+  publisher: Publisher<
+    AsyncPayloadRecord,
+    IncrementalResult,
+    SubsequentIncrementalExecutionResult
+  >;
 }
 
 /**
@@ -365,7 +369,7 @@ function executeImpl(
                 ...initialResult,
                 hasNext: true,
               },
-              subsequentResults: publisher.yieldSubsequentPayloads(),
+              subsequentResults: publisher.subscribe(),
             };
           }
           return initialResult;
@@ -384,7 +388,7 @@ function executeImpl(
           ...initialResult,
           hasNext: true,
         },
-        subsequentResults: publisher.yieldSubsequentPayloads(),
+        subsequentResults: publisher.subscribe(),
       };
     }
     return initialResult;
@@ -2187,9 +2191,7 @@ class StreamRecord {
   }
 }
 
-export type { StreamRecord };
-
-export type AsyncPayloadRecord = DeferredFragmentRecord | StreamRecord;
+type AsyncPayloadRecord = DeferredFragmentRecord | StreamRecord;
 
 function isStreamPayload(
   asyncPayload: AsyncPayloadRecord,

src/execution/publisher.ts

@@ -1,100 +1,90 @@
-import type {
-  AsyncPayloadRecord,
-  StreamRecord,
-  IncrementalResult,
-  SubsequentIncrementalExecutionResult,
-} from './execute.js';
-
-type ToIncrementalResult = (
-  asyncPayloadRecord: AsyncPayloadRecord,
-) => IncrementalResult;
-
-type ToPayload = (
-  incremental: ReadonlyArray<IncrementalResult>,
+interface Source {
+  promise: Promise<void>;
+  isCompleted: boolean;
+  isCompletedIterator?: boolean | undefined;
+  iterator?: AsyncIterator<unknown> | undefined;
+}
+
+type ToIncrementalResult<TSource extends Source, TIncremental> = (
+  source: TSource,
+) => TIncremental;
+
+type ToPayload<TIncremental, TPayload> = (
+  incremental: ReadonlyArray<TIncremental>,
   hasNext: boolean,
-) => SubsequentIncrementalExecutionResult;
+) => TPayload;
 
 /**
  * @internal
  */
-export class Publisher {
-  subsequentPayloads: Set<AsyncPayloadRecord>;
-  toIncrementalResult: ToIncrementalResult;
-  toPayload: ToPayload;
-
-  constructor(toIncrementalResult: ToIncrementalResult, toPayload: ToPayload) {
-    this.subsequentPayloads = new Set();
+export class Publisher<TSource extends Source, TIncremental, TPayload> {
+  sources: Set<TSource>;
+  toIncrementalResult: ToIncrementalResult<TSource, TIncremental>;
+  toPayload: ToPayload<TIncremental, TPayload>;
+
+  constructor(
+    toIncrementalResult: ToIncrementalResult<TSource, TIncremental>,
+    toPayload: ToPayload<TIncremental, TPayload>,
+  ) {
+    this.sources = new Set();
     this.toIncrementalResult = toIncrementalResult;
     this.toPayload = toPayload;
   }
 
-  add(payload: AsyncPayloadRecord) {
-    this.subsequentPayloads.add(payload);
+  add(source: TSource) {
+    this.sources.add(source);
   }
 
   hasNext(): boolean {
-    return this.subsequentPayloads.size > 0;
+    return this.sources.size > 0;
   }
 
-  filter(predicate: (payload: AsyncPayloadRecord) => boolean): void {
-    this.subsequentPayloads.forEach((asyncRecord) => {
-      if (predicate(asyncRecord)) {
+  filter(predicate: (source: TSource) => boolean): void {
+    this.sources.forEach((source) => {
+      if (predicate(source)) {
         return;
       }
-      // asyncRecord path points to nulled error field
-      if (isStreamPayload(asyncRecord) && asyncRecord.iterator?.return) {
-        asyncRecord.iterator.return().catch(() => {
+      if (source.iterator?.return) {
+        source.iterator.return().catch(() => {
           // ignore error
         });
       }
-      this.subsequentPayloads.delete(asyncRecord);
+      this.sources.delete(source);
     });
   }
 
-  getCompletedIncrementalResults(): Array<IncrementalResult> {
-    const incrementalResults: Array<IncrementalResult> = [];
-    for (const asyncPayloadRecord of this.subsequentPayloads) {
-      if (!asyncPayloadRecord.isCompleted) {
+  getCompletedIncrementalResults(): Array<TIncremental> {
+    const incrementalResults: Array<TIncremental> = [];
+    for (const source of this.sources) {
+      if (!source.isCompleted) {
         continue;
       }
-      this.subsequentPayloads.delete(asyncPayloadRecord);
-      if (isStreamPayload(asyncPayloadRecord)) {
-        if (asyncPayloadRecord.isCompletedIterator) {
-          // async iterable resolver just finished but there may be pending payloads
-          continue;
-        }
+      this.sources.delete(source);
+      if (source.isCompletedIterator) {
+        continue;
       }
-      incrementalResults.push(this.toIncrementalResult(asyncPayloadRecord));
+      incrementalResults.push(this.toIncrementalResult(source));
     }
     return incrementalResults;
   }
 
-  yieldSubsequentPayloads(): AsyncGenerator<
-    SubsequentIncrementalExecutionResult,
-    void,
-    void
-  > {
+  subscribe(): AsyncGenerator<TPayload, void, void> {
     let isDone = false;
     const publisher = this;
 
-    async function next(): Promise<
-      IteratorResult<SubsequentIncrementalExecutionResult, void>
-    > {
+    async function next(): Promise<IteratorResult<TPayload, void>> {
       if (isDone) {
         return { value: undefined, done: true };
       }
 
-      await Promise.race(
-        Array.from(publisher.subsequentPayloads).map((p) => p.promise),
-      );
+      await Promise.race(Array.from(publisher.sources).map((p) => p.promise));
 
       if (isDone) {
-        // a different call to next has exhausted all payloads
         return { value: undefined, done: true };
       }
 
       const incremental = publisher.getCompletedIncrementalResults();
-      const hasNext = publisher.subsequentPayloads.size > 0;
+      const hasNext = publisher.sources.size > 0;
 
       if (!incremental.length && hasNext) {
         return next();
@@ -110,14 +100,11 @@ export class Publisher {
       };
     }
 
-    function returnStreamIterators() {
+    function returnIterators() {
       const promises: Array<Promise<IteratorResult<unknown>>> = [];
-      publisher.subsequentPayloads.forEach((asyncPayloadRecord) => {
-        if (
-          isStreamPayload(asyncPayloadRecord) &&
-          asyncPayloadRecord.iterator?.return
-        ) {
-          promises.push(asyncPayloadRecord.iterator.return());
+      publisher.sources.forEach((source) => {
+        if (source.iterator?.return) {
+          promises.push(source.iterator.return());
         }
       });
       return Promise.all(promises);
@@ -128,26 +115,16 @@ export class Publisher {
         return this;
       },
       next,
-      async return(): Promise<
-        IteratorResult<SubsequentIncrementalExecutionResult, void>
-      > {
-        await returnStreamIterators();
+      async return(): Promise<IteratorResult<TPayload, void>> {
+        await returnIterators();
         isDone = true;
         return { value: undefined, done: true };
       },
-      async throw(
-        error?: unknown,
-      ): Promise<IteratorResult<SubsequentIncrementalExecutionResult, void>> {
-        await returnStreamIterators();
+      async throw(error?: unknown): Promise<IteratorResult<TPayload, void>> {
+        await returnIterators();
         isDone = true;
         return Promise.reject(error);
       },
     };
   }
 }
-
-function isStreamPayload(
-  asyncPayload: AsyncPayloadRecord,
-): asyncPayload is StreamRecord {
-  return asyncPayload.type === 'stream';
-}