document: reject writes during rejoin

9 files changed, 107 insertions, 36 deletions

diff --git a/src/api/mod.rs b/src/api/mod.rs
index 528df35..31f4fea 100644
--- a/src/api/mod.rs
+++ b/src/api/mod.rs
@@ -202,6 +202,7 @@ pub(crate) fn store_status(error: StoreError, context: &str) -> ApiError {
             StatusCode::CONFLICT
         }
         StoreError::NotFound => StatusCode::NOT_FOUND,
+        StoreError::Rejoining => StatusCode::SERVICE_UNAVAILABLE,
         other => {
             error!(error = %format!("{other:#}"), "{context}");
             return StatusCode::INTERNAL_SERVER_ERROR.into();
diff --git a/src/document/replica.rs b/src/document/replica.rs
index 27c81ad..61f2b3b 100644
--- a/src/document/replica.rs
+++ b/src/document/replica.rs
@@ -113,6 +113,12 @@ pub(crate) struct Replica {
     /// run: messages carry it, so a late one can't flip a peer's view of our epoch back.
     pub(crate) epoch_seq: u32,
     pub(crate) doc: Document,
+    /// Whether writes go in (`AppState::mutate_doc`). False while a rejoin syncs a fresh
+    /// document: a write would land in a document the epoch's snapshot hasn't reached
+    /// yet, or be thrown away with it if the rejoin is abandoned. Cleared and set again
+    /// in the same lock as the swap (`Host::switch_epoch`), so no write checked
+    /// before a swap lands after it.
+    pub(crate) writable: bool,
     heads: Vec<ChangeHash>,
     stamp: watch::Sender<Stamp>,
 }
@@ -138,6 +144,7 @@ impl Replica {
             epoch,
             epoch_seq: 0,
             doc,
+            writable: true,
             heads: Vec::new(),
             stamp,
         };
diff --git a/src/document/store.rs b/src/document/store.rs
index 0a3c8ad..45d46cd 100644
--- a/src/document/store.rs
+++ b/src/document/store.rs
@@ -94,6 +94,10 @@ pub(crate) enum StoreError {
     /// An index past a list's end (`length` items).
     #[error("index {index} is past the list's end ({length} items)")]
     PastTheEnd { index: usize, length: usize },
+    /// The node is rejoining a newer epoch: no writes until it is done
+    /// (`Replica::writable`).
+    #[error("the node is rejoining a newer epoch; writes are refused until it is done")]
+    Rejoining,
     #[error(transparent)]
     Doc(#[from] automerge::AutomergeError),
     #[error(transparent)]
diff --git a/src/remake/catchup.rs b/src/remake/catchup.rs
index 695e4d7..7269d8a 100644
--- a/src/remake/catchup.rs
+++ b/src/remake/catchup.rs
@@ -29,6 +29,7 @@
 
 use std::collections::{BTreeMap, HashMap};
 use std::convert::Infallible;
+use std::sync::Arc;
 use std::time::Duration;
 
 use crate::document::{ActorId, ChangeHash, Document, ReadDoc};
@@ -43,6 +44,7 @@ use super::quarantine;
 use super::snapshot::{self, EpochMetadata};
 use crate::clock::UnixMillis;
 use crate::document::replica::Epoch;
+use crate::document::store::StoreError;
 
 /// Pause before trying the same target again after a failed rejoin.
 const RETRY: Duration = Duration::from_secs(30);
@@ -179,8 +181,10 @@ async fn start<H: Host>(host: &H, local: Epoch, target: Epoch, attempts: u32) {
     let was_ready = host.suspend().await;
     let actor = ActorId::random();
     host.remake().add_own_actor(actor.clone());
+    // The fresh document takes no writes until the rejoin is done: one checked as
+    // ready just before the suspension above would land in it (`Replica::writable`).
     let old = match host
-        .switch_epoch(target, move |_live| {
+        .switch_epoch(target, false, move |_live| {
             Ok::<_, Infallible>(Document::new().with_actor(actor))
         })
         .await
@@ -318,6 +322,7 @@ async fn finish<H: Host>(host: &H) {
         Err(join_error) => {
             // The old document is gone with the panicked task; nothing to restore.
             error!(error = %join_error, "Rejoin task failed; staying on the new epoch");
+            host.accept_writes().await;
             if was_ready {
                 host.resume().await;
             }
@@ -331,26 +336,39 @@ async fn finish<H: Host>(host: &H) {
     let mut keep_archive = true;
     let mut old = Some(old);
 
+    // The new epoch is verified: writes go in again, the replay first.
+    host.accept_writes().await;
+
     match plan {
         Plan::Nothing => {}
         Plan::Replay(paths, translator) => {
-            let document = old.take().expect("old document");
-            let (replayed, document) = host
-                .commit_doc(move |doc| {
-                    let stats = delta::replay_paths(&document, doc, &paths, &translator);
-                    Ok::<_, Infallible>((stats, document))
+            let document = Arc::new(old.take().expect("old document"));
+            let replayed = {
+                let document = Arc::clone(&document);
+                let paths = paths.clone();
+                host.commit_doc(move |doc| {
+                    Ok::<_, StoreError>(delta::replay_paths(&document, doc, &paths, &translator))
                 })
                 .await
-                .unwrap_or_else(|never| match never {});
-            stats = replayed;
-            // Writes the remakes since had no place for.
-            quarantine_id = quarantine::add(
-                host,
-                previous,
-                target,
-                &document,
-                std::mem::take(&mut stats.quarantined),
-            );
+            };
+            // The commit's task held the only other reference, and it is done.
+            let document = Arc::into_inner(document).expect("the replay's reference is gone");
+            let unplaced = match replayed {
+                Ok(replayed) => {
+                    stats = replayed;
+                    // Writes the remakes since had no place for.
+                    std::mem::take(&mut stats.quarantined)
+                }
+                // Not after `accept_writes`; but if so, nothing is lost.
+                Err(error) => {
+                    let reason = format!("couldn't replay: {error}");
+                    paths
+                        .into_iter()
+                        .map(|path| (path, reason.clone()))
+                        .collect()
+                }
+            };
+            quarantine_id = quarantine::add(host, previous, target, &document, unplaced);
             old = Some(document);
         }
         Plan::Quarantine { paths, reason } => {
@@ -525,7 +543,7 @@ async fn abandon<H: Host>(host: &H, rejoin: Rejoin, reason: String) {
     warn!(target = %target, reason = %reason, "Rejoin failed; restoring the previous epoch");
 
     if let Err(never) = host
-        .switch_epoch(previous, move |_live| Ok::<_, Infallible>(old))
+        .switch_epoch(previous, true, move |_live| Ok::<_, Infallible>(old))
         .await
     {
         match never {}
diff --git a/src/remake/mod.rs b/src/remake/mod.rs
index e9d3021..04e9d8a 100644
--- a/src/remake/mod.rs
+++ b/src/remake/mod.rs
@@ -53,7 +53,7 @@ use axum::{
 use tokio::sync::broadcast;
 
 use crate::document::replica::Epoch;
-use crate::document::store::ChangeSet;
+use crate::document::store::{ChangeSet, StoreError};
 use crate::document::{Document, Transaction};
 
 /// Remakes and their aftermath on this node, in numbers (`/admin/stats`).
@@ -184,12 +184,18 @@ pub(crate) trait Host: Clone + Send + Sync + 'static {
         F: FnOnce(&Document) -> T + Send + 'static,
         T: Send + 'static;
 
-    /// One committed change (rolled back on error), published and replicated.
+    /// One committed change (rolled back on error), published and replicated. Refused
+    /// (`StoreError::Rejoining`) while the document isn't writable: from a rejoin's
+    /// swap until [`Host::accept_writes`].
     fn commit_doc<T, E, F>(&self, f: F) -> impl Future<Output = Result<T, E>> + Send
     where
         F: FnOnce(&mut Transaction<'_>) -> Result<T, E> + Send + 'static,
         T: Send + 'static,
-        E: Send + 'static;
+        E: From<StoreError> + Send + 'static;
+
+    /// Lets writes in again, after a rejoin swapped in a document that took none
+    /// ([`Host::switch_epoch`] with `writable` false).
+    fn accept_writes(&self) -> impl Future<Output = ()> + Send;
 
     /// Applies a change downloaded from peers (`fetch_change`), then publishes it like
     /// a change received by sync.
@@ -199,11 +205,13 @@ pub(crate) trait Host: Clone + Send + Sync + 'static {
     ) -> impl Future<Output = Result<(), String>> + Send;
 
     /// Replaces the document and its epoch atomically. Under the write lock, `f` gets
-    /// the live document and returns the new one. Afterwards the node resets every
-    /// peer's sync state and publishes the change. Returns the old document.
+    /// the live document and returns the new one, which takes writes if `writable`
+    /// (a rejoin's fresh document doesn't). Afterwards the node resets every peer's
+    /// sync state and publishes the change. Returns the old document.
     fn switch_epoch<E, F>(
         &self,
         epoch: Epoch,
+        writable: bool,
         f: F,
     ) -> impl Future<Output = Result<Document, E>> + Send
     where
diff --git a/src/remake/quarantine.rs b/src/remake/quarantine.rs
index 3ad5005..bb34d73 100644
--- a/src/remake/quarantine.rs
+++ b/src/remake/quarantine.rs
@@ -12,7 +12,7 @@
 //!
 //! Quarantine lives in memory on the node that holds it.
 
-use std::convert::Infallible;
+use std::sync::Arc;
 
 use crate::document::Document;
 use axum::{
@@ -33,7 +33,7 @@ use super::local::Quarantined;
 use crate::api::error_response;
 use crate::clock::UnixMillis;
 use crate::document::replica::Epoch;
-use crate::document::store;
+use crate::document::store::{self, StoreError};
 
 #[derive(Debug, Serialize)]
 #[serde(rename_all = "camelCase")]
@@ -167,15 +167,30 @@ pub(crate) async fn replay<H: Host>(State(host): State<H>, Path(id): Path<String
         return not_found();
     };
 
-    let (stats, entry) = host
-        .commit_doc(move |doc| {
+    let entry = Arc::new(entry);
+    let replayed = {
+        let entry = Arc::clone(&entry);
+        host.commit_doc(move |doc| {
             // As they are: the operator decided they belong at their old paths.
-            let stats =
-                delta::replay_paths(&entry.document, doc, &entry.paths, &Translator::default());
-            Ok::<_, Infallible>((stats, entry))
+            Ok::<_, StoreError>(delta::replay_paths(
+                &entry.document,
+                doc,
+                &entry.paths,
+                &Translator::default(),
+            ))
         })
         .await
-        .unwrap_or_else(|never| match never {});
+    };
+    // The commit's task held the only other reference, and it is done.
+    let entry = Arc::into_inner(entry).expect("the replay's reference is gone");
+    let stats = match replayed {
+        Ok(stats) => stats,
+        Err(error) => {
+            // Rejoining since the check above: keep the entry for later.
+            host.remake().restore_quarantined(entry);
+            return error_response(StatusCode::SERVICE_UNAVAILABLE, &error.to_string(), None);
+        }
+    };
 
     if stats.failed > 0 && stats.replayed == 0 {
         // Nothing went in; keep the entry for another look.
diff --git a/src/remake/switch.rs b/src/remake/switch.rs
index e60c942..9b03789 100644
--- a/src/remake/switch.rs
+++ b/src/remake/switch.rs
@@ -303,7 +303,7 @@ pub(super) async fn swap_in<H: Host>(
     let late_stats = std::sync::Arc::new(std::sync::Mutex::new(ReplayStats::default()));
     let late_out = std::sync::Arc::clone(&late_stats);
     let old = host
-        .switch_epoch(target, move |live| {
+        .switch_epoch(target, true, move |live| {
             // Only one worker switches, but be sure the copy still describes the live document.
             if !crate::document::changes::missing(live, &copy_heads).is_empty() {
                 return Err("the document was replaced while preparing the switch".to_owned());
diff --git a/src/state/host.rs b/src/state/host.rs
index 24a989e..6bb9d3b 100644
--- a/src/state/host.rs
+++ b/src/state/host.rs
@@ -12,7 +12,7 @@ use tracing::debug;
 
 use super::{AppState, HelloState, NodeLifecycle, StaleView, Standing};
 use crate::document::replica::Epoch;
-use crate::document::store::ChangeSet;
+use crate::document::store::{ChangeSet, StoreError};
 use crate::document::{ChangeHash, Document, ReadDoc, Transaction};
 use crate::remake::{Host, PeerSummary, RemakeState};
 
@@ -220,12 +220,22 @@ impl Host for AppState {
     where
         F: FnOnce(&mut Transaction<'_>) -> Result<T, E> + Send + 'static,
         T: Send + 'static,
-        E: Send + 'static,
+        E: From<StoreError> + Send + 'static,
     {
         let state = self.clone();
         async move { state.mutate_doc(f).await }
     }
 
+    fn accept_writes(&self) -> impl Future<Output = ()> + Send {
+        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)
+                .await
+                .expect("document write task panicked");
+        }
+    }
+
     fn apply_change(
         &self,
         change: crate::document::Change,
@@ -255,6 +265,7 @@ impl Host for AppState {
     fn switch_epoch<E, F>(
         &self,
         epoch: Epoch,
+        writable: bool,
         f: F,
     ) -> impl Future<Output = Result<Document, E>> + Send
     where
@@ -271,6 +282,7 @@ impl Host for AppState {
                     let old_doc = std::mem::replace(&mut replica.doc, new_doc);
                     replica.epoch = epoch;
                     replica.epoch_seq += 1;
+                    replica.writable = writable;
                     replica.refresh_heads();
                     // Updated under the same lock, so the cached epoch never disagrees
                     // with the document for anyone holding the lock.
diff --git a/src/state/mod.rs b/src/state/mod.rs
index d8bdac5..9ee08b7 100644
--- a/src/state/mod.rs
+++ b/src/state/mod.rs
@@ -38,7 +38,7 @@ use tracing::{debug, error, warn};
 
 use crate::document::lock::DocumentLock;
 use crate::document::replica::{Epoch, Stamp};
-use crate::document::store::ChangeSet;
+use crate::document::store::{ChangeSet, StoreError};
 use crate::document::{Document, Transaction};
 use crate::locks::MutexExt;
 use crate::remake::{Host, RemakeState};
@@ -393,20 +393,26 @@ 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.
+    /// cancelled request can't leave a committed write unannounced. Refused
+    /// (`StoreError::Rejoining`, `f` not called) while the document isn't writable.
     pub(crate) async fn mutate_doc<T, E>(
         &self,
         f: impl FnOnce(&mut Transaction<'_>) -> Result<T, E> + Send + 'static,
     ) -> Result<T, E>
     where
         T: Send + 'static,
-        E: Send + 'static,
+        E: From<StoreError> + Send + 'static,
     {
         let state = self.clone();
 
         tokio::task::spawn_blocking(move || {
             let result = {
                 let mut replica = state.doc.blocking_write();
+                // 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));
+                }
                 // One API write = one change; on error, its partial writes (e.g. the old
                 // value deleted, the new one not written) are dropped.
                 let result = replica.transact(state.clock.now_millis(), f);