peerspeak/tests/reconnect_eviction.rs

//! Regression tests for the reconnect-eviction logic in `core`'s conn-event
//! handler. The bug these guard: after a peer reconnects, a *second* sustained
//! outage never evicted it, because the grace timer was armed only by the gossip
//! `PeerConnectionLost` path — which goes silent once gossip's neighbor state is
//! stale — while the transport keeps reliably reporting the drop. The fix arms the
//! grace timer from the transport `ConnEvent::Connecting` path too. These tests
//! drive `ConnEventHandler` directly with a short grace so the second-outage
//! eviction (and the cases that must NOT evict) are verified without a live link.

use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::Duration;

use tokio::sync::mpsc;

use iroh::endpoint::presets;
use iroh::{Endpoint, EndpointId, RelayMode, SecretKey};

use peerspeak::core::ConnEventHandler;
use peerspeak::core::messages::UiEvent;
use peerspeak::network::ConnEvent;
use peerspeak::network::iroh_impl::IrohTransport;

/// A real (but unconnected) transport. The handler's eviction path calls
/// `disconnect_peer`, which is a safe no-op for a peer we never linked to.
async fn make_transport() -> Arc<IrohTransport> {
    let endpoint = Endpoint::builder(presets::Minimal)
        .secret_key(SecretKey::generate())
        .relay_mode(RelayMode::Disabled)
        .bind()
        .await
        .expect("bind endpoint");
    Arc::new(IrohTransport::new(endpoint))
}

/// A peer identity with no live endpoint — all we need to key the handler's maps.
fn fake_peer() -> EndpointId {
    SecretKey::generate().public()
}

fn make_handler(
    ui_tx: mpsc::Sender<UiEvent>,
    transport: Arc<IrohTransport>,
    grace: Duration,
) -> ConnEventHandler {
    let grace_timers = Arc::new(std::sync::Mutex::new(HashMap::new()));
    let seen_connected = Arc::new(std::sync::Mutex::new(HashSet::new()));
    let jitter = Arc::new(tokio::sync::Mutex::new(HashMap::new()));
    ConnEventHandler::new(ui_tx, grace_timers, seen_connected, transport, jitter).with_grace(grace)
}

/// Drain UI events until a `PeerConnectionFailed` for `peer` (eviction) arrives,
/// or `within` elapses. Returns whether the eviction fired.
async fn evicted_within(
    ui_rx: &mut mpsc::Receiver<UiEvent>,
    peer: EndpointId,
    within: Duration,
) -> bool {
    let deadline = tokio::time::Instant::now() + within;
    loop {
        match tokio::time::timeout_at(deadline, ui_rx.recv()).await {
            Ok(Some(UiEvent::PeerConnectionFailed { id })) if id == peer => return true,
            Ok(Some(_)) => continue, // ignore PeerConnecting / PeerConnected / PeerLeft
            Ok(None) => return false, // channel closed
            Err(_) => return false,   // timed out — no eviction
        }
    }
}

/// Drain UI events until a `PeerLeft` for `peer` arrives, or `within` elapses.
/// Returns whether the leave fired.
async fn left_within(
    ui_rx: &mut mpsc::Receiver<UiEvent>,
    peer: EndpointId,
    within: Duration,
) -> bool {
    let deadline = tokio::time::Instant::now() + within;
    loop {
        match tokio::time::timeout_at(deadline, ui_rx.recv()).await {
            Ok(Some(UiEvent::PeerLeft { id })) if id == peer => return true,
            Ok(Some(_)) => continue, // ignore PeerConnecting / PeerConnected
            Ok(None) => return false, // channel closed
            Err(_) => return false,   // timed out — no leave
        }
    }
}

const GRACE: Duration = Duration::from_millis(200);

/// The regression: connect, drop, reconnect, then drop again — the second outage
/// must still arm a grace timer and evict. (Before the fix, only the first did.)
#[tokio::test]
async fn second_outage_after_reconnect_still_evicts() {
    let (ui_tx, mut ui_rx) = mpsc::channel(100);
    let h = make_handler(ui_tx, make_transport().await, GRACE);
    let peer = fake_peer();

    h.handle(ConnEvent::Connected(peer)).await; // initial link up
    h.handle(ConnEvent::Connecting(peer)).await; // first drop -> arms
    h.handle(ConnEvent::Connected(peer)).await; // reconnects -> cancels
    h.handle(ConnEvent::Connecting(peer)).await; // SECOND drop -> must re-arm

    assert!(
        evicted_within(&mut ui_rx, peer, Duration::from_secs(2)).await,
        "second outage after a reconnect must evict the peer"
    );
}

