stochastic-routing-extended 1.0.2

//! Transport manager: multiplexes frames across active transports.
//!
//! Responsible for:
//! - Maintaining the pool of active transports
//! - Per-transport health monitoring (RTT, failure count, blockage detection)
//! - Dispatching frames according to routing masks
//! - Automatic fallback when DPI blocks a transport

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

use bytes::Bytes;

use super::{Transport, TransportKind};
use crate::channel::fallback::FallbackController;
use crate::channel::health::HealthMonitor;
use crate::error::{SrxError, TransportError};
use crate::routing::RoutingMask;
use crate::seed::SeedRng;

/// Default consecutive failure count before declaring a transport blocked.
const DEFAULT_BLOCK_THRESHOLD: u32 = 5;

/// Manages a pool of active transports and dispatches frames.
pub struct TransportManager {
    transports: HashMap<TransportKind, Arc<dyn Transport>>,
    /// Per-transport health monitors (created when a transport is added).
    health: HashMap<TransportKind, HealthMonitor>,
    /// Fallback controller for automatic transport switching.
    fallback: Option<FallbackController>,
}

impl TransportManager {
    pub fn new() -> Self {
        Self {
            transports: HashMap::new(),
            health: HashMap::new(),
            fallback: None,
        }
    }

    /// Register a transport in the pool (also creates a health monitor).
    pub fn add_transport(&mut self, transport: Arc<dyn Transport>) {
        let kind = transport.kind();
        self.transports.insert(kind, transport);
        self.health
            .entry(kind)
            .or_insert_with(|| HealthMonitor::new(DEFAULT_BLOCK_THRESHOLD));
    }

    /// Remove a transport from the pool.
    pub fn remove_transport(&mut self, kind: TransportKind) {
        self.transports.remove(&kind);
        self.health.remove(&kind);
    }

    /// Get a reference to a specific transport.
    pub fn get(&self, kind: TransportKind) -> Option<&Arc<dyn Transport>> {
        self.transports.get(&kind)
    }

    /// List all active transport kinds.
    pub fn active_kinds(&self) -> Vec<TransportKind> {
        self.transports.keys().copied().collect()
    }

    /// List healthy (non-blocked) transport kinds, sorted by descending health score.
    pub fn healthy_kinds(&self) -> Vec<TransportKind> {
        let mut scored: Vec<_> = self
            .transports
            .keys()
            .filter_map(|&k| {
                let monitor = self.health.get(&k)?;
                if monitor.is_blocked() {
                    None
                } else {
                    Some((k, monitor.score()))
                }
            })
            .collect();
        scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        scored.into_iter().map(|(k, _)| k).collect()
    }

    /// Get the health monitor for a transport (if registered).
    pub fn health_monitor(&self, kind: TransportKind) -> Option<&HealthMonitor> {
        self.health.get(&kind)
    }

    /// Check whether a specific transport is currently blocked.
    pub fn is_blocked(&self, kind: TransportKind) -> bool {
        self.health.get(&kind).is_some_and(|m| m.is_blocked())
    }

    /// Get health scores for all active transports (blocked transports get 0.0).
    pub fn health_scores(&self) -> Vec<(TransportKind, f64)> {
        self.transports
            .keys()
            .map(|&kind| {
                let monitor = self.health.get(&kind);
                let score = monitor
                    .filter(|m| !m.is_blocked())
                    .map(|m| m.score())
                    .unwrap_or(0.0);
                (kind, score)
            })
            .collect()
    }

    /// Set a fallback controller for automatic transport switching.
    pub fn set_fallback(&mut self, controller: FallbackController) {
        self.fallback = Some(controller);
    }

    /// If the current preferred transport is blocked, advance the fallback
    /// controller to the next healthy transport and return it.
    pub fn try_fallback(&mut self) -> Option<TransportKind> {
        let fb = self.fallback.as_mut()?;
        let current = fb.current()?;
        let blocked = self.health.get(&current).is_some_and(|m| m.is_blocked());
        if blocked {
            fb.fallback()
        } else {
            Some(current)
        }
    }

