path: root/src/sync.rs

//! Automerge synchronization: the sync worker, delivery with retries, inbound sync
//! messages, and the change publisher that feeds incremental streams.

use std::collections::HashSet;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;

use anyhow::{anyhow, bail, Context, Result};
use automerge::sync::SyncDoc;
use tokio::sync::mpsc;
use tokio::time::Instant;
use tracing::{debug, info, instrument, trace, warn};

use crate::peer::{peer_cluster, send_hello_once};
use crate::state::{AppState, Liveness, PeerMembership, PendingOutbound};
use crate::transport::{self, ParsedMessage};

const INITIAL_SYNC_RETRY: Duration = Duration::from_millis(500);
const MAX_SYNC_RETRY: Duration = Duration::from_secs(30);

#[derive(Debug)]
pub(crate) enum SyncEvent {
    PeerValidated(String),
    DocumentChanged,
    MessageReceived(String),
    PeerNeedsSync(String),
    RetryDelivery(String),
}

/// Store paths touched by one batch of document changes.
#[derive(Debug)]
pub(crate) struct ChangeSet {
    pub(crate) paths: Vec<Vec<String>>,
}

#[instrument(skip(state), fields(peer = %peer, retry_delivery))]
async fn deliver_sync_to_peer(state: &AppState, peer: &str, retry_delivery: bool) -> Result<()> {
    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
        trace!("Peer was forgotten; skipping synchronization");
        return Ok(());
    };

    let validated_cluster = {
        let peer_state = peer_state_mutex.lock().await;

        if matches!(peer_state.liveness, Liveness::Dormant { .. }) {
            // Reactivation re-handshakes, and PeerValidated then restarts sync.
            trace!("Skipping synchronization with dormant peer");
            return Ok(());
        }

        if peer_state
            .epoch
            .is_some_and(|epoch| epoch != state.local_epoch())
        {
            // Unrelated document histories must never be synced.
            trace!("Skipping synchronization with peer on another epoch");
            return Ok(());
        }

        match &peer_state.membership {
            PeerMembership::Unknown => return Ok(()),
            PeerMembership::Member { cluster_id } => cluster_id.clone(),
        }
    };

    let cluster_id = state
        .cluster_id()
        .await
        .ok_or_else(|| anyhow!("Cannot sync without an established cluster ID"))?;

    if validated_cluster != cluster_id {
        bail!("Peer {peer} belongs to a different cluster");
    }

    let existing = {
        let mut guard = peer_state_mutex.lock().await;
        let peer_state = &mut *guard;
        match &peer_state.pending_outbound_sync {
            Some(entry) if retry_delivery => Some((entry.message.clone(), entry.next_attempt)),
            Some(_) => {
                peer_state.needs_generate = true;
                return Ok(());
            }
            None => None,
        }
    };

    let encoded = match existing {
        Some((message, next_attempt)) => {
            if Instant::now() < next_attempt {
                return Ok(());
            }
            message
        }
        None if retry_delivery => return Ok(()),
        None => {
            let doc = Arc::clone(&state.doc);
            let psm = Arc::clone(&peer_state_mutex);
            let transport = Arc::clone(&state.transport);
            let peer_owned = peer.to_owned();

            let generated = tokio::task::spawn_blocking(
                move || -> Result<Option<transport::EncodedMessage>> {
                    // Label the message with the epoch of the document it was generated
                    // from, read under the same lock.
                    let (epoch, message) = {
                        let mut replica = doc.blocking_write();
                        let mut peer_state = psm.blocking_lock();
                        peer_state.needs_generate = false;
                        let epoch = replica.epoch;
                        let sync_doc = replica.doc.sync();
                        (
                            epoch,
                            sync_doc.generate_sync_message(&mut peer_state.sync_state),
                        )
                    };
                    let Some(message) = message else {
                        return Ok(None);
                    };
                    let envelope = transport::SyncEnvelope {
                        cluster_id,
                        epoch,
                        message,
                    };
                    transport
                        .encode(&peer_owned, &ParsedMessage::Sync(envelope))
                        .map(Some)
                },
            )
            .await
            .context("Sync generation task panicked")??;

            let Some(encoded) = generated else {
                return Ok(());
            };

            peer_state_mutex.lock().await.pending_outbound_sync = Some(PendingOutbound {
                message: encoded.clone(),
                retry_delay: INITIAL_SYNC_RETRY,
                next_attempt: Instant::now(),
            });

            encoded
        }
    };

    // No state locks are held across the send.
    match state.transport.send_encoded(&encoded).await {
        Ok(()) => {
            let regenerate = {
                let mut peer_state = peer_state_mutex.lock().await;
                peer_state.pending_outbound_sync = None;
                std::mem::take(&mut peer_state.needs_generate)
            };

            debug!(regenerate, "Outbound sync message accepted by Mycelium");

            if regenerate {
                state
                    .sync_events
                    .send(SyncEvent::PeerNeedsSync(peer.to_owned()))
                    .map_err(|error| anyhow!("Failed to notify sync worker: {error}"))?;
            }

            Ok(())
        }
        Err(error) => {
            let delay = {
                let mut peer_state = peer_state_mutex.lock().await;

                let Some(ref mut entry) = peer_state.pending_outbound_sync else {
                    debug!("Pending envelope discarded by a sync reset during delivery");
                    return Ok(());
                };

                entry.retry_delay = entry.retry_delay.saturating_mul(2).min(MAX_SYNC_RETRY);
                entry.next_attempt = Instant::now() + entry.retry_delay;
                entry.retry_delay
            };

            warn!(
                retry_ms = delay.as_millis(),
                error = %format!("{error:#}"),
                "Sync delivery failed; retaining the same outbound envelope"
            );

            let tx = state.sync_events.clone();
            let peer = peer.to_owned();
            tokio::spawn(async move {
                tokio::time::sleep(delay).await;
                let _ = tx.send(SyncEvent::RetryDelivery(peer));
            });

            Ok(())
        }
    }
}

