laminar-core 0.26.0

//! Static seed-list discovery with TCP heartbeats.

#![allow(clippy::disallowed_types)] // cold path: static discovery coordination
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

use parking_lot::RwLock;
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;

use super::{Discovery, DiscoveryError, NodeId, NodeInfo, NodeMetadata, NodeState};

/// TCP connect timeout for heartbeat connections.
const CONNECT_TIMEOUT: Duration = Duration::from_secs(2);

/// TCP read/write timeout per I/O operation.
const IO_TIMEOUT: Duration = Duration::from_secs(5);

/// Maximum concurrent handler tasks in the listener.
const MAX_HANDLER_TASKS: usize = 64;

/// Maximum message size (1 MB).
const MAX_MESSAGE_SIZE: usize = 1_048_576;

/// Grace period after which `Left` peers are reaped from the peer map.
const LEFT_REAP_THRESHOLD: u32 = 30;

/// Configuration for static discovery.
#[derive(Debug, Clone)]
pub struct StaticDiscoveryConfig {
    /// This node's info.
    pub local_node: NodeInfo,
    /// Seed addresses to connect to.
    pub seeds: Vec<String>,
    /// Heartbeat interval.
    pub heartbeat_interval: Duration,
    /// Number of missed heartbeats before marking as `Suspected`.
    pub suspect_threshold: u32,
    /// Number of missed heartbeats before marking as `Left`.
    pub dead_threshold: u32,
    /// Address to bind the heartbeat listener.
    pub listen_address: String,
}

impl Default for StaticDiscoveryConfig {
    fn default() -> Self {
        Self {
            local_node: NodeInfo {
                id: NodeId(1),
                name: "node-1".into(),
                rpc_address: "127.0.0.1:9000".into(),
                raft_address: "127.0.0.1:9001".into(),
                state: NodeState::Active,
                metadata: NodeMetadata::default(),
                last_heartbeat_ms: 0,
            },
            seeds: Vec::new(),
            heartbeat_interval: Duration::from_secs(1),
            suspect_threshold: 3,
            dead_threshold: 10,
            listen_address: "127.0.0.1:9002".into(),
        }
    }
}

/// Internal per-peer tracking state.
#[derive(Debug)]
struct PeerState {
    info: NodeInfo,
    /// Missed *outbound* heartbeats (managed exclusively by the heartbeater).
    missed_heartbeats: u32,
    /// Counter that keeps incrementing after `Left` state. Used for reaping.
    left_ticks: u32,
}

/// Static discovery implementation with TCP heartbeats.
#[derive(Debug)]
pub struct StaticDiscovery {
    config: StaticDiscoveryConfig,
    peers: Arc<RwLock<HashMap<u64, PeerState>>>,
    membership_tx: watch::Sender<Vec<NodeInfo>>,
    membership_rx: watch::Receiver<Vec<NodeInfo>>,
    cancel: CancellationToken,
    listener_handle: Option<tokio::task::JoinHandle<Result<(), DiscoveryError>>>,
    heartbeater_handle: Option<tokio::task::JoinHandle<()>>,
    started: bool,
}

impl StaticDiscovery {
    /// Create a new static discovery instance.
    #[must_use]
    pub fn new(config: StaticDiscoveryConfig) -> Self {
        debug_assert!(
            config.suspect_threshold < config.dead_threshold,
            "suspect_threshold ({}) must be less than dead_threshold ({})",
            config.suspect_threshold,
            config.dead_threshold,
        );
        let (tx, rx) = watch::channel(Vec::new());
        Self {
            config,
            peers: Arc::new(RwLock::new(HashMap::new())),
            membership_tx: tx,
            membership_rx: rx,
            cancel: CancellationToken::new(),
            listener_handle: None,
            heartbeater_handle: None,
            started: false,
        }
    }

    /// Serialize a `NodeInfo` for transmission.
    fn serialize_node_info(info: &NodeInfo) -> Result<Vec<u8>, DiscoveryError> {
        rkyv::to_bytes::<rkyv::rancor::Error>(info)
            .map(|v| v.to_vec())
            .map_err(|e| DiscoveryError::Serialization(e.to_string()))
    }

