peerspeak/src/core/recovery.rs

use crate::network::{RoomState, gossip::IrohGossipState};
use iroh::{EndpointAddr, EndpointId};
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use tokio::time::Instant;

const RECOVERY_COMMAND_CAPACITY: usize = 64;
const RECOVERY_DELAYS: [Duration; 7] = [
    Duration::from_secs(1),
    Duration::from_secs(2),
    Duration::from_secs(4),
    Duration::from_secs(8),
    Duration::from_secs(15),
    Duration::from_secs(30),
    Duration::from_secs(60),
];

fn recovery_delay(attempt: usize) -> Duration {
    RECOVERY_DELAYS[attempt.min(RECOVERY_DELAYS.len() - 1)]
}

/// Terminal retry budget for background recovery. After this many failed attempts
/// the coordinator gives up: it drops the entry, frees the active slot, and signals
/// the event task to forget the retained address (Tier C recovery-identity cap).
///
/// With the [`RECOVERY_DELAYS`] backoff this is roughly seven minutes of dialing
/// (1+2+4+8+15+30+60s, then 60s steps), far beyond any normal transient outage. A
/// genuine peer returning after a longer outage still rejoins on its own via a
/// gossip announce, so giving up only stops us from dialing a peer that is not
/// coming back — it does not break legitimate reconnect-after-outage.
const RECOVERY_TERMINAL_ATTEMPTS: usize = 12;

/// Capacity of the terminal-eviction notification channel. Bounded; on the rare
/// event of saturation the entry is still removed (the dial work stops) and only
/// the retained-address forget is skipped, which the per-topic retain cap bounds.
const RECOVERY_TERMINAL_CAPACITY: usize = 64;

/// Whether `attempt` completed recoveries have exhausted the terminal budget.
fn recovery_is_terminal(attempt: usize, max_attempts: usize) -> bool {
    attempt >= max_attempts
}

enum RecoveryCommand {
    Start {
        peer_id: EndpointId,
        addr: EndpointAddr,
    },
    Cancel(EndpointId),
}

struct RecoveryEntry {
    addr: EndpointAddr,
    attempt: usize,
    next_attempt: Instant,
}

#[async_trait::async_trait]
trait RecoveryRoom: Send + Sync {
    async fn rebootstrap_peers(&self, peers: Vec<EndpointAddr>) -> Result<(), String>;
}

#[async_trait::async_trait]
impl RecoveryRoom for IrohGossipState {
    async fn rebootstrap_peers(&self, peers: Vec<EndpointAddr>) -> Result<(), String> {
        RoomState::rebootstrap_peers(self, peers)
            .await
            .map_err(|error| error.to_string())
    }
}

/// Cloneable command side of the single per-session recovery coordinator.
/// `active` is shared with transport/event handlers so cancellation is visible
/// immediately even while the coordinator is awaiting an in-flight gossip call.
#[derive(Clone)]
pub(super) struct RecoveryCoordinator {
    tx: mpsc::Sender<RecoveryCommand>,
    active: Arc<Mutex<HashSet<EndpointId>>>,
}

impl RecoveryCoordinator {
    pub(super) fn spawn(
        room_state: Arc<IrohGossipState>,
    ) -> (Self, JoinHandle<()>, mpsc::Receiver<EndpointId>) {
        Self::spawn_inner(room_state)
    }

    fn spawn_inner(
        room_state: Arc<dyn RecoveryRoom>,
    ) -> (Self, JoinHandle<()>, mpsc::Receiver<EndpointId>) {
        let (tx, rx) = mpsc::channel(RECOVERY_COMMAND_CAPACITY);
        let (terminal_tx, terminal_rx) = mpsc::channel(RECOVERY_TERMINAL_CAPACITY);
        let active = Arc::new(Mutex::new(HashSet::new()));
        let handle = Self {
            tx,
            active: active.clone(),
        };
        let task = tokio::spawn(run_coordinator(room_state, active, rx, terminal_tx));
        (handle, task, terminal_rx)
    }