#[instrument(skip(state, envelope), fields(peer = %peer, cluster_id = %envelope.cluster_id))]
pub(crate) async fn handle_sync(
    state: &AppState,
    peer: &str,
    envelope: transport::SyncEnvelope,
) -> Result<()> {
    debug!("Processing sync envelope");

    let peer_state_mutex = state.get_or_create_peer_state(peer).await;

    let needs_hello = {
        let mut peer_state = peer_state_mutex.lock().await;

        match &peer_state.membership {
            PeerMembership::Unknown => {
                let duplicate = peer_state.pending_inbound_sync.iter().any(|pending| {
                    pending.message == envelope.message
                        && pending.cluster_id == envelope.cluster_id
                        && pending.epoch == envelope.epoch
                });
                if !duplicate {
                    peer_state.pending_inbound_sync.push(envelope.clone());
                    info!(
                        pending_count = peer_state.pending_inbound_sync.len(),
                        "Queued sync envelope pending peer membership"
                    );
                } else {
                    trace!("Ignoring duplicate pending sync envelope");
                }
                true
            }
            PeerMembership::Member { cluster_id } if cluster_id != &envelope.cluster_id => {
                bail!(
                    "Sync cluster mismatch for peer {peer}: validated_cluster={cluster_id}, envelope_cluster={}",
                    envelope.cluster_id
                );
            }
            PeerMembership::Member { .. } => false,
        }
    };

    if needs_hello {
        // The peer treats us as a member but we don't know it, typically because we
        // restarted. Our HELLO makes it reset its sync state and reply.
        if let Err(error) = send_hello_once(state, peer).await {
            warn!(peer = %peer, error = %format!("{error:#}"), "HELLO to unknown sync sender failed; will retry");
        }
        return Ok(());
    }

    let expected_cluster = state
        .cluster_id()
        .await
        .ok_or_else(|| anyhow!("Missing expected cluster ID after peer membership"))?;
    if envelope.cluster_id != expected_cluster {
        bail!(
            "Incompatible cluster from {peer}: expected {}, received {}",
            expected_cluster,
            envelope.cluster_id
        );
    }

    let epoch_label = envelope.epoch;

    let applied = {
        let doc = Arc::clone(&state.doc);
        let peer_state_mutex = Arc::clone(&peer_state_mutex);
        let check_caught_up = !state.initial_sync_done.load(Ordering::Acquire);
        let epoch = envelope.epoch;
        let message = envelope.message;

        // Some(caught_up) when applied, None when the message belongs to another epoch.
        tokio::task::spawn_blocking(move || -> Result<Option<bool>> {
            let mut replica = doc.blocking_write();

            // Authoritative check, under the same lock as the apply: the document may
            // have switched epochs since the message was accepted.
            if replica.epoch != epoch {
                return Ok(None);
            }

            let mut peer_state = peer_state_mutex.blocking_lock();

            replica
                .doc
                .sync()
                .receive_sync_message(&mut peer_state.sync_state, message)
                .context("Automerge synchronization failed")?;
            peer_state.last_sync_rx = Some(Instant::now());

            Ok(Some(
                check_caught_up
                    && peer_state
                        .sync_state
                        .their_heads
                        .as_ref()
                        .is_some_and(|heads| {
                            heads
                                .iter()
                                .all(|h| replica.doc.get_change_by_hash(h).is_some())
                        }),
            ))
        })
        .await
        .context("Sync task panicked")??
    };

    let Some(caught_up) = applied else {
        debug!(epoch = %epoch_label, "Dropped sync message from another epoch");
        return Ok(());
    };

    if caught_up && !state.initial_sync_done.swap(true, Ordering::AcqRel) {
        info!(peer = %peer, "Initial sync complete");
    }

    state.notify_store_changed();

    state
        .sync_events
        .send(SyncEvent::MessageReceived(peer.to_owned()))
        .map_err(|error| anyhow!("Failed to notify sync worker: {error}"))?;

    Ok(())
}