    /// Deserialize a `NodeInfo` from received bytes.
    fn deserialize_node_info(data: &[u8]) -> Result<NodeInfo, DiscoveryError> {
        rkyv::from_bytes::<NodeInfo, rkyv::rancor::Error>(data)
            .map_err(|e| DiscoveryError::Serialization(e.to_string()))
    }

    /// Update the membership watch channel from current peer state.
    fn broadcast_membership(&self) {
        let peers = self.peers.read();
        let peer_list: Vec<NodeInfo> = peers.values().map(|p| p.info.clone()).collect();
        publish_if_changed(&self.membership_tx, peer_list);
    }

    /// Send a heartbeat to a single seed address with connect + I/O timeouts.
    #[allow(clippy::cast_possible_truncation)]
    async fn send_heartbeat(address: &str, data: &[u8]) -> Result<Option<Vec<u8>>, DiscoveryError> {
        use tokio::io::{AsyncReadExt, AsyncWriteExt};

        // Connect with timeout (C2 fix)
        let mut stream = tokio::time::timeout(CONNECT_TIMEOUT, TcpStream::connect(address))
            .await
            .map_err(|_| DiscoveryError::Connection {
                address: address.into(),
                reason: "connect timeout".into(),
            })?
            .map_err(|e| DiscoveryError::Connection {
                address: address.into(),
                reason: e.to_string(),
            })?;

        // Validate message size before sending (W7 symmetric limit)
        if data.len() > MAX_MESSAGE_SIZE {
            return Err(DiscoveryError::Serialization(
                "message too large to send".into(),
            ));
        }

        // Write length-prefixed message with I/O timeout (W1 fix)
        let len = data.len() as u32;
        tokio::time::timeout(IO_TIMEOUT, async {
            stream.write_all(&len.to_be_bytes()).await?;
            stream.write_all(data).await
        })
        .await
        .map_err(|_| DiscoveryError::Connection {
            address: address.into(),
            reason: "write timeout".into(),
        })?
        .map_err(|e| DiscoveryError::Connection {
            address: address.into(),
            reason: e.to_string(),
        })?;

        // Read response with I/O timeout (W1 fix)
        let resp = tokio::time::timeout(IO_TIMEOUT, async {
            let mut len_buf = [0u8; 4];
            if stream.read_exact(&mut len_buf).await.is_err() {
                return Ok(None);
            }

            let resp_len = u32::from_be_bytes(len_buf) as usize;
            if resp_len > MAX_MESSAGE_SIZE {
                return Err(DiscoveryError::Serialization("response too large".into()));
            }
            let mut resp = vec![0u8; resp_len];
            stream.read_exact(&mut resp).await?;
            Ok(Some(resp))
        })
        .await
        .map_err(|_| DiscoveryError::Connection {
            address: address.into(),
            reason: "read timeout".into(),
        })?;

        resp.map_err(|e: DiscoveryError| e)
    }

