Retain queries if data in store

Author: PatrykWalach

libs/isograph-react/src/core/IsographEnvironment.ts

@@ -10,7 +10,7 @@ import {
   type StableIdForFragmentReference,
   type UnknownTReadFromStore,
 } from './FragmentReference';
-import { RetainedQuery } from './garbageCollection';
+import { RetainedQuery, type ReadyRetainedQuery } from './garbageCollection';
 import { LogFunction, WrappedLogFunction } from './logging';
 import { PromiseWrapper, wrapPromise } from './PromiseWrapper';
 import { WithEncounteredRecords } from './read';
@@ -71,7 +71,7 @@ export type IsographEnvironment = {
     PromiseWrapper<IsographEntrypoint<any, any, any>>
   >;
   readonly retainedQueries: Set<RetainedQuery>;
-  readonly gcBuffer: Array<RetainedQuery>;
+  readonly gcBuffer: Array<ReadyRetainedQuery>;
   readonly gcBufferSize: number;
   readonly loggers: Set<WrappedLogFunction>;
 };

libs/isograph-react/src/core/garbageCollection.ts

@@ -1,5 +1,5 @@
 import { getParentRecordKey } from './cache';
-import { NormalizationAstNodes } from './entrypoint';
+import { NormalizationAstNodes, type NormalizationAst } from './entrypoint';
 import { Variables } from './FragmentReference';
 import {
   assertLink,
@@ -12,20 +12,40 @@ import {
 } from './IsographEnvironment';
 
 export type RetainedQuery = {
-  readonly normalizationAst: NormalizationAstNodes;
+  normalizationAst:
+    | {
+        kind: 'Loading';
+      }
+    | {
+        kind: 'Ready';
+        value: NormalizationAst;
+      };
   readonly variables: {};
   readonly root: StoreLink;
 };
 
+export interface ReadyRetainedQuery extends RetainedQuery {
+  readonly normalizationAst: {
+    kind: 'Ready';
+    value: NormalizationAst;
+  };
+}
+
+function isRetainedQueryReady(
+  query: RetainedQuery,
+): query is ReadyRetainedQuery {
+  return query.normalizationAst.kind === 'Ready';
+}
+
 export type DidUnretainSomeQuery = boolean;
 export function unretainQuery(
   environment: IsographEnvironment,
   retainedQuery: RetainedQuery,
 ): DidUnretainSomeQuery {
   environment.retainedQueries.delete(retainedQuery);
-  environment.gcBuffer.push(retainedQuery);
-
-  if (environment.gcBuffer.length > environment.gcBufferSize) {
+  if (isRetainedQueryReady(retainedQuery)) {
+    environment.gcBuffer.push(retainedQuery);
+  } else if (environment.gcBuffer.length > environment.gcBufferSize) {
     environment.gcBuffer.shift();
     return true;
   }
@@ -47,7 +67,12 @@ export function garbageCollectEnvironment(environment: IsographEnvironment) {
   const retainedIds: RetainedIds = {};
 
   for (const query of environment.retainedQueries) {
-    recordReachableIds(environment.store, query, retainedIds);
+    if (isRetainedQueryReady(query)) {
+      recordReachableIds(environment.store, query, retainedIds);
+    } else {
+      // if we have any queries with loading normalizationAst, we can't garbage collect
+      return;
+    }
   }
   for (const query of environment.gcBuffer) {
     recordReachableIds(environment.store, query, retainedIds);
@@ -82,7 +107,7 @@ interface RetainedIds {
 
 function recordReachableIds(
   store: IsographStore,
-  retainedQuery: RetainedQuery,
+  retainedQuery: ReadyRetainedQuery,
   mutableRetainedIds: RetainedIds,
 ) {
   const record =
@@ -98,7 +123,7 @@ function recordReachableIds(
       store,
       record,
       mutableRetainedIds,
-      retainedQuery.normalizationAst,
+      retainedQuery.normalizationAst.value.selections,
       retainedQuery.variables,
     );
   }

libs/isograph-react/src/core/makeNetworkRequest.ts

@@ -60,7 +60,11 @@ export function maybeMakeNetworkRequest<
       );
     }
     case 'No': {
-      return [wrapResolvedValue(undefined), () => {}];
+      return loadNormalizationAstAndRetainQuery(
+        environment,
+        artifact,
+        variables,
+      );
     }
     case 'IfNecessary': {
       if (
@@ -82,7 +86,11 @@ export function maybeMakeNetworkRequest<
       );
 
       if (result.kind === 'EnoughData') {
-        return [wrapResolvedValue(undefined), () => {}];
+        return loadNormalizationAstAndRetainQuery(
+          environment,
+          artifact,
+          variables,
+        );
       } else {
         return makeNetworkRequest(
           environment,
@@ -109,6 +117,90 @@ function loadNormalizationAst(
   }
 }
 
+type NormalizationAstRequestStatus =
+  | {
+      readonly kind: 'Disposed';
+    }
+  | {
+      readonly kind: 'Undisposed';
+      readonly retainedQuery: RetainedQuery;
+    };
+
+function loadNormalizationAstAndRetainQuery<
+  TReadFromStore extends UnknownTReadFromStore,
+  TClientFieldValue,
+  TArtifact extends
+    | RefetchQueryNormalizationArtifact
+    | IsographEntrypoint<TReadFromStore, TClientFieldValue, TNormalizationAst>,
+  TNormalizationAst extends NormalizationAst | NormalizationAstLoader,
+>(
+  environment: IsographEnvironment,
+  artifact: TArtifact,
+  variables: ExtractParameters<TReadFromStore>,
+): ItemCleanupPair<PromiseWrapper<void, AnyError>> {
+  const root = { __link: ROOT_ID, __typename: artifact.concreteType };
+
+  const retainedQuery: RetainedQuery = {
+    normalizationAst: {
+      kind: 'Loading',
+    },
+    variables,
+    root,
+  };
+  let status: NormalizationAstRequestStatus = {
+    kind: 'Undisposed',
+    retainedQuery,
+  };
+  retainQuery(environment, retainedQuery);
+
+  switch (artifact.networkRequestInfo.normalizationAst.kind) {
+    case 'NormalizationAst': {
+      retainedQuery.normalizationAst = {
+        kind: 'Ready',
+        value: artifact.networkRequestInfo.normalizationAst,
+      };
+      break;
+    }
+    case 'NormalizationAstLoader': {
+      artifact.networkRequestInfo.normalizationAst
+        .loader()
+        .then((normalizationAst) => {
+          retainedQuery.normalizationAst = {
+            kind: 'Ready',
+            value: normalizationAst,
+          };
+        })
+        .catch(() => {
+          const didUnretainSomeQuery = unretainQuery(
+            environment,
+            retainedQuery,
+          );
+          if (didUnretainSomeQuery) {
+            garbageCollectEnvironment(environment);
+          }
+        });
+    }
+  }
+
+  return [
+    wrapResolvedValue(undefined),
+    () => {
+      if (status.kind === 'Undisposed') {
+        const didUnretainSomeQuery = unretainQuery(
+          environment,
+          status.retainedQuery,
+        );
+        if (didUnretainSomeQuery) {
+          garbageCollectEnvironment(environment);
+        }
+      }
+      status = {
+        kind: 'Disposed',
+      };
+    },
+  ];
+}
+
 export function makeNetworkRequest<
   TReadFromStore extends UnknownTReadFromStore,
   TClientFieldValue,
@@ -173,8 +265,11 @@ export function makeNetworkRequest<
           variables,
           root,
         );
-        const retainedQuery = {
-          normalizationAst: normalizationAst.selections,
+        const retainedQuery: RetainedQuery = {
+          normalizationAst: {
+            kind: 'Ready',
+            value: normalizationAst,
+          },
           variables,
           root,
         };

libs/isograph-react/src/tests/garbageCollection.test.ts

@@ -2,6 +2,7 @@ import { describe, expect, test } from 'vitest';
 import {
   garbageCollectEnvironment,
   retainQuery,
+  type RetainedQuery,
 } from '../core/garbageCollection';
 import {
   createIsographEnvironment,
@@ -55,9 +56,11 @@ export const meNameField = iso(`
 `)(() => {});
 
 const meNameEntrypoint = iso(`entrypoint Query.meName`);
-const meNameRetainedQuery = {
-  normalizationAst:
-    meNameEntrypoint.networkRequestInfo.normalizationAst.selections,
+const meNameRetainedQuery: RetainedQuery = {
+  normalizationAst: {
+    kind: 'Ready',
+    value: meNameEntrypoint.networkRequestInfo.normalizationAst,
+  },
   variables: {},
   root: { __link: ROOT_ID, __typename: 'Query' },
 };

libs/isograph-react/src/tests/meNameSuccessor.ts

@@ -1,3 +1,4 @@
+import type { RetainedQuery } from '../core/garbageCollection';
 import { ROOT_ID } from '../core/IsographEnvironment';
 import { iso } from './__isograph/iso';
 
@@ -14,9 +15,11 @@ export const meNameField = iso(`
   }
 `)(() => {});
 const meNameSuccessorEntrypoint = iso(`entrypoint Query.meNameSuccessor`);
-export const meNameSuccessorRetainedQuery = {
-  normalizationAst:
-    meNameSuccessorEntrypoint.networkRequestInfo.normalizationAst.selections,
+export const meNameSuccessorRetainedQuery: RetainedQuery = {
+  normalizationAst: {
+    kind: 'Ready',
+    value: meNameSuccessorEntrypoint.networkRequestInfo.normalizationAst,
+  },
   variables: {},
   root: {
     __link: ROOT_ID,