fix persistedQuery memory leak in SSR scenarios (#10805)

* fix `persistedQuery` memory leak in SSR scenarios * Update src/link/persisted-queries/index.ts Co-authored-by: Jerel Miller <[email protected]> * use jest spy to make test more error-resilient --------- Co-authored-by: Jerel Miller <[email protected]>
apollographql · Apr 27, 2023 · a550366 · a550366
1 parent fdd0c5b
commit a550366
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 23 deletions.
diff --git a/.changeset/long-yaks-clean.md b/.changeset/long-yaks-clean.md
@@ -0,0 +1,5 @@
+---
+"@apollo/client": patch
+---
+
+Fix a potential memory leak in SSR scenarios when many `persistedQuery` instances were created over time.
diff --git a/src/link/persisted-queries/__tests__/persisted-queries.test.ts b/src/link/persisted-queries/__tests__/persisted-queries.test.ts
@@ -99,17 +99,14 @@ describe('happy path', () => {
   itAsync('memoizes between requests', (resolve, reject) => {
     fetchMock.post("/graphql", () => new Promise(resolve => resolve({ body: response})), { repeat: 1 });
     fetchMock.post("/graphql", () => new Promise(resolve => resolve({ body: response})), { repeat: 1 });
-    const link = createPersistedQuery({ sha256 }).concat(createHttpLink());
+    const hashSpy = jest.fn(sha256)
+    const link = createPersistedQuery({ sha256: hashSpy }).concat(createHttpLink());
 
-    let start = new Date();
     execute(link, { query, variables }).subscribe(result => {
-      const firstRun = new Date().valueOf() - start.valueOf();
+      expect(hashSpy).toHaveBeenCalledTimes(1);
       expect(result.data).toEqual(data);
-      // this one should go faster becuase of memoization
-      let secondStart = new Date();
       execute(link, { query, variables }).subscribe(result2 => {
-        const secondRun = new Date().valueOf() - secondStart.valueOf();
-        expect(firstRun).toBeGreaterThan(secondRun);
+        expect(hashSpy).toHaveBeenCalledTimes(1);
         expect(result2.data).toEqual(data);
         resolve();
       }, reject);

diff --git a/src/link/persisted-queries/index.ts b/src/link/persisted-queries/index.ts
@@ -93,18 +93,13 @@ function operationDefinesMutation(operation: Operation) {
     d => d.kind === 'OperationDefinition' && d.operation === 'mutation');
 }
 
-const { hasOwnProperty } = Object.prototype;
-
-const hashesByQuery = new WeakMap<
-  DocumentNode,
-  Record<string, Promise<string>>
->();
-
-let nextHashesChildKey = 0;
-
 export const createPersistedQueryLink = (
   options: PersistedQueryLink.Options,
 ) => {
+  const hashesByQuery = new WeakMap<
+    DocumentNode,
+    Promise<string>
+  >();
   // Ensure a SHA-256 hash function is provided, if a custom hash
   // generation function is not provided. We don't supply a SHA-256 hash
   // function by default, to avoid forcing one as a dependency. Developers
@@ -136,8 +131,6 @@ export const createPersistedQueryLink = (
 
   let supportsPersistedQueries = true;
 
-  const hashesChildKey = 'forLink' + nextHashesChildKey++;
-
   const getHashPromise = (query: DocumentNode) =>
     new Promise<string>(resolve => resolve(generateHash(query)));
 
@@ -148,11 +141,9 @@ export const createPersistedQueryLink = (
       // what to do with the bogus query.
       return getHashPromise(query);
     }
-    let hashes = hashesByQuery.get(query)!;
-    if (!hashes) hashesByQuery.set(query, hashes = Object.create(null));
-    return hasOwnProperty.call(hashes, hashesChildKey)
-      ? hashes[hashesChildKey]
-      : hashes[hashesChildKey] = getHashPromise(query);
+    let hash = hashesByQuery.get(query)!;
+    if (!hash) hashesByQuery.set(query, hash = getHashPromise(query));
+    return hash;
   }
 
   return new ApolloLink((operation, forward) => {