From 798fbe9ed8136d808b16fc0a34a85c5e828814db Mon Sep 17 00:00:00 2001
From: Geoffrey Litt <gklitt@gmail.com>
Date: Tue, 12 Dec 2023 12:48:04 -0500
Subject: [PATCH] fix async bug in useDocument

previously useDocument did not behave well for documents that took a long time to load:

- while a load was blocking the UI, an old stale doc would be shown
- loads that finished out of order could result in an inconsistent final state

This change fixes those two issues:

- Upon a handle changing, the old doc contents are cleared out to avoid staleness inconsistencies.
- Out-of-order loads are accounted for; when a load finishes it's ignored if the handle has since changed.
---
 .../src/useDocument.ts                        | 29 ++++++++++++++-----
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/packages/automerge-repo-react-hooks/src/useDocument.ts b/packages/automerge-repo-react-hooks/src/useDocument.ts
index de51ae6fe..a2f335498 100644
--- a/packages/automerge-repo-react-hooks/src/useDocument.ts
+++ b/packages/automerge-repo-react-hooks/src/useDocument.ts
@@ -1,6 +1,6 @@
 import { ChangeFn, ChangeOptions, Doc } from "@automerge/automerge/next"
 import { AutomergeUrl, DocHandleChangePayload } from "@automerge/automerge-repo"
-import { useEffect, useState } from "react"
+import { useLayoutEffect, useRef, useState } from "react"
 import { useRepo } from "./useRepo.js"
 
 /** A hook which returns a document identified by a URL and a function to change the document.
@@ -19,16 +19,31 @@ export function useDocument<T>(
 
   const handle = documentUrl ? repo.find<T>(documentUrl) : null
 
-  useEffect(() => {
-    if (!handle) {
-      if (doc) {
-        setDoc(undefined)
-      }
+  const handleRef = useRef(null)
+
+  // Doing this in a useLayoutEffect seems to help make sure that
+  // a loading state gets rendered before we block the UI thread
+  // with a slow load.
+  useLayoutEffect(() => {
+    // When the handle has changed, reset the doc to an empty state.
+    // This ensures that if loading the doc takes a long time, the UI
+    // shows a loading state during that time rather than a stale doc.
+    setDoc(undefined)
 
+    if (!handle) {
       return
     }
 
-    handle.doc().then(v => setDoc(v))
+    handleRef.current = handle
+    handle.doc().then(v => {
+      // Bail out on updating the doc if the handle has changed since we started loading.
+      // This avoids problem with out-of-order loads when the handle is changing faster
+      // than documents are loading.
+      if (handleRef.current !== handle) {
+        return
+      }
+      setDoc(v)
+    })
 
     const onChange = (h: DocHandleChangePayload<T>) => setDoc(h.doc)
     handle.on("change", onChange)