document: reject writes during rejoin

2 files changed, 23 insertions, 5 deletions

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);