    /// Run the heartbeat listener (accepts incoming heartbeats).
    ///
    /// The listener records that remote peers exist and updates their
    /// address/metadata, but does **not** reset the heartbeater's failure
    /// counter. See module-level docs for the design rationale.
    #[allow(clippy::cast_possible_truncation)]
    async fn run_listener(
        listen_address: String,
        local_info: NodeInfo,
        peers: Arc<RwLock<HashMap<u64, PeerState>>>,
        membership_tx: watch::Sender<Vec<NodeInfo>>,
        cancel: CancellationToken,
    ) -> Result<(), DiscoveryError> {
        use tokio::io::{AsyncReadExt, AsyncWriteExt};

        let listener = TcpListener::bind(&listen_address)
            .await
            .map_err(|e| DiscoveryError::Bind(e.to_string()))?;

        // Bound concurrent handler tasks (W3 fix)
        let semaphore = Arc::new(tokio::sync::Semaphore::new(MAX_HANDLER_TASKS));

        loop {
            tokio::select! {
                () = cancel.cancelled() => break,
                accept = listener.accept() => {
                    let (mut stream, _) = accept?;
                    let local_info = local_info.clone();
                    let peers = Arc::clone(&peers);
                    let membership_tx = membership_tx.clone();
                    let permit = Arc::clone(&semaphore);

                    tokio::spawn(async move {
                        // Acquire semaphore permit — drop-guard releases on exit
                        let Ok(_permit) = permit.try_acquire() else {
                            return; // at capacity, drop connection
                        };

                        // Wrap all handler I/O in a timeout (W1 fix)
                        let result = tokio::time::timeout(IO_TIMEOUT, async {
                            let mut len_buf = [0u8; 4];
                            if stream.read_exact(&mut len_buf).await.is_err() {
                                return;
                            }
                            let msg_len = u32::from_be_bytes(len_buf) as usize;
                            if msg_len > MAX_MESSAGE_SIZE {
                                return;
                            }
                            let mut data = vec![0u8; msg_len];
                            if stream.read_exact(&mut data).await.is_err() {
                                return;
                            }

                            if let Ok(remote_info) = Self::deserialize_node_info(&data) {
                                // Skip self — don't add ourselves to the peer list
                                if remote_info.id == local_info.id {
                                    // Still respond so the heartbeater gets a reply
                                    if let Ok(resp) = Self::serialize_node_info(&local_info) {
                                        let len = resp.len() as u32;
                                        let _ = stream.write_all(&len.to_be_bytes()).await;
                                        let _ = stream.write_all(&resp).await;
                                    }
                                    return;
                                }

                                let peer_list = {
                                    let mut guard = peers.write();
                                    let now = chrono::Utc::now().timestamp_millis();
                                    let peer =
                                        guard.entry(remote_info.id.0).or_insert_with(|| {
                                            // First time seeing this peer via the listener.
                                            // We know it can reach us, but we haven't
                                            // confirmed outbound yet — start as Joining.
                                            PeerState {
                                                info: NodeInfo {
                                                    last_heartbeat_ms: now,
                                                    state: NodeState::Joining,
                                                    ..remote_info.clone()
                                                },
                                                missed_heartbeats: 0,
                                                left_ticks: 0,
                                            }
                                        });
                                    // Update addresses/metadata from the remote.
                                    // Do NOT touch missed_heartbeats or state — the
                                    // heartbeater is the sole authority (C1 fix).
                                    peer.info.rpc_address.clone_from(&remote_info.rpc_address);
                                    peer.info.raft_address.clone_from(&remote_info.raft_address);
                                    peer.info.name.clone_from(&remote_info.name);
                                    peer.info.metadata = remote_info.metadata.clone();
                                    peer.info.last_heartbeat_ms = now;
                                    guard.values().map(|p| p.info.clone()).collect::<Vec<_>>()
                                };
                                publish_if_changed(&membership_tx, peer_list);
                            }

                            if let Ok(resp) = Self::serialize_node_info(&local_info) {
                                let len = resp.len() as u32;
                                let _ = stream.write_all(&len.to_be_bytes()).await;
                                let _ = stream.write_all(&resp).await;
                            }
                        })
                        .await;

                        if result.is_err() {
                            // Handler timed out — connection dropped
                        }
                    });
                }
            }
        }

        Ok(())
    }

    /// Run the periodic heartbeat sender.
    ///
    /// Sends heartbeats concurrently to all seeds and uses the responses
    /// to track failure state. The heartbeater is the sole authority on
    /// `missed_heartbeats` and state transitions.
    async fn run_heartbeater(config: StaticDiscoveryConfig, ctx: HeartbeatContext) {
        let local_id = config.local_node.id;
        let mut interval = tokio::time::interval(config.heartbeat_interval);
        // Don't burst missed ticks — skip them to avoid thundering herd (W5 fix)
        interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);

        // seed_to_peer tracks which seed address maps to which node ID.
        // Protected by a Mutex since concurrent heartbeat tasks need it.
        let seed_to_peer = Arc::new(parking_lot::Mutex::new(HashMap::<String, u64>::new()));

