1 files changed, 307 insertions, 0 deletions
diff --git a/src/compaction/catchup.rs b/src/compaction/catchup.rs
new file mode 100644
index 0000000..f166f56
--- /dev/null
+++ b/src/compaction/catchup.rs
@@ -0,0 +1,307 @@
+//! Catching up with a newer epoch from a peer.
+//!
+//! Peers on another epoch don't sync, so a node that missed a switch (new, late, or
+//! one whose own build failed) can't learn the new epoch through ordinary sync. Instead,
+//! once an established peer advertises a newer generation, the node asks it for its
+//! whole document (`SnapshotRequest`), checks that the document starts with the
+//! advertised snapshot change, and switches to it:
+//!
+//! * a node that wrote nothing the snapshot lacks (e.g. a brand-new node) adopts the
+//!   document as is, whatever the gap;
+//! * a node exactly one epoch behind replays its own changes the snapshot lacks, found
+//!   with the snapshot's per-actor clock;
+//! * a node further behind with such changes stays where it is and reports it; that is
+//!   the returning-node case (replay window or quarantine, step 7).
+//!
+//! Only one fetch runs at a time, and the worker retries after `RETRY`. The reply is the
+//! peer's full document in one message; chunking comes later.
+
+use std::time::Duration;
+
+use automerge::{ActorId, AutoCommit};
+use tokio::time::Instant;
+use tracing::{debug, error, info, warn};
+
+use super::delta::{self, ReplayStats};
+use super::local::{CatchUp, LocalStatus};
+use super::records::unix_millis;
+use super::snapshot;
+use super::switch::swap_in;
+use super::Host;
+use crate::replica::Epoch;
+use crate::transport::{ParsedMessage, SnapshotReply, SnapshotRequest};
+
+const RETRY: Duration = Duration::from_secs(10);
+const FETCH_TIMEOUT: Duration = Duration::from_secs(60);
+
+/// Called by the switch worker when there is nothing to switch to locally.
+pub(super) async fn check<H: Host>(host: &H) {
+    let local = host.current_epoch();
+
+    let ahead: Vec<(String, Epoch)> = host
+        .peer_summaries()
+        .await
+        .into_iter()
+        .filter(|peer| peer.active && peer.established)
+        .filter_map(|peer| peer.epoch.map(|epoch| (peer.id, epoch)))
+        .filter(|(_, epoch)| epoch.generation > local.generation)
+        .collect();
+
+    let Some(generation) = ahead.iter().map(|(_, epoch)| epoch.generation).max() else {
+        if host.compaction().catch_up().is_some() {
+            host.compaction().set_catch_up(None);
+        }
+        return;
+    };
+
+    let mut targets: Vec<Epoch> = ahead
+        .iter()
+        .filter(|(_, epoch)| epoch.generation == generation)
+        .map(|(_, epoch)| *epoch)
+        .collect();
+    targets.sort_by_key(|epoch| epoch.id);
+    targets.dedup();
+    let [target] = targets.as_slice() else {
+        let listed: Vec<String> = targets.iter().map(ToString::to_string).collect();
+        warn!(epochs = ?listed, "Peers disagree about the newest epoch; not catching up");
+        return;
+    };
+    let target = *target;
+
+    // Back off between attempts for the same target.
+    let previous = host
+        .compaction()
+        .catch_up()
+        .filter(|catch_up| catch_up.target == target);
+    if previous
+        .as_ref()
+        .is_some_and(|catch_up| catch_up.last_attempt.elapsed() < RETRY)
+    {
+        return;
+    }
+    let attempts = previous.map_or(0, |catch_up| catch_up.attempts) + 1;
+
+    // Rotate through the peers on the target epoch.
+    let sources: Vec<&String> = ahead
+        .iter()
+        .filter(|(_, epoch)| *epoch == target)
+        .map(|(peer, _)| peer)
+        .collect();
+    let source = sources[(attempts as usize - 1) % sources.len()].clone();
+
+    host.compaction().set_catch_up(Some(CatchUp {
+        target,
+        attempts,
+        last_attempt: Instant::now(),
+        error: None,
+    }));
+
+    info!(local = %local, target = %target, peer = %source, attempt = attempts, "Fetching newer epoch from peer");
+
+    let result = match fetch(host, &source, local, target).await {
+        Ok(document) => adopt(host, local, target, document).await,
+        Err(error) => Err(error),
+    };
+
+    match result {
+        Ok(()) => host.compaction().set_catch_up(None),
+        Err(error) => {
+            warn!(target = %target, peer = %source, error = %error, "Catching up failed; will retry");
+            host.compaction().set_catch_up(Some(CatchUp {
+                target,
+                attempts,
+                last_attempt: Instant::now(),
+                error: Some(error),
+            }));
+        }
+    }
+}
+
+async fn fetch<H: Host>(
+    host: &H,
+    peer: &str,
+    local: Epoch,
+    target: Epoch,
+) -> Result<Vec<u8>, String> {
+    let reply = host.compaction().await_reply(peer, target);
+    let request = ParsedMessage::SnapshotRequest(SnapshotRequest {
+        epoch: local,
+        wanted: target,
+    });
+    if let Err(error) = host.send_message(peer, request).await {
+        host.compaction().stop_awaiting();
+        return Err(format!("sending the request failed: {error:#}"));
+    }
+
+    let reply = match tokio::time::timeout(FETCH_TIMEOUT, reply).await {
+        Ok(Ok(reply)) => reply,
+        Ok(Err(_)) => return Err("request superseded".to_owned()),
+        Err(_) => {
+            host.compaction().stop_awaiting();
+            return Err(format!("no reply within {}s", FETCH_TIMEOUT.as_secs()));
+        }
+    };
+
+    if reply.epoch != target || reply.document.is_empty() {
+        return Err(format!(
+            "peer is on {} and did not send {target}",
+            reply.epoch
+        ));
+    }
+    Ok(reply.document)
+}
+
+/// Verifies the fetched document, carries over this node's own missing changes, and
+/// switches to it.
+async fn adopt<H: Host>(
+    host: &H,
+    local: Epoch,
+    target: Epoch,
+    document: Vec<u8>,
+) -> Result<(), String> {
+    // Copy under the read lock only; everything else happens without a lock.
+    let copy = host.read_doc(|doc| doc.clone()).await;
+    let bytes = document.len();
+
+    let prepared = tokio::task::spawn_blocking(move || -> Result<Prepared, String> {
+        let mut fetched =
+            AutoCommit::load(&document).map_err(|error| format!("invalid document: {error}"))?;
+        let metadata = snapshot::verify_document(&mut fetched, target)?;
+
+        let mut copy = copy;
+        let own = copy.get_actor().clone();
+        let copy_heads = copy.get_heads();
+        let had_history = !copy_heads.is_empty();
+        let missing = delta::missing_from_snapshot(&mut copy, &metadata.clock, &own);
+
+        let mut doc = fetched.with_actor(ActorId::random());
+        let stats = if missing.own == 0 {
+            ReplayStats::default()
+        } else if metadata.previous == local {
+            delta::replay(&mut copy, &missing.before, &copy_heads, &mut doc)
+        } else {
+            return Err(format!(
+                "{} local changes are not in {target}, and this node is more than one epoch behind (on {local}, {target} follows {}); not replaying them automatically",
+                missing.own, metadata.previous
+            ));
+        };
+
+        Ok(Prepared {
+            doc,
+            copy_heads,
+            stats,
+            own_missing: missing.own,
+            had_history,
+            metadata,
+        })
+    })
+    .await
+    .map_err(|error| format!("catch-up task failed: {error}"))??;
+
+    let Prepared {
+        doc,
+        copy_heads,
+        stats,
+        own_missing,
+        had_history,
+        metadata,
+    } = prepared;
+
+    let old = swap_in(host, target, doc, copy_heads).await?;
+
+    if had_history {
+        host.compaction().archive(
+            local,
+            old,
+            metadata.policy.archive_retention.clone(),
+            target.generation,
+        );
+    }
+
+    // Visible in the compaction's report, if this node has its record.
+    host.compaction().set_status(
+        &metadata.compaction_id,
+        LocalStatus {
+            switched_at: Some(unix_millis()),
+            replayed_paths: Some(stats.replayed),
+            failed_paths: Some(stats.failed),
+            error: None,
+        },
+    );
+
+    info!(
+        from = %local,
+        to = %target,
+        document_bytes = bytes,
+        own_changes = own_missing,
+        replayed = stats.replayed,
+        failed = stats.failed,
+        "Caught up with newer epoch"
+    );
+    Ok(())
+}
+
+struct Prepared {
+    doc: AutoCommit,
+    copy_heads: Vec<automerge::ChangeHash>,
+    stats: ReplayStats,
+    own_missing: usize,
+    had_history: bool,
+    metadata: snapshot::EpochMetadata,
+}
+
+/// Answers a peer's snapshot request with this node's document, if it is on the wanted
+/// epoch. Only for members of our cluster; one at a time.
+pub(crate) async fn serve_snapshot<H: Host>(host: H, peer: String, request: SnapshotRequest) {
+    if !host.is_cluster_member(&peer).await {
+        debug!(peer = %peer, "Ignoring snapshot request from a non-member");
+        return;
+    }
+    if !host.compaction().try_start_serving() {
+        debug!(peer = %peer, "Already serving a snapshot; the peer will retry");
+        return;
+    }
+
+    let current = host.current_epoch();
+    let wanted = request.wanted;
+
+    let document = if current == wanted {
+        // Clone under the read lock, save without one. The history check guards against
+        // a switch between reading the epoch and cloning.
+        let mut doc = host.read_doc(|doc| doc.clone()).await;
+        tokio::task::spawn_blocking(move || {
+            if doc
+                .get_change_by_hash(&automerge::ChangeHash(wanted.id))
+                .is_some()
+            {
+                doc.save()
+            } else {
+                Vec::new()
+            }
+        })
+        .await
+        .unwrap_or_else(|error| {
+            error!(error = %error, "Saving the document for a peer failed");
+            Vec::new()
+        })
+    } else {
+        Vec::new()
+    };
+    host.compaction().done_serving();
+
+    info!(peer = %peer, wanted = %wanted, document_bytes = document.len(), "Serving snapshot");
+    let reply = ParsedMessage::SnapshotReply(SnapshotReply {
+        epoch: if document.is_empty() { current } else { wanted },
+        document,
+    });
+    if let Err(error) = host.send_message(&peer, reply).await {
+        warn!(peer = %peer, error = %format!("{error:#}"), "Sending snapshot failed");
+    }
+}
+
+/// Hands a peer's snapshot reply to the waiting fetch.
+pub(crate) fn receive_snapshot_reply<H: Host>(host: &H, peer: &str, reply: SnapshotReply) {
+    if !host.compaction().deliver_reply(peer, reply) {
+        debug!(peer = %peer, "Unexpected snapshot reply; dropped");
+    }
+}