read stamps sync and lock free

8 files changed, 68 insertions, 24 deletions

diff --git a/src/document/replica.rs b/src/document/replica.rs
index db47223..0cab0be 100644
--- a/src/document/replica.rs
+++ b/src/document/replica.rs
@@ -8,8 +8,10 @@
 
 use std::fmt;
 use std::ops::{Deref, DerefMut};
+use std::sync::Arc;
 
 use automerge::{AutoCommit, ChangeHash};
+use tokio::sync::watch;
 
 /// One generation of the swarm document.
 ///
@@ -84,6 +86,10 @@ impl fmt::Debug for Epoch {
 /// [`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`].
+///
+/// The refresh also publishes a [`Stamp`], for what must not wait for the lock: a
+/// tokio read lock waits behind a writer, and applying a large sync message holds the
+/// write lock for seconds.
 pub struct Replica {
     pub epoch: Epoch,
     /// How many times this node changed epoch (switch, rejoin, abandoned rejoin) in this
@@ -91,15 +97,32 @@ pub struct Replica {
     pub epoch_seq: u32,
     pub doc: AutoCommit,
     heads: Vec<ChangeHash>,
+    stamp: watch::Sender<Stamp>,
+}
+
+/// The epoch label and the heads of the replica, as of its last refresh, published
+/// together: keepalives and HELLOs read them without the document's lock.
+#[derive(Clone, Debug)]
+pub struct Stamp {
+    pub epoch: Epoch,
+    pub epoch_seq: u32,
+    /// Sorted.
+    pub heads: Arc<[ChangeHash]>,
 }
 
 impl Replica {
     pub fn new(epoch: Epoch, doc: AutoCommit) -> Self {
+        let (stamp, _) = watch::channel(Stamp {
+            epoch,
+            epoch_seq: 0,
+            heads: Arc::from([]),
+        });
         let mut replica = Self {
             epoch,
             epoch_seq: 0,
             doc,
             heads: Vec::new(),
+            stamp,
         };
         replica.refresh_heads();
         replica
@@ -110,10 +133,21 @@ impl Replica {
         &self.heads
     }
 
-    /// Call after changing the document (under the write lock).
+    /// Call after changing the document or its epoch (under the write lock). Publishes
+    /// the new [`Stamp`].
     pub fn refresh_heads(&mut self) {
         self.heads = self.doc.get_heads();
         self.heads.sort_unstable();
+        self.stamp.send_replace(Stamp {
+            epoch: self.epoch,
+            epoch_seq: self.epoch_seq,
+            heads: Arc::from(self.heads.as_slice()),
+        });
+    }
+
+    /// The published stamps, from the last refresh on.
+    pub fn stamps(&self) -> watch::Receiver<Stamp> {
+        self.stamp.subscribe()
     }
 }
 
diff --git a/src/node.rs b/src/node.rs
index 5b55c8d..ed0832c 100644
--- a/src/node.rs
+++ b/src/node.rs
@@ -67,12 +67,12 @@ impl Node {
         let traffic = Arc::new(Traffic::default());
         let initial_actor = config.initial_doc.get_actor().clone();
 
+        let replica = Replica::new(config.initial_epoch, config.initial_doc);
+        let stamp = replica.stamps();
         let state = AppState {
-            doc: Arc::new(RwLock::new(Replica::new(
-                config.initial_epoch,
-                config.initial_doc,
-            ))),
+            doc: Arc::new(RwLock::new(replica)),
             epoch_tx,
+            stamp,
             lifecycle: Arc::new(RwLock::new(NodeLifecycle::Joining {
                 discovered_swarm_id: None,
             })),
diff --git a/src/simtest/timing.rs b/src/simtest/timing.rs
index 06a2f12..ce1dfd5 100644
--- a/src/simtest/timing.rs
+++ b/src/simtest/timing.rs
@@ -255,7 +255,7 @@ async fn resent_hello_keeps_the_session() {
 
     let (run, handshake) = peer(a, "b", |p| p.last_hello).await.unwrap().unwrap();
     assert_eq!(run, b.run_id);
-    let (epoch, epoch_seq) = b.epoch_stamp().await;
+    let (epoch, epoch_seq) = b.epoch_stamp();
     let swarm_id = b.swarm_id().await;
     let hello = |handshake| {
         ParsedMessage::Hello(HelloRequest {
diff --git a/src/state/mod.rs b/src/state/mod.rs
index a51339c..962d62a 100644
--- a/src/state/mod.rs
+++ b/src/state/mod.rs
@@ -29,7 +29,7 @@ use tokio::sync::{Mutex, RwLock, broadcast, mpsc, watch};
 use tokio::time::Instant;
 use tracing::{debug, error, warn};
 
-use crate::document::replica::{Epoch, Replica};
+use crate::document::replica::{Epoch, Replica, Stamp};
 use crate::locks::MutexExt;
 use crate::remake::{Host, RemakeState};
 use crate::swarm::sync::{ChangeSet, SyncEvent};
@@ -52,6 +52,8 @@ pub(crate) struct AppState {
     /// Copy of `doc.epoch` for cheap, lock-free reads (handshakes, keepalives, /peers).
     /// Only ever changed while holding `doc`'s write lock, together with `doc.epoch`.
     pub(crate) epoch_tx: watch::Sender<Epoch>,
+    /// `doc`'s epoch label and heads, readable without its lock (`Replica::stamps`).
+    pub(crate) stamp: watch::Receiver<Stamp>,
 
     pub(crate) lifecycle: Arc<RwLock<NodeLifecycle>>,
 
@@ -186,10 +188,11 @@ impl AppState {
         *self.epoch_tx.borrow()
     }
 
-    /// The local epoch and `Replica::epoch_seq`, read together.
-    pub(crate) async fn epoch_stamp(&self) -> (Epoch, u32) {
-        let replica = self.doc.read().await;
-        (replica.epoch, replica.epoch_seq)
+    /// The local epoch and `Replica::epoch_seq`, read together, without waiting for the
+    /// document's lock.
+    pub(crate) fn epoch_stamp(&self) -> (Epoch, u32) {
+        let stamp = self.stamp.borrow();
+        (stamp.epoch, stamp.epoch_seq)
     }
 
     pub(crate) async fn is_ready(&self) -> bool {
diff --git a/src/swarm/blob.rs b/src/swarm/blob.rs
index c153a17..4b69562 100644
--- a/src/swarm/blob.rs
+++ b/src/swarm/blob.rs
@@ -166,7 +166,7 @@ async fn prepare(state: &AppState, id: BlobId) -> Option<Arc<Prepared>> {
 }
 
 async fn reply(state: &AppState, peer: &str, id: BlobId, body: BlobBody) {
-    let (epoch, epoch_seq) = state.epoch_stamp().await;
+    let (epoch, epoch_seq) = state.epoch_stamp();
     let message = ParsedMessage::Blob(BlobMessage {
         epoch,
         epoch_seq,
@@ -220,7 +220,7 @@ impl Drop for Registration<'_> {
 }
 
 async fn request(state: &AppState, peer: &str, id: BlobId, chunks: Vec<u32>) {
-    let (epoch, epoch_seq) = state.epoch_stamp().await;
+    let (epoch, epoch_seq) = state.epoch_stamp();
     let message = ParsedMessage::Blob(BlobMessage {
         epoch,
         epoch_seq,
diff --git a/src/swarm/peer/handshake.rs b/src/swarm/peer/handshake.rs
index f08929a..7e494c6 100644
--- a/src/swarm/peer/handshake.rs
+++ b/src/swarm/peer/handshake.rs
@@ -118,7 +118,7 @@ async fn transmit_hello(state: &AppState, peer: &str) -> Result<()> {
         .as_ref()
         .map(|key| key.hello_tag(swarm_id.as_deref().unwrap_or(""), state.local_id(), peer))
         .unwrap_or_default();
-    let (epoch, epoch_seq) = state.epoch_stamp().await;
+    let (epoch, epoch_seq) = state.epoch_stamp();
     // Retries repeat the handshake's number.
     let handshake = match state.get_peer_state(peer).await {
         Some(peer_state) => peer_state.lock().await.handshake,
@@ -209,7 +209,7 @@ pub(super) async fn send_hello_reply(
         .map(|key| key.hello_reply_tag(&swarm_id, state.local_id(), peer))
         .unwrap_or_default();
     let ready = state.is_ready().await;
-    let (epoch, epoch_seq) = state.epoch_stamp().await;
+    let (epoch, epoch_seq) = state.epoch_stamp();
     let reply = transport::HelloReply {
         swarm_id,
         peers,
diff --git a/src/swarm/peer/keepalive.rs b/src/swarm/peer/keepalive.rs
index 4fcad7e..4ce326e 100644
--- a/src/swarm/peer/keepalive.rs
+++ b/src/swarm/peer/keepalive.rs
@@ -25,14 +25,14 @@ 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 {
-    // Heads and epoch come from the same locked replica, so they always match. The heads
-    // are cached (`Replica`), so the read lock is enough.
+    // Heads and epoch are published together (`Replica::stamps`), so they always match,
+    // and reading them doesn't wait for a sync message being applied.
     let (heads, epoch, epoch_seq) = {
-        let replica = state.doc.read().await;
+        let stamp = state.stamp.borrow();
         (
-            transport::HeadsSummary::of(replica.heads().to_vec()),
-            replica.epoch,
-            replica.epoch_seq,
+            transport::HeadsSummary::of(stamp.heads.to_vec()),
+            stamp.epoch,
+            stamp.epoch_seq,
         )
     };
     let syncs_with_you = match state.get_peer_state(peer).await {
@@ -62,7 +62,13 @@ pub(super) async fn note_keepalive(state: &AppState, peer: &str, remote: &transp
         peer_state.syncs_with_us = Some(remote.syncs_with_you);
         peer_state.reported_heads = remote.heads.heads().map(<[_]>::len);
         drop(peer_state);
-        crate::swarm::relay::pull_if_behind(state, peer, remote).await;
+        // Off the consumer: telling which heads we lack takes the document's lock.
+        {
+            let (state, peer, remote) = (state.clone(), peer.to_owned(), remote.clone());
+            tokio::spawn(async move {
+                crate::swarm::relay::pull_if_behind(&state, &peer, &remote).await
+            });
+        }
         if became_ready {
             // We may have left the opening to it while it wasn't ready, and nothing else
             // would make us look again.
diff --git a/src/swarm/relay.rs b/src/swarm/relay.rs
index 4bfc2b6..34451c9 100644
--- a/src/swarm/relay.rs
+++ b/src/swarm/relay.rs
@@ -103,8 +103,9 @@ 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 {
+    // The common case, the same heads, needs only the published heads.
+    let ours = HeadsSummary::of(state.stamp.borrow().heads.to_vec());
+    if ours == remote.heads {
         return;
     }
     let missing = match &remote.heads {