cache heads

To reduce the number of write locks, heads are now cached.

1 files changed, 20 insertions, 13 deletions

diff --git a/src/state.rs b/src/state.rs
index 9383c89..0dd3a81 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -815,9 +815,14 @@ impl AppState {
         T: Send + 'static,
     {
         let doc = Arc::clone(&self.doc);
-        tokio::task::spawn_blocking(move || f(&mut doc.blocking_write().doc))
-            .await
-            .expect("document write task panicked")
+        tokio::task::spawn_blocking(move || {
+            let mut replica = doc.blocking_write();
+            let out = f(&mut replica.doc);
+            replica.refresh_heads();
+            out
+        })
+        .await
+        .expect("document write task panicked")
     }
 
     /// Mutates the document on the blocking pool as one Automerge change, then
@@ -852,6 +857,7 @@ impl AppState {
                         replica.doc.rollback();
                     }
                 }
+                replica.refresh_heads();
                 result
             }; // lock released before notifying
 
@@ -944,8 +950,8 @@ impl AppState {
     /// anyway with `unready_live`.
     ///
     /// With `with_heads`, a live read also returns the document's heads at the moment of
-    /// the read (for clients that base later requests on what they read). That needs the
-    /// write lock, so it is opt-in. Stale reads have no heads.
+    /// the read (for clients that base later requests on what they read). Stale reads
+    /// have no heads.
     pub(crate) async fn render_read<T, R>(
         &self,
         unready_live: bool,
@@ -995,14 +1001,14 @@ impl AppState {
         if !with_heads {
             return (self.render(read, finish).await, None);
         }
-        // AutoCommit::get_heads takes &mut self (it closes a pending transaction).
-        let (data, mut heads) = self
-            .write_doc(move |doc| {
-                let heads = doc.get_heads();
-                (read(doc), heads)
-            })
-            .await;
-        heads.sort();
+        // Heads and data under the same read lock (the heads are cached: `Replica`).
+        let doc = Arc::clone(&self.doc);
+        let (data, heads) = tokio::task::spawn_blocking(move || {
+            let replica = doc.blocking_read();
+            (read(&replica.doc), replica.heads().to_vec())
+        })
+        .await
+        .expect("document read task panicked");
         let rendered = tokio::task::spawn_blocking(move || finish(data))
             .await
             .expect("document read task panicked");
@@ -1198,6 +1204,7 @@ impl Host for AppState {
                     let new_doc = f(&mut replica.doc)?;
                     let old_doc = std::mem::replace(&mut replica.doc, new_doc);
                     replica.epoch = epoch;
+                    replica.refresh_heads();
                     // Updated under the same lock, so the cached epoch never disagrees
                     // with the document for anyone holding the lock.
                     state.epoch_tx.send_replace(epoch);