feat: batchedQuery

src/data/index.ts

@@ -1,5 +1,5 @@
 export { createAsync, createAsyncStore, type AccessorWithLatest } from "./createAsync.js";
 export { action, useSubmission, useSubmissions, useAction, type Action } from "./action.js";
-export { query, revalidate, cache, type CachedFunction } from "./query.js";
+export { query, revalidate, cache, type CachedFunction, batchedQuery, type BatchOptions } from "./query.js";
 export { redirect, reload, json } from "./response.js";
 

src/data/query.ts

@@ -275,3 +275,101 @@ function isPlainObject(obj: object) {
     (!(proto = Object.getPrototypeOf(obj)) || proto === Object.prototype)
   );
 }
+
+export interface BatchOptions<Query> {
+  /** Decides when two queries are the same read. Defaults to `hashKey([query])`. */
+  key?: (query: Query) => unknown;
+  /** The most unique queries passed to one `callback` call. Extra queries go to more calls. */
+  limit?: number;
+  /** How long calls are collected before reading, in milliseconds. */
+  wait?: number;
+}
+
+interface BatchCaller<Return> {
+  resolve: (value: Return) => void;
+  reject: (reason: unknown) => void;
+}
+
+interface BatchSlot<Query, Return> {
+  query: Query;
+  callers: BatchCaller<Return>[];
+}
+
+interface BatchWaiting<Query, Return> extends BatchCaller<Return> {
+  query: Query;
+}
+
+/**
+ * Collects calls made at the same time and reads them with one `callback` call.
+ * `lookup` picks each caller's result out of the data. `index` is the query's
+ * position in the list given to `callback`.
+ *
+ * On the server every call reads alone. The queue lives in the module, so
+ * batching there would mix queries from separate requests.
+ */
+export function batchedQuery<Query, Data, Return>(
+  callback: (queries: Query[]) => Promise<Data>,
+  lookup: (data: Data, query: Query, index: number) => Return,
+  options: BatchOptions<Query> = {}
+): (query: Query) => Promise<Return> {
+  const keyOf = options.key ?? ((query: Query): unknown => hashKey([query]));
+  const size = Math.max(1, Math.floor(options.limit ?? Number.POSITIVE_INFINITY));
+  let waiting: BatchWaiting<Query, Return>[] = [];
+  let timer: ReturnType<typeof setTimeout> | undefined;
+
+  // A failed read rejects only the callers in its chunk.
+  // A failed lookup rejects only the callers of that query.
+  const settle = async (chunk: BatchSlot<Query, Return>[]) => {
+    const queries: Query[] = [];
+    for (const slot of chunk) queries.push(slot.query);
+
+    let data: Data;
+    try {
+      data = await callback(queries);
+    } catch (error) {
+      for (const slot of chunk) {
+        for (const caller of slot.callers) caller.reject(error);
+      }
+      return;
+    }
+    for (let index = 0; index < chunk.length; index++) {
+      const slot = chunk[index];
+      try {
+        const value = lookup(data, slot.query, index);
+        for (const caller of slot.callers) caller.resolve(value);
+      } catch (error) {
+        for (const caller of slot.callers) caller.reject(error);
+      }
+    }
+  };
+
+  const flush = () => {
+    const gathered = waiting;
+    const slots = new Map<unknown, BatchSlot<Query, Return>>();
+    // Reset before reading, so calls made during `callback` start a new batch.
+    waiting = [];
+    timer = undefined;
+    for (const { query, resolve, reject } of gathered) {
+      const key = keyOf(query);
+      const slot = slots.get(key);
+      if (slot) slot.callers.push({ resolve, reject });
+      else slots.set(key, { query, callers: [{ resolve, reject }] });
+    }
+    const unique = Array.from(slots.values());
+    for (let start = 0; start < unique.length; start += size) {
+      settle(unique.slice(start, start + size));
+    }
+  };
+
+  return (query: Query) => {
+    if (isServer) {
+      return Promise.resolve([query])
+        .then(callback)
+        .then(data => lookup(data, query, 0));
+    }
+    return new Promise<Return>((resolve, reject) => {
+      waiting.push({ query, resolve, reject });
+      if (!timer) timer = setTimeout(flush, options.wait ?? 0);
+    });
+  };
+}

test/data/batched-query.spec.ts

