@graphql-yoga/nestjs/__tests__/utils/pubsub.ts

GraphQL Yoga driver for NestJS GraphQL.

export interface Generator<T> {
  gen: AsyncGenerator<T, void, T>;
  produce(val: T): void;
}

function createGenerator<T>(): Generator<T> {
  const pending: T[] = [];

  const deferred = {
    done: false,
    error: null as unknown,
    resolve: () => {
      // noop
    },
  };

  const gen = (async function* gen() {
    for (;;) {
      if (!pending.length) {
        // only wait if there are no pending messages available
        await new Promise<void>(resolve => (deferred.resolve = resolve));
      }
      // first flush
      while (pending.length) {
        yield pending.shift()!;
      }
      // then error
      if (deferred.error) {
        throw deferred.error;
      }
      // or complete
      if (deferred.done) {
        return;
      }
    }
  })();

  gen.throw = async err => {
    if (!deferred.done) {
      deferred.done = true;
      deferred.error = err;
      deferred.resolve();
    }
    return { done: true, value: undefined };
  };

  gen.return = async () => {
    if (!deferred.done) {
      deferred.done = true;
      deferred.resolve();
    }
    return { done: true, value: undefined };
  };

  return {
    gen,
    produce(val) {
      pending.push(val);
      deferred.resolve();
    },
  };
}

export function createPubSub<T>() {
  const producers: Generator<T>['produce'][] = [];
  return {
    pub(val: T) {
      for (const next of producers) next(val);
    },
    sub() {
      const { gen, produce } = createGenerator<T>();
      producers.push(produce);
      const origReturn = gen.return;
      gen.return = () => {
        producers.splice(producers.indexOf(produce), 1);
        return origReturn();
      };
      return gen;
    },
  };
}