openrtc 1.0.4

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

//! Iroh-level per-connection heartbeat.
//!
//! Wire protocol on iroh unidirectional streams:
//!   Ping sender opens a Uni stream and writes: [u32 len BE][typeId=3][payload]
//!   Responder opens a new Uni stream and writes: [u32 len BE][typeId=4][pong payload]
//!
//! The host layer (native_node.rs stream dispatcher) peeks at the first 5 bytes of
//! every incoming Uni stream: if typeId == 3 (ping), it parses the ping and immediately
//! replies via a new Uni stream, then drops the frame (does not forward to the app).
//! If typeId == 4 (pong), it parses the pong and delivers it to the active HeartbeatManager
//! for that endpoint, then drops the frame.

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

use bytes::Bytes;
use iroh::endpoint::Connection;
use tokio::sync::{mpsc, RwLock};
use tokio_util::sync::CancellationToken;

use super::codec::{self, Pong};
use super::state_machine::{aggregate_peer_health, TransportHealthEvaluator};
use super::transport_heartbeat::{HealthTransition, HeartbeatConfig};
use crate::connection_manager::ConnectionHealth;

/// Manages heartbeats for all active iroh connections.
#[derive(Default, Clone)]
pub struct IrohHeartbeatManager {
    inner: Arc<RwLock<Inner>>,
}

impl std::fmt::Debug for IrohHeartbeatManager {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("IrohHeartbeatManager")
            .finish_non_exhaustive()
    }
}

#[derive(Default)]
struct Inner {
    /// endpoint_id (string) -> heartbeat state for that connection
    connections: HashMap<String, ConnectionHeartbeatState>,
}

struct ConnectionHeartbeatState {
    pong_tx: mpsc::Sender<Pong>,
    health_rx: Arc<RwLock<ConnectionHealth>>,
    _cancel: CancellationToken,
}

impl IrohHeartbeatManager {
    pub fn new() -> Self {
        Self::default()
    }

    /// Start a heartbeat task for a new iroh connection.
    /// `connection` is used to open Uni streams for sending pings.
    /// `endpoint_id` is the string key for this connection.
    /// `health_tx` receives health transitions.
    pub async fn start_connection(
        &self,
        endpoint_id: String,
        connection: Connection,
        config: HeartbeatConfig,
        health_tx: mpsc::Sender<HealthTransition>,
    ) {
        let cancel = CancellationToken::new();
        let (pong_tx, pong_rx) = mpsc::channel::<Pong>(32);
        let health_store = Arc::new(RwLock::new(ConnectionHealth::Unknown));

        let task_cancel = cancel.clone();
        let task_health = health_store.clone();
        let task_endpoint = endpoint_id.clone();

        tokio::spawn(run_iroh_heartbeat(
            task_endpoint,
            connection,
            health_tx,
            pong_rx,
            config,
            task_cancel,
            task_health,
        ));

        let mut guard = self.inner.write().await;
        guard.connections.insert(
            endpoint_id,
            ConnectionHeartbeatState {
                pong_tx,
                health_rx: health_store,
                _cancel: cancel,
            },
        );
    }

    /// Deliver a pong received on an incoming Uni stream to the correct connection.
    pub async fn deliver_pong(&self, endpoint_id: &str, pong: Pong) {
        let guard = self.inner.read().await;
        if let Some(state) = guard.connections.get(endpoint_id) {
            let _ = state.pong_tx.try_send(pong);
        }
    }

    /// Stop heartbeat for a connection (called when connection closes).
    pub async fn stop_connection(&self, endpoint_id: &str) {
        let mut guard = self.inner.write().await;
        guard.connections.remove(endpoint_id);
        // CancellationToken dropped here, cancels the task
    }

