perf(query-core): optimize find/findAll lookups in QueryCache and MutationCache

.changeset/perf-cache-lookup-optimization.md

@@ -0,0 +1,5 @@
+---
+"@tanstack/query-core": patch
+---
+
+Optimize `find`/`findAll` lookups in `QueryCache` and `MutationCache` by using index maps for O(1) key-based access when `exact: true` and a key filter are provided, instead of iterating all entries.

packages/query-core/src/mutationCache.ts

@@ -1,6 +1,6 @@
 import { notifyManager } from './notifyManager'
 import { Mutation } from './mutation'
-import { matchMutation, noop } from './utils'
+import { hashKey, matchMutation, noop } from './utils'
 import { Subscribable } from './subscribable'
 import type { MutationObserver } from './mutationObserver'
 import type {
@@ -93,12 +93,14 @@ type MutationCacheListener = (event: MutationCacheNotifyEvent) => void
 export class MutationCache extends Subscribable<MutationCacheListener> {
   #mutations: Set<Mutation<any, any, any, any>>
   #scopes: Map<string, Array<Mutation<any, any, any, any>>>
+  #byMutationKey: Map<string, Array<Mutation<any, any, any, any>>>
   #mutationId: number
 
   constructor(public config: MutationCacheConfig = {}) {
     super()
     this.#mutations = new Set()
     this.#scopes = new Map()
+    this.#byMutationKey = new Map()
     this.#mutationId = 0
   }
 
@@ -131,6 +133,13 @@ export class MutationCache extends Subscribable<MutationCacheListener> {
         this.#scopes.set(scope, [mutation])
       }
     }
+    const { mutationKey } = mutation.options
+    if (mutationKey) {
+      const hash = hashKey(mutationKey)
+      const keyed = this.#byMutationKey.get(hash)
+      if (keyed) keyed.push(mutation)
+      else this.#byMutationKey.set(hash, [mutation])
+    }
     this.notify({ type: 'added', mutation })
   }
 
@@ -150,6 +159,19 @@ export class MutationCache extends Subscribable<MutationCacheListener> {
           }
         }
       }
+      const { mutationKey } = mutation.options
+      if (mutationKey) {
+        const hash = hashKey(mutationKey)
+        const keyed = this.#byMutationKey.get(hash)
+        if (keyed) {
+          if (keyed.length > 1) {
+            const index = keyed.indexOf(mutation)
+            if (index !== -1) keyed.splice(index, 1)
+          } else {
+            this.#byMutationKey.delete(hash)
+          }
+        }
+      }
     }
 
     // Currently we notify the removal even if the mutation was already removed.
@@ -194,6 +216,7 @@ export class MutationCache extends Subscribable<MutationCacheListener> {
       })
       this.#mutations.clear()
       this.#scopes.clear()
+      this.#byMutationKey.clear()
     })
   }
 
