document: allow writes during rejoin

3 files changed, 207 insertions, 110 deletions

diff --git a/src/state/host.rs b/src/state/host.rs
index eb7db35..e7921e4 100644
--- a/src/state/host.rs
+++ b/src/state/host.rs
@@ -11,12 +11,78 @@ use tokio::time::Instant;
 use tracing::debug;
 
 use super::{AppState, HelloState, NodeLifecycle, StaleView, Standing};
-use crate::document::replica::Epoch;
+use crate::document::replica::{Epoch, Replica};
 use crate::document::store::{ChangeSet, StoreError};
 use crate::document::{ChangeHash, Document, ReadDoc, Transaction};
-use crate::remake::{Host, PeerSummary, RemakeState};
+use crate::locks::RwLockExt;
+use crate::remake::{BegunRejoin, Host, PeerSummary, RemakeState};
 
 impl AppState {
+    /// Under the document's write lock: `doc` of `epoch` replaces the live document,
+    /// which is returned. Every peer's sync state is reset before the lock is released,
+    /// so no message is generated from or applied to a state that describes the old
+    /// document. (Lock order doc -> peer, as in sync generation.)
+    fn swap_document(&self, replica: &mut Replica, epoch: Epoch, doc: Document) -> Document {
+        let old = std::mem::replace(&mut replica.doc, doc);
+        replica.epoch = epoch;
+        replica.epoch_seq += 1;
+        // Publishes the epoch with the heads (`Replica::stamps`), under the same lock,
+        // so it never disagrees with the document for anyone holding the lock.
+        replica.refresh_heads();
+        let now = Instant::now();
+        for peer in self.peer_states.blocking_read().values() {
+            let mut peer = peer.blocking_lock();
+            peer.reset_sync();
+            // A fresh start, not an idle sync: no heal re-handshake for it.
+            peer.last_sync_rx = Some(now);
+        }
+        old
+    }
+
+    /// Stops serving requests (back to `Joining`, keeping the swarm ID), as while a
+    /// rejoin syncs a fresh document; on the blocking pool. Whether the node served.
+    fn suspend(&self) -> bool {
+        let mut lifecycle = self.lifecycle.blocking_write();
+        let NodeLifecycle::Ready { swarm_id } = &*lifecycle else {
+            return false;
+        };
+        let swarm_id = swarm_id.clone();
+        debug!(swarm_id = %swarm_id, "Node suspended (rejoining)");
+        *lifecycle = NodeLifecycle::Joining {
+            discovered_swarm_id: Some(swarm_id),
+        };
+        true
+    }
+
+    /// Serves requests again after [`suspend`](Self::suspend); on the blocking pool.
+    fn resume(&self) {
+        let mut lifecycle = self.lifecycle.blocking_write();
+        if let NodeLifecycle::Joining {
+            discovered_swarm_id: Some(swarm_id),
+        } = &*lifecycle
+        {
+            let swarm_id = swarm_id.clone();
+            debug!(swarm_id = %swarm_id, "Node resumed");
+            *lifecycle = NodeLifecycle::Ready { swarm_id };
+        }
+    }
+
+    /// The document before a rejoin, out of the stale view (which ends) if the node
+    /// keeps it there, else `old`; and whether it was kept there (the node was ready).
+    fn take_stale(&self, old: Option<Document>) -> (Document, bool) {
+        let Some(view) = self.stale.write_or_recover().take() else {
+            return (old.unwrap_or_default(), false);
+        };
+        // A read holding it a moment longer gets a copy of its own.
+        let doc = Arc::try_unwrap(view.doc)
+            .map(|doc| {
+                doc.into_inner()
+                    .unwrap_or_else(|poisoned| poisoned.into_inner())
+            })
+            .unwrap_or_else(|shared| shared.read_or_recover().clone());
+        (doc, true)
+    }
+
     /// After the document was replaced: every peer's sync state refers to the old
     /// document, so start over with all of them. Peers still on the old epoch are
     /// skipped by the sync worker until they switch too.
@@ -143,54 +209,6 @@ impl Host for AppState {
         }
     }
 