    /// Reserve one recovery slot before grace-expiry teardown begins. Returns
    /// false when the peer is already recovering, preventing duplicate work.
    pub(super) fn begin(&self, peer_id: EndpointId) -> bool {
        self.active.lock().unwrap().insert(peer_id)
    }

    /// Activate the reserved slot with its retained authenticated address.
    /// Uses a bounded non-blocking send while holding the active-set lock so a
    /// concurrent cancellation is ordered before or after this command.
    pub(super) fn activate(&self, peer_id: EndpointId, addr: EndpointAddr) -> Result<bool, ()> {
        let mut active = self.active.lock().unwrap();
        if !active.contains(&peer_id) {
            return Ok(false);
        }
        if self
            .tx
            .try_send(RecoveryCommand::Start { peer_id, addr })
            .is_err()
        {
            active.remove(&peer_id);
            return Err(());
        }
        Ok(true)
    }

    pub(super) fn cancel(&self, peer_id: EndpointId) {
        self.active.lock().unwrap().remove(&peer_id);
        // Cancellation is governed by the shared active set, so it remains
        // immediate even if the bounded command queue is temporarily full.
        let _ = self.tx.try_send(RecoveryCommand::Cancel(peer_id));
    }

    pub(super) fn is_active(&self, peer_id: &EndpointId) -> bool {
        self.active.lock().unwrap().contains(peer_id)
    }
}