    /// Send on the transport registered for `kind`.
    ///
    /// Measures wall-clock time and records success or failure on the
    /// per-transport [`HealthMonitor`].
    pub async fn send(&mut self, kind: TransportKind, data: Bytes) -> crate::error::Result<()> {
        let t = self
            .transports
            .get(&kind)
            .ok_or_else(|| SrxError::Transport(TransportError::NotRegistered(format!("{kind:?}"))))?
            .clone();
        let start = Instant::now();
        match t.send(data).await {
            Ok(()) => {
                if let Some(m) = self.health.get_mut(&kind) {
                    m.record_success(start.elapsed());
                }
                Ok(())
            }
            Err(e) => {
                if let Some(m) = self.health.get_mut(&kind) {
                    m.record_failure();
                }
                Err(e)
            }
        }
    }

    /// Receive from the transport registered for `kind`.
    ///
    /// Measures wall-clock time and records success or failure on the
    /// per-transport [`HealthMonitor`].
    pub async fn recv(&mut self, kind: TransportKind) -> crate::error::Result<Bytes> {
        let t = self
            .transports
            .get(&kind)
            .ok_or_else(|| SrxError::Transport(TransportError::NotRegistered(format!("{kind:?}"))))?
            .clone();
        let start = Instant::now();
        match t.recv().await {
            Ok(data) => {
                if let Some(m) = self.health.get_mut(&kind) {
                    m.record_success(start.elapsed());
                }
                Ok(data)
            }
            Err(e) => {
                if let Some(m) = self.health.get_mut(&kind) {
                    m.record_failure();
                }
                Err(e)
            }
        }
    }

    /// Send on the first transport from [`RoutingMask::generate`] that is registered in this manager.
    ///
    /// Prefers healthy transports: if a mask-selected transport is blocked, it
    /// is skipped in favour of the next candidate.
    pub async fn send_with_routing_mask(
        &mut self,
        rng: &mut SeedRng,
        available: &[TransportKind],
        frame_counter: u64,
        data: Bytes,
    ) -> crate::error::Result<TransportKind> {
        let mask = RoutingMask::generate(rng, available, frame_counter);
        for k in &mask.transports {
            if self.transports.contains_key(k) && !self.is_blocked(*k) {
                self.send(*k, data).await?;
                return Ok(*k);
            }
        }
        // Fallback: try any non-blocked registered transport.
        for k in &mask.transports {
            if self.transports.contains_key(k) {
                self.send(*k, data).await?;
                return Ok(*k);
            }
        }
        Err(SrxError::Transport(TransportError::NotRegistered(
            "no transport from routing mask is registered".into(),
        )))
    }

    /// Receive from the first transport from [`RoutingMask::generate`] that is registered in this manager.
    pub async fn recv_with_routing_mask(
        &mut self,
        rng: &mut SeedRng,
        available: &[TransportKind],
        frame_counter: u64,
    ) -> crate::error::Result<(TransportKind, Bytes)> {
        let mask = RoutingMask::generate(rng, available, frame_counter);
        for k in &mask.transports {
            if self.transports.contains_key(k) {
                let data = self.recv(*k).await?;
                return Ok((*k, data));
            }
        }
        Err(SrxError::Transport(TransportError::NotRegistered(
            "no transport from routing mask is registered".into(),
        )))
    }

    /// Health-check all transports and remove unhealthy ones.
    pub async fn health_check(&mut self) {
        let mut to_remove = Vec::new();
        for (kind, transport) in &self.transports {
            if !transport.is_healthy().await {
                to_remove.push(*kind);
            }
        }
        for kind in to_remove {
            self.transports.remove(&kind);
            self.health.remove(&kind);
        }
    }
}

impl Default for TransportManager {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::error::{SrxError, TransportError};
    use crate::seed::SeedRng;
    use crate::transport::{TcpTransport, UdpTransport};