/// A peer that recovers within the grace window must NOT be evicted.
#[tokio::test]
async fn reconnect_within_grace_is_not_evicted() {
    let (ui_tx, mut ui_rx) = mpsc::channel(100);
    let h = make_handler(ui_tx, make_transport().await, GRACE);
    let peer = fake_peer();

    h.handle(ConnEvent::Connected(peer)).await;
    h.handle(ConnEvent::Connecting(peer)).await; // drop -> arms
    h.handle(ConnEvent::Connected(peer)).await; // recovers -> cancels

    assert!(
        !evicted_within(&mut ui_rx, peer, GRACE * 3).await,
        "a peer that reconnects within the grace window must not be evicted"
    );
}

/// A first-ever dial (never connected) must not be handed an eviction timer.
#[tokio::test]
async fn first_ever_dial_is_not_evicted() {
    let (ui_tx, mut ui_rx) = mpsc::channel(100);
    let h = make_handler(ui_tx, make_transport().await, GRACE);
    let peer = fake_peer();

    // Connecting with no prior Connected: the peer isn't in `seen_connected`.
    h.handle(ConnEvent::Connecting(peer)).await;

    assert!(
        !evicted_within(&mut ui_rx, peer, GRACE * 3).await,
        "a first-ever dial must not arm an eviction timer"
    );
}

/// A graceful leave must scrub `seen_connected`, so if that identity later rejoins,
/// its initial dial isn't mistaken for a reconnect and given an eviction clock.
#[tokio::test]
async fn graceful_leave_lets_a_later_rejoin_dial_cleanly() {
    let (ui_tx, mut ui_rx) = mpsc::channel(100);
    let h = make_handler(ui_tx, make_transport().await, GRACE);
    let peer = fake_peer();

    h.handle(ConnEvent::Connected(peer)).await; // linked once
    h.handle(ConnEvent::Left(peer)).await; // graceful leave -> scrub seen
    h.handle(ConnEvent::Connecting(peer)).await; // later rejoin's initial dial

    assert!(
        !evicted_within(&mut ui_rx, peer, GRACE * 3).await,
        "a rejoin's initial dial after a graceful leave must not arm eviction"
    );
}

/// A graceful leave evicts *immediately* — it emits `PeerLeft` right away rather
/// than leaving the peer "reconnecting" until the grace timer (or the slow gossip
/// Leave) catches up.
#[tokio::test]
async fn graceful_leave_evicts_immediately() {
    let (ui_tx, mut ui_rx) = mpsc::channel(100);
    let h = make_handler(ui_tx, make_transport().await, GRACE);
    let peer = fake_peer();

    h.handle(ConnEvent::Connected(peer)).await;
    h.handle(ConnEvent::Left(peer)).await;

    // Well inside the grace window: the leave must not be gated on the timer.
    assert!(
        left_within(&mut ui_rx, peer, GRACE / 2).await,
        "a graceful leave must emit PeerLeft immediately, not wait for the grace timer"
    );
}

/// After a peer is evicted by the grace timer, its identity is scrubbed from
/// `seen_connected`, so a later rejoin's initial dial is a fresh first-dial and
/// must not be mistaken for a reconnect (which would arm a new eviction clock).
#[tokio::test]
async fn rejoin_after_grace_eviction_dials_cleanly() {
    let (ui_tx, mut ui_rx) = mpsc::channel(100);
    let h = make_handler(ui_tx, make_transport().await, GRACE);
    let peer = fake_peer();

    // Connect, drop, and let the grace timer fire (evicts + scrubs seen_connected).
    h.handle(ConnEvent::Connected(peer)).await;
    h.handle(ConnEvent::Connecting(peer)).await;
    assert!(
        evicted_within(&mut ui_rx, peer, GRACE * 4).await,
        "the first outage should evict once the grace window elapses"
    );

    // A later rejoin's initial dial: scrubbed on eviction, so this is a first-ever
    // dial again and must not arm a new eviction timer.
    h.handle(ConnEvent::Connecting(peer)).await;
    assert!(
        !evicted_within(&mut ui_rx, peer, GRACE * 3).await,
        "an initial dial after a grace eviction must not be treated as a reconnect"
    );
}