openrtc 1.0.1

#![cfg(not(target_arch = "wasm32"))]

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

use bytes::Bytes;
use tokio::sync::{mpsc, Mutex};
use tokio_util::sync::CancellationToken;

use super::codec::{self, Pong};
use super::state_machine::TransportHealthEvaluator;
use crate::connection_manager::ConnectionHealth;

#[derive(Debug, Clone)]
pub struct HeartbeatConfig {
    pub tick_interval: Duration,
    pub suspect_after: Duration,
    pub stale_after: Duration,
    pub send_timeout: Duration,
}

impl Default for HeartbeatConfig {
    fn default() -> Self {
        Self {
            tick_interval: Duration::from_secs(5),
            suspect_after: Duration::from_secs(10),
            stale_after: Duration::from_secs(30),
            send_timeout: Duration::from_secs(2),
        }
    }
}

#[derive(Debug)]
pub struct HealthTransition {
    pub transport: String,
    pub previous: ConnectionHealth,
    pub current: ConnectionHealth,
}

/// Per-transport heartbeat handle. Drop or call `cancel()` to stop.
pub struct TransportHeartbeat {
    transport: String,
    cancel: CancellationToken,
    pong_tx: mpsc::Sender<Pong>,
    current_health: Arc<Mutex<ConnectionHealth>>,
}

impl TransportHeartbeat {
    /// Spawn a heartbeat task for a transport.
    ///
    /// `send_fn` should send a raw byte slice over the transport (returns error on failure).
    /// `health_tx` receives health transitions whenever the health changes.
    /// `webrtc_mode` selects the wire format: false = framed iroh, true = WebRTC prefix.
    pub fn spawn(
        transport: String,
        send_fn: Arc<
            dyn Fn(Bytes) -> futures::future::BoxFuture<'static, anyhow::Result<()>> + Send + Sync,
        >,
        health_tx: mpsc::Sender<HealthTransition>,
        config: HeartbeatConfig,
        webrtc_mode: bool,
    ) -> Self {
        let cancel = CancellationToken::new();
        let (pong_tx, pong_rx) = mpsc::channel::<Pong>(32);
        let current_health = Arc::new(Mutex::new(ConnectionHealth::Unknown));

        let task_cancel = cancel.clone();
        let task_transport = transport.clone();
        let task_health = current_health.clone();

        tokio::spawn(run_heartbeat(
            task_transport,
            send_fn,
            health_tx,
            pong_rx,
            config,
            webrtc_mode,
            task_cancel,
            task_health,
        ));

        Self {
            transport,
            cancel,
            pong_tx,
            current_health,
        }
    }

    pub async fn handle_pong(&self, pong: Pong) {
        let _ = self.pong_tx.send(pong).await;
    }

    pub fn pong_sender(&self) -> mpsc::Sender<Pong> {
        self.pong_tx.clone()
    }

    pub async fn health(&self) -> ConnectionHealth {
        self.current_health.lock().await.clone()
    }

    pub fn cancel(self) {
        self.cancel.cancel();
    }

    pub fn transport(&self) -> &str {
        &self.transport
    }
}

async fn run_heartbeat(
    transport: String,
    send_fn: Arc<
        dyn Fn(Bytes) -> futures::future::BoxFuture<'static, anyhow::Result<()>> + Send + Sync,
    >,
    health_tx: mpsc::Sender<HealthTransition>,
    mut pong_rx: mpsc::Receiver<Pong>,
    config: HeartbeatConfig,
    webrtc_mode: bool,
    cancel: CancellationToken,
    current_health: Arc<Mutex<ConnectionHealth>>,
) {
    let evaluator = TransportHealthEvaluator::new(
        config.suspect_after.as_millis() as i64,
        config.stale_after.as_millis() as i64,
    );
    let mut ticker = tokio::time::interval(config.tick_interval);
    let mut seq: u64 = 0;
    let mut pending_pings: HashMap<u64, i64> = HashMap::new(); // seq -> sent_at_ms
    let mut last_pong_at_ms: Option<i64> = None;
    let mut first_ping_sent_at_ms: Option<i64> = None;

    loop {
        tokio::select! {
            biased;
            _ = cancel.cancelled() => break,

            pong = pong_rx.recv() => {
                let Some(pong) = pong else { break };
                if pending_pings.remove(&pong.seq).is_some() {
                    let now = codec::now_unix_ms();
                    last_pong_at_ms = Some(now);
                    let new_health = evaluator.evaluate(last_pong_at_ms, first_ping_sent_at_ms, now);
                    update_health(&current_health, &health_tx, &transport, new_health).await;
                }
            }

            _ = ticker.tick() => {
                let now = codec::now_unix_ms();

                // Garbage-collect pings older than stale_after
                let stale_cutoff = now - config.stale_after.as_millis() as i64;
                pending_pings.retain(|_, sent_at| *sent_at > stale_cutoff);

                // Send ping
                seq = seq.wrapping_add(1);
                let frame = if webrtc_mode {
                    Bytes::from(codec::build_webrtc_ping(seq, now))
                } else {
                    Bytes::from(codec::build_framed_ping(seq, now))
                };

                let sent = tokio::time::timeout(config.send_timeout, send_fn(frame))
                    .await
                    .ok()
                    .and_then(|r| r.ok())
                    .is_some();

                if sent {
                    if first_ping_sent_at_ms.is_none() {
                        first_ping_sent_at_ms = Some(now);
                    }
                    pending_pings.insert(seq, now);
                }

                let new_health = evaluator.evaluate(last_pong_at_ms, first_ping_sent_at_ms, now);
                update_health(&current_health, &health_tx, &transport, new_health).await;
            }
        }
    }
}

