fix: do not clear dialog state immediately on useDialog unmount (#2584)

### 🎯 Goal

When the `useDialog` hook's component unmounts, it immediately clears
dialog state in an effect cleanup. However, in some situations an effect
cleanup can run even if the hook's component is still mounted and
effect's dependencies didn't change - e.g., when `<StrictMode />` is
enabled.

So it's safer to keep dialog state for a short time after cleanup runs.

### 🛠 Implementation details

Instead of immediately removing dialog state, it's marked to be removed
after a short timeout. Referencing the dialog again quick cancels state
removal.

Fixes #2583.

Author: myandrienko

examples/vite/src/main.tsx

@@ -1,5 +1,10 @@
+import { StrictMode } from 'react';
 import ReactDOM from 'react-dom/client';
 import App from './App.tsx';
 import './index.scss';
 
-ReactDOM.createRoot(document.getElementById('root')!).render(<App />);
+ReactDOM.createRoot(document.getElementById('root')!).render(
+  <StrictMode>
+    <App />
+  </StrictMode>,
+);

examples/vite/yarn.lock

@@ -257,7 +257,7 @@
     "react-dom": "^18.1.0",
     "react-test-renderer": "^18.1.0",
     "semantic-release": "^19.0.5",
-    "stream-chat": "^8.46.1",
+    "stream-chat": "^8.47.1",
     "ts-jest": "^29.1.4",
     "typescript": "^5.4.5"
   },

package.json

@@ -11,6 +11,7 @@ export type Dialog = {
   id: DialogId;
   isOpen: boolean | undefined;
   open: (zIndex?: number) => void;
+  removalTimeout: NodeJS.Timeout | undefined;
   remove: () => void;
   toggle: (closeAll?: boolean) => void;
 };
@@ -64,6 +65,7 @@ export class DialogManager {
         open: () => {
           this.open({ id });
         },
+        removalTimeout: undefined,
         remove: () => {
           this.remove(id);
         },
@@ -76,6 +78,23 @@ export class DialogManager {
         ...{ dialogsById: { ...current.dialogsById, [id]: dialog } },
       }));
     }
+
+    if (dialog.removalTimeout) {
+      clearTimeout(dialog.removalTimeout);
+      this.state.next((current) => ({
+        ...current,
+        ...{
+          dialogsById: {
+            ...current.dialogsById,
+            [id]: {
+              ...dialog,
+              removalTimeout: undefined,
+            },
+          },
+        },
+      }));
+    }
+
     return dialog;
   }
 
@@ -117,6 +136,10 @@ export class DialogManager {
     const dialog = state.dialogsById[id];
     if (!dialog) return;
 
+    if (dialog.removalTimeout) {
+      clearTimeout(dialog.removalTimeout);
+    }
+
     this.state.next((current) => {
       const newDialogs = { ...current.dialogsById };
       delete newDialogs[id];
@@ -126,4 +149,30 @@ export class DialogManager {
       };
     });
   }
+
+  /**
+   * Marks the dialog state as unused. If the dialog id is referenced again quickly,
+   * the state will not be removed. Otherwise, the state will be removed after
+   * a short timeout.
+   */
+  markForRemoval(id: DialogId) {
+    const dialog = this.state.getLatestValue().dialogsById[id];
+
+    if (!dialog) {
+      return;
+    }
+
+    this.state.next((current) => ({
+      ...current,
+      dialogsById: {
+        ...current.dialogsById,
+        [id]: {
+          ...dialog,
+          removalTimeout: setTimeout(() => {
+            this.remove(id);
+          }, 16),
+        },
+      },
+    }));
+  }
 }

src/components/Dialog/DialogManager.ts

@@ -8,13 +8,15 @@ describe('DialogManager', () => {
     const dialogManager = new DialogManager({ id });
     expect(dialogManager.id).toBe(id);
   });
