Use provider hooks for single hook (#970)

hackerwins · web-flow · commit 324af4f4cb07 · 2025-04-07T15:55:43.000+09:00
This commit adds support for StrictMode in React's standalone mode,
updating the interface to be on par with the Provider mode.
Additionally, it removes existing duplicate logic to streamline the
codebase and enhance maintainability.
diff --git a/examples/react-flow/src/main.tsx b/examples/react-flow/src/main.tsx
@@ -2,6 +2,7 @@ import { createRoot } from 'react-dom/client';
 import { DocumentProvider, YorkieProvider } from '@yorkie-js/react';
 import './index.css';
 import App from './App';
+import { StrictMode } from 'react';
 
 const initialNodes = [
   {
@@ -42,18 +43,21 @@ const initialGraph = {
 };
 
 createRoot(document.getElementById('root')!).render(
-  <YorkieProvider
-    apiKey={import.meta.env.VITE_YORKIE_API_KEY}
-    rpcAddr={import.meta.env.VITE_YORKIE_API_ADDR}
-  >
-    <DocumentProvider
-      docKey={`react-flow-${new Date()
-        .toISOString()
-        .substring(0, 10)
-        .replace(/-/g, '')}`}
-      initialRoot={initialGraph}
+  <StrictMode>
+    <YorkieProvider
+      apiKey={import.meta.env.VITE_YORKIE_API_KEY}
+      rpcAddr={import.meta.env.VITE_YORKIE_API_ADDR}
     >
-      <App />
-    </DocumentProvider>
-  </YorkieProvider>,
+      <DocumentProvider
+        docKey={`react-flow-${new Date()
+          .toISOString()
+          .substring(0, 10)
+          .replace(/-/g, '')}`}
+        initialRoot={initialGraph}
+      >
+        <App />
+      </DocumentProvider>
+    </YorkieProvider>
+    ,
+  </StrictMode>,
 );
diff --git a/packages/react/src/DocumentProvider.tsx b/packages/react/src/DocumentProvider.tsx
@@ -26,43 +26,34 @@ import {
   Presence,
   Indexable,
   StreamConnectionStatus,
+  Client,
 } from '@yorkie-js/sdk';
 import { useYorkie } from './YorkieProvider';
 
-type DocumentContextType<R, P extends Indexable = Indexable> = {
-  root: R;
-  presences: { clientID: string; presence: P }[];
-  connection: StreamConnectionStatus;
-  update: (callback: (root: R) => void) => void;
-  loading: boolean;
-  error: Error | undefined;
-};
-
-const DocumentContext = createContext<DocumentContextType<any> | null>(null);
-
 /**
- * `DocumentProvider` is a component that provides a document to its children.
- * This component must be under a `YorkieProvider` component to initialize the
- * Yorkie client properly.
+ * `useYorkieDocument` is a custom hook that initializes a Yorkie document.
+ * @param client
+ * @param clientLoading
+ * @param clientError
+ * @param docKey
+ * @param initialRoot
+ * @param initialPresence
+ * @returns
  */
-export const DocumentProvider = <R, P extends Indexable = Indexable>({
-  docKey,
-  initialRoot = {} as R,
-  initialPresence = {} as P,
-  children,
-}: {
-  docKey: string;
-  initialRoot?: R;
-  initialPresence?: P;
-  children?: React.ReactNode;
-}) => {
-  const { client, loading: clientLoading, error: clientError } = useYorkie();
-  const [doc, setDoc] = useState<Document<R, P> | undefined>(undefined);
-  const [loading, setLoading] = useState<boolean>(true);
-  const [error, setError] = useState<Error | undefined>(undefined);
+export function useYorkieDocument<R, P extends Indexable = Indexable>(
+  client: Client | undefined,
+  clientLoading: boolean,
+  clientError: Error | undefined,
+  docKey: string,
+  initialRoot: R,
+  initialPresence: P,
+) {
+  const [doc, setDoc] = useState<Document<R, P>>();
+  const [loading, setLoading] = useState(true);
+  const [error, setError] = useState<Error>();
   const [root, setRoot] = useState(initialRoot);
   const [presences, setPresences] = useState<
-    { clientID: string; presence: any }[]
+    Array<{ clientID: string; presence: P }>
   >([]);
   const [connection, setConnection] = useState<StreamConnectionStatus>(
     StreamConnectionStatus.Disconnected,
@@ -152,6 +143,54 @@ export const DocumentProvider = <R, P extends Indexable = Indexable>({
     },
     [doc],
   );
+  return {
+    doc,
+    root,
+    presences,
+    connection,
+    update,
+    loading,
+    error,
+  };
+}
+
+type DocumentContextType<R, P extends Indexable = Indexable> = {
+  root: R;
+  presences: { clientID: string; presence: P }[];
+  connection: StreamConnectionStatus;
+  update: (callback: (root: R, presence: Presence<P>) => void) => void;
+  loading: boolean;
+  error: Error | undefined;
+};
+
+const DocumentContext = createContext<DocumentContextType<any> | null>(null);
+
+/**
+ * `DocumentProvider` is a component that provides a document to its children.
+ * This component must be under a `YorkieProvider` component to initialize the
+ * Yorkie client properly.
+ */
+export const DocumentProvider = <R, P extends Indexable = Indexable>({
+  docKey,
+  initialRoot = {} as R,
+  initialPresence = {} as P,
+  children,
+}: {
+  docKey: string;
+  initialRoot?: R;
+  initialPresence?: P;
+  children?: React.ReactNode;
+}) => {
+  const { client, loading: clientLoading, error: clientError } = useYorkie();
+  const { root, presences, connection, update, loading, error } =
+    useYorkieDocument(
+      client,
+      clientLoading,
+      clientError,
+      docKey,
+      initialRoot,
+      initialPresence,
+    );
 
   return (
     <DocumentContext.Provider
@@ -176,7 +215,7 @@ export const useDocument = <R, P extends Indexable = Indexable>() => {
     presences: context.presences as { clientID: string; presence: P }[],
     connection: context.connection,
     update: context.update as (
-      callback: (root: R, presence: P) => void,
+      callback: (root: R, presence: Presence<P>) => void,
     ) => void,
     loading: context.loading,
     error: context.error,
diff --git a/packages/react/src/YorkieProvider.tsx b/packages/react/src/YorkieProvider.tsx
@@ -45,37 +45,38 @@ interface YorkieProviderProps {
 }
 
 /**
- * `YorkieProvider` is a component that provides the Yorkie client to its children.
- * It initializes the Yorkie client with the given API key and RPC address.
+ * `useYorkieClient` is a custom hook that initializes a Yorkie client.
+ * @param apiKey
+ * @param rpcAddr
+ * @returns
  */
-export const YorkieProvider: React.FC<
-  PropsWithChildren<YorkieProviderProps>
-> = ({ apiKey, rpcAddr = 'https://api.yorkie.dev', children }) => {
+export function useYorkieClient(apiKey: string, rpcAddr: string) {
   const [client, setClient] = useState<Client | undefined>(undefined);
   const [loading, setLoading] = useState<boolean>(true);
   const [error, setError] = useState<Error | undefined>(undefined);
+  const [didMount, setDidMount] = useState(false);
 
-  // NOTE(hackerwins): We use `useRef` to keep the client instance to prevent
-  // re-creating the client instance in StrictMode.
-  const clientRef = useRef<Client | undefined>(undefined);
+  // NOTE(hackerwins): In StrictMode, the component will call twice
+  // useEffect in development mode. To prevent creating a new client
+  // twice, create a client after the mounting.
+  useEffect(() => {
+    setDidMount(true);
+  }, []);
 
   useEffect(() => {
     setLoading(true);
     setError(undefined);
 
     async function activateClient() {
-      try {
-        if (clientRef.current?.isActive()) {
-          setClient(clientRef.current);
-          setLoading(false);
-          return;
-        }
+      if (!didMount) {
+        return;
+      }
 
+      try {
         const newClient = new Client(rpcAddr, {
           apiKey,
         });
         await newClient.activate();
-        clientRef.current = newClient;
         setClient(newClient);
       } catch (e) {
         setError(
@@ -88,12 +89,23 @@ export const YorkieProvider: React.FC<
     activateClient();
 
     return () => {
-      if (clientRef.current?.isActive()) {
-        clientRef.current.deactivate({ keepalive: true });
-        clientRef.current = undefined;
+      if (client?.isActive()) {
+        client.deactivate({ keepalive: true });
       }
     };
-  }, [apiKey, rpcAddr]);
+  }, [apiKey, rpcAddr, didMount]);
+
+  return { client, loading, error };
+}
+
+/**
+ * `YorkieProvider` is a component that provides the Yorkie client to its children.
+ * It initializes the Yorkie client with the given API key and RPC address.
+ */
+export const YorkieProvider: React.FC<
+  PropsWithChildren<YorkieProviderProps>
+> = ({ apiKey, rpcAddr = 'https://api.yorkie.dev', children }) => {
+  const { client, loading, error } = useYorkieClient(apiKey, rpcAddr);
 
   return (
     <YorkieContext.Provider value={{ client, loading, error }}>
@@ -102,6 +114,10 @@ export const YorkieProvider: React.FC<
   );
 };
 
+/**
+ * `useYorkie` is a custom hook that returns the Yorkie client and its loading state.
+ * @returns
+ */
 export const useYorkie = () => {
   const context = useContext(YorkieContext);
   if (!context) {
diff --git a/packages/react/src/useYorkieDoc.ts b/packages/react/src/useYorkieDoc.ts
@@ -14,96 +14,58 @@
  * limitations under the License.
  */
 
-import { useCallback, useEffect, useState } from 'react';
-import { Client, Document } from '@yorkie-js/sdk';
+import { Indexable, Presence, StreamConnectionStatus } from '@yorkie-js/sdk';
+import { useYorkieClient } from './YorkieProvider';
+import { useYorkieDocument } from './DocumentProvider';
 
 /**
- * `useYorkieDoc` is a custom hook that initializes a Yorkie client and attaches a document.
+ * `useYorkieDoc` is a custom hook that initializes a Yorkie Client and a
+ * document in a single hook.
  *
  * @param apiKey
  * @param docKey
- * @param initialRoot
  * @returns
  */
-export function useYorkieDoc<T>(
+export function useYorkieDoc<R, P extends Indexable>(
   apiKey: string,
   docKey: string,
-  initialRoot: T,
   options?: {
+    initialRoot?: R;
+    initialPresence?: P;
     rpcAddr?: string;
   },
 ): {
-  root: T;
-  update: (callback: (root: T) => void) => void;
+  root: R;
+  presences: Array<{ clientID: string; presence: P }>;
+  connection: StreamConnectionStatus;
+  update: (callback: (root: R, presence: Presence<P>) => void) => void;
   loading: boolean;
   error: Error | undefined;
 } {
-  const [client, setClient] = useState<Client | undefined>(undefined);
-  const [doc, setDoc] = useState<Document<T> | undefined>(undefined);
-  const [root, setRoot] = useState<T>(initialRoot);
-  const [loading, setLoading] = useState<boolean>(true);
-  const [error, setError] = useState<Error | undefined>(undefined);
+  const rpcAddr = options?.rpcAddr || 'https://api.yorkie.dev';
+  const initialRoot = options?.initialRoot || ({} as R);
+  const initialPresence = options?.initialPresence || ({} as P);
 
-  useEffect(() => {
-    setLoading(true);
-    setError(undefined);
+  const {
+    client,
+    loading: clientLoading,
+    error: clientError,
+  } = useYorkieClient(apiKey, rpcAddr);
 
-    /**
-     * `setupYorkie` initializes the Yorkie client and attaches the document.
-     */
-    async function setupYorkie() {
-      try {
-        const client = new Client(
-          options?.rpcAddr || 'https://api.yorkie.dev',
-          { apiKey },
-        );
-        await client.activate();
-
-        const doc = new Document<T>(docKey);
-        await client.attach(doc, {
-          initialPresence: {},
-          initialRoot,
-        });
-
-        doc.subscribe((event) => {
-          if (event.type === 'remote-change' || event.type === 'local-change') {
-            setRoot(doc.getRoot());
-          }
-        });
-
-        setClient(client);
-        setDoc(doc);
-        setRoot(doc.getRoot());
-      } catch (err) {
-        setError(err instanceof Error ? err : new Error('Unknown error'));
-      } finally {
-        setLoading(false);
-      }
-    }
-
-    setupYorkie();
-
-    return () => {
-      client?.deactivate({ keepalive: true });
-    };
-  }, [apiKey, docKey, JSON.stringify(initialRoot), JSON.stringify(options)]);
-
-  const update = useCallback(
-    (callback: (root: T) => void) => {
-      if (!doc) {
-        console.warn('Attempted to update document before it was initialized');
-        return;
-      }
-
-      doc.update((root) => {
-        callback(root);
-      });
-    },
-    [doc],
-  );
+  const { root, presences, connection, update, loading, error } =
+    useYorkieDocument(
+      client,
+      clientLoading,
+      clientError,
+      docKey,
+      initialRoot,
+      initialPresence,
+    );
 
   return {
     root,
+    presences,
+    connection,
     update,
     loading,
     error,
