fix: server crash and memory leak after switching between databases

client/src/components/layout/GlobalEffect.tsx

@@ -1,54 +1,78 @@
-import React, { useContext } from 'react';
+import React, { useContext, useEffect } from 'react';
 import axiosInstance from '@/http/Axios';
-import { rootContext, authContext } from '@/context';
+import { rootContext, authContext, dataContext } from '@/context';
 import { HTTP_STATUS_CODE } from '@server/utils/Const';
 
 let axiosResInterceptor: number | null = null;
-// let timer: Record<string, ReturnType<typeof setTimeout> | number>[] = [];
-// we only take side effect here, nothing else
-const GlobalEffect = (props: { children: React.ReactNode }) => {
+
+const GlobalEffect = ({ children }: { children: React.ReactNode }) => {
   const { openSnackBar } = useContext(rootContext);
   const { logout } = useContext(authContext);
+  const { database } = useContext(dataContext);
 
-  // catch axios error here
-  if (axiosResInterceptor === null) {
-    axiosResInterceptor = axiosInstance.interceptors.response.use(
-      function (res: any) {
-        if (res.statusCode && res.statusCode !== HTTP_STATUS_CODE.OK) {
-          openSnackBar(res.data.message, 'warning');
-          return Promise.reject(res.data);
-        }
-
-        return res;
+  useEffect(() => {
+    // Add database header to all axios requests
+    const requestInterceptor = axiosInstance.interceptors.request.use(
+      config => {
+        config.headers['x-attu-database'] = database;
+        return config;
       },
-      function (error: any) {
-        const { response = {} } = error;
-
-        switch (response.status) {
-          case HTTP_STATUS_CODE.UNAUTHORIZED:
-          case HTTP_STATUS_CODE.FORBIDDEN:
-            setTimeout(logout, 1000);
-            break;
-          default:
-            break;
-        }
-        if (response.data) {
-          const { message: errMsg } = response.data;
-          // We need check status 401 in login page
-          // So server will return 500 when change the user password.
-          errMsg && openSnackBar(errMsg, 'error');
-          return Promise.reject(error);
-        }
-        if (error.message) {
+      error => Promise.reject(error)
+    );
+
+    // Clean up interceptor on unmount
+    return () => {
+      axiosInstance.interceptors.request.eject(requestInterceptor);
+    };
+  }, [database]);
+
+  useEffect(() => {
+    if (axiosResInterceptor === null) {
+      axiosResInterceptor = axiosInstance.interceptors.response.use(
+        (response: any) => {
+          if (
+            response.statusCode &&
+            response.statusCode !== HTTP_STATUS_CODE.OK
+          ) {
+            openSnackBar(response.data.message, 'warning');
+            return Promise.reject(response.data);
+          }
+          return response;
+        },
+        error => {
+          const { response } = error;
+          if (response) {
+            switch (response.status) {
+              case HTTP_STATUS_CODE.UNAUTHORIZED:
+              case HTTP_STATUS_CODE.FORBIDDEN:
+                setTimeout(() => logout(true), 1000);
+                break;
+              default:
+                break;
+            }
+            const errorMessage = response.data?.message;
+            if (errorMessage) {
+              openSnackBar(errorMessage, 'error');
+              return Promise.reject(error);
+            }
+          }
+          // Handle other error cases
           openSnackBar(error.message, 'error');
+          return Promise.reject(error);
         }
-        return Promise.reject(error);
+      );
+    }
+
+    // Clean up response interceptor on unmount
+    return () => {
+      if (axiosResInterceptor !== null) {
+        axiosInstance.interceptors.response.eject(axiosResInterceptor);
+        axiosResInterceptor = null;
       }
-    );
-  }
-  // get global data
+    };
+  }, [logout, openSnackBar]);
 
-  return <>{props.children}</>;
+  return <>{children}</>;
 };
 
 export default GlobalEffect;

client/src/components/layout/Header.tsx

@@ -106,7 +106,7 @@ const Header: FC = () => {
   };
 
   const handleLogout = async () => {
-    logout();
+    logout(false);
   };
 
   const useDatabase = async (database: string) => {

client/src/context/Auth.tsx

@@ -69,7 +69,11 @@ export const AuthProvider = (props: { children: React.ReactNode }) => {
     return res;
   };
   // logout API
-  const logout = () => {
+  const logout = async (pass?: boolean) => {
+    if (!pass) {
+      // close connetion
+      await MilvusService.closeConnection();
+    }
     // clear client id
     setClientId('');
     // remove client id from local storage

client/src/context/Data.tsx

@@ -136,14 +136,23 @@ export const DataProvider = (props: { children: React.ReactNode }) => {
 
   // Websocket Callback: update single collection
   const updateCollections = useCallback(
-    (updateCollections: CollectionFullObject[]) => {
+    (props: { collections: CollectionFullObject[]; database?: string }) => {
+      const { collections, database: remote } = props;
+      if (
+        remote !== database &&
+        database !== undefined &&
+        remote !== undefined
+      ) {
+        // console.log('database not matched', remote, database);
+        return;
+      }
       // check state to see if it is loading or building index, if so, start server cron job
-      detectLoadingIndexing(updateCollections);
+      detectLoadingIndexing(collections);
       // update single collection
       setCollections(prev => {
         // update exist collection
         const newCollections = prev.map(v => {
-          const collectionToUpdate = updateCollections.find(c => c.id === v.id);
+          const collectionToUpdate = collections.find(c => c.id === v.id);
 
           if (collectionToUpdate) {
             return collectionToUpdate;
@@ -215,7 +224,7 @@ export const DataProvider = (props: { children: React.ReactNode }) => {
     const res = await CollectionService.getCollection(name);
 
     // update collection
-    updateCollections([res]);
+    updateCollections({ collections: [res] });
 
     return res;
   };
@@ -260,7 +269,7 @@ export const DataProvider = (props: { children: React.ReactNode }) => {
     }
 
     // update collection, and trigger cron job
-    updateCollections([collection]);
+    updateCollections({ collections: [collection] });
   };
 
   // API: release collection
@@ -275,7 +284,7 @@ export const DataProvider = (props: { children: React.ReactNode }) => {
     const newCollection = await CollectionService.renameCollection(name, {
       new_collection_name: newName,
     });
-    updateCollections([newCollection]);
+    updateCollections({ collections: [newCollection] });
 
     return newCollection;
   };
@@ -307,7 +316,7 @@ export const DataProvider = (props: { children: React.ReactNode }) => {
     // create index
     const newCollection = await CollectionService.createIndex(param);
     // update collection
-    updateCollections([newCollection]);
+    updateCollections({ collections: [newCollection] });
 
     return newCollection;
   };
@@ -317,7 +326,7 @@ export const DataProvider = (props: { children: React.ReactNode }) => {
     // drop index
     const { data } = await CollectionService.dropIndex(params);
     // update collection
-    updateCollections([data]);
+    updateCollections({ collections: [data] });
 
     return data;
   };
@@ -329,7 +338,7 @@ export const DataProvider = (props: { children: React.ReactNode }) => {
       alias,
     });
     // update collection
-    updateCollections([newCollection]);
+    updateCollections({ collections: [newCollection] });
 
     return newCollection;
   };
@@ -342,7 +351,7 @@ export const DataProvider = (props: { children: React.ReactNode }) => {
     });
 
     // update collection
-    updateCollections([data]);
+    updateCollections({ collections: [data] });
 
     return data;
   };
@@ -359,7 +368,7 @@ export const DataProvider = (props: { children: React.ReactNode }) => {
     });
 
     // update existing collection
-    updateCollections([newCollection]);
+    updateCollections({ collections: [newCollection] });
 
     return newCollection;
   };

client/src/context/Types.ts

@@ -70,7 +70,7 @@ export type AuthContextType = {
   clientId: string;
   isManaged: boolean;
   isAuth: boolean;
-  logout: () => void;
+  logout: (pass?: boolean) => void;
   login: (params: AuthReq) => Promise<AuthObject>;
 };
 

server/package.json

@@ -13,7 +13,7 @@
   },
   "dependencies": {
     "@json2csv/plainjs": "^7.0.3",
-    "@zilliz/milvus2-sdk-node": "^2.4.8",
+    "@zilliz/milvus2-sdk-node": "^2.4.9",
     "axios": "^1.7.4",
     "chalk": "4.1.2",
     "class-sanitizer": "^1.0.1",