Ensure Proper Cleanup of Yorkie Initialization Process (#210)

* Clean up yorkie client if component unmounts before initialization is done

* Fix typo

frontend/src/hooks/useYorkieDocument.ts

@@ -10,7 +10,7 @@ import { selectAuth } from "../store/authSlice";
 yorkie.setLogLevel(4);
 
 export const useYorkieDocument = (
-	yorkieDocuentId?: string | null,
+	yorkieDocumentId?: string | null,
 	presenceName?: string | null
 ) => {
 	const [searchParams] = useSearchParams();
@@ -28,7 +28,8 @@ export const useYorkieDocument = (
 	}, [client, doc]);
 
 	useEffect(() => {
-		if (!yorkieDocuentId || !presenceName || doc || client) return;
+		let mounted = true;
+		if (!yorkieDocumentId || !presenceName || doc || client) return;
 
 		let yorkieToken = `default:${authStore.accessToken}`;
 
@@ -43,7 +44,10 @@ export const useYorkieDocument = (
 			});
 			await newClient.activate();
 
-			const newDoc = new yorkie.Document(yorkieDocuentId as string);
+			const newDoc = new yorkie.Document<
+				YorkieCodeMirrorDocType,
+				YorkieCodeMirrorPresenceType
+			>(yorkieDocumentId);
 
 			await newClient.attach(newDoc, {
 				initialPresence: {
@@ -53,13 +57,28 @@ export const useYorkieDocument = (
 				},
 			});
 
+			// Clean up if the component is unmounted before the initialization is done
+			if (!mounted) {
+				await newClient.detach(newDoc);
+				await newClient.deactivate();
+				return;
+			}
+
 			setClient(newClient);
-			setDoc(
-				newDoc as yorkie.Document<YorkieCodeMirrorDocType, YorkieCodeMirrorPresenceType>
-			);
+			setDoc(newDoc);
 		};
 		initializeYorkie();
-	}, [presenceName, yorkieDocuentId, doc, client, authStore.accessToken, searchParams]);
 
-	return { client, doc, cleanUpYorkieDocument };
+		return () => {
+			mounted = false;
+		};
+	}, [presenceName, yorkieDocumentId, doc, client, authStore.accessToken, searchParams]);
+
+	useEffect(() => {
+		return () => {
+			cleanUpYorkieDocument();
+		};
+	}, [cleanUpYorkieDocument]);
+
+	return { client, doc };
 };

frontend/src/pages/workspace/document/Index.tsx

@@ -16,10 +16,7 @@ function DocumentIndex() {
 	const userStore = useSelector(selectUser);
 	const { data: workspace } = useGetWorkspaceQuery(params.workspaceSlug);
 	const { data: document } = useGetDocumentQuery(workspace?.id, params.documentId);
-	const { doc, client, cleanUpYorkieDocument } = useYorkieDocument(
-		document?.yorkieDocumentId,
-		userStore.data?.nickname
-	);
+	const { doc, client } = useYorkieDocument(document?.yorkieDocumentId, userStore.data?.nickname);
 
 	useEffect(() => {
 		if (!doc || !client) return;
@@ -28,11 +25,10 @@ function DocumentIndex() {
 		dispatch(setClient(client));
 
 		return () => {
-			cleanUpYorkieDocument();
 			dispatch(setDoc(null));
 			dispatch(setClient(null));
 		};
-	}, [cleanUpYorkieDocument, dispatch, client, doc]);
+	}, [dispatch, client, doc]);
 
 	return (
 		<Box height="calc(100% - 64px)">

frontend/src/pages/workspace/document/share/Index.tsx

@@ -13,10 +13,7 @@ function DocumentShareIndex() {
 	const [searchParams] = useSearchParams();
 	const shareToken = useMemo(() => searchParams.get("token"), [searchParams]);
 	const { data: sharedDocument } = useGetDocumentBySharingTokenQuery(shareToken);
-	const { doc, client, cleanUpYorkieDocument } = useYorkieDocument(
-		sharedDocument?.yorkieDocumentId,
-		"Anonymous"
-	);
+	const { doc, client } = useYorkieDocument(sharedDocument?.yorkieDocumentId, "Anonymous");
 
 	useEffect(() => {
 		if (!sharedDocument?.role) return;
@@ -35,11 +32,10 @@ function DocumentShareIndex() {
 		dispatch(setClient(client));
 
 		return () => {
-			cleanUpYorkieDocument();
 			dispatch(setDoc(null));
 			dispatch(setClient(null));
 		};
-	}, [cleanUpYorkieDocument, dispatch, client, doc]);
+	}, [dispatch, client, doc]);
 
 	if (!shareToken) return <Navigate to="/" state={{ from: location }} replace />;