Suspense cache: use Trie directly (#10969)

Co-authored-by: Jerel Miller <jerelmiller@gmail.com>
apollographql · Jun 13, 2023 · 525a931 · 525a931
1 parent e187866
commit 525a931
Show file tree

Hide file tree

Showing 7 changed files with 118 additions and 33 deletions.
diff --git a/.changeset/tasty-wasps-relate.md b/.changeset/tasty-wasps-relate.md
@@ -0,0 +1,5 @@
+---
+'@apollo/client': patch
+---
+
+Slightly decrease bundle size and memory footprint of `SuspenseCache` by changing how cache entries are stored internally.
diff --git a/src/react/cache/SuspenseCache.ts b/src/react/cache/SuspenseCache.ts
@@ -19,12 +19,7 @@ interface SuspenseCacheOptions {
 }
 
 export class SuspenseCache {
-  private cacheKeys = new Trie<CacheKey>(
-    canUseWeakMap,
-    (cacheKey: CacheKey) => cacheKey
-  );
-
-  private queryRefs = new Map<CacheKey, QueryReference>();
+  private queryRefs = new Trie<{ current?: QueryReference }>(canUseWeakMap);
   private options: SuspenseCacheOptions;
 
   constructor(options: SuspenseCacheOptions = Object.create(null)) {
@@ -35,19 +30,18 @@ export class SuspenseCache {
     cacheKey: CacheKey,
     createObservable: () => ObservableQuery<TData>
   ) {
-    const stableCacheKey = this.cacheKeys.lookupArray(cacheKey);
+    const ref = this.queryRefs.lookupArray(cacheKey);
 
-    if (!this.queryRefs.has(stableCacheKey)) {
-      this.queryRefs.set(
-        stableCacheKey,
-        new QueryReference(createObservable(), {
-          key: stableCacheKey,
-          autoDisposeTimeoutMs: this.options.autoDisposeTimeoutMs,
-          onDispose: () => this.queryRefs.delete(stableCacheKey),
-        })
-      );
+    if (!ref.current) {
+      ref.current = new QueryReference(createObservable(), {
+        key: cacheKey,
+        autoDisposeTimeoutMs: this.options.autoDisposeTimeoutMs,
+        onDispose: () => {
+          delete ref.current;
+        },
+      });
     }
 
-    return this.queryRefs.get(stableCacheKey)! as QueryReference<TData>;
+    return ref.current as QueryReference<TData>;
   }
 }
diff --git a/src/react/hooks/__tests__/useBackgroundQuery.test.tsx b/src/react/hooks/__tests__/useBackgroundQuery.test.tsx
@@ -697,8 +697,11 @@ describe('useBackgroundQuery', () => {
       }
     );
 
-    expect(directSuspenseCache['queryRefs'].size).toBe(1);
-    expect(contextSuspenseCache['queryRefs'].size).toBe(0);
+    expect(directSuspenseCache).toHaveSuspenseCacheEntryUsing(client, query);
+    expect(contextSuspenseCache).not.toHaveSuspenseCacheEntryUsing(
+      client,
+      query
+    );
   });
 
   it('passes context to the link', async () => {

diff --git a/src/react/hooks/__tests__/useSuspenseQuery.test.tsx b/src/react/hooks/__tests__/useSuspenseQuery.test.tsx
@@ -378,8 +378,11 @@ describe('useSuspenseQuery', () => {
       }
     );
 
-    expect(directSuspenseCache['queryRefs'].size).toBe(1);
-    expect(contextSuspenseCache['queryRefs'].size).toBe(0);
+    expect(directSuspenseCache).toHaveSuspenseCacheEntryUsing(client, query);
+    expect(contextSuspenseCache).not.toHaveSuspenseCacheEntryUsing(
+      client,
+      query
+    );
   });
 
   it('ensures a valid fetch policy is used', () => {
@@ -678,7 +681,7 @@ describe('useSuspenseQuery', () => {
     );
 
     expect(client.getObservableQueries().size).toBe(1);
-    expect(suspenseCache['queryRefs'].size).toBe(1);
+    expect(suspenseCache).toHaveSuspenseCacheEntryUsing(client, query);
 
     unmount();
 
@@ -687,7 +690,7 @@ describe('useSuspenseQuery', () => {
     await wait(0);
 
     expect(client.getObservableQueries().size).toBe(0);
-    expect(suspenseCache['queryRefs'].size).toBe(0);
+    expect(suspenseCache).not.toHaveSuspenseCacheEntryUsing(client, query);
   });
 
   it('tears down all queries when rendering with multiple variable sets', async () => {
@@ -716,7 +719,12 @@ describe('useSuspenseQuery', () => {
     });
 
     expect(client.getObservableQueries().size).toBe(2);
-    expect(suspenseCache['queryRefs'].size).toBe(2);
+    expect(suspenseCache).toHaveSuspenseCacheEntryUsing(client, query, {
+      variables: { id: '1' },
+    });
+    expect(suspenseCache).toHaveSuspenseCacheEntryUsing(client, query, {
+      variables: { id: '2' },
+    });
 
     unmount();
 
@@ -725,7 +733,13 @@ describe('useSuspenseQuery', () => {
     await wait(0);
 
     expect(client.getObservableQueries().size).toBe(0);
-    expect(suspenseCache['queryRefs'].size).toBe(0);
+
+    expect(suspenseCache).not.toHaveSuspenseCacheEntryUsing(client, query, {
+      variables: { id: '1' },
+    });
+    expect(suspenseCache).not.toHaveSuspenseCacheEntryUsing(client, query, {
+      variables: { id: '2' },
+    });
   });
 
   it('tears down all queries when multiple clients are used', async () => {
@@ -773,9 +787,16 @@ describe('useSuspenseQuery', () => {
       });
     });
 
+    const variables = { id: '1' };
+
     expect(client1.getObservableQueries().size).toBe(1);
     expect(client2.getObservableQueries().size).toBe(1);
-    expect(suspenseCache['queryRefs'].size).toBe(2);
+    expect(suspenseCache).toHaveSuspenseCacheEntryUsing(client1, query, {
+      variables,
+    });
+    expect(suspenseCache).toHaveSuspenseCacheEntryUsing(client2, query, {
+      variables,
+    });
 
     unmount();
 
@@ -785,7 +806,12 @@ describe('useSuspenseQuery', () => {
 
     expect(client1.getObservableQueries().size).toBe(0);
     expect(client2.getObservableQueries().size).toBe(0);
-    expect(suspenseCache['queryRefs'].size).toBe(0);
+    expect(suspenseCache).not.toHaveSuspenseCacheEntryUsing(client1, query, {
+      variables,
+    });
+    expect(suspenseCache).not.toHaveSuspenseCacheEntryUsing(client2, query, {
+      variables,
+    });
   });
 
   it('tears down the query if the component never renders again after suspending', async () => {
@@ -834,12 +860,12 @@ describe('useSuspenseQuery', () => {
     link.simulateComplete();
 
     expect(client.getObservableQueries().size).toBe(1);
-    expect(suspenseCache['queryRefs'].size).toBe(1);
+    expect(suspenseCache).toHaveSuspenseCacheEntryUsing(client, query);
 
     jest.advanceTimersByTime(30_000);
 
     expect(client.getObservableQueries().size).toBe(0);
-    expect(suspenseCache['queryRefs'].size).toBe(0);
+    expect(suspenseCache).not.toHaveSuspenseCacheEntryUsing(client, query);
 
     jest.useRealTimers();
 
@@ -895,12 +921,12 @@ describe('useSuspenseQuery', () => {
     link.simulateComplete();
 
     expect(client.getObservableQueries().size).toBe(1);
-    expect(suspenseCache['queryRefs'].size).toBe(1);
+    expect(suspenseCache).toHaveSuspenseCacheEntryUsing(client, query);
 
     jest.advanceTimersByTime(5_000);
 
     expect(client.getObservableQueries().size).toBe(0);
-    expect(suspenseCache['queryRefs'].size).toBe(0);
+    expect(suspenseCache).not.toHaveSuspenseCacheEntryUsing(client, query);
 
     jest.useRealTimers();
 
@@ -957,7 +983,7 @@ describe('useSuspenseQuery', () => {
     jest.advanceTimersByTime(30_000);
 
     expect(client.getObservableQueries().size).toBe(1);
-    expect(suspenseCache['queryRefs'].size).toBe(1);
+    expect(suspenseCache).toHaveSuspenseCacheEntryUsing(client, query);
 
     jest.useRealTimers();
   });

diff --git a/src/testing/matchers/index.d.ts b/src/testing/matchers/index.d.ts
@@ -1,11 +1,27 @@
-import type { DocumentNode } from '../../core';
+import type {
+  ApolloClient,
+  DocumentNode,
+  OperationVariables,
+} from '../../core';
 
 interface ApolloCustomMatchers<R = void> {
   /**
    * Used to determine if two GraphQL query documents are equal to each other by
    * comparing their printed values. The document must be parsed by `gql`.
    */
   toMatchDocument(document: DocumentNode): R;
+
+  /**
+   * Used to determine if the Suspense cache has a cache entry.
+   */
+  toHaveSuspenseCacheEntryUsing(
+    client: ApolloClient<unknown>,
+    query: DocumentNode,
+    options?: {
+      variables?: OperationVariables;
+      queryKey?: string | number | any[];
+    }
+  ): R;
 }
 
 declare global {

diff --git a/src/testing/matchers/index.ts b/src/testing/matchers/index.ts
@@ -1,6 +1,8 @@
 import { expect } from '@jest/globals';
 import { toMatchDocument } from './toMatchDocument';
+import { toHaveSuspenseCacheEntryUsing } from './toHaveSuspenseCacheEntryUsing';
 
 expect.extend({
+  toHaveSuspenseCacheEntryUsing,
   toMatchDocument,
 });
diff --git a/src/testing/matchers/toHaveSuspenseCacheEntryUsing.ts b/src/testing/matchers/toHaveSuspenseCacheEntryUsing.ts
@@ -0,0 +1,39 @@
+import type { MatcherFunction } from 'expect';
+import type { DocumentNode } from 'graphql';
+import type { ApolloClient, OperationVariables } from '../../core';
+import { SuspenseCache } from '../../react';
+import { canonicalStringify } from '../../cache';
+
+export const toHaveSuspenseCacheEntryUsing: MatcherFunction<
+  [
+    client: ApolloClient<unknown>,
+    query: DocumentNode,
+    options: {
+      variables?: OperationVariables;
+      queryKey?: string | number | any[];
+    }
+  ]
+> = function (
+  suspenseCache,
+  client,
+  query,
+  { variables, queryKey = [] } = Object.create(null)
+) {
+  if (!(suspenseCache instanceof SuspenseCache)) {
+    throw new Error('Actual must be an instance of `SuspenseCache`');
+  }
+
+  const cacheKey = (
+    [client, query, canonicalStringify(variables)] as any[]
+  ).concat(queryKey);
+  const queryRef = suspenseCache['queryRefs'].lookupArray(cacheKey)?.current;
+
+  return {
+    pass: !!queryRef,
+    message: () => {
+      return `Expected suspense cache ${
+        queryRef ? 'not ' : ''
+      }to have cache entry using key`;
+    },
+  };
+};