    /// Get current health for a connection (Unknown if not tracked).
    pub async fn connection_health(&self, endpoint_id: &str) -> ConnectionHealth {
        let guard = self.inner.read().await;
        if let Some(state) = guard.connections.get(endpoint_id) {
            state.health_rx.read().await.clone()
        } else {
            ConnectionHealth::Unknown
        }
    }

    /// Aggregate health across all tracked connections.
    pub async fn aggregate_health(&self) -> ConnectionHealth {
        let guard = self.inner.read().await;
        let mut healths = HashMap::new();
        for (key, state) in &guard.connections {
            healths.insert(key.clone(), state.health_rx.read().await.clone());
        }
        aggregate_peer_health(&healths)
    }
}

async fn run_iroh_heartbeat(
    endpoint_id: String,
    connection: Connection,
    health_tx: mpsc::Sender<HealthTransition>,
    mut pong_rx: mpsc::Receiver<Pong>,
    config: HeartbeatConfig,
    cancel: CancellationToken,
    current_health: Arc<RwLock<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();
    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,
            _ = connection.closed() => 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, &endpoint_id, "iroh", new_health).await;
                }
            }

            _ = ticker.tick() => {
                let now = codec::now_unix_ms();
                let stale_cutoff = now - config.stale_after.as_millis() as i64;
                pending_pings.retain(|_, sent_at| *sent_at > stale_cutoff);

                seq = seq.wrapping_add(1);
                let frame = Bytes::from(codec::build_framed_ping(seq, now));

                let sent = send_uni_frame(&connection, frame, config.send_timeout).await;
                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, &endpoint_id, "iroh", new_health).await;
            }
        }
    }
}

/// Upper bound on **heartbeat-only** `open_uni` calls for a **symmetric** idle pair
/// (both peers run `IrohHeartbeatManager` with the same `HeartbeatConfig`).
///
/// Model: each `tick_interval`, each side sends one ping (one uni stream). The peer
/// answers each ping with one pong (another uni stream). With two peers, that is at
/// most **four** `open_uni` calls per tick (two pings + two pongs). This matches what
/// a browser logs as repeated `accept_uni` for inbound heartbeat traffic from the peer.
///
/// Real apps may exceed this bound when **non-heartbeat** code also opens uni streams
/// (signaling, framed sends, etc.).
pub fn idle_symmetric_heartbeat_max_send_uni_opens_upper_bound(
    observation_window: std::time::Duration,
    tick_interval: std::time::Duration,
) -> u64 {
    let w_ms = observation_window.as_millis().max(1) as u64;
    let t_ms = tick_interval.as_millis().max(1) as u64;
    let ticks = (w_ms + t_ms - 1) / t_ms;
    // Four uni opens per tick (symmetric ping + pong) plus slack for interval alignment.
    4 * ticks + 12
}

#[cfg(test)]
pub mod test_counters {
    use std::sync::atomic::{AtomicU64, Ordering};

    pub static HEARTBEAT_SEND_UNI_COUNT: AtomicU64 = AtomicU64::new(0);

    pub fn reset_heartbeat_send_uni_count() {
        HEARTBEAT_SEND_UNI_COUNT.store(0, Ordering::SeqCst);
    }
}

async fn send_uni_frame(
    connection: &Connection,
    frame: Bytes,
    timeout: std::time::Duration,
) -> bool {
    let open = tokio::time::timeout(timeout, connection.open_uni()).await;
    let Ok(Ok(mut send)) = open else { return false };
    #[cfg(test)]
    test_counters::HEARTBEAT_SEND_UNI_COUNT.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
    let write = tokio::time::timeout(timeout, async {
        tokio::io::AsyncWriteExt::write_all(&mut send, &frame).await?;
        let _ = send.finish();
        Ok::<(), std::io::Error>(())
    })
    .await;
    write.is_ok()
}

