Suspense cache: use Trie directly

.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.

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>;
   }
 }

src/react/hooks/__tests__/useSuspenseQuery.test.tsx

@@ -47,6 +47,51 @@ import { ApolloProvider } from '../../context';
 import { SuspenseCache } from '../../cache';
 import { SuspenseQueryHookFetchPolicy } from '../../../react';
 import { useSuspenseQuery } from '../useSuspenseQuery';
+import { canonicalStringify } from '../../../cache';
+
+expect.extend({
+  toHaveSuspenseCacheEntryUsing(
+    suspenseCache: SuspenseCache,
+    client: ApolloClient<unknown>,
+    query: DocumentNode,
+    {
+      variables,
+      queryKey = [],
+    }: {
+      variables?: OperationVariables;
+      queryKey?: string | number | any[];
+    } = Object.create(null)
+  ) {
+    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`;
+      },
+    };
+  },
+});
+
+declare global {
+  namespace jest {
+    interface Matchers<R = void> {
+      toHaveSuspenseCacheEntryUsing(
+        client: ApolloClient<unknown>,
+        query: DocumentNode,
+        options?: {
+          variables?: OperationVariables;
+          queryKey?: string | number | any[];
+        }
+      ): R;
+    }
+  }
+}
 
 type RenderSuspenseHookOptions<Props, TSerializedCache = {}> = Omit<
   RenderHookOptions<Props>,
@@ -378,8 +423,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 +726,7 @@ describe('useSuspenseQuery', () => {
     );
 
     expect(client.getObservableQueries().size).toBe(1);
-    expect(suspenseCache['queryRefs'].size).toBe(1);
+    expect(suspenseCache).toHaveSuspenseCacheEntryUsing(client, query);
 
     unmount();
 
@@ -687,7 +735,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 +764,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 +778,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 +832,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 +851,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 +905,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 +966,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 +1028,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();
   });