diff --git a/src/data/index.ts b/src/data/index.ts index c6209164..805d3e22 100644 --- a/src/data/index.ts +++ b/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"; diff --git a/src/data/query.ts b/src/data/query.ts index bf3183bb..137bdbc8 100644 --- a/src/data/query.ts +++ b/src/data/query.ts @@ -275,3 +275,101 @@ function isPlainObject(obj: object) { (!(proto = Object.getPrototypeOf(obj)) || proto === Object.prototype) ); } + +export interface BatchOptions { + /** 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 { + resolve: (value: Return) => void; + reject: (reason: unknown) => void; +} + +interface BatchSlot { + query: Query; + callers: BatchCaller[]; +} + +interface BatchWaiting extends BatchCaller { + 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( + callback: (queries: Query[]) => Promise, + lookup: (data: Data, query: Query, index: number) => Return, + options: BatchOptions = {} +): (query: Query) => Promise { + const keyOf = options.key ?? ((query: Query): unknown => hashKey([query])); + const size = Math.max(1, Math.floor(options.limit ?? Number.POSITIVE_INFINITY)); + let waiting: BatchWaiting[] = []; + let timer: ReturnType | 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[]) => { + 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>(); + // 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((resolve, reject) => { + waiting.push({ query, resolve, reject }); + if (!timer) timer = setTimeout(flush, options.wait ?? 0); + }); + }; +} diff --git a/test/data/batched-query.spec.ts b/test/data/batched-query.spec.ts new file mode 100644 index 00000000..5dfe51b8 --- /dev/null +++ b/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(); + 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 | undefined; + const read: (query: number) => Promise = 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"]]); + }); +});