From 5640bacf735c6728acf4bc22217b0c07ac4bf7f9 Mon Sep 17 00:00:00 2001 From: tv Date: Sun, 1 Nov 2026 19:35:20 +0100 Subject: document: allow writes during rejoin --- src/state/host.rs | 269 ++++++++++++++++++++++++++++++++++-------------------- 1 file changed, 169 insertions(+), 100 deletions(-) (limited to 'src/state/host.rs') 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, 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 + 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 + 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(&self, f: F) -> impl Future> + Send - where - F: FnOnce(&mut Document) -> T + Send + 'static, - T: Send + 'static, - { + fn begin_rejoin( + &self, + target: Epoch, + actor: crate::document::ActorId, + ) -> impl Future + 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> + Send { + fn rejoin_document(&self) -> impl Future> + 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 + Send { + fn finish_rejoin( + &self, + old: Option, + f: F, + ) -> impl Future + 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, + ) -> impl Future + 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( &self, epoch: Epoch, - writable: bool, f: F, ) -> impl Future> + Send where @@ -321,25 +408,7 @@ impl Host for AppState { tokio::task::spawn_blocking(move || -> Result { 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") -- cgit v1.3.1