@@ -0,0 +1,126 @@
+import { batchedQuery } from "../../src/data/query.js";
+
+// A read that squares every number it gets and records each batch.
+function squares(fail: (queries: number[]) => boolean = () => false) {
+  const batches: number[][] = [];
+  const read = batchedQuery(
+    async (queries: number[]) => {
+      // Settle on a later tick, like a real read.
+      await Promise.resolve();
+      batches.push(queries);
+      if (fail(queries)) throw new Error(`refused ${queries.join(",")}`);
+      const answers = new Map<number, number>();
+      for (const query of queries) answers.set(query, query * query);
+      return answers;
+    },
+    (answers, query) => answers.get(query) ?? Number.NaN,
+    { limit: 2 }
+  );
+
+  return { read, batches };
+}
+
+describe("batchedQuery should", () => {
+  test("read calls made together with one read", async () => {
+    const { read, batches } = squares();
+
+    expect(await Promise.all([read(2), read(3)])).toEqual([4, 9]);
+    expect(batches).toEqual([[2, 3]]);
+  });
+
+  test("read a repeated query only once", async () => {
+    const { read, batches } = squares();
+
+    expect(await Promise.all([read(2), read(2), read(3)])).toEqual([4, 4, 9]);
+    expect(batches).toEqual([[2, 3]]);
+  });
+
+  test("split queries past the limit across reads", async () => {
+    const { read, batches } = squares();
+
+    expect(await Promise.all([read(1), read(2), read(3)])).toEqual([1, 4, 9]);
+    expect(batches).toEqual([[1, 2], [3]]);
+  });
+
+  test("reject only the callers of the failed read", async () => {
+    const { read } = squares(queries => queries.includes(3));
+    const answers = await Promise.allSettled([read(1), read(2), read(3)]);
+
+    expect(answers[0]).toEqual({ status: "fulfilled", value: 1 });
+    expect(answers[1]).toEqual({ status: "fulfilled", value: 4 });
+    expect(answers[2].status).toBe("rejected");
+  });
+
+  test("reject only the callers of a failed lookup", async () => {
+    const read = batchedQuery(
+      async (queries: number[]) => queries,
+      (_data, query) => {
+        if (query === 2) throw new Error("missing");
+        return query;
+      }
+    );
+    const answers = await Promise.allSettled([read(1), read(2)]);
+
+    expect(answers[0]).toEqual({ status: "fulfilled", value: 1 });
+    expect(answers[1].status).toBe("rejected");
+  });
+
+  test("start a new read for calls made after a read went out", async () => {
+    const { read, batches } = squares();
+
+    expect(await read(2)).toBe(4);
+    expect(await read(2)).toBe(4);
+    expect(batches).toEqual([[2], [2]]);
+  });
+
+  test("start a new read for calls made during a read", async () => {
+    const batches: number[][] = [];
+    let inner: Promise<number> | undefined;
+    const read: (query: number) => Promise<number> = batchedQuery(
+      async (queries: number[]) => {
+        batches.push(queries);
+        if (!inner) inner = read(99);
+        return queries.map(query => query * 10);
+      },
+      (data, _query, index) => data[index]
+    );
+
+    expect(await read(1)).toBe(10);
+    expect(await inner).toBe(990);
+    expect(batches).toEqual([[1], [99]]);
+  });
+
+  test("match object queries by value", async () => {
+    const seen: string[][] = [];
+    const read = batchedQuery(
+      async (queries: { id: string }[]) => {
+        const ids = queries.map(query => query.id);
+        seen.push(ids);
+        return ids;
+      },
+      (ids, _query, index) => ids[index]
+    );
+
+    expect(await Promise.all([read({ id: "a" }), read({ id: "a" })])).toEqual(["a", "a"]);
+    expect(seen).toEqual([["a"]]);
+  });
+
+  test("match queries by a custom key", async () => {
+    const seen: string[][] = [];
+    const read = batchedQuery(
+      async (queries: { id: string; extra: number }[]) => {
+        const ids = queries.map(query => query.id);
+        seen.push(ids);
+        return ids;
+      },
+      (ids, _query, index) => ids[index],
+      { key: query => query.id }
+    );
+
+    expect(await Promise.all([read({ id: "a", extra: 1 }), read({ id: "a", extra: 2 })])).toEqual([
+      "a",
+      "a"
+    ]);
+    expect(seen).toEqual([["a"]]);
+  });
+});