-    fn suspend(&self) -> impl Future<Output = bool> + Send {
-        let state = self.clone();
-        async move {
-            let mut lifecycle = state.lifecycle.write().await;
-            let NodeLifecycle::Ready { swarm_id } = &*lifecycle else {
-                return false;
-            };
-            let swarm_id = swarm_id.clone();
-            debug!(swarm_id = %swarm_id, "Node suspended (rejoining)");
-            *lifecycle = NodeLifecycle::Joining {
-                discovered_swarm_id: Some(swarm_id),
-            };
-            true
-        }
-    }
-
-    fn resume(&self) -> impl Future<Output = ()> + Send {
-        let state = self.clone();
-        async move {
-            let resumed = {
-                let mut lifecycle = state.lifecycle.write().await;
-                let swarm_id = match &*lifecycle {
-                    NodeLifecycle::Joining {
-                        discovered_swarm_id: Some(swarm_id),
-                    } => Some(swarm_id.clone()),
-                    _ => None,
-                };
-                if let Some(swarm_id) = &swarm_id {
-                    debug!(swarm_id = %swarm_id, "Node resumed");
-                    *lifecycle = NodeLifecycle::Ready {
-                        swarm_id: swarm_id.clone(),
-                    };
-                }
-                swarm_id.is_some()
-            };
-            if resumed {
-                crate::swarm::peer::announce(&state);
-            }
-        }
-    }
-
-    fn set_stale_view(&self, view: Option<(Epoch, Document)>) {
-        self.set_stale(view.map(|(epoch, doc)| StaleView {
-            epoch,
-            doc: Arc::new(doc),
-        }));
-    }
-
     fn initial_sync_done(&self) -> bool {
         self.initial_sync_done
             .load(std::sync::atomic::Ordering::Acquire)
@@ -230,51 +248,121 @@ impl Host for AppState {
         self.superseded_writes
     }
 
-    fn commit_changes<T, F>(&self, f: F) -> impl Future<Output = Result<T, StoreError>> + Send
-    where
-        F: FnOnce(&mut Document) -> T + Send + 'static,
-        T: Send + 'static,
-    {
+    fn begin_rejoin(
+        &self,
+        target: Epoch,
+        actor: crate::document::ActorId,
+    ) -> impl Future<Output = BegunRejoin> + Send {
         let state = self.clone();
         async move {
-            tokio::task::spawn_blocking(move || {
-                let result = {
+            let begun = {
+                let state = state.clone();
+                tokio::task::spawn_blocking(move || {
                     let mut replica = state.doc.blocking_write();
-                    if !replica.writable {
-                        return Err(StoreError::Rejoining);
+                    let previous = replica.epoch;
+                    let fresh = Document::new().with_actor(actor);
+                    let old = state.swap_document(&mut replica, target, fresh);
+                    replica.writable = false;
+                    let heads = old.get_heads();
+                    // Under the document's lock: no read or write in between finds the
+                    // node ready on the fresh document.
+                    let was_ready = state.suspend();
+                    let old = if was_ready {
+                        state.set_stale(Some(StaleView {
+                            epoch: previous,
+                            doc: Arc::new(std::sync::RwLock::new(old)),
+                        }));
+                        None
+                    } else {
+                        Some(old)
+                    };
+                    BegunRejoin {
+                        was_ready,
+                        heads,
+                        old,
                     }
-                    let result = f(&mut replica.doc);
-                    replica.refresh_heads();
-                    result
-                };
-                state.publish_change();
-                Ok(result)
-            })
-            .await
-            .expect("document write task panicked")
+                })
+                .await
+                .expect("rejoin task panicked")
+            };
+            state.after_epoch_switch().await;
+            begun
         }
     }
 
-    fn refuse_writes(&self, copy: bool) -> impl Future<Output = Option<Document>> + Send {
+    fn rejoin_document(&self) -> impl Future<Output = Option<Document>> + Send {
         let state = self.clone();
         async move {
-            tokio::task::spawn_blocking(move || {
-                let mut replica = state.doc.blocking_write();
-                replica.writable = false;
-                copy.then(|| replica.doc.clone())
-            })
-            .await
-            .expect("document write task panicked")
+            let (_, stale) = state.stale_document()?;
+            tokio::task::spawn_blocking(move || stale.read_or_recover().clone())
+                .await
+                .ok()
         }
     }
 
-    fn accept_writes(&self) -> impl Future<Output = ()> + Send {
+    fn finish_rejoin<T, F>(
+        &self,
+        old: Option<Document>,
+        f: F,
+    ) -> impl Future<Output = (T, Document)> + Send
+    where
+        F: FnOnce(&mut Document, &Document) -> T + Send + 'static,
+        T: Send + 'static,
+    {
         let state = self.clone();
         async move {
-            // On the blocking pool, as every write lock is taken.
-            tokio::task::spawn_blocking(move || state.doc.blocking_write().writable = true)
+            let (out, old, resumed) = {
+                let state = state.clone();
+                tokio::task::spawn_blocking(move || {
+                    let mut replica = state.doc.blocking_write();
+                    // No write goes into the document before the rejoin any more: they
+                    // take the document's lock first (`mutate_doc`).
+                    let (old, resumed) = state.take_stale(old);
+                    let out = f(&mut replica.doc, &old);
+                    replica.writable = true;
+                    replica.refresh_heads();
+                    if resumed {
+                        state.resume();
+                    }
+                    (out, old, resumed)
+                })
                 .await
-                .expect("document write task panicked");
+                .expect("rejoin task panicked")
+            };
+            state.publish_change();
+            if resumed {
+                crate::swarm::peer::announce(&state);
+            }
+            (out, old)
+        }
+    }
+
+    fn abandon_rejoin(
+        &self,
+        previous: Epoch,
+        old: Option<Document>,
+    ) -> impl Future<Output = ()> + Send {
+        let state = self.clone();
+        async move {
+            let resumed = {
+                let state = state.clone();
+                tokio::task::spawn_blocking(move || {
+                    let mut replica = state.doc.blocking_write();
+                    let (old, resumed) = state.take_stale(old);
+                    state.swap_document(&mut replica, previous, old);
+                    replica.writable = true;
+                    if resumed {
+                        state.resume();
+                    }
+                    resumed
+                })
+                .await
+                .expect("rejoin task panicked")
+            };
+            state.after_epoch_switch().await;
+            if resumed {
+                crate::swarm::peer::announce(&state);
+            }
         }
     }
 
@@ -307,7 +395,6 @@ impl Host for AppState {
     fn switch_epoch<E, F>(
         &self,
         epoch: Epoch,
-        writable: bool,
         f: F,
     ) -> impl Future<Output = Result<Document, E>> + Send
     where
@@ -321,25 +408,7 @@ impl Host for AppState {
                 tokio::task::spawn_blocking(move || -> Result<Document, E> {
                     let mut replica = state.doc.blocking_write();
                     let new_doc = f(&replica.doc)?;
-                    let old_doc = std::mem::replace(&mut replica.doc, new_doc);
-                    replica.epoch = epoch;
-                    replica.epoch_seq += 1;
-                    replica.writable = writable;
-                    // Publishes the epoch with the heads (`Replica::stamps`), under the
-                    // same lock, so it never disagrees with the document for anyone
-                    // holding the lock.
-                    replica.refresh_heads();
-                    // Reset every peer's sync state before the lock is released, so no
-                    // message is generated from or applied to a state that describes the
-                    // old document. (Lock order doc -> peer, as in sync generation.)
-                    let now = Instant::now();
-                    for peer in state.peer_states.blocking_read().values() {
-                        let mut peer = peer.blocking_lock();
-                        peer.reset_sync();
-                        // A fresh start, not an idle sync: no heal re-handshake for it.
-                        peer.last_sync_rx = Some(now);
-                    }
-                    Ok(old_doc)
+                    Ok(state.swap_document(&mut replica, epoch, new_doc))
                 })
                 .await
                 .expect("epoch switch task panicked")
diff --git a/src/state/mod.rs b/src/state/mod.rs
index 6e6e48c..7c1299f 100644
--- a/src/state/mod.rs
+++ b/src/state/mod.rs
@@ -40,7 +40,7 @@ use crate::document::lock::DocumentLock;
 use crate::document::replica::{Epoch, Stamp};
 use crate::document::store::{ChangeSet, StoreError};
 use crate::document::{Document, Transaction};
-use crate::locks::MutexExt;
+use crate::locks::{MutexExt, RwLockExt};
 use crate::remake::{Host, RemakeState};
 use crate::swarm::sync::SyncEvent;
 use crate::swarm::throughput::ThroughputTable;
@@ -394,8 +394,10 @@ impl AppState {
 
     /// Mutates the document on the blocking pool as one change, then
     /// publishes it. Write and notification run in the same blocking task, so a
-    /// cancelled request can't leave a committed write unannounced. Refused
-    /// (`StoreError::Rejoining`, `f` not called) while the document isn't writable.
+    /// cancelled request can't leave a committed write unannounced. While a rejoin runs,
+    /// the write goes into the document before it (`StaleView`), replayed when the
+    /// rejoin ends; on a node that keeps none, it is refused (`StoreError::Rejoining`,
+    /// `f` not called).
     pub(crate) async fn mutate_doc<T, E>(
         &self,
         f: impl FnOnce(&mut Transaction<'_>) -> Result<T, E> + Send + 'static,
@@ -412,7 +414,18 @@ impl AppState {
                 // Checked here, under the lock: the caller's readiness check (if any)
                 // was before a rejoin could swap in a fresh document.
                 if !replica.writable {
-                    return Err(E::from(StoreError::Rejoining));
+                    // Rejoining: into the document before the rejoin, which the rejoin
+                    // replays when it ends (`StaleView`); refused by a node that keeps
+                    // none (it wasn't ready).
+                    let Some((_, stale)) = state.stale_document() else {
+                        return Err(E::from(StoreError::Rejoining));
+                    };
+                    let mut doc = stale.write_or_recover();
+                    return crate::document::replica::transact(
+                        &mut doc,
+                        state.clock.now_millis(),
+                        f,
+                    );
                 }
                 // One API write = one change; on error, its partial writes (e.g. the old
                 // value deleted, the new one not written) are dropped.
diff --git a/src/state/read.rs b/src/state/read.rs
index a9b3116..e325e61 100644
--- a/src/state/read.rs
+++ b/src/state/read.rs
@@ -1,4 +1,5 @@
-//! Serving reads: from the live document, or from the pre-rejoin copy while rejoining.
+//! Serving reads: from the live document, or from the pre-rejoin document while
+//! rejoining.
 
 use std::sync::Arc;
 
@@ -7,10 +8,13 @@ use crate::document::replica::Epoch;
 use crate::document::{ChangeHash, Document};
 use crate::locks::RwLockExt;
 
-/// A read-only view of the document the node served before it started rejoining.
+/// The document the node served before it started rejoining, of the epoch it left:
+/// reads are answered from it (marked stale), and writes go into it
+/// (`AppState::mutate_doc`), until the rejoin replays them into the new epoch or
+/// restores this document (`Host::finish_rejoin`, `Host::abandon_rejoin`).
 pub(crate) struct StaleView {
     pub(crate) epoch: Epoch,
-    pub(crate) doc: Arc<Document>,
+    pub(crate) doc: Arc<std::sync::RwLock<Document>>,
 }
 
 /// Where a read was served from.
@@ -57,9 +61,12 @@ impl AppState {
             return Err(axum::http::StatusCode::SERVICE_UNAVAILABLE);
         };
 
-        let rendered = tokio::task::spawn_blocking(move || finish(read(&doc)))
-            .await
-            .expect("document read task panicked");
+        let rendered = tokio::task::spawn_blocking(move || {
+            let data = read(&doc.read_or_recover());
+            finish(data)
+        })
+        .await
+        .expect("document read task panicked");
         Ok((rendered, ReadSource::Stale(epoch), None))
     }
 
@@ -93,4 +100,12 @@ impl AppState {
     pub(super) fn set_stale(&self, view: Option<StaleView>) {
         *self.stale.write_or_recover() = view;
     }
+
+    /// The stale view's document and epoch, if there is one.
+    pub(crate) fn stale_document(&self) -> Option<(Epoch, Arc<std::sync::RwLock<Document>>)> {
+        self.stale
+            .read_or_recover()
+            .as_ref()
+            .map(|view| (view.epoch, Arc::clone(&view.doc)))
+    }
 }