async fn update_health(
    current: &Arc<Mutex<ConnectionHealth>>,
    tx: &mpsc::Sender<HealthTransition>,
    transport: &str,
    new: ConnectionHealth,
) {
    let mut guard = current.lock().await;
    if *guard == new {
        return;
    }
    let previous = guard.clone();
    *guard = new.clone();
    drop(guard);

    let _ = tx
        .send(HealthTransition {
            transport: transport.to_string(),
            previous,
            current: new,
        })
        .await;
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicU64, Ordering};
    use std::sync::Arc as StdArc;

    #[tokio::test]
    async fn transitions_to_suspect_after_missed_pings() {
        let send_count = StdArc::new(AtomicU64::new(0));
        let sc = send_count.clone();
        let send_fn: Arc<
            dyn Fn(Bytes) -> futures::future::BoxFuture<'static, anyhow::Result<()>> + Send + Sync,
        > = Arc::new(move |_| {
            sc.fetch_add(1, Ordering::SeqCst);
            Box::pin(async { Ok(()) })
        });

        let (health_tx, mut health_rx) = mpsc::channel(16);
        // `run_heartbeat` uses wall-clock ms (`codec::now_unix_ms`); paused Tokio time does not advance it.
        let config = HeartbeatConfig {
            tick_interval: Duration::from_millis(20),
            suspect_after: Duration::from_millis(60),
            stale_after: Duration::from_millis(200),
            send_timeout: Duration::from_millis(5),
        };
        let _hb = TransportHeartbeat::spawn("iroh".to_string(), send_fn, health_tx, config, false);

        // Real sleep until silence since first ping exceeds suspect_after.
        tokio::time::sleep(Duration::from_millis(150)).await;

        // Drain Unknown/Healthy transitions; the first recv is often still Healthy.
        let mut saw_suspect_or_stale = false;
        for _ in 0..16 {
            let t = tokio::time::timeout(Duration::from_secs(2), health_rx.recv())
                .await
                .expect("timeout waiting for health transition")
                .expect("channel closed");
            if matches!(
                t.current,
                ConnectionHealth::Suspect | ConnectionHealth::Stale
            ) {
                saw_suspect_or_stale = true;
                break;
            }
        }
        assert!(
            saw_suspect_or_stale,
            "expected Suspect or Stale after missed pongs"
        );
    }

    #[tokio::test]
    async fn recovers_to_healthy_on_pong() {
        tokio::time::pause();

        let send_fn: Arc<
            dyn Fn(Bytes) -> futures::future::BoxFuture<'static, anyhow::Result<()>> + Send + Sync,
        > = Arc::new(|_| Box::pin(async { Ok(()) }));

        let (health_tx, mut health_rx) = mpsc::channel(16);
        let config = HeartbeatConfig {
            tick_interval: Duration::from_secs(5),
            suspect_after: Duration::from_secs(10),
            stale_after: Duration::from_secs(30),
            send_timeout: Duration::from_secs(2),
        };
        let hb = TransportHeartbeat::spawn("iroh".to_string(), send_fn, health_tx, config, false);

        // Let it go suspect
        tokio::time::advance(Duration::from_secs(15)).await;
        tokio::task::yield_now().await;

        // Send a pong with seq=1 (first ping)
        hb.handle_pong(Pong {
            seq: 1,
            sender_unix_ms: 0,
            recv_unix_ms: 0,
        })
        .await;
        tokio::task::yield_now().await;

        // Drain events and look for Healthy
        let mut saw_healthy = false;
        while let Ok(Some(t)) =
            tokio::time::timeout(Duration::from_millis(50), health_rx.recv()).await
        {
            if matches!(t.current, ConnectionHealth::Healthy) {
                saw_healthy = true;
                break;
            }
        }
        assert!(saw_healthy, "expected recovery to Healthy after pong");
    }
}