@@ -211,13 +234,50 @@ export class MutationCache extends Subscribable<MutationCacheListener> {
   ): Mutation<TData, TError, TVariables, TOnMutateResult> | undefined {
     const defaultedFilters = { exact: true, ...filters }
 
-    return this.getAll().find((mutation) =>
-      matchMutation(defaultedFilters, mutation),
-    ) as Mutation<TData, TError, TVariables, TOnMutateResult> | undefined
+    if (defaultedFilters.exact && defaultedFilters.mutationKey) {
+      const candidates = this.#byMutationKey.get(
+        hashKey(defaultedFilters.mutationKey),
+      )
+      if (!candidates) return undefined
+      const { mutationKey: _m, ...filtersWithoutKey } = defaultedFilters
+      for (const mutation of candidates) {
+        if (matchMutation(filtersWithoutKey as MutationFilters, mutation)) {
+          return mutation as Mutation<
+            TData,
+            TError,
+            TVariables,
+            TOnMutateResult
+          >
+        }
+      }
+      return undefined
+    }
+
+    for (const mutation of this.#mutations) {
+      if (matchMutation(defaultedFilters, mutation)) {
+        return mutation as Mutation<TData, TError, TVariables, TOnMutateResult>
+      }
+    }
+    return undefined
   }
 
   findAll(filters: MutationFilters = {}): Array<Mutation> {
-    return this.getAll().filter((mutation) => matchMutation(filters, mutation))
+    if (filters.exact && filters.mutationKey) {
+      const candidates = this.#byMutationKey.get(hashKey(filters.mutationKey))
+      if (!candidates) return []
+      const { mutationKey: _m, ...filtersWithoutKey } = filters
+      return candidates.filter((m) =>
+        matchMutation(filtersWithoutKey as MutationFilters, m),
+      )
+    }
+
+    const result: Array<Mutation> = []
+    for (const mutation of this.#mutations) {
+      if (matchMutation(filters, mutation)) {
+        result.push(mutation)
+      }
+    }
+    return result
   }
 
   notify(event: MutationCacheNotifyEvent) {
@@ -229,7 +289,12 @@ export class MutationCache extends Subscribable<MutationCacheListener> {
   }
 
   resumePausedMutations(): Promise<unknown> {
-    const pausedMutations = this.getAll().filter((x) => x.state.isPaused)
+    const pausedMutations: Array<Mutation> = []
+    for (const mutation of this.#mutations) {
+      if (mutation.state.isPaused) {
+        pausedMutations.push(mutation)
+      }
+    }
 
     return notifyManager.batch(() =>
       Promise.all(

packages/query-core/src/queryCache.ts

@@ -1,4 +1,4 @@
-import { hashQueryKeyByOptions, matchQuery } from './utils'
+import { hashKey, hashQueryKeyByOptions, matchQuery } from './utils'
 import { Query } from './query'
 import { notifyManager } from './notifyManager'
 import { Subscribable } from './subscribable'
@@ -185,16 +185,50 @@ export class QueryCache extends Subscribable<QueryCacheListener> {
   ): Query<TQueryFnData, TError, TData> | undefined {
     const defaultedFilters = { exact: true, ...filters }
 
-    return this.getAll().find((query) =>
-      matchQuery(defaultedFilters, query),
-    ) as Query<TQueryFnData, TError, TData> | undefined
+    let found: Query | undefined
+
+    if (defaultedFilters.exact) {
+      const candidate = this.#queries.get(hashKey(defaultedFilters.queryKey))
+      if (candidate) {
+        const { queryKey: _q, ...filtersWithoutKey } = defaultedFilters
+        found = matchQuery(filtersWithoutKey as QueryFilters, candidate)
+          ? candidate
+          : undefined
+        return found as Query<TQueryFnData, TError, TData> | undefined
+      }
+    }
+
+    for (const query of this.#queries.values()) {
+      if (matchQuery(defaultedFilters, query)) {
+        found = query
+        break
+      }
+    }
+    return found as Query<TQueryFnData, TError, TData> | undefined
   }
 
   findAll(filters: QueryFilters<any> = {}): Array<Query> {
-    const queries = this.getAll()
-    return Object.keys(filters).length > 0
-      ? queries.filter((query) => matchQuery(filters, query))
-      : queries
+    if (Object.keys(filters).length === 0) {
+      return [...this.#queries.values()]
+    }
+
+    if (filters.exact && filters.queryKey) {
+      const candidate = this.#queries.get(hashKey(filters.queryKey))
+      if (candidate) {
+        const { queryKey: _q, ...filtersWithoutKey } = filters
+        return matchQuery(filtersWithoutKey as QueryFilters, candidate)
+          ? [candidate]
+          : []
+      }
+    }
+
+    const result: Array<Query> = []
+    for (const query of this.#queries.values()) {
+      if (matchQuery(filters, query)) {
+        result.push(query)
+      }
+    }
+    return result
   }
 
   notify(event: QueryCacheNotifyEvent): void {
@@ -207,17 +241,17 @@ export class QueryCache extends Subscribable<QueryCacheListener> {
 
   onFocus(): void {
     notifyManager.batch(() => {
-      this.getAll().forEach((query) => {
+      for (const query of this.#queries.values()) {
         query.onFocus()
-      })
+      }
     })
   }
 
   onOnline(): void {
     notifyManager.batch(() => {
-      this.getAll().forEach((query) => {
+      for (const query of this.#queries.values()) {
         query.onOnline()
-      })
+      }
     })
   }
 }