    #[tokio::test]
    async fn send_recv_dispatches_tcp() {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let server = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let t = TcpTransport::from_stream(stream);
            let got = t.recv().await.unwrap();
            assert_eq!(got.as_ref(), b"mgr-ping");
            t.send(Bytes::from_static(b"mgr-pong")).await.unwrap();
        });

        let client = TcpTransport::connect(addr).await.unwrap();
        let mut mgr = TransportManager::new();
        mgr.add_transport(Arc::new(client));

        mgr.send(TransportKind::Tcp, Bytes::from_static(b"mgr-ping"))
            .await
            .unwrap();
        let reply = mgr.recv(TransportKind::Tcp).await.unwrap();
        assert_eq!(reply.as_ref(), b"mgr-pong");

        server.await.unwrap();
    }

    #[tokio::test]
    async fn send_with_routing_mask_uses_registered_kind() {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let server = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let t = TcpTransport::from_stream(stream);
            let got = t.recv().await.unwrap();
            assert_eq!(got.as_ref(), b"mask-ping");
            t.send(Bytes::from_static(b"mask-pong")).await.unwrap();
        });

        let client = TcpTransport::connect(addr).await.unwrap();
        let mut mgr = TransportManager::new();
        mgr.add_transport(Arc::new(client));

        let mut rng = SeedRng::new([0xBBu8; 32]);
        let available = [TransportKind::Tcp];
        let used = mgr
            .send_with_routing_mask(&mut rng, &available, 7, Bytes::from_static(b"mask-ping"))
            .await
            .unwrap();
        assert_eq!(used, TransportKind::Tcp);
        let reply = mgr.recv(used).await.unwrap();
        assert_eq!(reply.as_ref(), b"mask-pong");

        server.await.unwrap();
    }

    #[tokio::test]
    async fn recv_with_routing_mask_uses_registered_kind() {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let server = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let mut mgr = TransportManager::new();
            mgr.add_transport(Arc::new(TcpTransport::from_stream(stream)));
            let mut rng = SeedRng::new([0xCCu8; 32]);
            mgr.recv_with_routing_mask(&mut rng, &[TransportKind::Tcp], 3)
                .await
        });

        let client = TcpTransport::connect(addr).await.unwrap();
        client
            .send(Bytes::from_static(b"recv-mask-ping"))
            .await
            .unwrap();

        let got = server.await.unwrap().unwrap();
        assert_eq!(got.0, TransportKind::Tcp);
        assert_eq!(got.1.as_ref(), b"recv-mask-ping");
    }

    #[tokio::test]
    async fn missing_kind_returns_not_registered() {
        let mut mgr = TransportManager::new();
        let e = mgr
            .send(TransportKind::Udp, Bytes::from_static(b"x"))
            .await
            .expect_err("expected error");
        match e {
            SrxError::Transport(TransportError::NotRegistered(name)) => {
                assert!(name.contains("Udp"));
            }
            other => panic!("unexpected: {other:?}"),
        }
    }

    #[tokio::test]
    async fn send_records_health_success() {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let server = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let t = TcpTransport::from_stream(stream);
            t.recv().await.unwrap();
        });

        let client = TcpTransport::connect(addr).await.unwrap();
        let mut mgr = TransportManager::new();
        mgr.add_transport(Arc::new(client));

        mgr.send(TransportKind::Tcp, Bytes::from_static(b"health-ok"))
            .await
            .unwrap();

        let monitor = mgr.health_monitor(TransportKind::Tcp).unwrap();
        assert!(!monitor.is_blocked());
        assert!(monitor.score() > 0.0);
        assert_eq!(monitor.consecutive_failures, 0);

        server.await.unwrap();
    }

    #[test]
    fn healthy_kinds_excludes_blocked() {
        let mut mgr = TransportManager::new();
        mgr.health.insert(TransportKind::Tcp, HealthMonitor::new(2));
        mgr.health.insert(TransportKind::Udp, HealthMonitor::new(2));

        assert!(!mgr.is_blocked(TransportKind::Tcp));

        // Block TCP.
        mgr.health
            .get_mut(&TransportKind::Tcp)
            .unwrap()
            .record_failure();
        mgr.health
            .get_mut(&TransportKind::Tcp)
            .unwrap()
            .record_failure();
        assert!(mgr.is_blocked(TransportKind::Tcp));
        assert!(!mgr.is_blocked(TransportKind::Udp));
    }

    #[test]
    fn fallback_advances_on_block() {
        let mut mgr = TransportManager::new();
        mgr.health.insert(TransportKind::Tcp, HealthMonitor::new(2));
        mgr.health.insert(TransportKind::Udp, HealthMonitor::new(2));
        mgr.set_fallback(FallbackController::new(vec![
            TransportKind::Tcp,
            TransportKind::Udp,
        ]));

        // TCP healthy — fallback stays on TCP.
        assert_eq!(mgr.try_fallback(), Some(TransportKind::Tcp));

        // Block TCP.
        mgr.health
            .get_mut(&TransportKind::Tcp)
            .unwrap()
            .record_failure();
        mgr.health
            .get_mut(&TransportKind::Tcp)
            .unwrap()
            .record_failure();

        // Fallback should advance to UDP.
        assert_eq!(mgr.try_fallback(), Some(TransportKind::Udp));
    }

    // -----------------------------------------------------------------------
    // Multi-transport tests (TCP + UDP simultaneously)
    // -----------------------------------------------------------------------

    #[tokio::test]
    async fn register_tcp_and_udp_simultaneously() {
        // Set up TCP listener.
        let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let tcp_addr = tcp_listener.local_addr().unwrap();

        // Set up UDP pair (connected sockets).
        let udp_server_sock = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let udp_client_sock = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let udp_srv_addr = udp_server_sock.local_addr().unwrap();
        let udp_cli_addr = udp_client_sock.local_addr().unwrap();
        udp_server_sock.connect(udp_cli_addr).await.unwrap();
        udp_client_sock.connect(udp_srv_addr).await.unwrap();

        let server = tokio::spawn(async move {
            // Accept TCP and handle UDP echo in parallel.
            let (tcp_stream, _) = tcp_listener.accept().await.unwrap();
            let tcp = TcpTransport::from_stream(tcp_stream);
            let udp = UdpTransport::from_socket(udp_server_sock);

            // Echo on both transports.
            let tcp_data = tcp.recv().await.unwrap();
            tcp.send(Bytes::from_static(b"tcp-pong")).await.unwrap();
            assert_eq!(tcp_data.as_ref(), b"tcp-ping");

            let udp_data = udp.recv().await.unwrap();
            udp.send(Bytes::from_static(b"udp-pong")).await.unwrap();
            assert_eq!(udp_data.as_ref(), b"udp-ping");
        });

        let tcp_client = TcpTransport::connect(tcp_addr).await.unwrap();
        let udp_client = UdpTransport::from_socket(udp_client_sock);

        let mut mgr = TransportManager::new();
        mgr.add_transport(Arc::new(tcp_client));
        mgr.add_transport(Arc::new(udp_client));

        assert_eq!(
            mgr.active_kinds().len(),
            2,
            "should have 2 active transports"
        );
        assert!(mgr.active_kinds().contains(&TransportKind::Tcp));
        assert!(mgr.active_kinds().contains(&TransportKind::Udp));

        // Send/recv on TCP.
        mgr.send(TransportKind::Tcp, Bytes::from_static(b"tcp-ping"))
            .await
            .unwrap();
        let reply = mgr.recv(TransportKind::Tcp).await.unwrap();
        assert_eq!(reply.as_ref(), b"tcp-pong");

        // Send/recv on UDP.
        mgr.send(TransportKind::Udp, Bytes::from_static(b"udp-ping"))
            .await
            .unwrap();
        let reply = mgr.recv(TransportKind::Udp).await.unwrap();
        assert_eq!(reply.as_ref(), b"udp-pong");

        server.await.unwrap();
    }

    #[tokio::test]
    async fn send_with_routing_mask_multiple_transports() {
        let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let tcp_addr = tcp_listener.local_addr().unwrap();

        let udp_srv = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let udp_cli = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let udp_srv_addr = udp_srv.local_addr().unwrap();
        let udp_cli_addr = udp_cli.local_addr().unwrap();
        udp_srv.connect(udp_cli_addr).await.unwrap();
        udp_cli.connect(udp_srv_addr).await.unwrap();

        let server = tokio::spawn(async move {
            let (tcp_stream, _) = tcp_listener.accept().await.unwrap();
            let tcp = TcpTransport::from_stream(tcp_stream);
            let udp = UdpTransport::from_socket(udp_srv);

            // Receive on whichever transport the client picks first.
            tokio::select! {
                result = tcp.recv() => {
                    let data = result.unwrap();
                    tcp.send(Bytes::from_static(b"tcp-ack")).await.unwrap();
                    assert_eq!(data.as_ref(), b"mask-tcp");
                }
                result = udp.recv() => {
                    let data = result.unwrap();
                    udp.send(Bytes::from_static(b"udp-ack")).await.unwrap();
                    assert_eq!(data.as_ref(), b"mask-udp");
                }
            }
        });

        let tcp_client = TcpTransport::connect(tcp_addr).await.unwrap();
        let udp_client = UdpTransport::from_socket(udp_cli);

        let mut mgr = TransportManager::new();
        mgr.add_transport(Arc::new(tcp_client));
        mgr.add_transport(Arc::new(udp_client));

        let mut rng = SeedRng::new([0xDDu8; 32]);
        let available = [TransportKind::Tcp, TransportKind::Udp];

        // Send with routing mask — should pick one of the two.
        let used = mgr
            .send_with_routing_mask(&mut rng, &available, 42, Bytes::from_static(b"mask-tcp"))
            .await
            .unwrap();

        // The used kind must be one of the available transports.
        assert!(
            used == TransportKind::Tcp || used == TransportKind::Udp,
            "routing mask selected unexpected kind {used:?}"
        );

        // Receive the reply on the same transport.
        let reply = mgr.recv(used).await.unwrap();
        assert!(
            reply.as_ref() == b"tcp-ack" || reply.as_ref() == b"udp-ack",
            "unexpected reply: {:?}",
            reply
        );

        server.await.unwrap();
    }

    #[tokio::test]
    async fn recv_with_routing_mask_multiple_transports() {
        let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let tcp_addr = tcp_listener.local_addr().unwrap();

        let udp_srv = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let udp_cli = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let udp_srv_addr = udp_srv.local_addr().unwrap();
        let udp_cli_addr = udp_cli.local_addr().unwrap();
        udp_srv.connect(udp_cli_addr).await.unwrap();
        udp_cli.connect(udp_srv_addr).await.unwrap();

        let server = tokio::spawn(async move {
            let (tcp_stream, _) = tcp_listener.accept().await.unwrap();
            let mut mgr = TransportManager::new();
            mgr.add_transport(Arc::new(TcpTransport::from_stream(tcp_stream)));
            mgr.add_transport(Arc::new(UdpTransport::from_socket(udp_srv)));
            let mut rng = SeedRng::new([0xEEu8; 32]);
            mgr.recv_with_routing_mask(&mut rng, &[TransportKind::Tcp, TransportKind::Udp], 10)
                .await
        });

        // Client sends on TCP.
        let tcp_client = TcpTransport::connect(tcp_addr).await.unwrap();
        tcp_client
            .send(Bytes::from_static(b"multi-recv-tcp"))
            .await
            .unwrap();

        let (kind, data) = server.await.unwrap().unwrap();
        assert!(
            kind == TransportKind::Tcp || kind == TransportKind::Udp,
            "unexpected recv kind {kind:?}"
        );
        assert_eq!(data.as_ref(), b"multi-recv-tcp");
    }

    #[tokio::test]
    async fn health_check_removes_unhealthy_transport() {
        let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let tcp_addr = tcp_listener.local_addr().unwrap();

        // Create UDP socket but close it immediately to make it unhealthy.
        let udp_sock = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let udp = UdpTransport::from_socket(udp_sock);
        // Close the UDP transport so is_healthy returns false.
        udp.close().await.unwrap();

        let server = tokio::spawn(async move {
            let (stream, _) = tcp_listener.accept().await.unwrap();
            let _tcp = TcpTransport::from_stream(stream);
            // TCP stays healthy (socket open).
            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
        });

        let tcp_client = TcpTransport::connect(tcp_addr).await.unwrap();

        let mut mgr = TransportManager::new();
        mgr.add_transport(Arc::new(tcp_client));
        mgr.add_transport(Arc::new(udp));

        assert_eq!(mgr.active_kinds().len(), 2);

        // Run health check — UDP should be removed.
        mgr.health_check().await;

        let kinds = mgr.active_kinds();
        assert_eq!(kinds.len(), 1, "only TCP should remain after health check");
        assert!(kinds.contains(&TransportKind::Tcp));
        assert!(
            !kinds.contains(&TransportKind::Udp),
            "UDP should have been removed"
        );

        server.await.unwrap();
    }

    #[test]
    fn healthy_kinds_sorts_by_score_descending() {
        // Test the score comparison logic directly on HealthMonitor.
        let mut tcp_mon = HealthMonitor::new(5);
        let mut udp_mon = HealthMonitor::new(5);
        let mut quic_mon = HealthMonitor::new(5);

        tcp_mon.record_success(std::time::Duration::from_millis(5));
        udp_mon.record_success(std::time::Duration::from_millis(50));
        quic_mon.record_success(std::time::Duration::from_millis(200));

        // Verify score ordering.
        let tcp_score = tcp_mon.score();
        let udp_score = udp_mon.score();
        let quic_score = quic_mon.score();

        assert!(
            tcp_score > udp_score,
            "TCP score {tcp_score} should be > UDP score {udp_score}"
        );
        assert!(
            udp_score > quic_score,
            "UDP score {udp_score} should be > QUIC score {quic_score}"
        );

        // Verify blocked transports get 0.0.
        let mut blocked_mon = HealthMonitor::new(2);
        blocked_mon.record_failure();
        blocked_mon.record_failure();
        assert!(blocked_mon.is_blocked());
        assert_eq!(blocked_mon.score(), 0.0);
    }

    #[tokio::test]
    async fn healthy_kinds_excludes_blocked_includes_healthy() {
        let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let tcp_addr = tcp_listener.local_addr().unwrap();

        let udp_srv = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let udp_cli = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let udp_srv_addr = udp_srv.local_addr().unwrap();
        let udp_cli_addr = udp_cli.local_addr().unwrap();
        udp_srv.connect(udp_cli_addr).await.unwrap();
        udp_cli.connect(udp_srv_addr).await.unwrap();

        let server = tokio::spawn(async move {
            let (tcp_stream, _) = tcp_listener.accept().await.unwrap();
            let _tcp = TcpTransport::from_stream(tcp_stream);
            let _udp = UdpTransport::from_socket(udp_srv);
        });

        let tcp_client = TcpTransport::connect(tcp_addr).await.unwrap();
        let udp_client = UdpTransport::from_socket(udp_cli);

        let mut mgr = TransportManager::new();
        mgr.add_transport(Arc::new(tcp_client));
        mgr.add_transport(Arc::new(udp_client));

        // Block TCP by recording failures up to threshold.
        for _ in 0..DEFAULT_BLOCK_THRESHOLD {
            mgr.health
                .get_mut(&TransportKind::Tcp)
                .unwrap()
                .record_failure();
        }

        // Give UDP a good score with low RTT.
        let udp_mon = mgr.health.get_mut(&TransportKind::Udp).unwrap();
        udp_mon.record_success(std::time::Duration::from_millis(10));

        let kinds = mgr.healthy_kinds();
        assert_eq!(kinds.len(), 1, "blocked TCP should be excluded");
        assert!(!kinds.contains(&TransportKind::Tcp));
        assert!(kinds.contains(&TransportKind::Udp));

        server.await.unwrap();
    }

    #[tokio::test]
    async fn fallback_when_one_transport_blocked() {
        let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let tcp_addr = tcp_listener.local_addr().unwrap();

        let udp_srv = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let udp_cli = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let udp_srv_addr = udp_srv.local_addr().unwrap();
        let udp_cli_addr = udp_cli.local_addr().unwrap();
        udp_srv.connect(udp_cli_addr).await.unwrap();
        udp_cli.connect(udp_srv_addr).await.unwrap();

        let server = tokio::spawn(async move {
            let (tcp_stream, _) = tcp_listener.accept().await.unwrap();
            let tcp = TcpTransport::from_stream(tcp_stream);
            let udp = UdpTransport::from_socket(udp_srv);

            // TCP: receive the initial message (before we block it on client side).
            let _ = tcp.recv().await;

            // UDP: handle the fallback message.
            let data = udp.recv().await.unwrap();
            assert_eq!(data.as_ref(), b"fallback-via-udp");
            udp.send(Bytes::from_static(b"udp-fallback-ok"))
                .await
                .unwrap();
        });

        let tcp_client = TcpTransport::connect(tcp_addr).await.unwrap();
        let udp_client = UdpTransport::from_socket(udp_cli);

        let mut mgr = TransportManager::new();
        mgr.add_transport(Arc::new(tcp_client));
        mgr.add_transport(Arc::new(udp_client));
        mgr.set_fallback(FallbackController::new(vec![
            TransportKind::Tcp,
            TransportKind::Udp,
        ]));

        // Send one message on TCP to establish health.
        mgr.send(TransportKind::Tcp, Bytes::from_static(b"pre-block"))
            .await
            .unwrap();

        // Now artificially block TCP by recording failures.
        let tcp_mon = mgr.health.get_mut(&TransportKind::Tcp).unwrap();
        tcp_mon.record_failure();
        tcp_mon.record_failure();
        tcp_mon.record_failure();
        tcp_mon.record_failure();
        tcp_mon.record_failure();
        assert!(mgr.is_blocked(TransportKind::Tcp));

        // try_fallback should advance to UDP.
        let fallback_kind = mgr.try_fallback();
        assert_eq!(
            fallback_kind,
            Some(TransportKind::Udp),
            "fallback should switch to UDP when TCP is blocked"
        );

        // Send on UDP via fallback.
        let used = fallback_kind.unwrap();
        mgr.send(used, Bytes::from_static(b"fallback-via-udp"))
            .await
            .unwrap();
        let reply = mgr.recv(used).await.unwrap();
        assert_eq!(reply.as_ref(), b"udp-fallback-ok");

        server.await.unwrap();
    }

    #[tokio::test]
    async fn send_with_routing_mask_skips_blocked_uses_healthy() {
        let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let tcp_addr = tcp_listener.local_addr().unwrap();

        let udp_srv = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let udp_cli = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let udp_srv_addr = udp_srv.local_addr().unwrap();
        let udp_cli_addr = udp_cli.local_addr().unwrap();
        udp_srv.connect(udp_cli_addr).await.unwrap();
        udp_cli.connect(udp_srv_addr).await.unwrap();

        // Server: handle both TCP and UDP.
        let server = tokio::spawn(async move {
            let (tcp_stream, _) = tcp_listener.accept().await.unwrap();
            let tcp = TcpTransport::from_stream(tcp_stream);
            let udp = UdpTransport::from_socket(udp_srv);

            // Try to receive on either transport.
            tokio::select! {
                result = tcp.recv() => {
                    let data = result.unwrap();
                    tcp.send(Bytes::from_static(b"tcp-ok")).await.unwrap();
                    assert_eq!(data.as_ref(), b"skip-blocked-tcp");
                }
                result = udp.recv() => {
                    let data = result.unwrap();
                    udp.send(Bytes::from_static(b"udp-ok")).await.unwrap();
                    assert_eq!(data.as_ref(), b"skip-blocked-tcp");
                }
            }
        });

        let tcp_client = TcpTransport::connect(tcp_addr).await.unwrap();
        let udp_client = UdpTransport::from_socket(udp_cli);

        let mut mgr = TransportManager::new();
        mgr.add_transport(Arc::new(tcp_client));
        mgr.add_transport(Arc::new(udp_client));

        // Block TCP.
        for _ in 0..DEFAULT_BLOCK_THRESHOLD {
            mgr.health
                .get_mut(&TransportKind::Tcp)
                .unwrap()
                .record_failure();
        }
        assert!(mgr.is_blocked(TransportKind::Tcp));

        // Try multiple seeds to find one where the mask includes UDP.
        // The routing mask is deterministic, so we find a seed that works.
        let mut found_seed = None;
        for seed_byte in 0u8..255 {
            let mut probe_rng = SeedRng::new([seed_byte; 32]);
            let available = [TransportKind::Tcp, TransportKind::Udp];
            let mask = RoutingMask::generate(&mut probe_rng, &available, 99);
            // We need the mask to include UDP so the first loop can pick it.
            if mask.transports.contains(&TransportKind::Udp) {
                found_seed = Some(seed_byte);
                break;
            }
        }
        let seed_byte = found_seed.expect("no seed found where mask includes UDP");

        let mut rng = SeedRng::new([seed_byte; 32]);
        let available = [TransportKind::Tcp, TransportKind::Udp];

        // send_with_routing_mask should skip blocked TCP and use UDP.
        let used = mgr
            .send_with_routing_mask(
                &mut rng,
                &available,
                99,
                Bytes::from_static(b"skip-blocked-tcp"),
            )
            .await
            .unwrap();
        assert_eq!(
            used,
            TransportKind::Udp,
            "should have fallen back to UDP since TCP is blocked"
        );

        let reply = mgr.recv(used).await.unwrap();
        assert_eq!(reply.as_ref(), b"udp-ok");

        server.await.unwrap();
    }

    #[tokio::test]
    async fn health_scores_reflect_multiple_transports() {
        let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let tcp_addr = tcp_listener.local_addr().unwrap();

        let udp_srv = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let udp_cli = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let udp_srv_addr = udp_srv.local_addr().unwrap();
        let udp_cli_addr = udp_cli.local_addr().unwrap();
        udp_srv.connect(udp_cli_addr).await.unwrap();
        udp_cli.connect(udp_srv_addr).await.unwrap();

        let server = tokio::spawn(async move {
            let (tcp_stream, _) = tcp_listener.accept().await.unwrap();
            let tcp = TcpTransport::from_stream(tcp_stream);
            let udp = UdpTransport::from_socket(udp_srv);
            // Receive on both to record health.
            let _ = tcp.recv().await;
            let _ = udp.recv().await;
        });

        let tcp_client = TcpTransport::connect(tcp_addr).await.unwrap();
        let udp_client = UdpTransport::from_socket(udp_cli);

        let mut mgr = TransportManager::new();
        mgr.add_transport(Arc::new(tcp_client));
        mgr.add_transport(Arc::new(udp_client));

        // Send on both to record health.
        mgr.send(TransportKind::Tcp, Bytes::from_static(b"health-tcp"))
            .await
            .unwrap();
        mgr.send(TransportKind::Udp, Bytes::from_static(b"health-udp"))
            .await
            .unwrap();

        let scores = mgr.health_scores();
        assert_eq!(scores.len(), 2);

        // Both should have positive scores.
        for (kind, score) in &scores {
            assert!(
                *score > 0.0,
                "transport {kind:?} should have positive score, got {score}"
            );
        }

        server.await.unwrap();
    }

    #[tokio::test]
    async fn remove_transport_removes_health_too() {
        let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let tcp_addr = tcp_listener.local_addr().unwrap();

        let udp_srv = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let udp_cli = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let udp_srv_addr = udp_srv.local_addr().unwrap();
        let udp_cli_addr = udp_cli.local_addr().unwrap();
        udp_srv.connect(udp_cli_addr).await.unwrap();
        udp_cli.connect(udp_srv_addr).await.unwrap();

        let server = tokio::spawn(async move {
            let (tcp_stream, _) = tcp_listener.accept().await.unwrap();
            let _tcp = TcpTransport::from_stream(tcp_stream);
            let _udp = UdpTransport::from_socket(udp_srv);
        });

        let tcp_client = TcpTransport::connect(tcp_addr).await.unwrap();
        let udp_client = UdpTransport::from_socket(udp_cli);

        let mut mgr = TransportManager::new();
        mgr.add_transport(Arc::new(tcp_client));
        mgr.add_transport(Arc::new(udp_client));
        assert_eq!(mgr.active_kinds().len(), 2);

        // Remove TCP.
        mgr.remove_transport(TransportKind::Tcp);
        assert_eq!(mgr.active_kinds().len(), 1);
        assert!(mgr.active_kinds().contains(&TransportKind::Udp));
        assert!(
            mgr.health_monitor(TransportKind::Tcp).is_none(),
            "TCP health monitor should also be removed"
        );
        assert!(
            mgr.health_monitor(TransportKind::Udp).is_some(),
            "UDP health monitor should still exist"
        );

        server.await.unwrap();
    }
}