async fn update_health(
    current: &Arc<RwLock<ConnectionHealth>>,
    tx: &mpsc::Sender<HealthTransition>,
    _endpoint_id: &str,
    transport: &str,
    new: ConnectionHealth,
) {
    let previous = {
        let guard = current.read().await;
        guard.clone()
    };
    if previous == new {
        return;
    }
    {
        let mut guard = current.write().await;
        *guard = new.clone();
    }
    let _ = tx
        .send(HealthTransition {
            transport: format!("iroh:{}", transport),
            previous,
            current: new,
        })
        .await;
}

/// Call this when an incoming Uni stream header reveals a heartbeat ping (typeId=3).
/// Parses the ping from `frame_body` (the bytes after typeId) and sends a pong
/// back on a new Uni stream opened on `connection`.
pub async fn handle_incoming_ping(connection: &Connection, frame_body: &[u8]) {
    let now = codec::now_unix_ms();
    let Some(ping) = codec::decode_ping(frame_body) else {
        return;
    };
    let pong_frame = Bytes::from(codec::build_framed_pong(&ping, now));
    let _ = send_uni_frame(connection, pong_frame, std::time::Duration::from_secs(2)).await;
}

/// Call this when an incoming Uni stream header reveals a heartbeat pong (typeId=4).
/// Returns the decoded Pong so the caller can deliver it to `IrohHeartbeatManager`.
pub fn parse_incoming_pong(frame_body: &[u8]) -> Option<Pong> {
    codec::decode_pong(frame_body)
}

/// Peek-parse an incoming control-stream header without matching application
/// uni-stream bytes that happen to resemble a short length/type header.
/// Returns `Some((typeId, remaining payload offset))` if it's a heartbeat frame.
pub fn classify_framed_stream(buf: &[u8]) -> Option<(u8, usize)> {
    if buf.len() < codec::IROH_CONTROL_HEADER_LEN
        || &buf[..codec::IROH_CONTROL_MAGIC.len()] != codec::IROH_CONTROL_MAGIC
    {
        return None;
    }
    let type_id = buf[codec::IROH_CONTROL_MAGIC.len() + 4];
    if type_id == codec::TYPE_PING || type_id == codec::TYPE_PONG {
        Some((type_id, codec::IROH_CONTROL_HEADER_LEN))
    } else {
        None
    }
}

/// Read a full heartbeat frame from a unidirectional recv stream.
/// Returns (typeId, payload) or None on error/non-heartbeat.
pub async fn try_read_heartbeat_frame(
    recv: &mut iroh::endpoint::RecvStream,
) -> Option<(u8, Vec<u8>)> {
    let mut header = [0u8; codec::IROH_CONTROL_HEADER_LEN];
    recv.read_exact(&mut header).await.ok()?;

    if &header[..codec::IROH_CONTROL_MAGIC.len()] != codec::IROH_CONTROL_MAGIC {
        return None;
    }
    let offset = codec::IROH_CONTROL_MAGIC.len();
    let len = u32::from_be_bytes(header[offset..offset + 4].try_into().ok()?) as usize;
    let type_id = header[offset + 4];

    if type_id != codec::TYPE_PING
        && type_id != codec::TYPE_PONG
        && type_id != codec::TYPE_MANUAL_DISCONNECT
    {
        return None;
    }

    if len == 0 {
        return None;
    }

    let payload_len = len - 1; // len includes type byte
    let mut payload = vec![0u8; payload_len];
    recv.read_exact(&mut payload).await.ok()?;
    Some((type_id, payload))
}

#[cfg(test)]
mod model_tests {
    use super::idle_symmetric_heartbeat_max_send_uni_opens_upper_bound;
    use std::time::Duration;

    #[test]
    fn idle_symmetric_max_send_uni_upper_bound_formula() {
        // ceil(900/200) = 5 ticks → 4*5 + 12 = 32
        assert_eq!(
            idle_symmetric_heartbeat_max_send_uni_opens_upper_bound(
                Duration::from_millis(900),
                Duration::from_millis(200),
            ),
            32
        );
    }
}