+
   it('initiates with default options', () => {
     const mockedId = '12345';
     const spy = jest.spyOn(Date.prototype, 'getTime').mockReturnValueOnce(mockedId);
     const dialogManager = new DialogManager();
     expect(dialogManager.id).toBe(mockedId);
     spy.mockRestore();
   });
+
   it('creates a new closed dialog', () => {
     const dialogManager = new DialogManager();
     expect(Object.keys(dialogManager.state.getLatestValue().dialogsById)).toHaveLength(0);
@@ -142,4 +144,33 @@ describe('DialogManager', () => {
     expect(dialogManager.openDialogCount).toBe(1);
     expect(Object.keys(dialogManager.state.getLatestValue().dialogsById)).toHaveLength(1);
   });
+
+  it('marks dialog for removal', () => {
+    jest.useFakeTimers();
+
+    const dialogManager = new DialogManager();
+    dialogManager.getOrCreate({ id: dialogId });
+    dialogManager.open({ id: dialogId });
+    dialogManager.markForRemoval(dialogId);
+
+    jest.runAllTimers();
+
+    expect(dialogManager.openDialogCount).toBe(0);
+    expect(Object.keys(dialogManager.state.getLatestValue().dialogsById)).toHaveLength(0);
+  });
+
+  it('cancels dialog removal if it is referenced again quickly', () => {
+    jest.useFakeTimers();
+
+    const dialogManager = new DialogManager();
+    dialogManager.getOrCreate({ id: dialogId });
+    dialogManager.open({ id: dialogId });
+    dialogManager.markForRemoval(dialogId);
+    dialogManager.getOrCreate({ id: dialogId });
+
+    jest.runAllTimers();
+
+    expect(dialogManager.openDialogCount).toBe(1);
+    expect(Object.keys(dialogManager.state.getLatestValue().dialogsById)).toHaveLength(1);
+  });
 });

src/components/Dialog/__tests__/DialogsManager.test.js

@@ -9,7 +9,11 @@ export const useDialog = ({ id }: GetOrCreateDialogParams) => {
 
   useEffect(
     () => () => {
-      dialogManager.remove(id);
+      // Since this cleanup can run even if the component is still mounted
+      // and dialog id is unchanged (e.g. in <StrictMode />), it's safer to
+      // mark state as unused and only remove it after a timeout, rather than
+      // to remove it immediately.
+      dialogManager.markForRemoval(id);
     },
     [dialogManager, id],
   );

src/components/Dialog/hooks/useDialog.ts

@@ -12231,10 +12231,10 @@ [email protected]:
   resolved "https://registry.yarnpkg.com/statuses/-/statuses-2.0.1.tgz#55cb000ccf1d48728bd23c685a063998cf1a1b63"
   integrity sha512-RwNA9Z/7PrK06rYLIzFMlaF+l73iwpzsqRIFgbMLbTcLD6cOao82TaWefPXQvB2fOC4AjuYSEndS7N/mTCbkdQ==
 
-stream-chat@^8.46.1:
-  version "8.46.1"
-  resolved "https://registry.yarnpkg.com/stream-chat/-/stream-chat-8.46.1.tgz#9624bbd4e8e357414389e4b0dcb5f0134487f679"
-  integrity sha512-jVg148tZDCAmZa6b31cnamiUvt58IQB4YTRIGFt3+4zp6jI2tKeHvEj24uEcVF2d7fhT1IXSuGhNfcrStqSZHQ==
+stream-chat@^8.47.1:
+  version "8.47.1"
+  resolved "https://registry.yarnpkg.com/stream-chat/-/stream-chat-8.47.1.tgz#5390c87cbb1929e7ca183aa1204dae3ab38469a2"
+  integrity sha512-raMAGYLT4UCVluMF0TMfdPKH9OUhDjH6e1HQdJIlllAFLaA8oxtG+e/7jyuPmVodLPzYCPqOt2eBH7soAkhV/A==
   dependencies:
     "@babel/runtime" "^7.16.3"
     "@types/jsonwebtoken" "~9.0.0"