pub(crate) async fn run_sync_worker(state: AppState, mut rx: mpsc::UnboundedReceiver<SyncEvent>) {
    let mut dirty_peers = HashSet::<String>::new();
    let mut retry_peers = HashSet::<String>::new();
    let mut document_dirty = false;

    info!("Starting event-driven synchronization worker");

    loop {
        if dirty_peers.is_empty() && retry_peers.is_empty() && !document_dirty {
            match rx.recv().await {
                Some(event) => match event {
                    SyncEvent::PeerValidated(peer)
                    | SyncEvent::MessageReceived(peer)
                    | SyncEvent::PeerNeedsSync(peer) => {
                        dirty_peers.insert(peer);
                    }
                    SyncEvent::RetryDelivery(peer) => {
                        retry_peers.insert(peer);
                    }
                    SyncEvent::DocumentChanged => {
                        document_dirty = true;
                    }
                },
                None => {
                    warn!("Synchronization event channel closed");
                    return;
                }
            }
        }

        // Coalesce event bursts.
        while let Ok(event) = rx.try_recv() {
            match event {
                SyncEvent::PeerValidated(peer)
                | SyncEvent::MessageReceived(peer)
                | SyncEvent::PeerNeedsSync(peer) => {
                    dirty_peers.insert(peer);
                }
                SyncEvent::RetryDelivery(peer) => {
                    retry_peers.insert(peer);
                }
                SyncEvent::DocumentChanged => {
                    document_dirty = true;
                }
            }
        }

        // Let pending HTTP task locks wake up and complete their lifecycle transitions gracefully.
        tokio::task::yield_now().await;

        if document_dirty {
            document_dirty = false;

            // Lock and extract only the cloned string keys into a small transient vector
            let peer_keys: Vec<String> = {
                let peers_guard = state.peer_states.read().await;
                peers_guard.keys().cloned().collect()
            };

            for peer in peer_keys {
                if peer_cluster(&state, &peer).await.is_some() {
                    dirty_peers.insert(peer);
                }
            }
        }

        // Retry existing envelopes separately from generating new ones.
        let retries: Vec<String> = retry_peers.drain().collect();

        for peer in retries {
            if let Err(error) = deliver_sync_to_peer(&state, &peer, true).await {
                warn!(
                    peer = %peer,
                    error = %format!("{error:#}"),
                    "Sync delivery retry failed"
                );
            }
        }

        let peers: Vec<String> = dirty_peers.drain().collect();

        for peer in peers {
            if let Err(error) = deliver_sync_to_peer(&state, &peer, false).await {
                warn!(
                    peer = %peer,
                    error = %format!("{error:#}"),
                    "Synchronization processing failed"
                );
            }
        }
    }
}

// --- change publisher --------------------------------------------------------

fn patch_store_path(patch: &automerge::Patch) -> Vec<String> {
    let mut path: Vec<String> = patch
        .path
        .iter()
        .map(|(_, prop)| match prop {
            automerge::Prop::Map(key) => key.clone(),
            automerge::Prop::Seq(index) => index.to_string(),
        })
        .collect();

    match &patch.action {
        automerge::PatchAction::PutMap { key, .. }
        | automerge::PatchAction::DeleteMap { key, .. } => path.push(key.clone()),
        // List, text and counter changes: the containing object changed.
        _ => {}
    }
    path
}

pub(crate) fn paths_related(a: &[String], b: &[String]) -> bool {
    a.starts_with(b) || b.starts_with(a)
}

pub(crate) async fn run_change_publisher(state: AppState) {
    let mut changed = state.change_tx.subscribe();
    let mut last_heads = state.write_doc(|doc| doc.get_heads()).await;

    info!("Starting change publisher");

    while changed.changed().await.is_ok() {
        let before = last_heads.clone();

        let (heads, paths, conflicts) = state
            .write_doc(move |doc| {
                let heads = doc.get_heads();
                if heads == before {
                    return (heads, Vec::new(), Vec::new());
                }

                // After an epoch switch the previous heads belong to the old document:
                // report the root as changed, which every stream is related to.
                if !before.iter().all(|h| doc.get_change_by_hash(h).is_some()) {
                    return (heads, vec![Vec::new()], Vec::new());
                }

                let mut paths = Vec::new();
                let mut conflicts = Vec::new();
                for patch in doc.diff(&before, &heads) {
                    let path = patch_store_path(&patch);
                    if matches!(
                        patch.action,
                        automerge::PatchAction::PutMap { conflict: true, .. }
                    ) {
                        conflicts.push(path.clone());
                    }
                    paths.push(path);
                }
                (heads, paths, conflicts)
            })
            .await;

        if heads == last_heads {
            continue;
        }
        last_heads = heads;

        for path in &conflicts {
            warn!(path = ?path, "Concurrent writes produced a conflict");
        }

        // Fails only when nobody is subscribed.
        let _ = state.change_events.send(Arc::new(ChangeSet { paths }));
    }
}