From c0b30631c11c8a6b2ceeda22c03e537d4f949c63 Mon Sep 17 00:00:00 2001
From: Sebastian Kranz <sebastian.kranz@nubix.de>
Date: Sat, 17 Feb 2024 13:32:47 +0100
Subject: [PATCH] perf(queryClient): use `queryCache.get` instead of `find({
 queryKey })`

Methods that access a query by its `queryKey` can bypass `find` to avoid performance issues in large caches.
---
 packages/query-core/src/queryClient.ts        | 24 ++++++-----
 .../query-core/src/tests/queryClient.test.tsx | 42 +++++++++++++++++++
 2 files changed, 56 insertions(+), 10 deletions(-)

diff --git a/packages/query-core/src/queryClient.ts b/packages/query-core/src/queryClient.ts
index 67bd65d534..960650b30d 100644
--- a/packages/query-core/src/queryClient.ts
+++ b/packages/query-core/src/queryClient.ts
@@ -119,7 +119,8 @@ export class QueryClient {
       : TQueryFnData,
   >(queryKey: TTaggedQueryKey): TInferredQueryFnData | undefined
   getQueryData(queryKey: QueryKey) {
-    return this.#queryCache.find({ queryKey })?.state.data
+    const options = this.defaultQueryOptions({ queryKey })
+    return this.#queryCache.get(options.queryHash)?.state.data
   }
 
   ensureQueryData<
@@ -165,14 +166,6 @@ export class QueryClient {
     >,
     options?: SetDataOptions,
   ): TInferredQueryFnData | undefined {
-    const query = this.#queryCache.find<TInferredQueryFnData>({ queryKey })
-    const prevData = query?.state.data
-    const data = functionalUpdate(updater, prevData)
-
-    if (typeof data === 'undefined') {
-      return undefined
-    }
-
     const defaultedOptions = this.defaultQueryOptions<
       any,
       any,
@@ -181,6 +174,16 @@ export class QueryClient {
       QueryKey
     >({ queryKey })
 
+    const query = this.#queryCache.get<TInferredQueryFnData>(
+      defaultedOptions.queryHash,
+    )
+    const prevData = query?.state.data
+    const data = functionalUpdate(updater, prevData)
+
+    if (typeof data === 'undefined') {
+      return undefined
+    }
+
     return this.#queryCache
       .build(this, defaultedOptions)
       .setData(data, { ...options, manual: true })
@@ -204,7 +207,8 @@ export class QueryClient {
   getQueryState<TQueryFnData = unknown, TError = DefaultError>(
     queryKey: QueryKey,
   ): QueryState<TQueryFnData, TError> | undefined {
-    return this.#queryCache.find<TQueryFnData, TError>({ queryKey })?.state
+    const options = this.defaultQueryOptions({ queryKey })
+    return this.#queryCache.get<TQueryFnData, TError>(options.queryHash)?.state
   }
 
   removeQueries(filters?: QueryFilters): void {
diff --git a/packages/query-core/src/tests/queryClient.test.tsx b/packages/query-core/src/tests/queryClient.test.tsx
index 79e77892ec..66a485d93e 100644
--- a/packages/query-core/src/tests/queryClient.test.tsx
+++ b/packages/query-core/src/tests/queryClient.test.tsx
@@ -356,6 +356,16 @@ describe('queryClient', () => {
         }),
       )
     })
+
+    test('should set 10k data in less than 500ms', () => {
+      const key = queryKey()
+      const start = performance.now()
+      for (let i = 0; i < 10000; i++) {
+        queryClient.setQueryData([key, i], i)
+      }
+      const end = performance.now()
+      expect(end - start).toBeLessThan(500)
+    })
   })
 
   describe('setQueriesData', () => {
@@ -419,6 +429,38 @@ describe('queryClient', () => {
       queryClient.setQueryData([key, 'id'], 'bar')
       expect(queryClient.getQueryData([key])).toBeUndefined()
     })
+
+    test('should get 10k queries in less than 500ms', () => {
+      const key = queryKey()
+      for (let i = 0; i < 10000; i++) {
+        queryClient.setQueryData([key, i], i)
+      }
+
+      const start = performance.now()
+      for (let i = 0; i < 10000; i++) {
+        queryClient.getQueryData([key, i])
+      }
+      const end = performance.now()
+
+      expect(end - start).toBeLessThan(500)
+    })
+  })
+
+  describe('getQueryState', () => {
+    test('should get 10k queries in less than 500ms', () => {
+      const key = queryKey()
+      for (let i = 0; i < 10000; i++) {
+        queryClient.setQueryData([key, i], i)
+      }
+
+      const start = performance.now()
+      for (let i = 0; i < 10000; i++) {
+        queryClient.getQueryState([key, i])
+      }
+      const end = performance.now()
+
+      expect(end - start).toBeLessThan(500)
+    })
   })
 
   describe('ensureQueryData', () => {