add epoch switch

1 files changed, 67 insertions, 1 deletions

diff --git a/src/state.rs b/src/state.rs
index 30efbc6..44cba2a 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -12,7 +12,7 @@ use tokio::sync::{Mutex, RwLock, broadcast, mpsc, watch};
 use tokio::time::Instant;
 use tracing::{debug, error, warn};
 
-use crate::compaction::{Host, PeerSummary};
+use crate::compaction::{Compaction, Host, PeerSummary};
 use crate::firewall::Firewall;
 use crate::replica::{Epoch, Replica};
 use crate::sync::{ChangeSet, SyncEvent};
@@ -279,6 +279,9 @@ pub(crate) struct AppState {
     pub(crate) change_events: broadcast::Sender<Arc<ChangeSet>>,
 
     pub(crate) timing: Timing,
+
+    /// Per-node compaction state (prebuilt snapshots, switch status, archive).
+    pub(crate) compaction: Arc<Compaction>,
 }
 
 impl AppState {
@@ -495,6 +498,31 @@ impl AppState {
     }
 }
 
+impl AppState {
+    /// 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.
+    async fn after_epoch_switch(&self) {
+        let peers: Vec<(String, Arc<Mutex<PeerState>>)> = self
+            .peer_states
+            .read()
+            .await
+            .iter()
+            .map(|(id, peer)| (id.clone(), Arc::clone(peer)))
+            .collect();
+
+        for (_, peer) in &peers {
+            peer.lock().await.reset_sync();
+        }
+
+        self.publish_change();
+
+        for (id, _) in peers {
+            let _ = self.sync_events.send(SyncEvent::PeerNeedsSync(id));
+        }
+    }
+}
+
 pub(crate) fn unix_millis() -> u64 {
     std::time::SystemTime::now()
         .duration_since(std::time::UNIX_EPOCH)
@@ -504,6 +532,10 @@ pub(crate) fn unix_millis() -> u64 {
 
 /// The node as seen by the compaction module.
 impl Host for AppState {
+    fn compaction(&self) -> &Compaction {
+        &self.compaction
+    }
+
     fn node_id(&self) -> String {
         self.local_id().to_owned()
     }
@@ -563,4 +595,38 @@ impl Host for AppState {
         let state = self.clone();
         async move { state.mutate_doc(f).await }
     }
+
+    fn switch_epoch<E, F>(
+        &self,
+        epoch: Epoch,
+        f: F,
+    ) -> impl Future<Output = Result<AutoCommit, E>> + Send
+    where
+        F: FnOnce(&mut AutoCommit) -> Result<AutoCommit, E> + Send + 'static,
+        E: Send + 'static,
+    {
+        let state = self.clone();
+        async move {
+            let swapped = {
+                let state = state.clone();
+                tokio::task::spawn_blocking(move || -> Result<AutoCommit, E> {
+                    let mut replica = state.doc.blocking_write();
+                    let new_doc = f(&mut replica.doc)?;
+                    let old_doc = std::mem::replace(&mut replica.doc, new_doc);
+                    replica.epoch = epoch;
+                    // Updated under the same lock, so the cached epoch never disagrees
+                    // with the document for anyone holding the lock.
+                    state.epoch_tx.send_replace(epoch);
+                    Ok(old_doc)
+                })
+                .await
+                .expect("epoch switch task panicked")
+            };
+
+            if swapped.is_ok() {
+                state.after_epoch_switch().await;
+            }
+            swapped
+        }
+    }
 }