useSubscribe that supports suspense

use-subscribe-suspense.ts

import { useEffect, useState } from 'react';
import { unstable_batchedUpdates } from 'react-dom';
import type {
  ReadonlyJSONValue,
  ReadTransaction,
  Replicache,
} from 'replicache';

// We wrap all the callbacks in a `unstable_batchedUpdates` call to ensure that
// we do not render things more than once over all of the changed subscriptions.

let hasPendingCallback = false;
let callbacks: (() => void)[] = [];

function doCallback() {
  const cbs = callbacks;
  callbacks = [];
  hasPendingCallback = false;
  unstable_batchedUpdates(() => {
    for (const callback of cbs) {
      callback();
    }
  });
}

function wrapPromise<T>(promise: Promise<T>) {
  let status = 'pending';
  let result: any;
  let suspender = promise.then(
    (r: T) => {
      status = 'success';
      result = r;
    },
    (e: any) => {
      status = 'error';
      result = e;
    }
  );
  return {
    read() {
      if (status === 'pending') {
        throw suspender;
      } else if (status === 'error') {
        throw result;
      } else if (status === 'success') {
        return result;
      }
    },
  };
}

const cache = new Map<string, any>();

function getSuspender<T>(
  rep: Replicache,
  query: (tx: ReadTransaction) => Promise<T>,
  keys: Array<any>
) {
  const key = JSON.stringify(keys);
  const existing = cache.get(key);
  if (existing) return existing;
  console.log('new :(');
  const suspender = wrapPromise(rep.query(query));
  cache.set(key, suspender);
  return suspender;
}

export function useSubscribe<R extends ReadonlyJSONValue>(
  rep: Replicache,
  query: (tx: ReadTransaction) => Promise<R>,
  keys: Array<any> = []
): R {
  const [snapshot, setSnapshot] = useState(() =>
    getSuspender(rep, query, keys).read()
  );

  useEffect(() => {
    return rep.subscribe(query, {
      onData: (data: R) => {
        if (data === snapshot) {
          return;
        }
        callbacks.push(() => setSnapshot(data));
        if (!hasPendingCallback) {
          void Promise.resolve().then(doCallback);
          hasPendingCallback = true;
        }
      },
    });
  }, [rep, query]);

  return snapshot;
}

I feel like you should be able to use useCallback and use the function instance in the map instead of stringifying the keys.

I still need to look deeper into suspense to get a better sense for how it is supposed to work.

I thought so too, but the function returned by useCallback is not the same as when the promise resolves. I suppose react is using the same storage for useState and it all gets discarded on suspend. It seems like quite a revere restriction for supporting suspense!

Thanks for sharing @dpeek, this is really great.

Any reason you're not updating the cache in onData? With the current code, when a component has its initial render after the cache has been set for that key, it gets stale data.

Seems to be fixed if you add this after line 82:

const updatedSuspender = wrapPromise<QueryRet>(Promise.resolve(data));
cache.set(suspenseKey, updatedSuspender);