        loop {
            tokio::select! {
                () = ctx.cancel.cancelled() => break,
                _ = interval.tick() => {
                    let Ok(data) = Self::serialize_node_info(&config.local_node) else {
                        continue;
                    };
                    let data = Arc::new(data);

                    // Send heartbeats to all seeds concurrently (W5 fix)
                    let mut tasks = Vec::with_capacity(config.seeds.len());
                    for seed in &config.seeds {
                        let seed = seed.clone();
                        let data = Arc::clone(&data);
                        tasks.push(tokio::spawn(async move {
                            let result = Self::send_heartbeat(&seed, &data).await;
                            (seed, result)
                        }));
                    }

                    // Collect results and update peer state
                    for task in tasks {
                        let Ok((seed, result)) = task.await else {
                            continue; // task panicked
                        };

                        if let Ok(Some(resp_data)) = result {
                            if let Ok(remote_info) = Self::deserialize_node_info(&resp_data) {
                                // Skip self
                                if remote_info.id == local_id {
                                    continue;
                                }

                                // Record seed → peer mapping
                                seed_to_peer.lock().insert(seed, remote_info.id.0);

                                let mut peers = ctx.peers.write();
                                let now = chrono::Utc::now().timestamp_millis();
                                let peer =
                                    peers.entry(remote_info.id.0).or_insert_with(|| {
                                        PeerState {
                                            info: remote_info.clone(),
                                            missed_heartbeats: 0,
                                            left_ticks: 0,
                                        }
                                    });
                                peer.info = NodeInfo {
                                    last_heartbeat_ms: now,
                                    state: NodeState::Active,
                                    ..remote_info
                                };
                                peer.missed_heartbeats = 0;
                                peer.left_ticks = 0;
                            }
                        } else {
                            // Heartbeat failed — increment missed counter
                            let map = seed_to_peer.lock();
                            if let Some(&peer_id) = map.get(seed.as_str()) {
                                drop(map);
                                let mut peers = ctx.peers.write();
                                if let Some(peer) = peers.get_mut(&peer_id) {
                                    peer.missed_heartbeats += 1;
                                    if peer.missed_heartbeats >= config.dead_threshold {
                                        peer.info.state = NodeState::Left;
                                    } else if peer.missed_heartbeats
                                        >= config.suspect_threshold
                                    {
                                        peer.info.state = NodeState::Suspected;
                                    }
                                }
                            }
                        }
                    }

                    // Reap peers stuck in Left state for too long (W2 fix)
                    {
                        let mut peers = ctx.peers.write();
                        peers.retain(|_id, peer| {
                            if peer.info.state == NodeState::Left {
                                peer.left_ticks += 1;
                                peer.left_ticks < LEFT_REAP_THRESHOLD
                            } else {
                                true
                            }
                        });
                    }

                    // Also clean up seed_to_peer for reaped peers (W2 fix)
                    {
                        let peers = ctx.peers.read();
                        let mut map = seed_to_peer.lock();
                        map.retain(|_, peer_id| peers.contains_key(peer_id));
                    }

                    let peer_list: Vec<NodeInfo> = {
                        let peers = ctx.peers.read();
                        peers.values().map(|p| p.info.clone()).collect()
                    };
                    publish_if_changed(&ctx.membership_tx, peer_list);
                }
            }
        }
    }
}

/// Update the membership watch only when the list actually changed.
/// `watch::Sender::send` notifies receivers on EVERY call, and the
/// heartbeater publishes each tick — unconditional sends starve any
/// consumer that debounces on `changed()` waiting for a quiet period
/// (the rebalance controller never saw 5s of quiet, so vnode rotation
/// never ran and a dead node's vnodes were never shed).
fn publish_if_changed(tx: &watch::Sender<Vec<NodeInfo>>, mut peer_list: Vec<NodeInfo>) {
    // Stable order: the peer map iterates in arbitrary order, and
    // equality must not depend on it.
    peer_list.sort_by_key(|n| n.id.0);
    tx.send_if_modified(|cur| {
        // Compare everything EXCEPT `last_heartbeat_ms` — it refreshes
        // every tick, so whole-struct equality would still notify
        // continuously. All other fields (state, addresses, name,
        // metadata/locality) are real membership changes watchers must
        // see.
        fn same_member(a: &NodeInfo, b: &NodeInfo) -> bool {
            let mut a = a.clone();
            a.last_heartbeat_ms = b.last_heartbeat_ms;
            a == *b
        }
        let same = cur.len() == peer_list.len()
            && cur.iter().zip(&peer_list).all(|(a, b)| same_member(a, b));
        if same {
            false
        } else {
            *cur = peer_list;
            true
        }
    });
}

