path: root/src/state.rs

//! Shared node state: the document, lifecycle, per-peer state and timing.

use std::collections::HashMap;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
use std::time::Duration;

use anyhow::{bail, Result};
use automerge::{sync, AutoCommit};
use tokio::sync::{broadcast, mpsc, watch, Mutex, RwLock};
use tokio::time::Instant;
use tracing::{debug, error, warn};

use crate::firewall::Firewall;
use crate::replica::{Epoch, Replica};
use crate::sync::{ChangeSet, SyncEvent};
use crate::transport;

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum NodeLifecycle {
    // TODO maybe rename discovered_cluster_id to cluster_id if that doesn't impair readability.
    Joining {
        discovered_cluster_id: Option<String>,
    },
    Ready {
        cluster_id: String,
    },
}

#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) enum PeerMembership {
    #[default]
    Unknown,
    Member {
        cluster_id: String,
    },
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum Liveness {
    Active,
    /// No inbound message for Timing::dormant_after. Kept indefinitely, probed slowly.
    Dormant {
        since: Instant,
    },
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub(crate) enum HelloState {
    #[default]
    NotSent,
    /// A send is in flight. `next_delay` is how long to wait after it before retrying.
    Sending { next_delay: Duration },
    /// HELLO was sent (or the send failed); no HELLO_REPLY yet.
    AwaitingReply {
        next_attempt: Instant,
        next_delay: Duration,
    },
    /// HELLO_REPLY received. No more retries.
    Established,
}

#[derive(Clone)]
pub(crate) struct PendingOutbound {
    pub(crate) message: transport::EncodedMessage,
    pub(crate) retry_delay: Duration,
    pub(crate) next_attempt: Instant,
}

#[derive(Clone)]
pub(crate) struct PeerState {
    pub(crate) membership: PeerMembership,
    pub(crate) sync_state: sync::State,
    pub(crate) hello_state: HelloState,

    /// The peer sent HELLO; we still owe it a HELLO_REPLY (e.g. we had no cluster ID yet,
    /// or sending failed).
    pub(crate) reply_owed: bool,

    /// Received from this peer, waiting to be processed.
    pub(crate) pending_inbound_sync: Vec<transport::SyncEnvelope>,

    /// Generated for this peer, waiting for delivery.
    pub(crate) pending_outbound_sync: Option<PendingOutbound>,

    /// A sync pass was requested while an envelope was pending. Once that envelope
    /// is delivered, the peer must be re-enqueued so the deferred changes go out.
    pub(crate) needs_generate: bool,

    pub(crate) liveness: Liveness,
    pub(crate) added_at: Instant,
    pub(crate) last_seen: Option<Instant>,
    pub(crate) last_ping_at: Option<Instant>,

    /// Last time a sync message from this peer was applied.
    pub(crate) last_sync_rx: Option<Instant>,

    /// Last self-healing re-handshake (cooldown).
    pub(crate) last_resync_at: Option<Instant>,

    /// Epoch the peer advertised in its most recent message.
    pub(crate) epoch: Option<Epoch>,
}

impl Default for PeerState {
    fn default() -> Self {
        Self {
            membership: PeerMembership::default(),
            sync_state: sync::State::new(),
            hello_state: HelloState::default(),
            reply_owed: false,
            pending_inbound_sync: Vec::new(),
            pending_outbound_sync: None,
            needs_generate: false,
            liveness: Liveness::Active,
            added_at: Instant::now(),
            last_seen: None,
            last_ping_at: None,
            last_sync_rx: None,
            last_resync_at: None,
            epoch: None,
        }
    }
}

impl PeerState {
    /// Start Automerge sync with this peer from scratch. Any envelope generated against
    /// the old state is discarded; a stale RetryDelivery event is then a no-op.
    pub(crate) fn reset_sync(&mut self) {
        self.sync_state = sync::State::new();
        self.pending_outbound_sync = None;
        self.needs_generate = false;
    }

    fn quiet_for(&self, now: Instant) -> Duration {
        now.saturating_duration_since(self.last_seen.unwrap_or(self.added_at))
    }

    /// Records inbound activity. Returns true if the peer was dormant.
    pub(crate) fn mark_seen(&mut self, now: Instant) -> bool {
        self.last_seen = Some(now);

        if matches!(self.liveness, Liveness::Dormant { .. }) {
            self.liveness = Liveness::Active;
            // Force a fresh handshake; its reply resets sync on both sides.
            if !matches!(self.hello_state, HelloState::Sending { .. }) {
                self.hello_state = HelloState::NotSent;
            }
            return true;
        }
        false
    }

    /// Moves a silent peer to Dormant. Returns true on the transition.
    pub(crate) fn check_dormancy(&mut self, timing: &Timing, now: Instant) -> bool {
        if self.liveness != Liveness::Active || self.quiet_for(now) < timing.dormant_after {
            return false;
        }

        self.liveness = Liveness::Dormant { since: now };
        self.reset_sync();

        if !matches!(self.hello_state, HelloState::Sending { .. }) {
            self.hello_state = HelloState::AwaitingReply {
                next_attempt: now + timing.dormant_probe,
                next_delay: timing.dormant_probe,
            };
        }
        true
    }

    /// True (and records the ping) if an established, quiet peer needs a keepalive.
    pub(crate) fn take_keepalive_due(&mut self, timing: &Timing, now: Instant) -> bool {
        let due = self.liveness == Liveness::Active
            && self.hello_state == HelloState::Established
            && self.quiet_for(now) >= timing.keepalive
            && self
                .last_ping_at
                .is_none_or(|t| now.saturating_duration_since(t) >= timing.keepalive);

        if due {
            self.last_ping_at = Some(now);
        }
        due
    }

    /// Self-healing check on PING/PONG. Returns true if a re-handshake should start.
    pub(crate) fn needs_heal(
        &mut self,
        local: &transport::HeadsDigest,
        remote: &transport::HeadsDigest,
        timing: &Timing,
        now: Instant,
    ) -> bool {
        if local == remote
            || self.liveness != Liveness::Active
            || self.hello_state != HelloState::Established
        {
            return false;
        }

        let sync_idle = now.saturating_duration_since(self.last_sync_rx.unwrap_or(self.added_at))
            >= timing.resync_after;
        let cooled_down = self
            .last_resync_at
            .is_none_or(|t| now.saturating_duration_since(t) >= timing.resync_after);

        if !(sync_idle && cooled_down) {
            return false;
        }

        self.last_resync_at = Some(now);
        self.hello_state = HelloState::NotSent; // establish_peer resets sync on the reply
        true
    }

    /// Heard from recently. A returning node counts as alive from its first message,
    /// before the handshake finishes, so a cluster-wide forget can't make it flap.
    pub(crate) fn is_alive(&self, timing: &Timing, now: Instant) -> bool {
        self.liveness == Liveness::Active
            && self
                .last_seen
                .is_some_and(|t| now.saturating_duration_since(t) < timing.dormant_after)
    }
}

#[derive(Clone, Copy, Debug)]
pub(crate) struct Timing {
    pub(crate) keepalive: Duration,
    pub(crate) dormant_after: Duration,
    pub(crate) resync_after: Duration,
    pub(crate) dormant_probe: Duration,
}

impl Timing {
    // Recommended values
    //  Environment                                 | Keepalive | Notes
    //  --------------------------------------------+-----------+------------------------------------------------------------
    //  Test harness                                | 1-2 s     | Dormancy after about 4–8 s and healing within seconds keep tests short.
    //  LAN or small cluster (up to ~50 nodes)      | 15-30 s   | Under 3.5 messages/s per node at 50 nodes and 30 s. Current setting.
    //  WAN, mobile or metered links, or 100+ nodes | 60-120 s  | At 100 nodes and 30 s it’s about 6–7 messages/s per node. Fine, but rarely worth it on battery or metered links.
    //  Several hundred nodes or more               | -         | The full mesh becomes the problem in itself, not the interval. Message counts grow O(N²) across the cluster, which calls for partial views (SWIM-style random probing and syncing with a subset of peers).
    pub(crate) fn from_keepalive(keepalive: Duration) -> Self {
        Self {
            keepalive,
            dormant_after: keepalive * 4, // dormant_after >= 4* keepalive to tolerate three lost PINGs in a row
            resync_after: keepalive * 2, // resync_after should be >= 2* keepalive to guarantee at least one digest exchange while sync is idle
            dormant_probe: keepalive * 10, // >= 10*: dormant peers cost almost nothing, but still rejoin within minutes
        }
    }
}

#[derive(Clone)]
pub(crate) struct AppState {
    /// The document and its epoch. Derefs to `AutoCommit`.
    pub(crate) doc: Arc<RwLock<Replica>>,

    /// Copy of `doc.epoch` for cheap, lock-free reads (handshakes, keepalives, /peers).
    /// Only ever changed while holding `doc`'s write lock, together with `doc.epoch`.
    pub(crate) epoch_tx: watch::Sender<Epoch>,

    pub(crate) lifecycle: Arc<RwLock<NodeLifecycle>>,

    // Serializes cluster-ID establishment, peer membership, and readiness transitions.
    //[concurrent]cluster_validation: Arc<Mutex<()>>,
    pub(crate) peer_states: Arc<RwLock<HashMap<String, Arc<Mutex<PeerState>>>>>,
    pub(crate) transport: Arc<dyn transport::Transport>,
    pub(crate) firewall: Arc<Firewall>,

    pub(crate) sync_events: mpsc::UnboundedSender<SyncEvent>,

    // TODO see key_watchers.md ("Use a per-key notification registry instead of the global `watch<u64>`")
    pub(crate) change_tx: watch::Sender<u64>,

    pub(crate) initial_sync_done: Arc<AtomicBool>,

    pub(crate) change_events: broadcast::Sender<Arc<ChangeSet>>,

    pub(crate) timing: Timing,
}

impl AppState {
    pub(crate) fn local_id(&self) -> &str {
        self.transport.local_id()
    }

    /// The local epoch. Code that labels or applies sync messages must use the epoch
    /// of the locked `Replica` instead, so label and document can't disagree.
    pub(crate) fn local_epoch(&self) -> Epoch {
        *self.epoch_tx.borrow()
    }

    pub(crate) async fn is_ready(&self) -> bool {
        matches!(*self.lifecycle.read().await, NodeLifecycle::Ready { .. })
    }

    pub(crate) async fn cluster_id(&self) -> Option<String> {
        match &*self.lifecycle.read().await {
            NodeLifecycle::Ready { cluster_id } => Some(cluster_id.clone()),
            NodeLifecycle::Joining {
                discovered_cluster_id,
            } => discovered_cluster_id.clone(),
        }
    }

    // Record the expected cluster without making the node Ready. Bootstrap owns readiness.
    pub(crate) async fn set_expected_cluster_id(&self, cluster_id: String) -> Result<()> {
        //[concurrent]let _guard = self.cluster_validation.lock().await;
        let mut lifecycle = self.lifecycle.write().await;
        match &*lifecycle {
            NodeLifecycle::Joining {
                discovered_cluster_id: Some(existing),
            } if existing != &cluster_id => {
                bail!(
                    "Conflicting expected cluster IDs: expected {}, received {}",
                    existing,
                    cluster_id
                );
            }
            NodeLifecycle::Joining {
                discovered_cluster_id: None,
            } => {
                let old_lifecycle = lifecycle.clone();
                *lifecycle = NodeLifecycle::Joining {
                    discovered_cluster_id: Some(cluster_id),
                };
                debug!("Node transitioned from {old_lifecycle:?} {lifecycle:?}");
                Ok(())
            }
            NodeLifecycle::Joining {
                discovered_cluster_id: Some(_),
            } => {
                warn!("Node already Joining with the same identity");
                Ok(())
            }
            NodeLifecycle::Ready {
                cluster_id: existing,
            } if existing == &cluster_id => {
                warn!("Node already Ready with the same identity");
                Ok(())
            }
            NodeLifecycle::Ready {
                cluster_id: existing,
            } => {
                bail!(
                    "Cannot change established cluster ID: existing {}, received {}",
                    existing,
                    cluster_id
                );
            }
        }
    }

    // Only bootstrap should call this, after validating document metadata.
    pub(crate) async fn mark_ready(&self, cluster_id: String) -> Result<()> {
        //[concurrent]let _guard = self.cluster_validation.lock().await;
        let mut lifecycle = self.lifecycle.write().await;
        match &*lifecycle {
            NodeLifecycle::Joining {
                discovered_cluster_id: Some(expected),
            } if expected != &cluster_id => {
                bail!(
                    "Cannot mark node ready: expected cluster {}, received {}",
                    expected,
                    cluster_id
                );
            }
            NodeLifecycle::Joining { .. } => {
                let old_lifecycle = lifecycle.clone();
                *lifecycle = NodeLifecycle::Ready { cluster_id };
                debug!("Node transitioned from {old_lifecycle:?} {lifecycle:?}");
            }
            NodeLifecycle::Ready {
                cluster_id: existing,
            } if existing == &cluster_id => {
                debug!("Node already ready with the same identity");
            }
            NodeLifecycle::Ready {
                cluster_id: existing,
            } => {
                bail!(
                    "Cannot change established cluster ID: existing {}, received {}",
                    existing,
                    cluster_id
                );
            }
        }
        Ok(())
    }

    /// Looks up a peer without creating it. Background work (sync, retries,
    /// keepalives) must use this so a forgotten peer isn't resurrected.
    pub(crate) async fn get_peer_state(&self, pk: &str) -> Option<Arc<Mutex<PeerState>>> {
        self.peer_states.read().await.get(pk).cloned()
    }

    pub(crate) async fn get_or_create_peer_state(&self, pk: &str) -> Arc<Mutex<PeerState>> {
        let mut peers = self.peer_states.write().await;

        Arc::clone(
            peers
                .entry(pk.to_owned())
                .or_insert_with(|| Arc::new(Mutex::new(PeerState::default()))),
        )
    }

    pub(crate) fn notify_store_changed(&self) {
        self.change_tx.send_modify(|version| {
            *version = version.wrapping_add(1);
        });
    }

    /// Reads the document on the blocking pool; `finish` runs after the lock is released.
    pub(crate) async fn render<T, R>(
        &self,
        read: impl FnOnce(&AutoCommit) -> T + Send + 'static,
        finish: impl FnOnce(T) -> R + Send + 'static,
    ) -> R
    where
        R: Send + 'static,
    {
        let doc = Arc::clone(&self.doc);
        tokio::task::spawn_blocking(move || {
            let data = {
                let guard = doc.blocking_read();
                read(&guard.doc)
            };
            finish(data)
        })
        .await
        .expect("document read task panicked")
    }

    /// Mutates the document on the blocking pool.
    pub(crate) async fn write_doc<T>(
        &self,
        f: impl FnOnce(&mut AutoCommit) -> T + Send + 'static,
    ) -> T
    where
        T: Send + 'static,
    {
        let doc = Arc::clone(&self.doc);
        tokio::task::spawn_blocking(move || f(&mut doc.blocking_write().doc))
            .await
            .expect("document write task panicked")
    }

    /// Mutates the document on the blocking pool as one Automerge change, then
    /// publishes it. Write and notification run in the same blocking task, so a
    /// cancelled request can't leave a committed write unannounced.
    pub(crate) async fn mutate_doc<T, E>(
        &self,
        f: impl FnOnce(&mut AutoCommit) -> Result<T, E> + Send + 'static,
    ) -> Result<T, E>
    where
        T: Send + 'static,
        E: Send + 'static,
    {
        let state = self.clone();

        tokio::task::spawn_blocking(move || {
            let result = {
                let mut replica = state.doc.blocking_write();
                let result = f(&mut replica.doc);
                match &result {
                    // One API write = one Automerge change.
                    Ok(_) => {
                        replica.doc.commit();
                    }
                    // Discard partial ops (e.g. old value deleted, new one not written).
                    Err(_) => {
                        replica.doc.rollback();
                    }
                }
                result
            }; // lock released before notifying

            if result.is_ok() {
                state.publish_change();
            }
            result
        })
        .await
        .expect("document write task panicked")
    }

    pub(crate) fn publish_change(&self) {
        self.notify_store_changed();
        if let Err(error) = self.sync_events.send(SyncEvent::DocumentChanged) {
            // The write is committed; failing the request would misreport it.
            error!(error = %error, "Sync worker is gone; change will not be replicated");
        }
    }
}

pub(crate) fn unix_millis() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_millis() as u64)
        .unwrap_or(0)
}