1 files changed, 78 insertions, 0 deletions
diff --git a/src/compaction/local.rs b/src/compaction/local.rs
index 6290a5e..869c94e 100644
--- a/src/compaction/local.rs
+++ b/src/compaction/local.rs
@@ -5,20 +5,25 @@
 //! archive and rebuilds snapshots when needed.
 
 use std::collections::HashMap;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Mutex, MutexGuard};
 
 use automerge::AutoCommit;
 use serde::Serialize;
+use tokio::sync::oneshot;
 use tokio::time::Instant;
 
 use super::records::{unix_millis, Retention};
 use super::snapshot::Snapshot;
 use crate::replica::Epoch;
+use crate::transport::SnapshotReply;
 
 #[derive(Default)]
 pub(crate) struct Compaction {
     // A std mutex: every access is short and never spans an await.
     inner: Mutex<Inner>,
+    /// Set while this node saves its document for a peer; one at a time.
+    serving: AtomicBool,
 }
 
 #[derive(Default)]
@@ -29,6 +34,27 @@ struct Inner {
     first_seen: HashMap<String, Instant>,
     status: HashMap<String, LocalStatus>,
     archive: Vec<Archived>,
+    /// The snapshot request this node is waiting for (peer, wanted epoch).
+    awaiting: Option<(String, Epoch, oneshot::Sender<SnapshotReply>)>,
+    catch_up: Option<CatchUp>,
+}
+
+/// Fetching a newer epoch from a peer.
+#[derive(Clone, Debug)]
+pub(crate) struct CatchUp {
+    pub(crate) target: Epoch,
+    pub(crate) attempts: u32,
+    pub(crate) last_attempt: Instant,
+    pub(crate) error: Option<String>,
+}
+
+#[derive(Clone, Debug, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub(crate) struct CatchUpReport {
+    pub(crate) target: String,
+    pub(crate) attempts: u32,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub(crate) error: Option<String>,
 }
 
 /// What happened to a compaction on this node.
@@ -118,6 +144,58 @@ impl Compaction {
         prune(&mut self.lock().archive, current_generation);
     }
 
+    /// Registers interest in a snapshot reply from `peer`, replacing any earlier one.
+    pub(crate) fn await_reply(
+        &self,
+        peer: &str,
+        wanted: Epoch,
+    ) -> oneshot::Receiver<SnapshotReply> {
+        let (sender, receiver) = oneshot::channel();
+        self.lock().awaiting = Some((peer.to_owned(), wanted, sender));
+        receiver
+    }
+
+    /// Hands a reply to the waiting fetch. False if nobody asked that peer.
+    pub(crate) fn deliver_reply(&self, peer: &str, reply: SnapshotReply) -> bool {
+        let mut inner = self.lock();
+        match inner.awaiting.take() {
+            Some((expected, _, sender)) if expected == peer => sender.send(reply).is_ok(),
+            other => {
+                inner.awaiting = other;
+                false
+            }
+        }
+    }
+
+    pub(crate) fn stop_awaiting(&self) {
+        self.lock().awaiting = None;
+    }
+
+    /// Claims the right to serve one snapshot; release with [`Compaction::done_serving`].
+    pub(crate) fn try_start_serving(&self) -> bool {
+        !self.serving.swap(true, Ordering::AcqRel)
+    }
+
+    pub(crate) fn done_serving(&self) {
+        self.serving.store(false, Ordering::Release);
+    }
+
+    pub(crate) fn catch_up(&self) -> Option<CatchUp> {
+        self.lock().catch_up.clone()
+    }
+
+    pub(crate) fn set_catch_up(&self, catch_up: Option<CatchUp>) {
+        self.lock().catch_up = catch_up;
+    }
+
+    pub(crate) fn catch_up_report(&self) -> Option<CatchUpReport> {
+        self.lock().catch_up.as_ref().map(|catch_up| CatchUpReport {
+            target: catch_up.target.to_string(),
+            attempts: catch_up.attempts,
+            error: catch_up.error.clone(),
+        })
+    }
+
     pub(crate) fn archived_epochs(&self) -> Vec<ArchivedEpoch> {
         self.lock()
             .archive