/// Shared context for the heartbeater background task.
struct HeartbeatContext {
    peers: Arc<RwLock<HashMap<u64, PeerState>>>,
    membership_tx: watch::Sender<Vec<NodeInfo>>,
    cancel: CancellationToken,
}

impl Discovery for StaticDiscovery {
    async fn start(&mut self) -> Result<(), DiscoveryError> {
        if self.started {
            return Ok(());
        }

        // Create a fresh cancellation token so restart after stop() works (W4 fix)
        self.cancel = CancellationToken::new();

        let peers = Arc::clone(&self.peers);
        let membership_tx = self.membership_tx.clone();
        let cancel = self.cancel.clone();
        let listen_address = self.config.listen_address.clone();
        let local_info = self.config.local_node.clone();

        // Spawn listener and keep the handle (W4 fix)
        self.listener_handle = Some(tokio::spawn(Self::run_listener(
            listen_address,
            local_info,
            Arc::clone(&peers),
            membership_tx.clone(),
            cancel.clone(),
        )));

        // Spawn heartbeater and keep the handle (W4 fix)
        self.heartbeater_handle = Some(tokio::spawn(Self::run_heartbeater(
            self.config.clone(),
            HeartbeatContext {
                peers,
                membership_tx,
                cancel,
            },
        )));

        self.started = true;
        Ok(())
    }

    async fn peers(&self) -> Result<Vec<NodeInfo>, DiscoveryError> {
        if !self.started {
            return Err(DiscoveryError::NotStarted);
        }
        let peers = self.peers.read();
        Ok(peers.values().map(|p| p.info.clone()).collect())
    }

    async fn announce(&self, info: NodeInfo) -> Result<(), DiscoveryError> {
        if !self.started {
            return Err(DiscoveryError::NotStarted);
        }
        {
            let mut peers = self.peers.write();
            peers.insert(
                info.id.0,
                PeerState {
                    info,
                    missed_heartbeats: 0,
                    left_ticks: 0,
                },
            );
        }
        self.broadcast_membership();
        Ok(())
    }

    fn membership_watch(&self) -> watch::Receiver<Vec<NodeInfo>> {
        self.membership_rx.clone()
    }