async fn run_coordinator(
    room_state: Arc<dyn RecoveryRoom>,
    active: Arc<Mutex<HashSet<EndpointId>>>,
    mut rx: mpsc::Receiver<RecoveryCommand>,
    terminal_tx: mpsc::Sender<EndpointId>,
) {
    let mut entries: HashMap<EndpointId, RecoveryEntry> = HashMap::new();

    loop {
        // The shared active set is the authoritative cancellation gate. Prune
        // here as well as on Cancel commands so a saturated command queue cannot
        // leave an inactive, past-due entry spinning the timer loop.
        let active_snapshot = active.lock().unwrap().clone();
        entries.retain(|peer_id, _| active_snapshot.contains(peer_id));
        let next_deadline = entries.values().map(|entry| entry.next_attempt).min();
        let command = match next_deadline {
            Some(deadline) => {
                tokio::select! {
                    command = rx.recv() => command,
                    _ = tokio::time::sleep_until(deadline) => {
                        let now = Instant::now();
                        let active_snapshot = active.lock().unwrap().clone();
                        let due: Vec<(EndpointId, EndpointAddr)> = entries
                            .iter()
                            .filter(|(id, entry)| {
                                entry.next_attempt <= now && active_snapshot.contains(*id)
                            })
                            .map(|(id, entry)| (*id, entry.addr.clone()))
                            .collect();

                        if !due.is_empty() {
                            let addrs = due.iter().map(|(_, addr)| addr.clone()).collect();
                            if let Err(error) = room_state.rebootstrap_peers(addrs).await {
                                crate::log_msg(&format!(
                                    "Background peer recovery attempt failed: {error}"
                                ));
                            }

                            let scheduled_at = Instant::now();
                            for (peer_id, _) in due {
                                if !active.lock().unwrap().contains(&peer_id) {
                                    entries.remove(&peer_id);
                                    continue;
                                }
                                // Advance the backoff, then check the terminal budget.
                                // `attempt` counts completed attempts, so the delay
                                // uses the current value before it is incremented.
                                let terminal = if let Some(entry) = entries.get_mut(&peer_id) {
                                    entry.next_attempt = scheduled_at + recovery_delay(entry.attempt);
                                    entry.attempt = entry.attempt.saturating_add(1);
                                    recovery_is_terminal(entry.attempt, RECOVERY_TERMINAL_ATTEMPTS)
                                } else {
                                    false
                                };
                                if terminal {
                                    // Give up on a peer that has not returned within the
                                    // budget: drop its entry, free the active slot, and
                                    // signal the event task to forget its retained
                                    // address so the per-topic retain table drains.
                                    entries.remove(&peer_id);
                                    active.lock().unwrap().remove(&peer_id);
                                    let _ = terminal_tx.try_send(peer_id);
                                }
                            }
                        }
                        continue;
                    }
                }
            }
            None => rx.recv().await,
        };

        match command {
            Some(RecoveryCommand::Start { peer_id, addr }) => {
                if active.lock().unwrap().contains(&peer_id) {
                    entries.entry(peer_id).or_insert(RecoveryEntry {
                        addr,
                        attempt: 0,
                        next_attempt: Instant::now(),
                    });
                }
            }
            Some(RecoveryCommand::Cancel(peer_id)) => {
                entries.remove(&peer_id);
            }
            None => break,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use iroh::SecretKey;

    struct RecordingRoom {
        attempts: mpsc::UnboundedSender<Vec<EndpointAddr>>,
    }

    #[async_trait::async_trait]
    impl RecoveryRoom for RecordingRoom {
        async fn rebootstrap_peers(&self, peers: Vec<EndpointAddr>) -> Result<(), String> {
            self.attempts.send(peers).map_err(|error| error.to_string())
        }
    }

    #[test]
    fn retry_backoff_reaches_and_stays_at_sixty_seconds() {
        let actual: Vec<u64> = (0..10)
            .map(|attempt| recovery_delay(attempt).as_secs())
            .collect();
        assert_eq!(actual, vec![1, 2, 4, 8, 15, 30, 60, 60, 60, 60]);
    }

    #[test]
    fn recovery_budget_is_terminal_only_at_or_past_the_cap() {
        assert!(!recovery_is_terminal(0, RECOVERY_TERMINAL_ATTEMPTS));
        assert!(!recovery_is_terminal(
            RECOVERY_TERMINAL_ATTEMPTS - 1,
            RECOVERY_TERMINAL_ATTEMPTS
        ));
        assert!(recovery_is_terminal(
            RECOVERY_TERMINAL_ATTEMPTS,
            RECOVERY_TERMINAL_ATTEMPTS
        ));
        assert!(recovery_is_terminal(
            RECOVERY_TERMINAL_ATTEMPTS + 5,
            RECOVERY_TERMINAL_ATTEMPTS
        ));
    }

    #[test]
    fn recovery_slots_are_deduplicated_and_cancel_immediately() {
        let (tx, mut rx) = mpsc::channel(4);
        let coordinator = RecoveryCoordinator {
            tx,
            active: Arc::new(Mutex::new(HashSet::new())),
        };
        let peer_id = SecretKey::generate().public();

        assert!(coordinator.begin(peer_id));
        assert!(
            !coordinator.begin(peer_id),
            "a peer gets only one recovery slot"
        );
        assert_eq!(
            coordinator.activate(peer_id, EndpointAddr::from(peer_id)),
            Ok(true)
        );
        assert!(matches!(
            rx.try_recv(),
            Ok(RecoveryCommand::Start { peer_id: id, .. }) if id == peer_id
        ));

        coordinator.cancel(peer_id);
        assert!(!coordinator.is_active(&peer_id));
        assert!(matches!(
            rx.try_recv(),
            Ok(RecoveryCommand::Cancel(id)) if id == peer_id
        ));
    }

    #[tokio::test]
    async fn coordinator_attempts_rebootstrap_immediately() {
        let (attempts_tx, mut attempts_rx) = mpsc::unbounded_channel();
        let (coordinator, task, _terminal_rx) =
            RecoveryCoordinator::spawn_inner(Arc::new(RecordingRoom {
                attempts: attempts_tx,
            }));
        let peer_id = SecretKey::generate().public();
        let addr = EndpointAddr::from(peer_id);

        assert!(coordinator.begin(peer_id));
        assert_eq!(coordinator.activate(peer_id, addr.clone()), Ok(true));
        let attempted = tokio::time::timeout(Duration::from_secs(1), attempts_rx.recv())
            .await
            .expect("first recovery attempt should be immediate")
            .expect("recording room remains subscribed");
        assert_eq!(attempted, vec![addr]);

        coordinator.cancel(peer_id);
        task.abort();
    }
}