cache heads

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

7 files changed, 61 insertions, 24 deletions

diff --git a/src/api/peers.rs b/src/api/peers.rs
index 7caaaf3..919f187 100644
--- a/src/api/peers.rs
+++ b/src/api/peers.rs
@@ -182,7 +182,7 @@ pub(crate) async fn list_peers(
     }
 
     peers.sort_by(|a, b| a.public_key.cmp(&b.public_key));
-    let local_heads = state.doc.write().await.doc.get_heads().len();
+    let local_heads = state.doc.read().await.heads().len();
     Ok(Json(PeerListResponse {
         local_alias: state.alias.clone(),
         local_epoch: local_epoch.into(),
diff --git a/src/api/stats.rs b/src/api/stats.rs
index bc925ff..838cc0f 100644
--- a/src/api/stats.rs
+++ b/src/api/stats.rs
@@ -134,8 +134,8 @@ async fn collect(state: &AppState) -> StatsResponse {
     let node = state.stats.snapshot();
 
     let (heads, tier) = {
-        let mut replica = state.doc.write().await;
-        let heads = replica.doc.get_heads().len();
+        let replica = state.doc.read().await;
+        let heads = replica.heads().len();
         (heads, members::tier(&replica.doc, state.local_id()))
     };
     let saved = state.stats.saved_size();
diff --git a/src/document/replica.rs b/src/document/replica.rs
index af807ea..e24f93c 100644
--- a/src/document/replica.rs
+++ b/src/document/replica.rs
@@ -9,7 +9,7 @@
 use std::fmt;
 use std::ops::{Deref, DerefMut};
 
-use automerge::AutoCommit;
+use automerge::{AutoCommit, ChangeHash};
 
 /// One generation of the swarm document.
 ///
@@ -77,14 +77,39 @@ impl fmt::Debug for Epoch {
 ///
 /// Derefs to the document, so code that only needs the `AutoCommit` keeps working
 /// with a guard of `Replica`.
+///
+/// The heads are cached: `AutoCommit::get_heads` takes `&mut self` (it closes a pending
+/// transaction), so without the cache every keepalive and status read would need the
+/// write lock. Every path that changes the document under the write lock calls
+/// [`Replica::refresh_heads`] before releasing it (`AppState::write_doc`, `mutate_doc`,
+/// `switch_epoch`, the sync apply); readers take the read lock and use
+/// [`Replica::heads`].
 pub struct Replica {
     pub epoch: Epoch,
     pub doc: AutoCommit,
+    heads: Vec<ChangeHash>,
 }
 
 impl Replica {
     pub fn new(epoch: Epoch, doc: AutoCommit) -> Self {
-        Self { epoch, doc }
+        let mut replica = Self {
+            epoch,
+            doc,
+            heads: Vec::new(),
+        };
+        replica.refresh_heads();
+        replica
+    }
+
+    /// The document's heads as of the last refresh, sorted.
+    pub fn heads(&self) -> &[ChangeHash] {
+        &self.heads
+    }
+
+    /// Call after changing the document (under the write lock).
+    pub fn refresh_heads(&mut self) {
+        self.heads = self.doc.get_heads();
+        self.heads.sort_unstable();
     }
 }
 
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);
diff --git a/src/swarm/peer.rs b/src/swarm/peer.rs
index dd70bbb..3d7da32 100644
--- a/src/swarm/peer.rs
+++ b/src/swarm/peer.rs
@@ -108,12 +108,12 @@ pub(crate) async fn send_ping(state: &AppState, peer: &str) -> Result<()> {
 
 /// Our keepalive for `peer`.
 pub(crate) async fn local_keepalive(state: &AppState, peer: &str) -> transport::Keepalive {
-    // AutoCommit::get_heads takes &mut self (it closes the pending transaction).
-    // Heads and epoch come from the same locked replica, so they always match.
+    // Heads and epoch come from the same locked replica, so they always match. The heads
+    // are cached (`Replica`), so the read lock is enough.
     let (heads, epoch) = {
-        let mut replica = state.doc.write().await;
+        let replica = state.doc.read().await;
         (
-            transport::HeadsSummary::of(replica.doc.get_heads()),
+            transport::HeadsSummary::of(replica.heads().to_vec()),
             replica.epoch,
         )
     };
diff --git a/src/swarm/relay.rs b/src/swarm/relay.rs
index 454a46a..492b59d 100644
--- a/src/swarm/relay.rs
+++ b/src/swarm/relay.rs
@@ -52,6 +52,10 @@ pub(crate) async fn pull_if_behind(state: &AppState, peer: &str, remote: &transp
     if remote.epoch != state.local_epoch() || !topology::is_neighbour(state, peer) {
         return;
     }
+    // The common case, the same heads, needs only the cached heads (read lock).
+    if HeadsSummary::of(state.doc.read().await.heads().to_vec()) == remote.heads {
+        return;
+    }
     let missing = match &remote.heads {
         HeadsSummary::Heads(heads) => {
             // AutoCommit's lookups take &mut self.
diff --git a/src/swarm/sync.rs b/src/swarm/sync.rs
index 17e25c4..80d6d56 100644
--- a/src/swarm/sync.rs
+++ b/src/swarm/sync.rs
@@ -343,7 +343,7 @@ pub(crate) async fn handle_sync(
 
             let mut peer_state = peer_state_mutex.blocking_lock();
             peer_state.reports_ready = Some(from_ready);
-            let heads_before = replica.doc.get_heads();
+            let heads_before = replica.heads().to_vec();
 
             replica
                 .doc
@@ -360,7 +360,8 @@ pub(crate) async fn handle_sync(
             peer_state.sync.reply_settled();
 
             resend_if_needed(&peer, &mut peer_state.sync.state);
-            let heads_after = replica.doc.get_heads();
+            replica.refresh_heads();
+            let heads_after = replica.heads().to_vec();
             stats.heads_seen(heads_after.len());
             let changed = heads_after != heads_before;
             if changed && measure_propagation {