From a5503666c2cc8220ac1d877e3296556e54e58ff6 Mon Sep 17 00:00:00 2001
From: Lenz Weber-Tronic <lorenz.weber-tronic@apollographql.com>
Date: Thu, 27 Apr 2023 11:10:09 +0200
Subject: [PATCH] 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 <jerelmiller@gmail.com>

* use jest spy to make test more error-resilient

---------

Co-authored-by: Jerel Miller <jerelmiller@gmail.com>
---
 .changeset/long-yaks-clean.md                 |  5 ++++
 .../__tests__/persisted-queries.test.ts       | 11 ++++-----
 src/link/persisted-queries/index.ts           | 23 ++++++-------------
 3 files changed, 16 insertions(+), 23 deletions(-)
 create mode 100644 .changeset/long-yaks-clean.md

diff --git a/.changeset/long-yaks-clean.md b/.changeset/long-yaks-clean.md
new file mode 100644
index 00000000000..f0fe9c8ee66
--- /dev/null
+++ 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
index 64db01a1897..e472c4dbf5d 100644
--- 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
index 795731daa5e..1aa661918dc 100644
--- 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) => {