diff --git a/packages/query-core/src/__tests__/queryCache.test.tsx b/packages/query-core/src/__tests__/queryCache.test.tsx index 6edb5598f7..746959dbf0 100644 --- a/packages/query-core/src/__tests__/queryCache.test.tsx +++ b/packages/query-core/src/__tests__/queryCache.test.tsx @@ -293,6 +293,69 @@ describe('queryCache', () => { expect(queryCache.findAll({ fetchStatus: 'fetching' })).toEqual([]) }) + test('only visits queries with the given prefix', async () => { + const prefix = ['posts'] + queryClient.prefetchQuery({ + queryKey: ['other'], + queryFn: () => sleep(100).then(() => 'other'), + }) + queryClient.prefetchQuery({ + queryKey: ['other', {}], + queryFn: () => sleep(100).then(() => 'otherObjectSuffix'), + }) + queryClient.prefetchQuery({ + queryKey: [{}, 'objectPrefix'], + queryFn: () => sleep(100).then(() => 'otherObjectPrefix'), + }) + + queryClient.prefetchQuery({ + queryKey: prefix, + queryFn: () => sleep(100).then(() => 'exactMatch'), + }) + const exactMatchQuery = queryCache.find({ queryKey: prefix }) + + queryClient.prefetchQuery({ + queryKey: [...prefix, 1], + queryFn: () => sleep(100).then(() => 'primitiveSuffix'), + }) + const primitiveSuffixQuery = queryCache.find({ queryKey: [...prefix, 1] }) + + queryClient.prefetchQuery({ + queryKey: [...prefix, 2, 'primitiveSuffix'], + queryFn: () => sleep(100).then(() => 'primitiveSuffixLength2'), + }) + const primitiveSuffixLength2Query = queryCache.find({ + queryKey: [...prefix, 2, 'primitiveSuffix'], + }) + + queryClient.prefetchQuery({ + queryKey: [...prefix, 3, {}, 'obj'], + queryFn: () => sleep(100).then(() => 'matchingObjectSuffix'), + }) + const matchingObjectSuffixQuery = queryCache.find({ + queryKey: [...prefix, 3, {}, 'obj'], + }) + + await vi.advanceTimersByTimeAsync(100) + + let predicateCallCount = 0 + const results = queryCache.findAll({ + queryKey: prefix, + predicate: () => { + predicateCallCount++ + return true + }, + }) + + expect(predicateCallCount).toBe(results.length) + expect(results).toEqual([ + exactMatchQuery, + primitiveSuffixQuery, + primitiveSuffixLength2Query, + matchingObjectSuffixQuery, + ]) + }) + test('should return all the queries when no filters are defined', async () => { const key1 = queryKey() const key2 = queryKey() diff --git a/packages/query-core/src/query.ts b/packages/query-core/src/query.ts index df5b7c030e..ba6f0c6174 100644 --- a/packages/query-core/src/query.ts +++ b/packages/query-core/src/query.ts @@ -202,8 +202,23 @@ export class Query< setOptions( options?: QueryOptions, ): void { + const oldOptions = (this.options as typeof this.options | undefined) + ? this.options + : undefined + this.options = { ...this.#defaultOptions, ...options } + if (oldOptions) { + // Do not do this update when first created. + // The QueryCache manages admission / removal itself. + // We only need to keep it up-to-date here. + const prevHashFn = oldOptions.queryKeyHashFn + const newHashFn = this.options.queryKeyHashFn + if (prevHashFn !== newHashFn) { + this.#cache.onQueryKeyHashFunctionChanged(prevHashFn, newHashFn) + } + } + this.updateGcTime(this.options.gcTime) } diff --git a/packages/query-core/src/queryCache.ts b/packages/query-core/src/queryCache.ts index dd7123eaac..34febf87ca 100644 --- a/packages/query-core/src/queryCache.ts +++ b/packages/query-core/src/queryCache.ts @@ -1,13 +1,19 @@ -import { hashQueryKeyByOptions, matchQuery } from './utils' +import { + hashKey, + hashQueryKeyByOptions, + matchQuery, + partialMatchKey, +} from './utils' import { Query } from './query' import { notifyManager } from './notifyManager' import { Subscribable } from './subscribable' -import type { QueryFilters } from './utils' import type { Action, QueryState } from './query' +import type { QueryFilters } from './utils' import type { DefaultError, NotifyEvent, QueryKey, + QueryKeyHashFunction, QueryOptions, WithRequired, } from './types' @@ -87,10 +93,389 @@ export interface QueryStore { values: () => IterableIterator } +class RefCountSet { + #refcounts = new Map(); + + [Symbol.iterator]() { + return this.#refcounts.keys() + } + + get size() { + return this.#refcounts.size + } + + add(value: T) { + const n = this.#refcounts.get(value) ?? 0 + this.#refcounts.set(value, n + 1) + } + + remove(value: T) { + let n = this.#refcounts.get(value) + if (n === undefined) { + return + } + n-- + if (n === 0) { + this.#refcounts.delete(value) + } else { + this.#refcounts.set(value, n) + } + } +} + +type Primitive = string | number | boolean | bigint | symbol | undefined | null +function isPrimitive(value: unknown): value is Primitive { + if (value === undefined || value === null) { + return true + } + const t = typeof value + switch (t) { + case 'object': + case 'function': + return false + case 'string': + case 'number': + case 'boolean': + case 'bigint': + case 'symbol': + case 'undefined': + return true + default: + t satisfies never + return false + } +} + +/** + * Like Map, but object keys have value semantics equality based + * on partialMatchKey, instead of reference equality. + * + * Lookups by object are O(NumberOfKeys) instead of O(1). + * + * ```ts + * const queryKeyMap = new QueryKeyElementMap<{ orderBy: 'likes', order: 'desc', limit: 30 }, string>() + * queryKeyMap.set({ orderBy: 'likes', order: 'desc', limit: 30 }, 'value') + * queryKeyMap.get({ orderBy: 'likes', order: 'desc', limit: 30 }) // 'value' + * + * const vanillaMap = new Map<{ orderBy: 'likes', order: 'desc', limit: 30 }, string>() + * vanillaMap.set({ orderBy: 'likes', order: 'desc', limit: 30 }, 'value') + * vanillaMap.get({ orderBy: 'likes', order: 'desc', limit: 30 }) // undefined + * ``` + */ +class QueryKeyElementMap { + #primitiveMap = new Map() + #objectMap = new Map() + + get size() { + return this.#primitiveMap.size + this.#objectMap.size + } + + get( + key: (TKeyElement & Primitive) | (TKeyElement & object), + ): TValue | undefined { + if (isPrimitive(key)) { + return this.#primitiveMap.get(key) + } + + const matchingKey = this.findMatchingObjectKey(key) + if (matchingKey) { + return this.#objectMap.get(matchingKey) + } + + return undefined + } + + set( + key: (TKeyElement & Primitive) | (TKeyElement & object), + value: TValue, + ): void { + if (isPrimitive(key)) { + this.#primitiveMap.set(key, value) + return + } + + const matchingKey = this.findMatchingObjectKey(key) + this.#objectMap.set(matchingKey ?? key, value) + } + + delete(key: (TKeyElement & Primitive) | (TKeyElement & object)): boolean { + if (isPrimitive(key)) { + return this.#primitiveMap.delete(key) + } + + const matchingKey = this.findMatchingObjectKey(key) + if (matchingKey) { + this.#objectMap.delete(matchingKey) + return true + } + return false + } + + values(): Iterable | undefined { + if (!this.#primitiveMap.size && !this.#objectMap.size) { + return undefined + } + + if (this.#primitiveMap.size && !this.#objectMap.size) { + return this.#primitiveMap.values() + } + + if (!this.#primitiveMap.size && this.#objectMap.size) { + return this.#objectMap.values() + } + + const primitiveValues = this.#primitiveMap.values() + const objectValues = this.#objectMap.values() + return (function* () { + yield* primitiveValues + yield* objectValues + })() + } + + private findMatchingObjectKey( + key: TKeyElement & object, + ): (TKeyElement & object) | undefined { + // Reference equality + if (this.#objectMap.has(key)) { + return key + } + + // Linear search for the matching key. + // This makes lookups in the trie O(NumberOfObjectKeys) + // but it also gives lookups in the trie like + // `map.get(['a', { obj: true }, 'c'])` the same semantics + // as `partialMatchKey` itself. + const keyArray = [key] + for (const candidateKey of this.#objectMap.keys()) { + if (partialMatchKey([candidateKey], keyArray)) { + return candidateKey + } + } + + return undefined + } +} + +type QueryKeyTrieNode = { + /** Element in the query key, QueryKey[number]. */ + key: (TKeyElement & Primitive) | (TKeyElement & object) + /** + * Value stored at the end of the path leading to this node. + * This holds: `key[key.length - 1] === node.key` + * ```ts + * map.set(['a', 'b', 'c'], '123') + * // -> + * const root = { + * children: { + * 'a': { + * key: 'a', + * children: { + * 'b': { + * key: 'b', + * children: { + * 'c': { + * key: 'c', + * value: '123', + * insertionOrder: 0, + * } + * } + * } + * } + * } + * } + * } + * ``` + */ + value?: TValue + /** + * Insertion order of the *value* being stored in the trie. + * Unfortunately the natural iteration order of values in the trie does not + * match the insertion order, as expected from a map-like data structure. + * + * We need to track it explicitly. + */ + insertionOrder?: number + /** + * Children map to the next index element in the queryKey. + */ + children?: QueryKeyElementMap< + TKeyElement, + QueryKeyTrieNode +> +} + +type QueryKeyTrieNodeWithValue = { + key: (TKeyElement & Primitive) | (TKeyElement & object) + value: TValue + insertionOrder: number + children?: QueryKeyElementMap< + TKeyElement, + QueryKeyTrieNode +> +} + +/** + * We only consider a value to be stored in a node when insertionOrder is defined. + * This allows storing `undefined` as a value. + */ +function nodeHasValue( + node: QueryKeyTrieNode, +): node is QueryKeyTrieNodeWithValue { + return node.insertionOrder !== undefined +} + +/** + * Path length is always 1 greater than the key length, as it includes the root + * node. + * ```ts + * map.set(['a', 'b', 'c'], '123') + * const path = [root, n1, n2, n3] + * const key = ['a', 'b', 'c'] + */ +function traverse< + TKey extends QueryKey, + TValue, + TLookup extends QueryKeyTrieNode | undefined, +>( + root: QueryKeyTrieNode, + key: TKey, + // May create a child node if needed + lookup: ( + parent: QueryKeyTrieNode, + key: (TKey[number] & Primitive) | (TKey[number] & object), + ) => TLookup, +): TLookup extends undefined ? undefined : Array { + const path: Array> = [root] + let node: QueryKeyTrieNode | undefined = root + // In hot code like this data structures, it is best to avoid creating + // Iterators with for-of loops. + // eslint-disable-next-line @typescript-eslint/prefer-for-of + for (let i = 0; i < key.length; i++) { + const keyPart = key[i] + node = lookup( + node, + keyPart as (TKey[number] & Primitive) | (TKey[number] & object), + ) + if (node) { + path.push(node) + } else { + return undefined as never + } + } + return path as never +} + +function* iterateSubtreeValueNodes( + node: QueryKeyTrieNode, +): Generator, void, undefined> { + if (nodeHasValue(node)) { + yield node + } + + const children = node.children?.values() + if (!children) { + return + } + + for (const child of children) { + yield* iterateSubtreeValueNodes(child) + } +} + +class QueryKeyTrie { + #root: QueryKeyTrieNode = { + key: undefined, + } + // Provides relative insertion ordering between values in the trie. + #insertionOrder = 0 + + set(key: TKey, value: TValue): void { + const path = traverse(this.#root, key, (parent, keyPart) => { + parent.children ??= new QueryKeyElementMap() + let child = parent.children.get(keyPart) + if (!child) { + // Note: insertionOrder is for values, not when nodes enter the trie. + child = { key: keyPart } + parent.children.set(keyPart, child) + } + return child + }) + + const lastNode = path[path.length - 1]! + if (!nodeHasValue(lastNode)) { + lastNode.insertionOrder = this.#insertionOrder++ + } + lastNode.value = value + } + + delete(key: TKey): void { + const path = traverse(this.#root, key, (parent, keyPart) => + parent.children?.get(keyPart), + ) + if (!path) { + return + } + + const lastNode = path[path.length - 1]! + if (lastNode.insertionOrder === undefined) { + // No value stored at key. + return + } + + // Drop. + lastNode.value = undefined + lastNode.insertionOrder = undefined + + // GC nodes in path that are no longer needed. + for (let i = path.length - 1; i> 0; i--) { + const node = path[i]! + if (nodeHasValue(node) || node.children?.size) { + // Has data. Do not GC. + return + } + + const parent = path[i - 1] + parent?.children?.delete(node.key) + } + } + + /** + * Returns all values that match the given key: + * Either the value has the same key and is all primitives, + * Or the value's key is a suffix of the given key and contains a non-primitive key. + */ + iteratePrefix(key: TKey): Array | undefined { + const path = traverse(this.#root, key, (parent, keyPart) => + parent.children?.get(keyPart), + ) + if (!path) { + return undefined + } + + const lastNode = path[path.length - 1]! + if (!lastNode.children?.size) { + // No children - either return value if we have one, or nothing. + if (nodeHasValue(lastNode)) { + return [lastNode.value] + } + return undefined + } + + const subtreeInDepthFirstOrder = Array.from( + iterateSubtreeValueNodes(lastNode), + ) + return subtreeInDepthFirstOrder + .sort((a, b) => a.insertionOrder - b.insertionOrder) + .map((node) => node.value) + } +} + // CLASS export class QueryCache extends Subscribable { #queries: QueryStore + #keyIndex = new QueryKeyTrie() + #knownHashFns = new RefCountSet>() constructor(public config: QueryCacheConfig = {}) { super() @@ -133,6 +518,11 @@ export class QueryCache extends Subscribable { add(query: Query): void { if (!this.#queries.has(query.queryHash)) { this.#queries.set(query.queryHash, query) + this.#keyIndex.set(query.queryKey, query) + const hashFn = query.options.queryKeyHashFn + if (hashFn) { + this.#knownHashFns.add(hashFn) + } this.notify({ type: 'added', @@ -149,6 +539,11 @@ export class QueryCache extends Subscribable { if (queryInMap === query) { this.#queries.delete(query.queryHash) + this.#keyIndex.delete(query.queryKey) + const hashFn = query.options.queryKeyHashFn + if (hashFn) { + this.#knownHashFns.remove(hashFn) + } } this.notify({ type: 'removed', query }) @@ -184,19 +579,78 @@ export class QueryCache extends Subscribable { filters: WithRequired, ): Query | undefined { const defaultedFilters = { exact: true, ...filters } + if (defaultedFilters.exact) { + return this.findExact(defaultedFilters) + } + + const candidates = this.#keyIndex.iteratePrefix(defaultedFilters.queryKey) + if (!candidates) { + return undefined + } - return this.getAll().find((query) => - matchQuery(defaultedFilters, query), - ) as Query | undefined + return candidates.find((query) => matchQuery(defaultedFilters, query)) as + | Query + | undefined } findAll(filters: QueryFilters = {}): Array { + if (filters.exact && filters.queryKey) { + const query = this.findExact(filters) + return query ? [query] : [] + } + + if (filters.queryKey) { + const candidates = this.#keyIndex.iteratePrefix(filters.queryKey) + if (!candidates) { + return [] + } + + return Object.keys(filters).length> 1 + ? candidates.filter((query) => matchQuery(filters, query)) + : candidates + } + const queries = this.getAll() return Object.keys(filters).length> 0 ? queries.filter((query) => matchQuery(filters, query)) : queries } + private findExact< + TQueryFnData = unknown, + TError = DefaultError, + TData = TQueryFnData, +>( + filters: QueryFilters, + ): Query | undefined { + const tryHashFn = (hashFn: QueryKeyHashFunction) => { + try { + const query = this.get(hashFn(filters.queryKey)) + if (query && matchQuery(filters, query)) { + // Confirmed the query actually uses the hash function we tried + // and matches the non-queryKey filters + return query + } else { + return undefined + } + } catch (error) { + return undefined + } + } + + let query = tryHashFn(hashKey) + if (!query) { + for (const hashFn of this.#knownHashFns) { + query = tryHashFn(hashFn) + if (query) { + break + } + } + } + + return query as unknown as Query | undefined + } + notify(event: QueryCacheNotifyEvent): void { notifyManager.batch(() => { this.listeners.forEach((listener) => { @@ -220,4 +674,16 @@ export class QueryCache extends Subscribable { }) }) } + + onQueryKeyHashFunctionChanged( + before: QueryKeyHashFunction | undefined, + after: QueryKeyHashFunction | undefined, + ): void { + if (before) { + this.#knownHashFns.remove(before) + } + if (after) { + this.#knownHashFns.add(after) + } + } }

AltStyle によって変換されたページ (->オリジナル) /