fix cleanup

Author: PatrykWalach

libs/isograph-react-disposable-state/src/useDisposableState.ts

@@ -22,7 +22,6 @@ export function useDisposableState<T = never>(
   const preCommitItem = useCachedResponsivePrecommitValue(
     parentCache,
     (pair) => {
-      itemCleanupPairRef.current?.[1]();
       itemCleanupPairRef.current = pair;
     },
   );
@@ -47,14 +46,23 @@ export function useDisposableState<T = never>(
     [stateFromDisposableStateHook],
   );
 
-  useEffect(function cleanupItemCleanupPairRefIfSetStateNotCalled() {
-    return () => {
-      if (itemCleanupPairRef.current !== null) {
-        itemCleanupPairRef.current[1]();
-        itemCleanupPairRef.current = null;
+  const lastCommittedParentCache = useRef<ParentCache<T> | null>(null);
+
+  useEffect(
+    function cleanupItemCleanupPairRefIfSetStateNotCalled() {
+      if (lastCommittedParentCache.current === parentCache) {
+        return;
       }
-    };
-  }, []);
+      lastCommittedParentCache.current = parentCache;
+      // capture last set pair in a variable
+      const current = itemCleanupPairRef.current;
+      return () => {
+        // current is a stale variable
+        current?.[1]();
+      };
+    },
+    [parentCache],
+  );
 
   // Safety: we can be in one of three states. Pre-commit, in which case
   // preCommitItem is assigned, post-commit but before setState has been

libs/isograph-react-disposable-state/src/useLazyDisposableState.test.tsx

@@ -1,5 +1,5 @@
 import { ItemCleanupPair } from '@isograph/disposable-types';
-import React, { useEffect, useState } from 'react';
+import React, { StrictMode, useEffect, useState } from 'react';
 import { create } from 'react-test-renderer';
 import { describe, expect, test, vi } from 'vitest';
 import { ParentCache } from './ParentCache';
@@ -55,7 +55,12 @@ describe('useLazyDisposableState', async () => {
       return null;
     }
 
-    const root = create(<TestComponent />, { unstable_isConcurrent: true });
+    const root = create(
+      <StrictMode>
+        <TestComponent />
+      </StrictMode>,
+      { unstable_isConcurrent: true },
+    );
     await committed.promise;
     expect(cache1.disposeItem).toHaveBeenCalled();
     expect(cache1.cache.factory).toHaveBeenCalledOnce();

libs/isograph-react-disposable-state/src/useLazyDisposableState.ts

@@ -23,18 +23,24 @@ export function useLazyDisposableState<T>(
   state: T;
 } {
   const itemCleanupPairRef = useRef<ItemCleanupPair<T> | null>(null);
-
   const preCommitItem = useCachedResponsivePrecommitValue(
     parentCache,
     (pair) => {
-      itemCleanupPairRef.current?.[1]();
       itemCleanupPairRef.current = pair;
     },
   );
 
+  const lastCommittedParentCache = useRef<ParentCache<T> | null>(null);
   useEffect(() => {
+    if (lastCommittedParentCache.current === parentCache) {
+      return;
+    }
+    lastCommittedParentCache.current = parentCache;
+    // capture last set pair in a variable
+    const current = itemCleanupPairRef.current;
     return () => {
-      const cleanupFn = itemCleanupPairRef.current?.[1];
+      // current is a stale variable
+      const cleanupFn = current?.[1];
       // TODO confirm useEffect is called in order.
       if (cleanupFn == null) {
         throw new Error(
@@ -43,7 +49,7 @@ export function useLazyDisposableState<T>(
       }
       return cleanupFn();
     };
-  }, []);
+  }, [parentCache]);
 
   const returnedItem = preCommitItem?.state ?? itemCleanupPairRef.current?.[0];
 

libs/isograph-react-disposable-state/src/useUpdatableDisposableState.test.tsx

@@ -1,5 +1,5 @@
 import { describe, test, vi, expect } from 'vitest';
-import React from 'react';
+import React, { StrictMode } from 'react';
 import { create } from 'react-test-renderer';
 import {
   useUpdatableDisposableState,
@@ -45,7 +45,8 @@ function promiseAndResolver() {
 // not is a bit worrisome.
 async function awaitableCreate(Component, isConcurrent) {
   const element = create(
-    Component,
+    <StrictMode>{Component}</StrictMode>,
+
     isConcurrent ? { unstable_isConcurrent: true } : undefined,
   );
   await shortPromise();