    async fn stop(&mut self) -> Result<(), DiscoveryError> {
        self.cancel.cancel();
        self.started = false;

        // Wait for spawned tasks to exit (W4 fix)
        if let Some(h) = self.listener_handle.take() {
            let _ = h.await;
        }
        if let Some(h) = self.heartbeater_handle.take() {
            let _ = h.await;
        }

        Ok(())
    }
}

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

    #[test]
    fn test_config_default() {
        let config = StaticDiscoveryConfig::default();
        assert_eq!(config.heartbeat_interval, Duration::from_secs(1));
        assert_eq!(config.suspect_threshold, 3);
        assert_eq!(config.dead_threshold, 10);
    }

    #[test]
    fn test_serialize_round_trip() {
        let info = NodeInfo {
            id: NodeId(42),
            name: "test".into(),
            rpc_address: "127.0.0.1:9000".into(),
            raft_address: "127.0.0.1:9001".into(),
            state: NodeState::Active,
            metadata: NodeMetadata::default(),
            last_heartbeat_ms: 1000,
        };

        let data = StaticDiscovery::serialize_node_info(&info).unwrap();
        let back = StaticDiscovery::deserialize_node_info(&data).unwrap();
        assert_eq!(back.id, NodeId(42));
        assert_eq!(back.name, "test");
    }

    #[test]
    fn test_deserialize_invalid() {
        let result = StaticDiscovery::deserialize_node_info(&[0xff, 0xff]);
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_not_started_errors() {
        let config = StaticDiscoveryConfig::default();
        let disc = StaticDiscovery::new(config);
        assert!(disc.peers().await.is_err());
    }

    #[tokio::test]
    async fn test_start_stop() {
        let config = StaticDiscoveryConfig {
            listen_address: "127.0.0.1:0".into(),
            ..StaticDiscoveryConfig::default()
        };
        let mut disc = StaticDiscovery::new(config);
        disc.start().await.unwrap();
        assert!(disc.started);
        disc.stop().await.unwrap();
        assert!(!disc.started);
    }

    #[tokio::test]
    async fn test_double_start_ok() {
        let config = StaticDiscoveryConfig {
            listen_address: "127.0.0.1:0".into(),
            ..StaticDiscoveryConfig::default()
        };
        let mut disc = StaticDiscovery::new(config);
        disc.start().await.unwrap();
        disc.start().await.unwrap(); // Should be idempotent
        disc.stop().await.unwrap();
    }

    #[tokio::test]
    async fn test_membership_watch() {
        let config = StaticDiscoveryConfig {
            listen_address: "127.0.0.1:0".into(),
            ..StaticDiscoveryConfig::default()
        };
        let disc = StaticDiscovery::new(config);
        let rx = disc.membership_watch();
        assert!(rx.borrow().is_empty());
    }

    #[tokio::test]
    async fn test_announce_adds_peer() {
        let config = StaticDiscoveryConfig {
            listen_address: "127.0.0.1:0".into(),
            ..StaticDiscoveryConfig::default()
        };
        let mut disc = StaticDiscovery::new(config);
        disc.start().await.unwrap();

        let peer = NodeInfo {
            id: NodeId(99),
            name: "peer".into(),
            rpc_address: "127.0.0.1:8000".into(),
            raft_address: "127.0.0.1:8001".into(),
            state: NodeState::Active,
            metadata: NodeMetadata::default(),
            last_heartbeat_ms: 0,
        };
        disc.announce(peer).await.unwrap();

        let peers = disc.peers().await.unwrap();
        assert_eq!(peers.len(), 1);
        assert_eq!(peers[0].id, NodeId(99));

        disc.stop().await.unwrap();
    }

    #[tokio::test]
    async fn test_two_node_heartbeat() {
        let listener1 = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr1 = listener1.local_addr().unwrap().to_string();
        drop(listener1);

        let listener2 = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr2 = listener2.local_addr().unwrap().to_string();
        drop(listener2);

        let config1 = StaticDiscoveryConfig {
            local_node: NodeInfo {
                id: NodeId(1),
                name: "node-1".into(),
                rpc_address: addr1.clone(),
                raft_address: addr1.clone(),
                state: NodeState::Active,
                metadata: NodeMetadata::default(),
                last_heartbeat_ms: 0,
            },
            seeds: vec![addr2.clone()],
            heartbeat_interval: Duration::from_millis(100),
            listen_address: addr1.clone(),
            ..StaticDiscoveryConfig::default()
        };

        let config2 = StaticDiscoveryConfig {
            local_node: NodeInfo {
                id: NodeId(2),
                name: "node-2".into(),
                rpc_address: addr2.clone(),
                raft_address: addr2.clone(),
                state: NodeState::Active,
                metadata: NodeMetadata::default(),
                last_heartbeat_ms: 0,
            },
            seeds: vec![addr1],
            heartbeat_interval: Duration::from_millis(100),
            listen_address: addr2,
            ..StaticDiscoveryConfig::default()
        };

        let mut disc1 = StaticDiscovery::new(config1);
        let mut disc2 = StaticDiscovery::new(config2);

        disc1.start().await.unwrap();
        disc2.start().await.unwrap();

        tokio::time::sleep(Duration::from_millis(500)).await;

        let peers1 = disc1.peers().await.unwrap();
        let peers2 = disc2.peers().await.unwrap();

        assert!(
            !peers1.is_empty() || !peers2.is_empty(),
            "at least one node should have discovered peers"
        );

        disc1.stop().await.unwrap();
        disc2.stop().await.unwrap();
    }

    #[tokio::test]
    async fn test_restart_after_stop() {
        let config = StaticDiscoveryConfig {
            listen_address: "127.0.0.1:0".into(),
            ..StaticDiscoveryConfig::default()
        };
        let mut disc = StaticDiscovery::new(config);

        // First start/stop cycle
        disc.start().await.unwrap();
        disc.stop().await.unwrap();

        // Second start should work (fresh CancellationToken)
        disc.start().await.unwrap();
        assert!(disc.started);
        disc.stop().await.unwrap();
    }
}