Simplifying Data Streaming (SSE, WS, FireStore)

[okalil]

I've seen a lot of different solutions for real time data with react query, but looks like there's not a consensus about the best way to deal with it. I made a simplified solution that is intended to be data source agnostic and focused on native APIs, like async generators, AbortSignal and the new Promise.withResolvers.
I appreciate if you leave a feedback if you think this approach makes sense.

Why

Real time subcriptions allow us to keep data in sync, which is one of the main objectives of react query. Sounds like one overlaps the other, but I still think using both together is useful because of the great state management that react query provides. I can easily show loading and error UIs for initial load "without touch any global state" or dealing with useState and useContext. Besides, if the next update event fails, in some cases I can keep showing the latest cached version.

How

// hook to build the stream handler
function useStreamHandler<T>(
  create: (signal: AbortSignal) => AsyncGenerator<T, never>
) {
  const client = useQueryClient();
  const abortControllerRef = useRef<AbortController>();

  const run = async (key: QueryKey) => {
    const controller = new AbortController();
    abortControllerRef.current?.abort(); // any previous signal
    abortControllerRef.current = controller;

    const stream = create(controller.signal);
    const { value: initialData } = await stream.next();
    startConsuming();
    abortOnRemove();

    return initialData;

    async function startConsuming() {
      for await (const nextValue of stream) {
        client.setQueryData(key, nextValue);
      }
    }
    function abortOnRemove() {
      const unsubscribe = client.getQueryCache().subscribe(e => {
        if (e.type === 'observerRemoved' && e.query.getObserversCount() === 0) {
          const matches =
            e.query.queryHash ===
            client.getQueryCache().find({ queryKey: key })?.queryHash;
          if (!matches) return;
          controller.abort();
          unsubscribe();
        }
      });
    }
  };
  return run;
}

// usage
function useCoinQuery() {
  const handleStream = useStreamHandler(signal => getCoinStream(signal));
  return useQuery({
    queryKey: ['realtime_coin_rate'],
    queryFn: ctx => handleStream(ctx.queryKey),
  });
}

Implementations

Sockets Live Data

interface Coin {
    price: number;
    volume: number;
}
async function* getCoinStream_socket(signal: AbortSignal) {
    let defer: PromiseWithResolvers<Coin>;
    const next = (data: Coin) => {
        defer.resolve(data)
    }
    socket.on('rate', next)
    signal.onabort = () => defer.reject(signal.reason ?? 'AbortError')
    try {
        while (true) {
            defer = Promise.withResolvers()
            yield await defer.promise
        }
    } finally {
        socket.off('rate', next)
    }
}

FireStore Live Data

interface FirebaseTask { name: string}
async function* getTasksStream_firestore(signal: AbortSignal) {
    let defer: PromiseWithResolvers<FirebaseTask[]>;
    const unsubscribe = onSnapshot(ref, docs => {
        const data = docs.map(it => it.data())
        defer.resolve(data)
    })
    signal.onabort = () => defer.reject(signal.reason ?? 'AbortError')
    try {
        while (true) {
            defer = Promise.withResolvers()
            yield await defer.promise
        }
    } finally {
        unsubscribe()
    }
}

SSE Partial Updates

async function* getNotificationsStream_sse(signal: AbortSignal) {
    const notifications = await getLatestNotifications() // initial data via http
    yield notifications;

    let defer: PromiseWithResolvers<Message>;
    const source = new EventSource('url')
    source.onmessage = (e) => {
        const data = JSON.parse(e.data) // new notification
        notifications.push(data)
        defer.resolve(notifications)
    }
    signal.onabort = () => defer.reject(signal.reason ?? 'AbortError')
    try {
        while (true) {
            defer = Promise.withResolvers()
            yield await defer.promise
        }
    } finally {
        source.close()
    }
}

Notes

1. I didn't use useEffect to clean the subscription to avoid multiple effects when using custom hook in multiple places at the same time.
2. Promise.withResolvers is a new API, supported in most modern devices and browsers, but can be easily polyfilled if necessary.

Some Related discussions:
#1061
#2621
#418