Merge pull request #74 from xmtp/rygine/performance-improvements

Fix `useClient` and `useAttachment` performance

.changeset/slow-waves-chew.md

@@ -0,0 +1,5 @@
+---
+"@xmtp/react-sdk": patch
+---
+
+Fix `useClient` and `useAttachment` performance

packages/react-sdk/src/helpers/caching/contentTypes/attachment.ts

@@ -61,7 +61,7 @@ export const updateAttachmentData = async (
  * @returns The attachment data, or `undefined` if the message is not an
  * attachment content type
  */
-export const getAttachmentData = (message: CachedMessage) => {
+export const getRemoteAttachmentData = (message: CachedMessage) => {
   if (message.contentType === ContentTypeRemoteAttachment.toString()) {
     const metadata = message.metadata?.[NAMESPACE] as Attachment | undefined;
     return metadata;

packages/react-sdk/src/hooks/useAttachment.ts

@@ -1,4 +1,4 @@
-import { useCallback, useEffect, useState } from "react";
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import {
   ContentTypeAttachment,
   ContentTypeRemoteAttachment,
@@ -12,7 +12,7 @@ import { useDb } from "./useDb";
 import { type CachedMessage } from "@/helpers/caching/messages";
 import { useClient } from "@/hooks/useClient";
 import {
-  getAttachmentData,
+  getRemoteAttachmentData,
   updateAttachmentData,
 } from "@/helpers/caching/contentTypes/attachment";
 import { useMessage } from "@/hooks/useMessage";
@@ -47,29 +47,39 @@ export const useAttachment = (
   const [attachment, setAttachment] = useState<Attachment | undefined>(
     undefined,
   );
+  const getAttachmentRef = useRef(false);
 
   const {
     disableAutoload = false,
     autoloadMaxFileSize = defaultAutoloadMaxFileSize,
   } = options ?? {};
 
+  // get remote attachment data from the cache
+  const remoteAttachmentData = useMemo(
+    () => getRemoteAttachmentData(message),
+    [message],
+  );
+
   const loadRemoteAttachment = useCallback(
     async (force: boolean = false) => {
+      // check if attachment data is already cached
+      if (remoteAttachmentData && status !== "loaded") {
+        setAttachment(remoteAttachmentData);
+        setStatus("loaded");
+        return;
+      }
+
       if (
         client &&
         message.contentType === ContentTypeRemoteAttachment.toString() &&
+        !remoteAttachmentData &&
+        status !== "loading" &&
+        status !== "loaded" &&
+        !attachment &&
         // if the attachment already failed to load, don't automatically reload
         // unless forced
         ((message.hasLoadError && force) || !message.hasLoadError)
       ) {
-        // check if attachment data is already cached
-        const attachmentData = getAttachmentData(message);
-        if (attachmentData) {
-          setAttachment(attachmentData);
-          setStatus("loaded");
-          return;
-        }
-
         try {
           setStatus("loading");
           const loadedAttachment = await RemoteAttachmentCodec.load<Attachment>(
@@ -86,8 +96,8 @@ export const useAttachment = (
           setAttachment(loadedAttachment);
           setStatus("loaded");
         } catch (e) {
+          await updateMessage(message, { hasLoadError: true });
           setError(new Error("Unable to load remote attachment"));
-          void updateMessage(message, { hasLoadError: true });
           setStatus("error");
         }
       } else {
@@ -97,7 +107,15 @@ export const useAttachment = (
         setStatus("error");
       }
     },
-    [client, db, message, updateMessage],
+    [
+      remoteAttachmentData,
+      status,
+      client,
+      message,
+      attachment,
+      db,
+      updateMessage,
+    ],
   );
 
   const forceLoadRemoteAttachment = useCallback(() => {
@@ -106,31 +124,62 @@ export const useAttachment = (
 
   useEffect(() => {
     const getAttachment = async () => {
+      // prevent running this hook multiple times in parallel
+      if (getAttachmentRef.current) {
+        return;
+      }
+
+      getAttachmentRef.current = true;
+
+      if (attachment || status === "loading" || status === "loaded") {
+        return;
+      }
+
       switch (message.contentType) {
         case ContentTypeAttachment.toString(): {
           setAttachment(message.content as Attachment);
           setStatus("loaded");
+          getAttachmentRef.current = false;
           break;
         }
         case ContentTypeRemoteAttachment.toString(): {
-          const remoteAttachmentData = message.content as RemoteAttachment;
-          if (remoteAttachmentData.contentLength > autoloadMaxFileSize) {
+          // check if attachment data is already cached
+          if (remoteAttachmentData) {
+            setAttachment(remoteAttachmentData);
+            setStatus("loaded");
+            getAttachmentRef.current = false;
+            return;
+          }
+          const attachmentData = message.content as RemoteAttachment;
+          if (attachmentData.contentLength > autoloadMaxFileSize) {
             setStatus("autoloadMaxFileSizeExceeded");
+            getAttachmentRef.current = false;
             return;
           }
           if (!disableAutoload) {
             await loadRemoteAttachment();
           }
+          getAttachmentRef.current = false;
           break;
         }
         default: {
           setError(new Error("Message is not an attachment content type"));
           setStatus("error");
+          getAttachmentRef.current = false;
         }
       }
     };
     void getAttachment();
-  }, [autoloadMaxFileSize, db, disableAutoload, loadRemoteAttachment, message]);
+  }, [
+    remoteAttachmentData,
+    autoloadMaxFileSize,
+    db,
+    disableAutoload,
+    loadRemoteAttachment,
+    message,
+    attachment,
+    status,
+  ]);
 
   return {
     attachment,

packages/react-sdk/src/hooks/useClient.test.tsx

@@ -71,7 +71,7 @@ describe("useClient", () => {
     expect(clientCreateSpy).not.toHaveBeenCalled();
 
     await waitFor(() => {
-      expect(processUnprocessedMessagesMock).toBeCalledTimes(1);
+      expect(processUnprocessedMessagesMock).not.toHaveBeenCalled();
     });
   });
 
@@ -136,11 +136,7 @@ describe("useClient", () => {
     processUnprocessedMessagesMock.mockRejectedValue(testError);
 
     const { result } = renderHook(() => useClient(onErrorMock), {
-      wrapper: ({ children }) => (
-        <TestWrapper client={mockClient as unknown as Client}>
-          {children}
-        </TestWrapper>
-      ),
+      wrapper: ({ children }) => <TestWrapper>{children}</TestWrapper>,
     });
 
     await act(async () => {

packages/react-sdk/src/hooks/useClient.ts

@@ -1,4 +1,4 @@
-import { useCallback, useContext, useEffect, useRef, useState } from "react";
+import { useCallback, useContext, useRef, useState } from "react";
 import type { ClientOptions, Signer } from "@xmtp/xmtp-js";
 import { Client } from "@xmtp/xmtp-js";
 import { XMTPContext } from "../contexts/XMTPContext";
@@ -31,10 +31,6 @@ export const useClient = (onError?: OnError["onError"]) => {
   const [error, setError] = useState<Error | null>(null);
   // client is initializing
   const initializingRef = useRef(false);
-  // unprocessed messages are being processed
-  const processingRef = useRef(false);
-  // unprocessed messages have been processed
-  const processedRef = useRef(false);
 
   const {
     client,
@@ -89,13 +85,37 @@ export const useClient = (onError?: OnError["onError"]) => {
         }
 
         setIsLoading(false);
-        initializingRef.current = false;
+
+        // process unprocessed messages
+        try {
+          await processUnprocessedMessages({
+            client: xmtpClient,
+            db,
+            processors,
+            namespaces,
+            validators,
+          });
+        } catch (e) {
+          onError?.(e as Error);
+        } finally {
+          initializingRef.current = false;
+        }
 
         return xmtpClient;
       }
       return client;
     },
-    [client, codecs, onError, setClient, setClientSigner],
+    [
+      client,
+      codecs,
+      db,
+      namespaces,
+      onError,
+      processors,
+      setClient,
+      setClientSigner,
+      validators,
+    ],
   );
 
   /**
@@ -109,35 +129,6 @@ export const useClient = (onError?: OnError["onError"]) => {
     }
   }, [client, setClient, setClientSigner]);
 
-  /**
-   * Process unprocessed messages when there's an available client, but only
-   * do it once
-   */
-  useEffect(() => {
-    if (client && !processingRef.current && !processedRef.current) {
-      processingRef.current = true;
-      setIsLoading(true);
-      const reprocess = async () => {
-        try {
-          await processUnprocessedMessages({
-            client,
-            db,
-            processors,
-            namespaces,
-            validators,
-          });
-          processedRef.current = true;
-        } catch (e) {
-          onError?.(e as Error);
-        } finally {
-          processingRef.current = false;
-          setIsLoading(false);
-        }
-      };
-      void reprocess();
-    }
-  }, [client, db, namespaces, onError, processors, validators]);
-
   return {
     client,
     disconnect,

packages/react-sdk/src/index.ts

@@ -51,7 +51,6 @@ export type {
 } from "./helpers/caching/messages";
 export {
   getMessageByXmtpID,
-  processUnprocessedMessages,
   toCachedMessage,
 } from "./helpers/caching/messages";