fix: memory leak during remount or user change (#2639)

* fix: memory leak during remount or user change

* fix: add additional unsubscribe on unmount

* chore: remove console.log

* chore: fix mounting hook and add tests

* fix: improve tests and fix bug with remounting

* chore: remove dead variable

* fix: linter errors

Author: isekovanic

package/src/components/Chat/Chat.tsx

@@ -204,14 +204,21 @@ const ChatWithContext = <
 
   useEffect(() => {
     if (userID && enableOfflineSupport) {
+      // This acts as a lock for some very rare occurrences of concurrency
+      // issues we've encountered before with the QuickSqliteClient being
+      // uninitialized before it's being invoked.
       setInitialisedDatabaseConfig({ initialised: false, userID });
       QuickSqliteClient.initializeDatabase();
-      DBSyncManager.init(client as unknown as StreamChat);
       setInitialisedDatabaseConfig({ initialised: true, userID });
+      DBSyncManager.init(client as unknown as StreamChat);
     }
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [userID, enableOfflineSupport]);
 
+  // In case something went wrong, make sure to also unsubscribe the listener
+  // on unmount if it exists to prevent a memory leak.
+  useEffect(() => () => DBSyncManager.connectionChangedListener?.unsubscribe(), []);
+
   const initialisedDatabase =
     initialisedDatabaseConfig.initialised && userID === initialisedDatabaseConfig.userID;
 

package/src/components/Chat/__tests__/Chat.test.js

@@ -9,7 +9,8 @@ import { useChatContext } from '../../../contexts/chatContext/ChatContext';
 import { useTranslationContext } from '../../../contexts/translationContext/TranslationContext';
 import dispatchConnectionChangedEvent from '../../../mock-builders/event/connectionChanged';
 import dispatchConnectionRecoveredEvent from '../../../mock-builders/event/connectionRecovered';
-import { getTestClient } from '../../../mock-builders/mock';
+import { getTestClient, getTestClientWithUser, setUser } from '../../../mock-builders/mock';
+import { DBSyncManager } from '../../../utils/DBSyncManager';
 import { Streami18n } from '../../../utils/i18n/Streami18n';
 import { Chat } from '../Chat';
 
@@ -24,7 +25,10 @@ const TranslationContextConsumer = ({ fn }) => {
 };
 
 describe('Chat', () => {
-  afterEach(cleanup);
+  afterEach(() => {
+    cleanup();
+    jest.clearAllMocks();
+  });
   const chatClient = getTestClient();
 
   it('renders children without crashing', async () => {
@@ -215,5 +219,79 @@ describe('Chat', () => {
         expect(context.tDateTimeParser).toBe(newI18nInstance.tDateTimeParser);
       });
     });
+
+    it('makes sure DBSyncManager listeners are cleaned up after Chat remount', async () => {
+      const chatClientWithUser = await getTestClientWithUser({ id: 'testID' });
+      jest.spyOn(DBSyncManager, 'init');
+
+      // initial mount and render
+      const { rerender } = render(
+        <Chat client={chatClientWithUser} enableOfflineSupport key={1} />,
+      );
+
+      // the unsubscribe fn changes during init(), so we keep a reference to the spy
+      const unsubscribeSpy = jest.spyOn(DBSyncManager.connectionChangedListener, 'unsubscribe');
+      const listenersAfterInitialMount = chatClientWithUser.listeners['connection.changed'];
+
+      // remount
+      rerender(<Chat client={chatClientWithUser} enableOfflineSupport key={2} />);
+
+      await waitFor(() => {
+        expect(DBSyncManager.init).toHaveBeenCalledTimes(2);
+        expect(unsubscribeSpy).toHaveBeenCalledTimes(2);
+        expect(chatClientWithUser.listeners['connection.changed'].length).toBe(
+          listenersAfterInitialMount.length,
+        );
+      });
+    });
+
+    it('makes sure DBSyncManager listeners are cleaned up if the user changes', async () => {
+      const chatClientWithUser = await getTestClientWithUser({ id: 'testID1' });
+      jest.spyOn(DBSyncManager, 'init');
+
+      // initial render
+      const { rerender } = render(<Chat client={chatClientWithUser} enableOfflineSupport />);
+
+      // the unsubscribe fn changes during init(), so we keep a reference to the spy
+      const unsubscribeSpy = jest.spyOn(DBSyncManager.connectionChangedListener, 'unsubscribe');
+      await act(async () => {
+        await setUser(chatClientWithUser, { id: 'testID2' });
+      });
+      const listenersAfterInitialMount = chatClientWithUser.listeners['connection.changed'];
+
+      // rerender with different user ID
+      rerender(<Chat client={chatClientWithUser} enableOfflineSupport />);
+
+      await waitFor(() => {
+        expect(DBSyncManager.init).toHaveBeenCalledTimes(2);
+        expect(unsubscribeSpy).toHaveBeenCalledTimes(1);
+        expect(chatClientWithUser.listeners['connection.changed'].length).toBe(
+          listenersAfterInitialMount.length,
+        );
+      });
+    });
+
+    it('makes sure DBSyncManager state stays intact during normal rerenders', async () => {
+      const chatClientWithUser = await getTestClientWithUser({ id: 'testID' });
+      jest.spyOn(DBSyncManager, 'init');
+
+      // initial render
+      const { rerender } = render(<Chat client={chatClientWithUser} enableOfflineSupport />);
+
+      // the unsubscribe fn changes during init(), so we keep a reference to the spy
+      const unsubscribeSpy = jest.spyOn(DBSyncManager.connectionChangedListener, 'unsubscribe');
+      const listenersAfterInitialMount = chatClientWithUser.listeners['connection.changed'];
+
+      // rerender
+      rerender(<Chat client={chatClientWithUser} enableOfflineSupport />);
+
+      await waitFor(() => {
+        expect(DBSyncManager.init).toHaveBeenCalledTimes(1);
+        expect(unsubscribeSpy).toHaveBeenCalledTimes(0);
+        expect(chatClientWithUser.listeners['connection.changed'].length).toBe(
+          listenersAfterInitialMount.length,
+        );
+      });
+    });
   });
 });

package/src/mock-builders/mock.js

@@ -6,7 +6,7 @@ import { StreamChat } from 'stream-chat';
 const apiKey = 'API_KEY';
 const token = 'dummy_token';
 
-const setUser = (client, user) =>
+export const setUser = (client, user) =>
   new Promise((resolve) => {
     client.connectionId = 'dumm_connection_id';
     client.user = user;

package/src/utils/DBSyncManager.ts

@@ -38,6 +38,7 @@ export class DBSyncManager {
   static syncStatus = false;
   static listeners: Array<(status: boolean) => void> = [];
   static client: StreamChat | null = null;
+  static connectionChangedListener: { unsubscribe: () => void } | null = null;
 
   /**
    * Returns weather channel states in local DB are synced with backend or not.
@@ -62,7 +63,16 @@ export class DBSyncManager {
       this.listeners.forEach((l) => l(true));
     }
 
-    this.client.on('connection.changed', async (event) => {
+    // If a listener has already been registered, unsubscribe from it so
+    // that it can be reinstated. This can happen if we reconnect with a
+    // different user or the component invoking the init() function gets
+    // unmounted and then remounted again. This part of the code makes
+    // sure the stale listener doesn't produce a memory leak.
+    if (this.connectionChangedListener) {
+      this.connectionChangedListener.unsubscribe();
+    }
+
+    this.connectionChangedListener = this.client.on('connection.changed', async (event) => {
       if (event.online) {
         await this.syncAndExecutePendingTasks();
         this.syncStatus = true;