hyperactor/channel/

net.rs

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! A simple socket channel implementation using a single-stream
//! framing protocol. Each frame is encoded as an 8-byte
//! **big-endian** length prefix (u64), followed by exactly that many
//! bytes of payload.
//!
//! Message frames carry a `serde_multipart::Message` (not raw
//! bincode). In compat mode (current default), this is encoded as a
//! sentinel `u64::MAX` followed by a single bincode payload. Response frames
//! are a bincode-serialized NetRxResponse enum, containing either the acked
//! sequence number, or the Reject value indicating that the server rejected
//! the connection.
//!
//! Message frame (compat/unipart) example:
//! ```text
//! +------------------ len: u64 (BE) ------------------+----------------------- data -----------------------+
//! | \x00\x00\x00\x00\x00\x00\x00\x10                  | \xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF | <bincode bytes> |
//! |                       16                          |           u64::MAX             |                   |
//! +---------------------------------------------------+-----------------------------------------------------+
//! ```
//!
//! Response frame (wire format):
//! ```text
//! +------------------ len: u64 (BE) ------------------+---------------- data ------------------+
//! | \x00\x00\x00\x00\x00\x00\x00\x??                  | <bincode acked sequence num or reject> |
//! +---------------------------------------------------+----------------------------------------+
//! ```
//!
//! I/O is handled by `FrameReader`/`FrameWrite`, which are
//! cancellation-safe and avoid extra copies. Helper fns
//! `serialize_response(NetRxResponse) -> Result<Bytes, bincode::Error>`
//! and `deserialize_response(Bytes) -> Result<NetRxResponse, bincode::Error>`
//! convert to/from the response payload.
//!
//! ### Limits & EOF semantics
//! * **Max frame size:** frames larger than
//!   `config::CODEC_MAX_FRAME_LENGTH` are rejected with
//!   `io::ErrorKind::InvalidData`.
//! * **EOF handling:** `FrameReader::next()` returns `Ok(None)` only
//!   when EOF occurs exactly on a frame boundary. If EOF happens
//!   mid-frame, it returns `Err(io::ErrorKind::UnexpectedEof)`.

use std::fmt;
use std::fmt::Debug;
use std::net::ToSocketAddrs;
use std::time::Duration;

use backoff::ExponentialBackoffBuilder;
use backoff::backoff::Backoff;
use bytes::Bytes;
use enum_as_inner::EnumAsInner;
use serde::Deserialize;
use serde::Serialize;
use serde::de::Error;
use tokio::io::AsyncRead;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWrite;
use tokio::io::AsyncWriteExt;
use tokio::sync::watch;
use tokio::time::Instant;

use super::*;
use crate::RemoteMessage;

pub mod duplex;
mod framed;
pub(super) mod server;
pub(super) mod session;
pub use server::ServerHandle;

pub(crate) trait Stream:
    AsyncRead + AsyncWrite + Unpin + Send + Sync + Debug + 'static
{
}
impl<S: AsyncRead + AsyncWrite + Unpin + Send + Sync + Debug + 'static> Stream for S {}

/// Opaque identifier for a session. Generated by the client,
/// sent to the server on each connect so the server can correlate
/// reconnections to the same logical session.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub(crate) struct SessionId(u64);

impl SessionId {
    /// Generate a new random session ID.
    pub fn random() -> Self {
        Self(rand::random())
    }
}

impl fmt::Display for SessionId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:016x}", self.0)
    }
}

/// Logical channel tag for initiator→acceptor traffic.
pub(crate) const INITIATOR_TO_ACCEPTOR: u8 = 0;

/// Logical channel tag for acceptor→initiator traffic.
pub(crate) const ACCEPTOR_TO_INITIATOR: u8 = 1;

/// Fixed-size header sent at the start of every physical connection.
/// This is written/read directly on the wire (not framed), before
/// any session framing begins.
///
/// Wire format (13 bytes, big-endian):
/// ```text
/// [magic: 4B "LNK\0"] [session_id: 8B u64 BE] [stream_id: 1B u8]
/// ```
const LINK_INIT_MAGIC: [u8; 4] = *b"LNK\0";
const LINK_INIT_SIZE: usize = 4 + 8 + 1;

/// Parsed LinkInit header.
#[derive(Debug, Clone, Copy)]
pub(crate) struct LinkInit {
    pub session_id: SessionId,
    pub stream_id: u8,
}

/// Write a LinkInit header to the stream.
async fn write_link_init<S: AsyncWrite + Unpin>(
    stream: &mut S,
    session_id: SessionId,
    stream_id: u8,
) -> Result<(), std::io::Error> {
    let mut buf = [0u8; LINK_INIT_SIZE];
    buf[0..4].copy_from_slice(&LINK_INIT_MAGIC);
    buf[4..12].copy_from_slice(&session_id.0.to_be_bytes());
    buf[12] = stream_id;
    stream.write_all(&buf).await
}

/// Read a LinkInit header from the stream.
async fn read_link_init<S: AsyncRead + Unpin>(stream: &mut S) -> Result<LinkInit, std::io::Error> {
    let mut buf = [0u8; LINK_INIT_SIZE];
    stream.read_exact(&mut buf).await?;
    if buf[0..4] != LINK_INIT_MAGIC {
        return Err(std::io::Error::new(
            std::io::ErrorKind::InvalidData,
            format!("invalid LinkInit magic: expected LNK, got {:?}", &buf[0..4]),
        ));
    }
    let session_id = SessionId(u64::from_be_bytes(buf[4..12].try_into().unwrap()));
    let stream_id = buf[12];
    Ok(LinkInit {
        session_id,
        stream_id,
    })
}

/// Link represents a network link through which connections may be
/// acquired. The session ID is baked in. Initiator links dial;
/// acceptor links wait for dispatched streams.
#[async_trait]
pub(crate) trait Link: Send + Sync + Debug + 'static {
    /// The underlying stream type.
    type Stream: Stream;

    /// The address of the link's destination.
    fn dest(&self) -> ChannelAddr;

    /// The session ID for this link.
    fn link_id(&self) -> SessionId;

    /// Acquire the next usable connection. For initiator links this
    /// dials; for acceptor links this waits on a dispatch channel.
    async fn next(&mut self) -> Result<Self::Stream, ClientError>;
}

use session::Session;

use crate::config;
use crate::metrics;

pub(crate) enum LinkStatus {
    NeverConnected,
    Connected(tokio::time::Instant),
    Disconnected {
        last_connected: tokio::time::Instant,
        since: tokio::time::Instant,
    },
}

impl LinkStatus {
    fn connected(&mut self) {
        *self = LinkStatus::Connected(tokio::time::Instant::now());
    }

    fn disconnected(&mut self) {
        match *self {
            LinkStatus::Connected(at) => {
                *self = LinkStatus::Disconnected {
                    last_connected: at,
                    since: tokio::time::Instant::now(),
                };
            }
            // Already disconnected or never connected — leave as is.
            _ => {}
        }
    }
}

impl std::fmt::Display for LinkStatus {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            LinkStatus::NeverConnected => write!(f, "never connected"),
            LinkStatus::Connected(at) => {
                write!(f, "connected for {:.1}s", at.elapsed().as_secs_f64())
            }
            LinkStatus::Disconnected {
                last_connected,
                since,
            } => {
                write!(
                    f,
                    "last connected {:.1}s ago, disconnected for {:.1}s",
                    last_connected.elapsed().as_secs_f64(),
                    since.elapsed().as_secs_f64(),
                )
            }
        }
    }
}

/// Log a send-loop error and return `true` if the error is terminal
/// (caller should exit), `false` if recoverable (caller should reconnect).
fn log_send_error(
    error: &session::SendLoopError,
    dest: &ChannelAddr,
    session_id: u64,
    mode: &str,
    link_status: &LinkStatus,
) -> bool {
    match error {
        session::SendLoopError::Io(err) => {
            tracing::info!(dest = %dest, session_id, error = %err, mode, "send error; {link_status}");
            metrics::CHANNEL_ERRORS.add(
                1,
                hyperactor_telemetry::kv_pairs!(
                    "dest" => dest.to_string(),
                    "session_id" => session_id.to_string(),
                    "error_type" => metrics::ChannelErrorType::SendError.as_str(),
                    "mode" => mode.to_string(),
                ),
            );
            false
        }
        session::SendLoopError::AppClosed => true,
        session::SendLoopError::Rejected(reason) => {
            tracing::error!(dest = %dest, session_id, mode, "server rejected connection: {reason}; {link_status}");
            true
        }
        session::SendLoopError::ServerClosed => {
            tracing::info!(dest = %dest, session_id, mode, "server closed the channel; {link_status}");
            true
        }
        session::SendLoopError::DeliveryTimeout => {
            let timeout = hyperactor_config::global::get(config::MESSAGE_DELIVERY_TIMEOUT);
            tracing::error!(
                dest = %dest, session_id, mode,
                "failed to receive ack within timeout {timeout:?}; link is currently connected; {link_status}"
            );
            true
        }
        session::SendLoopError::OversizedFrame(reason) => {
            tracing::error!(dest = %dest, session_id, mode, "oversized frame: {reason}; {link_status}");
            true
        }
    }
}

/// Establish a simplex (send-only) session over the given link. Returns a send handle.
pub(crate) fn spawn<M: RemoteMessage>(link: impl Link) -> NetTx<M> {
    spawn_inner(link)
}

/// Establish a multi-stream (unordered) simplex session over N
/// links sharing the same `SessionId`. Returns a single send handle
/// that distributes frames across streams.
pub(crate) fn spawn_unordered<M: RemoteMessage>(links: Vec<impl Link + 'static>) -> NetTx<M> {
    assert!(!links.is_empty());
    if links.len() == 1 {
        return spawn(links.into_iter().next().unwrap());
    }

    let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel();
    let dest = links[0].dest();
    let session_id = links[0].link_id();
    let (notify, status) = watch::channel(TxStatus::Active);
    let tx = NetTx {
        sender,
        dest: dest.clone(),
        status,
    };

    let num_streams = links.len();

    crate::init::get_runtime().spawn(async move {
        // Shared MPMC work queue. The dispatcher enqueues *unserialized*
        // messages; each writer serializes its own pulls before writing.
        // This spreads serialization cost across all N writer tasks
        // instead of bottlenecking on the dispatcher.
        let (queue_tx, queue_rx) =
            async_channel::bounded::<session::PendingMessage<M>>(num_streams * 8);

        // Shared unacked buffer: any writer's ack reader can prune it.
        // BTreeMap keyed by seq — writers insert out of order
        // (multiple streams, interleaved), so not a VecDeque.
        let unacked: Arc<
            tokio::sync::Mutex<std::collections::BTreeMap<u64, session::QueuedMessage<M>>>,
        > = Arc::new(tokio::sync::Mutex::new(std::collections::BTreeMap::new()));

        let mut writer_handles: Vec<tokio::task::JoinHandle<()>> = Vec::with_capacity(num_streams);
        let log_id = format!("session {}.{:016x}", dest, session_id.0);

        for (i, link) in links.into_iter().enumerate() {
            let dest = dest.clone();
            let unacked = unacked.clone();
            let queue_rx = queue_rx.clone();
            let log_id = log_id.clone();

            writer_handles.push(tokio::spawn(async move {
                let mut session = Session::new(link);
                let mut reconnect_backoff = ExponentialBackoffBuilder::new()
                    .with_initial_interval(Duration::from_millis(10))
                    .with_multiplier(2.0)
                    .with_randomization_factor(0.1)
                    .with_max_interval(Duration::from_secs(5))
                    .with_max_elapsed_time(None)
                    .build();

                loop {
                    let connected = match session.connect().await {
                        Ok(s) => s,
                        Err(_) => {
                            tracing::info!(
                                dest = %dest, stream = i,
                                "multi-stream writer {} connect failed", i
                            );
                            break;
                        }
                    };
                    tracing::info!(
                        dest = %dest, stream = i, "multi-stream writer {} connected", i
                    );

                    let stream = connected.stream(INITIATOR_TO_ACCEPTOR);
                    let connected_at = tokio::time::Instant::now();

                    // Pull from the shared queue; serialize locally; write; interleave ack reads.
                    let result: Result<(), session::SendLoopError> = async {
                        loop {
                            tokio::select! {
                                biased;

                                ack_result = stream.next() => {
                                    match ack_result {
                                        Ok(Some(buffer)) => {
                                            let response = deserialize_response(buffer)
                                                .map_err(|e| session::SendLoopError::Io(e.into()))?;
                                            match response {
                                                NetRxResponse::Ack(ack) => {
                                                    let mut guard = unacked.lock().await;
                                                    // Remove all entries with seq <= ack.
                                                    let retain: std::collections::BTreeMap<u64, session::QueuedMessage<M>> = guard.split_off(&(ack + 1));
                                                    drop(std::mem::replace(&mut *guard, retain));
                                                }
                                                NetRxResponse::Reject(reason) => {
                                                    return Err(session::SendLoopError::Rejected(reason));
                                                }
                                                NetRxResponse::Closed => {
                                                    return Err(session::SendLoopError::ServerClosed);
                                                }
                                            }
                                        }
                                        Ok(None) => return Ok(()),
                                        Err(e) => return Err(session::SendLoopError::Io(e.into())),
                                    }
                                }

                                msg = queue_rx.recv() => {
                                    let pending = match msg {
                                        Ok(m) => m,
                                        // Dispatcher closed the queue: clean shutdown.
                                        Err(_) => return Ok(()),
                                    };
                                    let session::PendingMessage {
                                        seq,
                                        message,
                                        received_at,
                                        return_channel,
                                    } = pending;
                                    let frame = Frame::Message(seq, message);
                                    let serialized = match serde_multipart::serialize_bincode(&frame) {
                                        Ok(m) => m,
                                        Err(e) => {
                                            tracing::error!(
                                                "{log_id}: serialization error: {e}"
                                            );
                                            // Drops return_channel; sender perceives success
                                            // (preserving prior behavior of the dispatcher-side
                                            // serialize path).
                                            continue;
                                        }
                                    };
                                    let mut queued = session::QueuedMessage {
                                        seq,
                                        message: serialized,
                                        received_at,
                                        sent_at: None,
                                        return_channel,
                                    };
                                    let framed = queued.message.clone().framed();
                                    stream.write(framed).drive().await.map_err(|e| {
                                        session::SendLoopError::Io(e.into())
                                    })?;
                                    queued.sent_at = Some(tokio::time::Instant::now());
                                    unacked.lock().await.insert(queued.seq, queued);
                                }
                            }
                        }
                    }
                    .await;

                    session = connected.release();

                    if connected_at.elapsed() > Duration::from_secs(1) {
                        reconnect_backoff.reset();
                    }

                    match result {
                        Ok(()) => {
                            if queue_rx.is_closed() {
                                // Dispatcher is gone and queue is drained.
                                break;
                            }
                            if let Some(delay) = reconnect_backoff.next_backoff() {
                                tokio::time::sleep(delay).await;
                            }
                        }
                        Err(ref e) => {
                            if log_send_error(e, &dest, session_id.0, "multi-stream", &LinkStatus::NeverConnected) {
                                break;
                            }
                            if let Some(delay) = reconnect_backoff.next_backoff() {
                                tokio::time::sleep(delay).await;
                            }
                        }
                    }
                }

                tracing::info!(
                    dest = %dest,
                    stream = i,
                    "multi-stream writer {} shutting down",
                    i,
                );
            }));
        }

        // Drop our local receiver clone so the queue closes once the
        // dispatcher's sender (queue_tx) is dropped at shutdown.
        drop(queue_rx);

        // Dispatcher: receive from app and enqueue for writers — no
        // serialization here; the writer that pulls the item serializes
        // it before writing.
        let mut next_seq = 0u64;

        tracing::info!(
            %dest, session = %log_id, num_streams,
            "multi-stream dispatcher started"
        );

        while let Some((message, return_channel, received_at)) = receiver.recv().await {
            let pending = session::PendingMessage {
                seq: next_seq,
                message,
                received_at,
                return_channel,
            };
            next_seq += 1;

            if queue_tx.send(pending).await.is_err() {
                // All writers are gone.
                break;
            }
        }

        // Shutdown: close the shared queue and wait for writers to drain.
        drop(queue_tx);
        for handle in writer_handles {
            let _ = handle.await;
        }

        let reason = format!("{log_id}: dispatcher closed");
        let _ = notify.send(TxStatus::Closed(reason.into()));
    });

    tx
}

fn spawn_inner<M: RemoteMessage>(link: impl Link) -> NetTx<M> {
    let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
    let dest = link.dest();
    let session_id = link.link_id();
    let (notify, status) = watch::channel(TxStatus::Active);
    let tx = NetTx {
        sender,
        dest: dest.clone(),
        status,
    };
    crate::init::get_runtime().spawn(async move {
        let mut session = Session::new(link);
        let log_id = format!("session {}.{:016x}", dest, session_id.0);
        let mut deliveries = session::Deliveries {
            outbox: session::Outbox::new(log_id.clone(), dest.clone(), session_id.0),
            unacked: session::Unacked::new(None, log_id.clone()),
        };
        let mut receiver = receiver;

        // Lazy connect: wait for first message.
        match receiver.recv().await {
            Some(msg) => {
                if let Err(err) = deliveries.outbox.push_back(msg) {
                    tracing::error!(
                        dest = %dest,
                        session_id = session_id.0,
                        error = %err,
                        "failed to push message to outbox"
                    );
                    let _ = notify.send(TxStatus::Closed("failed to push to outbox".into()));
                    return;
                }
            }
            None => {
                let _ = notify.send(TxStatus::Closed("sender dropped".into()));
                return;
            }
        }

        let mut reconnect_backoff = ExponentialBackoffBuilder::new()
            .with_initial_interval(Duration::from_millis(10))
            .with_multiplier(2.0)
            .with_randomization_factor(0.1)
            .with_max_interval(Duration::from_secs(5))
            .with_max_elapsed_time(None)
            .build();

        let mut link_status = LinkStatus::NeverConnected;

        let reason: String = 'outer: loop {
            let connected = match deliveries.expiry_time() {
                Some(deadline) => match session.connect_by(deadline).await {
                    Ok(s) => s,
                    Err(_) => {
                        let timeout =
                            hyperactor_config::global::get(config::MESSAGE_DELIVERY_TIMEOUT);
                        let error_msg = if deliveries.outbox.is_expired(timeout) {
                            format!("failed to deliver message within timeout {timeout:?}; {link_status}")
                        } else {
                            format!(
                                "failed to receive ack within timeout {timeout:?}; \
                                 link is currently broken; {link_status}",
                            )
                        };
                        tracing::error!(
                            dest = %dest, session_id = session_id.0, "{}", error_msg
                        );
                        break 'outer format!("{log_id}: {error_msg}");
                    }
                },
                None => match session.connect().await {
                    Ok(s) => s,
                    Err(_) => break 'outer "session shut down".into(),
                },
            };

            metrics::CHANNEL_CONNECTIONS.add(
                1,
                hyperactor_telemetry::kv_pairs!(
                    "transport" => dest.transport().to_string(),
                    "mode" => "simplex",
                    "reason" => "link connected",
                ),
            );

            if !deliveries.unacked.is_empty() {
                metrics::CHANNEL_RECONNECTIONS.add(
                    1,
                    hyperactor_telemetry::kv_pairs!(
                        "dest" => dest.to_string(),
                        "transport" => dest.transport().to_string(),
                        "mode" => "simplex",
                        "reason" => "reconnect_with_unacked",
                    ),
                );
            }
            deliveries.requeue_unacked();

            link_status.connected();
            let connected_at = tokio::time::Instant::now();

            let result = {
                let stream = connected.stream(INITIATOR_TO_ACCEPTOR);
                session::send_connected(&stream, &mut deliveries, &mut receiver).await
            };
            session = connected.release();

            link_status.disconnected();

            // Reset backoff if the connection was alive long enough to have
            // been useful (i.e. not an immediate EOF/error).
            if connected_at.elapsed() > Duration::from_secs(1) {
                reconnect_backoff.reset();
            }

            match result {
                Ok(()) => {
                    // EOF — connection closed normally, reconnect after backoff.
                    if let Some(delay) = reconnect_backoff.next_backoff() {
                        tracing::info!(
                            dest = %dest,
                            session_id = session_id.0,
                            delay_ms = delay.as_millis() as u64,
                            "send_connected returned EOF, reconnecting after backoff; {link_status}"
                        );
                        tokio::time::sleep(delay).await;
                    }
                }
                Err(ref e) => {
                    if log_send_error(e, &dest, session_id.0, "simplex", &link_status) {
                        break 'outer format!("{log_id}: {e}");
                    }
                    // Recoverable error — reconnect after backoff.
                    if let Some(delay) = reconnect_backoff.next_backoff() {
                        tracing::info!(
                            dest = %dest,
                            session_id = session_id.0,
                            delay_ms = delay.as_millis() as u64,
                            error = %e,
                            "send_connected returned recoverable error, reconnecting after backoff; {link_status}"
                        );
                        tokio::time::sleep(delay).await;
                    }
                }
            }
        };

        tracing::info!(
            dest = %dest, session_id = session_id.0, "NetTx closing: {reason}"
        );

        receiver.close();
        deliveries
            .unacked
            .deque
            .drain(..)
            .chain(deliveries.outbox.deque.drain(..))
            .for_each(|queued| queued.try_return(Some(reason.clone())));
        while let Ok((msg, return_channel, _)) = receiver.try_recv() {
            let _ = return_channel.send(SendError {
                error: ChannelError::Closed,
                message: msg,
                reason: Some(reason.clone()),
            });
        }

        let _ = notify.send(TxStatus::Closed(reason.into()));
    });
    tx
}

/// Transport-agnostic link that dispatches to the appropriate
/// transport based on the channel address.
#[derive(Debug)]
pub(crate) enum NetLink {
    Tcp(tcp::TcpLink),
    Unix(unix::UnixLink),
    Tls(tls::TlsLink),
}

/// Create a link for the given channel address with the given
/// `session_id` and `stream_id`. Single-stream callers pass a fresh
/// `SessionId::random()` and `stream_id = 0`.
pub(crate) fn link(
    addr: ChannelAddr,
    session_id: SessionId,
    stream_id: u8,
) -> Result<NetLink, ClientError> {
    match addr {
        ChannelAddr::Tcp(socket_addr) => {
            Ok(NetLink::Tcp(tcp::link(socket_addr, session_id, stream_id)))
        }
        ChannelAddr::Unix(unix_addr) => {
            Ok(NetLink::Unix(unix::link(unix_addr, session_id, stream_id)))
        }
        ChannelAddr::Tls(tls_addr) => Ok(NetLink::Tls(tls::link(tls_addr, session_id, stream_id)?)),
        ChannelAddr::MetaTls(meta_addr) => {
            Ok(NetLink::Tls(meta::link(meta_addr, session_id, stream_id)?))
        }
        other => Err(ClientError::Connect(
            other,
            std::io::Error::other("unsupported transport"),
            "unsupported transport".into(),
        )),
    }
}

#[async_trait]
impl Link for NetLink {
    type Stream = Box<dyn Stream>;

    fn dest(&self) -> ChannelAddr {
        match self {
            Self::Tcp(l) => l.dest(),
            Self::Unix(l) => l.dest(),
            Self::Tls(l) => l.dest(),
        }
    }

    fn link_id(&self) -> SessionId {
        match self {
            Self::Tcp(l) => l.link_id(),
            Self::Unix(l) => l.link_id(),
            Self::Tls(l) => l.link_id(),
        }
    }

    async fn next(&mut self) -> Result<Box<dyn Stream>, ClientError> {
        match self {
            Self::Tcp(l) => Ok(Box::new(l.next().await?)),
            Self::Unix(l) => Ok(Box::new(l.next().await?)),
            Self::Tls(l) => Ok(Box::new(l.next().await?)),
        }
    }
}

/// Listener represents the server side of a network link: it accepts inbound connections.
///
/// This is the counterpart to [`Link`]. Each transport module (tcp, unix, tls)
/// provides both a `Link` impl (for dialing) and a `Listener` impl (for accepting).
#[async_trait]
pub(crate) trait Listener: Send + Unpin + 'static {
    /// The underlying stream type produced by accepting a connection.
    type Stream: Stream;

    /// Accept the next inbound connection, returning the stream and the peer's address.
    async fn accept(&mut self) -> Result<(Self::Stream, ChannelAddr), ServerError>;
}

/// Transport-agnostic listener that dispatches to the appropriate
/// transport based on the channel address. TLS has no variant — it
/// uses `Tcp` under the hood (the TLS handshake happens in `prepare`,
/// not the listener).
#[derive(Debug)]
pub(crate) enum NetListener {
    Tcp(tcp::TcpSocketListener),
    Unix(unix::UnixSocketListener),
}

#[async_trait]
impl Listener for NetListener {
    type Stream = Box<dyn Stream>;

    async fn accept(&mut self) -> Result<(Box<dyn Stream>, ChannelAddr), ServerError> {
        match self {
            Self::Tcp(l) => {
                let (stream, addr) = l.accept().await?;
                Ok((Box::new(stream), addr))
            }
            Self::Unix(l) => {
                let (stream, addr) = l.accept().await?;
                Ok((Box::new(stream), addr))
            }
        }
    }
}

/// Bind a listener for the given channel address, optionally using a pre-opened TCP listener.
/// Returns the listener and the canonical address callers should advertise.
/// When `prebound` is `Some`, it is used for TCP/TLS transports instead of binding a new socket.
pub(crate) fn listen_with_prebound(
    addr: ChannelAddr,
    prebound: Option<std::net::TcpListener>,
) -> Result<(NetListener, ChannelAddr), ServerError> {
    match addr {
        ChannelAddr::Tcp(socket_addr) => {
            let std_listener = match prebound {
                Some(l) => l,
                None => std::net::TcpListener::bind(socket_addr)
                    .map_err(|err| ServerError::Listen(ChannelAddr::Tcp(socket_addr), err))?,
            };
            std_listener
                .set_nonblocking(true)
                .map_err(|e| ServerError::Listen(ChannelAddr::Tcp(socket_addr), e))?;
            let tokio_listener = tokio::net::TcpListener::from_std(std_listener)
                .map_err(|e| ServerError::Listen(ChannelAddr::Tcp(socket_addr), e))?;
            let local_addr = tokio_listener
                .local_addr()
                .map_err(|err| ServerError::Resolve(ChannelAddr::Tcp(socket_addr), err))?;
            let listener = tcp::TcpSocketListener {
                inner: tokio_listener,
                addr: local_addr,
            };
            Ok((NetListener::Tcp(listener), ChannelAddr::Tcp(local_addr)))
        }
        ChannelAddr::Unix(ref unix_addr) => {
            use std::os::unix::net::UnixDatagram as StdUnixDatagram;
            use std::os::unix::net::UnixListener as StdUnixListener;

            let caddr = addr.clone();
            let maybe_listener = match unix_addr {
                unix::SocketAddr::Bound(sock_addr) => StdUnixListener::bind_addr(sock_addr),
                unix::SocketAddr::Unbound => StdUnixDatagram::unbound()
                    .and_then(|u| u.local_addr())
                    .and_then(|uaddr| StdUnixListener::bind_addr(&uaddr)),
            };
            let std_listener =
                maybe_listener.map_err(|err| ServerError::Listen(caddr.clone(), err))?;
            std_listener
                .set_nonblocking(true)
                .map_err(|err| ServerError::Listen(caddr.clone(), err))?;
            let local_addr = std_listener
                .local_addr()
                .map_err(|err| ServerError::Resolve(caddr.clone(), err))?;
            let tokio_listener = tokio::net::UnixListener::from_std(std_listener)
                .map_err(|err| ServerError::Io(caddr, err))?;
            let bound_addr = unix::SocketAddr::new(local_addr);
            let listener = unix::UnixSocketListener {
                inner: tokio_listener,
                addr: bound_addr.clone(),
            };
            Ok((NetListener::Unix(listener), ChannelAddr::Unix(bound_addr)))
        }
        addr @ (ChannelAddr::Tls(_) | ChannelAddr::MetaTls(_)) => {
            let is_meta = matches!(addr, ChannelAddr::MetaTls(_));
            let tls_addr = match addr {
                ChannelAddr::Tls(a) | ChannelAddr::MetaTls(a) => a,
                _ => unreachable!(),
            };
            let TlsAddr { hostname, port } = tls_addr;
            let make_channel_addr = |h: &str, p: Port| {
                if is_meta {
                    ChannelAddr::MetaTls(TlsAddr::new(h, p))
                } else {
                    ChannelAddr::Tls(TlsAddr::new(h, p))
                }
            };

            let addrs: Vec<core::net::SocketAddr> = (hostname.as_ref(), port)
                .to_socket_addrs()
                .map_err(|err| ServerError::Resolve(make_channel_addr(&hostname, port), err))?
                .collect();

            if addrs.is_empty() {
                return Err(ServerError::Resolve(
                    make_channel_addr(&hostname, port),
                    std::io::Error::other("no available socket addr"),
                ));
            }

            let channel_addr = make_channel_addr(&hostname, port);
            let std_listener = match prebound {
                Some(l) => l,
                None => std::net::TcpListener::bind(&addrs[..])
                    .map_err(|err| ServerError::Listen(channel_addr.clone(), err))?,
            };
            std_listener
                .set_nonblocking(true)
                .map_err(|e| ServerError::Listen(channel_addr.clone(), e))?;
            let tokio_listener = tokio::net::TcpListener::from_std(std_listener)
                .map_err(|e| ServerError::Listen(channel_addr.clone(), e))?;
            let local_addr = tokio_listener
                .local_addr()
                .map_err(|err| ServerError::Resolve(channel_addr, err))?;
            let listener = tcp::TcpSocketListener {
                inner: tokio_listener,
                addr: local_addr,
            };
            Ok((
                NetListener::Tcp(listener),
                make_channel_addr(&hostname, local_addr.port()),
            ))
        }
        other => Err(ServerError::Listen(
            other.clone(),
            std::io::Error::other(format!("unsupported transport: {}", other)),
        )),
    }
}

/// Bind a listener for the given channel address. Returns the listener
/// and the canonical address callers should advertise (which encodes
/// the transport — e.g. `ChannelAddr::Tls` for TLS).
#[expect(
    dead_code,
    reason = "canonical listen() entry point; callers currently route through listen_with_prebound"
)]
pub(crate) fn listen(addr: ChannelAddr) -> Result<(NetListener, ChannelAddr), ServerError> {
    listen_with_prebound(addr, None)
}

/// Frames are the messages sent between clients and servers over sessions.
#[derive(Debug, Serialize, Deserialize, EnumAsInner, PartialEq)]
pub(super) enum Frame<M> {
    /// Send a message with the provided sequence number.
    Message(u64, M),
}

#[derive(Debug, Serialize, Deserialize, EnumAsInner)]
pub(super) enum NetRxResponse {
    Ack(u64),
    /// This session is rejected with the given reason. NetTx should stop reconnecting.
    Reject(String),
    /// This channel is closed.
    Closed,
}

pub(super) fn serialize_response(
    response: NetRxResponse,
) -> Result<Bytes, bincode::error::EncodeError> {
    bincode::serde::encode_to_vec(&response, bincode::config::legacy()).map(|bytes| bytes.into())
}

pub(super) fn deserialize_response(
    data: Bytes,
) -> Result<NetRxResponse, bincode::error::DecodeError> {
    bincode::serde::decode_from_slice(&data, bincode::config::legacy()).map(|(v, _)| v)
}

/// A Tx implemented on top of a Link. The Tx manages the link state,
/// reconnections, etc.
pub(crate) struct NetTx<M: RemoteMessage> {
    sender: mpsc::UnboundedSender<(M, oneshot::Sender<SendError<M>>, Instant)>,
    dest: ChannelAddr,
    status: watch::Receiver<TxStatus>,
}

#[async_trait]
impl<M: RemoteMessage> Tx<M> for NetTx<M> {
    fn addr(&self) -> ChannelAddr {
        self.dest.clone()
    }

    fn status(&self) -> &watch::Receiver<TxStatus> {
        &self.status
    }

    fn do_post(&self, message: M, return_channel: Option<oneshot::Sender<SendError<M>>>) {
        tracing::trace!(
            name = "post",
            dest = %self.dest,
            "sending message"
        );

        let return_channel = return_channel.unwrap_or_else(|| oneshot::channel().0);
        if let Err(mpsc::error::SendError((message, return_channel, _))) =
            self.sender
                .send((message, return_channel, tokio::time::Instant::now()))
        {
            let reason = self.status.borrow().as_closed().map(|r| r.to_string());
            let _ = return_channel.send(SendError {
                error: ChannelError::Closed,
                message,
                reason,
            });
        }
    }
}

pub struct NetRx<M: RemoteMessage>(mpsc::Receiver<M>, ChannelAddr, ServerHandle);

#[async_trait]
impl<M: RemoteMessage> Rx<M> for NetRx<M> {
    async fn recv(&mut self) -> Result<M, ChannelError> {
        tracing::trace!(
            name = "recv",
            dest = %self.1,
            "receiving message"
        );
        self.0.recv().await.ok_or(ChannelError::Closed)
    }

    fn addr(&self) -> ChannelAddr {
        self.1.clone()
    }

    /// Gracefully shut down the channel server, waiting for pending
    /// acks to be flushed before returning.
    async fn join(mut self) {
        self.2
            .stop(&format!("NetRx joined; channel address: {}", self.1));
        let _ = (&mut self.2).await;
        // Drop will call stop() again which is harmless (token already cancelled).
    }
}

impl<M: RemoteMessage> Drop for NetRx<M> {
    fn drop(&mut self) {
        self.2
            .stop(&format!("NetRx dropped; channel address: {}", self.1));
    }
}

/// Error returned during server operations.
#[derive(Debug, thiserror::Error)]
pub enum ServerError {
    /// An I/O error occurred while operating on the server at the given address.
    #[error("io: {1}")]
    Io(ChannelAddr, #[source] std::io::Error),
    /// Listening on the given address failed.
    #[error("listen: {0} {1}")]
    Listen(ChannelAddr, #[source] std::io::Error),
    /// Resolving the given address failed.
    #[error("resolve: {0} {1}")]
    Resolve(ChannelAddr, #[source] std::io::Error),
    /// An internal server error occurred for the given address.
    #[error("internal: {0} {1}")]
    Internal(ChannelAddr, #[source] anyhow::Error),
}

#[derive(thiserror::Error, Debug)]
pub enum ClientError {
    #[error("connection to {0} failed: {1}: {2}")]
    Connect(ChannelAddr, std::io::Error, String),
    #[error("unable to resolve address: {0}")]
    Resolve(ChannelAddr),
    #[error("io: {0} {1}")]
    Io(ChannelAddr, std::io::Error),
    #[error("send {0}: serialize: {1}")]
    Serialize(ChannelAddr, bincode::error::EncodeError),
    #[error("invalid address: {0}")]
    InvalidAddress(String),
}

/// Tells whether the address is a 'net' address. These currently have different semantics
/// from local transports.
#[cfg(test)]
pub(super) fn is_net_addr(addr: &ChannelAddr) -> bool {
    match addr.transport() {
        ChannelTransport::Tcp(_) => true,
        ChannelTransport::MetaTls(_) => true,
        ChannelTransport::Tls => true,
        ChannelTransport::Unix => true,
        _ => false,
    }
}

pub(crate) mod unix {

    use core::str;
    use std::os::unix::net::SocketAddr as StdSocketAddr;
    use std::os::unix::net::UnixStream as StdUnixStream;

    use rand::RngExt as _;
    use rand::distr::Alphanumeric;
    use tokio::net::UnixListener;
    use tokio::net::UnixStream;

    use super::*;

    #[derive(Debug)]
    pub(crate) struct UnixLink {
        pub(super) addr: SocketAddr,
        pub(super) session_id: SessionId,
        pub(super) stream_id: u8,
    }

    #[async_trait]
    impl Link for UnixLink {
        type Stream = UnixStream;

        fn dest(&self) -> ChannelAddr {
            ChannelAddr::Unix(self.addr.clone())
        }

        fn link_id(&self) -> SessionId {
            self.session_id
        }

        async fn next(&mut self) -> Result<Self::Stream, ClientError> {
            let session_id = self.session_id;
            let sock_addr = match &self.addr {
                SocketAddr::Bound(a) => a,
                SocketAddr::Unbound => return Err(ClientError::Resolve(self.dest())),
            };
            let mut backoff = ExponentialBackoffBuilder::new()
                .with_initial_interval(Duration::from_millis(1))
                .with_multiplier(2.0)
                .with_randomization_factor(0.1)
                .with_max_interval(Duration::from_millis(1000))
                .with_max_elapsed_time(None)
                .build();
            loop {
                match StdUnixStream::connect_addr(sock_addr) {
                    Ok(std_stream) => {
                        std_stream
                            .set_nonblocking(true)
                            .map_err(|err| ClientError::Io(self.dest(), err))?;
                        let mut stream = UnixStream::from_std(std_stream)
                            .map_err(|err| ClientError::Io(self.dest(), err))?;
                        write_link_init(&mut stream, session_id, self.stream_id)
                            .await
                            .map_err(|err| ClientError::Io(self.dest(), err))?;
                        return Ok(stream);
                    }
                    Err(err) => {
                        tracing::debug!(error = %err, "unix connect failed, backing off");
                        if let Some(delay) = backoff.next_backoff() {
                            tokio::time::sleep(delay).await;
                        }
                    }
                }
            }
        }
    }

    /// Server-side listener for Unix domain sockets.
    #[derive(Debug)]
    pub(crate) struct UnixSocketListener {
        pub(super) inner: UnixListener,
        pub(super) addr: SocketAddr,
    }

    #[async_trait]
    impl super::Listener for UnixSocketListener {
        type Stream = UnixStream;

        async fn accept(&mut self) -> Result<(Self::Stream, ChannelAddr), ServerError> {
            let (stream, peer_addr) = self
                .inner
                .accept()
                .await
                .map_err(|err| ServerError::Io(ChannelAddr::Unix(self.addr.clone()), err))?;
            // tokio::net::unix::SocketAddr -> std::os::unix::net::SocketAddr
            let std_addr: StdSocketAddr = peer_addr.into();
            Ok((stream, ChannelAddr::Unix(SocketAddr::new(std_addr))))
        }
    }

    /// Create a unix link to the given socket address.
    pub(crate) fn link(addr: SocketAddr, session_id: SessionId, stream_id: u8) -> UnixLink {
        UnixLink {
            addr,
            session_id,
            stream_id,
        }
    }

    /// Wrapper around std-lib's unix::SocketAddr that lets us implement equality functions
    #[derive(Clone, Debug)]
    pub enum SocketAddr {
        Bound(Box<StdSocketAddr>),
        Unbound,
    }

    impl PartialOrd for SocketAddr {
        fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
            Some(self.cmp(other))
        }
    }

    impl Ord for SocketAddr {
        fn cmp(&self, other: &Self) -> std::cmp::Ordering {
            self.to_string().cmp(&other.to_string())
        }
    }

    impl<'de> Deserialize<'de> for SocketAddr {
        fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
        where
            D: serde::Deserializer<'de>,
        {
            let s = String::deserialize(deserializer)?;
            Self::from_str(&s).map_err(D::Error::custom)
        }
    }

    impl Serialize for SocketAddr {
        fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
        where
            S: serde::Serializer,
        {
            serializer.serialize_str(String::from(self).as_str())
        }
    }

    impl From<&SocketAddr> for String {
        fn from(value: &SocketAddr) -> Self {
            match value {
                SocketAddr::Bound(addr) => match addr.as_pathname() {
                    Some(path) => path
                        .to_str()
                        .expect("unable to get str for path")
                        .to_string(),
                    #[cfg(target_os = "linux")]
                    _ => match addr.as_abstract_name() {
                        Some(name) => format!("@{}", String::from_utf8_lossy(name)),
                        _ => String::from("(unnamed)"),
                    },
                    #[cfg(not(target_os = "linux"))]
                    _ => String::from("(unnamed)"),
                },
                SocketAddr::Unbound => String::from("(unbound)"),
            }
        }
    }

    impl FromStr for SocketAddr {
        type Err = anyhow::Error;

        fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
            match s {
                "" => {
                    // TODO: ensure this socket doesn't already exist. 24 bytes of randomness should be good for now but is not perfect.
                    // We can't use annon sockets because those are not valid across processes that aren't in the same process hierarchy aka forked.
                    let random_string = rand::rng()
                        .sample_iter(&Alphanumeric)
                        .take(24)
                        .map(char::from)
                        .collect::<String>();
                    SocketAddr::from_abstract_name(&random_string)
                }
                // by convention, named sockets are displayed with an '@' prefix
                name if name.starts_with("@") => {
                    SocketAddr::from_abstract_name(name.strip_prefix("@").unwrap())
                }
                path => SocketAddr::from_pathname(path),
            }
        }
    }

    impl Eq for SocketAddr {}
    impl std::hash::Hash for SocketAddr {
        fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
            String::from(self).hash(state);
        }
    }
    impl PartialEq for SocketAddr {
        fn eq(&self, other: &Self) -> bool {
            match (self, other) {
                (Self::Bound(saddr), Self::Bound(oaddr)) => {
                    if saddr.is_unnamed() || oaddr.is_unnamed() {
                        return false;
                    }

                    #[cfg(target_os = "linux")]
                    {
                        saddr.as_pathname() == oaddr.as_pathname()
                            && saddr.as_abstract_name() == oaddr.as_abstract_name()
                    }
                    #[cfg(not(target_os = "linux"))]
                    {
                        // On non-Linux platforms, only compare pathname since no abstract names
                        saddr.as_pathname() == oaddr.as_pathname()
                    }
                }
                (Self::Unbound, _) | (_, Self::Unbound) => false,
            }
        }
    }

    impl fmt::Display for SocketAddr {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            match self {
                Self::Bound(addr) => match addr.as_pathname() {
                    Some(path) => {
                        write!(f, "{}", path.to_string_lossy())
                    }
                    #[cfg(target_os = "linux")]
                    _ => match addr.as_abstract_name() {
                        Some(name) => {
                            if name.starts_with(b"@") {
                                return write!(f, "{}", String::from_utf8_lossy(name));
                            }
                            write!(f, "@{}", String::from_utf8_lossy(name))
                        }
                        _ => write!(f, "(unnamed)"),
                    },
                    #[cfg(not(target_os = "linux"))]
                    _ => write!(f, "(unnamed)"),
                },
                Self::Unbound => write!(f, "(unbound)"),
            }
        }
    }

    impl SocketAddr {
        /// Wraps the stdlib socket address for use with this module
        pub fn new(addr: StdSocketAddr) -> Self {
            Self::Bound(Box::new(addr))
        }

        /// Abstract socket names start with a "@" by convention when displayed. If there is an
        /// "@" prefix, it will be stripped from the name before used.
        #[cfg(target_os = "linux")]
        pub fn from_abstract_name(name: &str) -> anyhow::Result<Self> {
            Ok(Self::new(StdSocketAddr::from_abstract_name(
                name.strip_prefix("@").unwrap_or(name),
            )?))
        }

        #[cfg(not(target_os = "linux"))]
        pub fn from_abstract_name(name: &str) -> anyhow::Result<Self> {
            // On non-Linux platforms, convert abstract names to filesystem paths
            let name = name.strip_prefix("@").unwrap_or(name);
            let path = Self::abstract_to_filesystem_path(name);
            Self::from_pathname(&path.to_string_lossy())
        }

        #[cfg(not(target_os = "linux"))]
        fn abstract_to_filesystem_path(abstract_name: &str) -> std::path::PathBuf {
            use std::collections::hash_map::DefaultHasher;
            use std::hash::Hash;
            use std::hash::Hasher;

            // Generate a stable hash of the abstract name for deterministic paths
            let mut hasher = DefaultHasher::new();
            abstract_name.hash(&mut hasher);
            let hash = hasher.finish();

            // Include process ID to prevent inter-process conflicts
            let process_id = std::process::id();

            // TODO: we just leak these. Should we do something smarter?
            std::path::PathBuf::from(format!("/tmp/hyperactor_{}_{:x}", process_id, hash))
        }

        /// Pathnames may be absolute or relative.
        pub fn from_pathname(name: &str) -> anyhow::Result<Self> {
            Ok(Self::new(StdSocketAddr::from_pathname(name)?))
        }
    }

    impl TryFrom<SocketAddr> for StdSocketAddr {
        type Error = anyhow::Error;

        fn try_from(value: SocketAddr) -> Result<Self, Self::Error> {
            match value {
                SocketAddr::Bound(addr) => Ok(*addr),
                SocketAddr::Unbound => Err(anyhow::anyhow!(
                    "std::os::unix::SocketAddr must be a bound address"
                )),
            }
        }
    }
}

pub(crate) mod tcp {
    use tokio::net::TcpListener;
    use tokio::net::TcpStream;

    use super::*;

    #[derive(Debug)]
    pub(crate) struct TcpLink {
        pub(super) addr: SocketAddr,
        pub(super) session_id: SessionId,
        pub(super) stream_id: u8,
    }

    #[async_trait]
    impl Link for TcpLink {
        type Stream = TcpStream;

        fn dest(&self) -> ChannelAddr {
            ChannelAddr::Tcp(self.addr)
        }

        fn link_id(&self) -> SessionId {
            self.session_id
        }

        async fn next(&mut self) -> Result<Self::Stream, ClientError> {
            let session_id = self.session_id;
            let mut backoff = ExponentialBackoffBuilder::new()
                .with_initial_interval(Duration::from_millis(1))
                .with_multiplier(2.0)
                .with_randomization_factor(0.1)
                .with_max_interval(Duration::from_millis(1000))
                .with_max_elapsed_time(None)
                .build();
            loop {
                match TcpStream::connect(&self.addr).await {
                    Ok(mut stream) => {
                        stream.set_nodelay(true).map_err(|err| {
                            ClientError::Connect(
                                self.dest(),
                                err,
                                "cannot disable Nagle algorithm".to_string(),
                            )
                        })?;
                        write_link_init(&mut stream, session_id, self.stream_id)
                            .await
                            .map_err(|err| ClientError::Io(self.dest(), err))?;
                        return Ok(stream);
                    }
                    Err(err) => {
                        tracing::debug!(error = %err, "tcp connect failed, backing off");
                        if let Some(delay) = backoff.next_backoff() {
                            tokio::time::sleep(delay).await;
                        }
                    }
                }
            }
        }
    }

    /// Server-side listener for TCP sockets.
    #[derive(Debug)]
    pub(crate) struct TcpSocketListener {
        pub(super) inner: TcpListener,
        pub(super) addr: SocketAddr,
    }

    #[async_trait]
    impl super::Listener for TcpSocketListener {
        type Stream = TcpStream;

        async fn accept(&mut self) -> Result<(Self::Stream, ChannelAddr), ServerError> {
            let (stream, peer_addr) = self
                .inner
                .accept()
                .await
                .map_err(|err| ServerError::Io(ChannelAddr::Tcp(self.addr), err))?;
            stream
                .set_nodelay(true)
                .map_err(|err| ServerError::Io(ChannelAddr::Tcp(self.addr), err))?;
            Ok((stream, ChannelAddr::Tcp(peer_addr)))
        }
    }

    /// Create a TCP link to the given socket address.
    pub(crate) fn link(addr: SocketAddr, session_id: SessionId, stream_id: u8) -> TcpLink {
        TcpLink {
            addr,
            session_id,
            stream_id,
        }
    }
}

// TODO: Try to simplify the TLS creation T208304433
pub(crate) mod meta {
    use std::io;
    use std::path::PathBuf;
    use std::sync::Arc;

    use anyhow::Result;
    use tokio_rustls::TlsAcceptor;
    use tokio_rustls::TlsConnector;

    use super::*;
    use crate::config::Pem;
    use crate::config::PemBundle;

    const THRIFT_TLS_SRV_CA_PATH_ENV: &str = "THRIFT_TLS_SRV_CA_PATH";
    const DEFAULT_SRV_CA_PATH: &str = "/var/facebook/rootcanal/ca.pem";
    const THRIFT_TLS_CL_CERT_PATH_ENV: &str = "THRIFT_TLS_CL_CERT_PATH";
    const THRIFT_TLS_CL_KEY_PATH_ENV: &str = "THRIFT_TLS_CL_KEY_PATH";
    const DEFAULT_SERVER_PEM_PATH: &str = "/var/facebook/x509_identities/server.pem";

    #[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `ChannelError`.
    pub(crate) fn parse(addr_string: &str) -> Result<ChannelAddr, ChannelError> {
        // Use right split to allow for ipv6 addresses where ":" is expected.
        let parts = addr_string.rsplit_once(":");
        match parts {
            Some((hostname, port_str)) => {
                let Ok(port) = port_str.parse() else {
                    return Err(ChannelError::InvalidAddress(addr_string.to_string()));
                };
                Ok(ChannelAddr::MetaTls(TlsAddr::new(hostname, port)))
            }
            _ => Err(ChannelError::InvalidAddress(addr_string.to_string())),
        }
    }

    /// Construct a PemBundle for server operations from Meta-specific paths.
    /// Server cert and key come from the same file (server.pem).
    pub(super) fn get_server_pem_bundle() -> PemBundle {
        let ca_path = std::env::var_os(THRIFT_TLS_SRV_CA_PATH_ENV)
            .map(PathBuf::from)
            .unwrap_or_else(|| PathBuf::from(DEFAULT_SRV_CA_PATH));
        let server_pem_path = PathBuf::from(DEFAULT_SERVER_PEM_PATH);
        PemBundle {
            ca: Pem::File(ca_path),
            cert: Pem::File(server_pem_path.clone()),
            key: Pem::File(server_pem_path),
        }
    }

    /// Construct a PemBundle for client operations from Meta-specific env vars.
    /// Returns None if client cert/key env vars are not set.
    fn get_client_pem_bundle() -> Option<PemBundle> {
        let cert_path = std::env::var_os(THRIFT_TLS_CL_CERT_PATH_ENV)?;
        let key_path = std::env::var_os(THRIFT_TLS_CL_KEY_PATH_ENV)?;
        let ca_path = std::env::var_os(THRIFT_TLS_SRV_CA_PATH_ENV)
            .map(PathBuf::from)
            .unwrap_or_else(|| PathBuf::from(DEFAULT_SRV_CA_PATH));
        Some(PemBundle {
            ca: Pem::File(ca_path),
            cert: Pem::File(PathBuf::from(cert_path)),
            key: Pem::File(PathBuf::from(key_path)),
        })
    }

    /// Creates a TLS acceptor by looking for necessary certs and keys in a Meta server environment.
    pub(crate) fn tls_acceptor(enforce_client_tls: bool) -> Result<TlsAcceptor> {
        let bundle = get_server_pem_bundle();
        tls::tls_acceptor_from_bundle(&bundle, enforce_client_tls)
    }

    /// Try to create a TLS connector for Meta environments.
    ///
    /// Returns `Ok` when the root CA is present (optional client certs
    /// are added when `THRIFT_TLS_CL_CERT_PATH` / `THRIFT_TLS_CL_KEY_PATH`
    /// are set).
    pub(super) fn try_tls_connector() -> Result<TlsConnector> {
        tls_connector()
    }

    /// Creates a TLS connector by looking for necessary certs and keys in a Meta server environment.
    /// Supports optional client authentication (unlike the tls module which always requires it).
    fn tls_connector() -> Result<TlsConnector> {
        // Ensure ring is installed as the process-level crypto provider.
        // No-op when already installed (e.g. under Buck with native-tls).
        let _ = rustls::crypto::ring::default_provider().install_default();

        let ca_path = std::env::var_os(THRIFT_TLS_SRV_CA_PATH_ENV)
            .map(PathBuf::from)
            .unwrap_or_else(|| PathBuf::from(DEFAULT_SRV_CA_PATH));
        let ca_pem = Pem::File(ca_path);
        let root_store = tls::build_root_store(&ca_pem)?;

        // If client certs are available, use mutual TLS; otherwise, no client auth
        let config = rustls::ClientConfig::builder().with_root_certificates(Arc::new(root_store));

        let config = if let Some(bundle) = get_client_pem_bundle() {
            let certs = tls::load_certs(&bundle.cert)?;
            let key = tls::load_key(&bundle.key)?;
            config
                .with_client_auth_cert(certs, key)
                .map_err(|e| anyhow::anyhow!("load client certs: {}", e))?
        } else {
            config.with_no_client_auth()
        };

        Ok(TlsConnector::from(Arc::new(config)))
    }

    /// Create a MetaTLS link to the given address.
    pub fn link(
        addr: TlsAddr,
        session_id: SessionId,
        stream_id: u8,
    ) -> Result<tls::TlsLink, ClientError> {
        let connector = tls_connector().map_err(|e| {
            ClientError::Connect(
                ChannelAddr::MetaTls(addr.clone()),
                io::Error::other(e.to_string()),
                "failed to create TLS connector".to_string(),
            )
        })?;
        let TlsAddr { hostname, port } = addr;
        Ok(tls::TlsLink {
            hostname,
            port,
            connector,
            addr_type: tls::TlsAddrType::MetaTls,
            session_id,
            stream_id,
        })
    }
}

/// TLS transport module using configurable certificates via hyperactor config attributes.
pub(crate) mod tls {
    use std::io;
    use std::io::BufReader;
    use std::sync::Arc;

    use anyhow::Context;
    use anyhow::Result;
    use rustls::RootCertStore;
    use rustls::pki_types::CertificateDer;
    use rustls::pki_types::PrivateKeyDer;
    use rustls::pki_types::ServerName;
    use tokio::net::TcpStream;
    use tokio_rustls::TlsAcceptor;
    use tokio_rustls::TlsConnector;
    use tokio_rustls::client::TlsStream;

    use super::*;
    use crate::channel::TlsAddr;
    use crate::config::Pem;
    use crate::config::PemBundle;
    use crate::config::TLS_CA;
    use crate::config::TLS_CERT;
    use crate::config::TLS_KEY;

    /// Distinguishes between Tls and MetaTls for address construction.
    #[derive(Debug, Clone, Copy)]
    pub(crate) enum TlsAddrType {
        Tls,
        MetaTls,
    }

    /// Parse an address string into a TlsAddr.
    #[allow(clippy::result_large_err)]
    pub(crate) fn parse(addr_string: &str) -> Result<ChannelAddr, ChannelError> {
        // Use right split to allow for ipv6 addresses where ":" is expected.
        let parts = addr_string.rsplit_once(":");
        match parts {
            Some((hostname, port_str)) => {
                let Ok(port) = port_str.parse() else {
                    return Err(ChannelError::InvalidAddress(addr_string.to_string()));
                };
                Ok(ChannelAddr::Tls(TlsAddr::new(hostname, port)))
            }
            _ => Err(ChannelError::InvalidAddress(addr_string.to_string())),
        }
    }

    /// Load certificates from a Pem value.
    pub(super) fn load_certs(pem: &Pem) -> Result<Vec<CertificateDer<'static>>> {
        let mut reader = BufReader::new(pem.reader()?);
        let certs = rustls_pemfile::certs(&mut reader)
            .filter_map(Result::ok)
            .collect();
        Ok(certs)
    }

    /// Load a private key from a Pem value.
    pub(super) fn load_key(pem: &Pem) -> Result<PrivateKeyDer<'static>> {
        let mut reader = BufReader::new(pem.reader()?);
        loop {
            break match rustls_pemfile::read_one(&mut reader)? {
                Some(rustls_pemfile::Item::Pkcs1Key(key)) => Ok(PrivateKeyDer::Pkcs1(key)),
                Some(rustls_pemfile::Item::Pkcs8Key(key)) => Ok(PrivateKeyDer::Pkcs8(key)),
                Some(rustls_pemfile::Item::Sec1Key(key)) => Ok(PrivateKeyDer::Sec1(key)),
                Some(_) => continue,
                None => anyhow::bail!("no private key found in TLS key file"),
            };
        }
    }

    /// Build root certificate store from the CA pem.
    pub(super) fn build_root_store(ca_pem: &Pem) -> Result<RootCertStore> {
        let mut root_store = RootCertStore::empty();
        let certs = load_certs(ca_pem)?;
        root_store.add_parsable_certificates(certs);
        Ok(root_store)
    }

    /// Get the PEM bundle from configuration.
    fn get_pem_bundle() -> PemBundle {
        PemBundle {
            ca: hyperactor_config::global::get_cloned(TLS_CA),
            cert: hyperactor_config::global::get_cloned(TLS_CERT),
            key: hyperactor_config::global::get_cloned(TLS_KEY),
        }
    }

    /// Creates a TLS acceptor using certificates from the provided PEM bundle.
    /// If `enforce_client_tls` is true, requires client certificates for mutual TLS.
    pub(super) fn tls_acceptor_from_bundle(
        bundle: &PemBundle,
        enforce_client_tls: bool,
    ) -> Result<TlsAcceptor> {
        // Ensure ring is installed as the process-level crypto provider.
        // No-op when already installed (e.g. under Buck with native-tls).
        let _ = rustls::crypto::ring::default_provider().install_default();

        let certs = load_certs(&bundle.cert).context("load TLS certificate")?;
        let key = load_key(&bundle.key).context("load TLS key")?;
        let root_store = build_root_store(&bundle.ca).context("build root cert store")?;

        let config = rustls::ServerConfig::builder();
        let config = if enforce_client_tls {
            // Build server config with mutual TLS (require client certs)
            let client_verifier =
                rustls::server::WebPkiClientVerifier::builder(Arc::new(root_store))
                    .build()
                    .map_err(|e| anyhow::anyhow!("failed to build client verifier: {}", e))?;
            config.with_client_cert_verifier(client_verifier)
        } else {
            config.with_no_client_auth()
        }
        .with_single_cert(certs, key)?;

        Ok(TlsAcceptor::from(Arc::new(config)))
    }

    /// Creates a TLS acceptor using certificates from config (always enforces mutual TLS).
    pub(crate) fn tls_acceptor() -> Result<TlsAcceptor> {
        tls_acceptor_from_bundle(&get_pem_bundle(), true)
    }

    /// Creates a TLS connector using certificates from the provided PEM bundle.
    pub(super) fn tls_connector_from_bundle(bundle: &PemBundle) -> Result<TlsConnector> {
        // Ensure ring is installed as the process-level crypto provider.
        // No-op when already installed (e.g. under Buck with native-tls).
        let _ = rustls::crypto::ring::default_provider().install_default();

        let certs = load_certs(&bundle.cert).context("load TLS certificate")?;
        let key = load_key(&bundle.key).context("load TLS key")?;
        let root_store = build_root_store(&bundle.ca).context("build root cert store")?;

        let config = rustls::ClientConfig::builder()
            .with_root_certificates(Arc::new(root_store))
            .with_client_auth_cert(certs, key)
            .context("configure client auth")?;

        Ok(TlsConnector::from(Arc::new(config)))
    }

    /// Creates a TLS connector using certificates from config.
    fn tls_connector() -> Result<TlsConnector> {
        tls_connector_from_bundle(&get_pem_bundle())
    }

    /// Shared TLS link implementation used by both tls and metatls transports.
    pub(crate) struct TlsLink {
        pub(crate) hostname: Hostname,
        pub(crate) port: Port,
        pub(crate) connector: TlsConnector,
        pub(crate) addr_type: TlsAddrType,
        pub(crate) session_id: SessionId,
        pub(crate) stream_id: u8,
    }

    impl std::fmt::Debug for TlsLink {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            f.debug_struct("TlsLink")
                .field("hostname", &self.hostname)
                .field("port", &self.port)
                .field("addr_type", &self.addr_type)
                .finish()
        }
    }

    #[async_trait]
    impl Link for TlsLink {
        type Stream = TlsStream<TcpStream>;

        fn dest(&self) -> ChannelAddr {
            let addr = TlsAddr::new(self.hostname.clone(), self.port);
            match self.addr_type {
                TlsAddrType::Tls => ChannelAddr::Tls(addr),
                TlsAddrType::MetaTls => ChannelAddr::MetaTls(addr),
            }
        }

        fn link_id(&self) -> SessionId {
            self.session_id
        }

        async fn next(&mut self) -> Result<Self::Stream, ClientError> {
            let session_id = self.session_id;
            let server_name = ServerName::try_from(self.hostname.clone()).map_err(|e| {
                ClientError::Connect(
                    self.dest(),
                    io::Error::other(e.to_string()),
                    "invalid server name".to_string(),
                )
            })?;
            let mut backoff = ExponentialBackoffBuilder::new()
                .with_initial_interval(Duration::from_millis(1))
                .with_multiplier(2.0)
                .with_randomization_factor(0.1)
                .with_max_interval(Duration::from_millis(1000))
                .with_max_elapsed_time(None)
                .build();
            loop {
                let mut addrs = (self.hostname.as_ref(), self.port)
                    .to_socket_addrs()
                    .map_err(|_| ClientError::Resolve(self.dest()))?;
                let addr = addrs.next().ok_or(ClientError::Resolve(self.dest()))?;
                match TcpStream::connect(&addr).await {
                    Ok(stream) => {
                        stream.set_nodelay(true).map_err(|err| {
                            ClientError::Connect(
                                self.dest(),
                                err,
                                "cannot disable Nagle algorithm".to_string(),
                            )
                        })?;
                        let mut tls_stream = self
                            .connector
                            .connect(server_name.clone(), stream)
                            .await
                            .map_err(|err| {
                                tracing::info!(
                                    dest = %self.dest(),
                                    error = %err,
                                    "TLS handshake failed"
                                );
                                ClientError::Connect(
                                    self.dest(),
                                    err,
                                    format!("cannot establish TLS connection to {:?}", server_name),
                                )
                            })?;
                        write_link_init(&mut tls_stream, session_id, self.stream_id)
                            .await
                            .map_err(|err| ClientError::Io(self.dest(), err))?;
                        return Ok(tls_stream);
                    }
                    Err(err) => {
                        tracing::debug!(error = %err, "tls connect failed, backing off");
                        if let Some(delay) = backoff.next_backoff() {
                            tokio::time::sleep(delay).await;
                        }
                    }
                }
            }
        }
    }

    /// Create a TLS link to the given address.
    pub fn link(
        addr: TlsAddr,
        session_id: SessionId,
        stream_id: u8,
    ) -> Result<TlsLink, ClientError> {
        let connector = tls_connector().map_err(|e| {
            ClientError::Connect(
                ChannelAddr::Tls(addr.clone()),
                io::Error::other(e.to_string()),
                "failed to create TLS connector".to_string(),
            )
        })?;
        let TlsAddr { hostname, port } = addr;
        Ok(TlsLink {
            hostname,
            port,
            connector,
            addr_type: TlsAddrType::Tls,
            session_id,
            stream_id,
        })
    }

    #[cfg(test)]
    mod tests {
        use timed_test::async_timed_test;

        use super::*;
        use crate::channel::Rx;
        use crate::channel::net::server;
        use crate::config::Pem;
        use crate::config::TLS_CA;
        use crate::config::TLS_CERT;
        use crate::config::TLS_KEY;

        // Dummy test certificates generated with openssl for testing only.
        // These certificates include Subject Alternative Names (SAN) for localhost, 127.0.0.1, and ::1
        // CA certificate
        const TEST_CA_CERT: &str = r#"-----BEGIN CERTIFICATE-----
MIIDBTCCAe2gAwIBAgIUaGNmboiIosG+8Up0vgDr/+cg+2IwDQYJKoZIhvcNAQEL
BQAwEjEQMA4GA1UEAwwHVGVzdCBDQTAeFw0yNjAxMjgxNzA4MzlaFw0yNzAxMjgx
NzA4MzlaMBIxEDAOBgNVBAMMB1Rlc3QgQ0EwggEiMA0GCSqGSIb3DQEBAQUAA4IB
DwAwggEKAoIBAQC9RBoMYXCajklswt8Vi1JI1lEYzic0WNOmz45vG/7H6jTWkgL3
K5Ri+Seg3MobDNc48YHWXYm4hP9wCzkx8ih3ntT5XiY1My/G3jLUuoIEE9pF/BoJ
YQwZVoPNFhA9WhXNRsINf1cXFf8NzRfXpxBfKWtQJxYXU4JiDBQ6rLnQQABo8JmQ
vYFhJbBaYip5jTSiVNn7mB1zNr5jsVxuoSF53Pb7xQ76bwBdOq4zd6PSxL5/lr4G
cHSoxwZQdZMG7PL6hbxDQ2S2YI2lYVET1zwc2WPKCfjbEXBC/jzx828CInQtuksk
18gJt6xHkTFEA8CSA29GM3lejnwYWf51xyyBAgMBAAGjUzBRMB0GA1UdDgQWBBRX
cbxSZ70NsUkAS3Hhy6irugywJDAfBgNVHSMEGDAWgBRXcbxSZ70NsUkAS3Hhy6ir
ugywJDAPBgNVHRMBAf8EBTADAQH/MA0GCSqGSIb3DQEBCwUAA4IBAQA7aAFfyW67
Z+uGSVYhpsT/uH/3Z3nr7X1smTz5CGEfq2czEcTC7gbYI2l8GZ47GPfnAvHTBZVm
V/XncBCsj7/thOh2jYEHFyCbPckoaSCRyCOnK7LPUlr4HN5uP9EFe45qBLCJDEoY
GTTw7MtzwdovfjchNfKQCTtkBJCXQ95WLCf6UOh02Sn28UTlgfXzF0X0FrcWqWa3
uJZd4XOo4O6hKKlHaBaQPiEr++1xc3SWPV7jZHbckI/vKBnDdEZ9JQX5fFZuypUI
sgomYHxvxrU2hWx+7k53CRdjfaIvT9Ie44z9sSdsU/+blw2S8f/ZTmuECoIAAXYO
0qpzlxZMdr7T
-----END CERTIFICATE-----"#;

        // Server certificate (signed by CA) with SAN for localhost, 127.0.0.1, ::1
        const TEST_SERVER_CERT: &str = r#"-----BEGIN CERTIFICATE-----
MIIDJDCCAgygAwIBAgIUaz66DsWaH5ZXM4hCFnbVbMsyN1cwDQYJKoZIhvcNAQEL
BQAwEjEQMA4GA1UEAwwHVGVzdCBDQTAeFw0yNjAxMjgxNzA4MzlaFw0yNzAxMjgx
NzA4MzlaMBQxEjAQBgNVBAMMCWxvY2FsaG9zdDCCASIwDQYJKoZIhvcNAQEBBQAD
ggEPADCCAQoCggEBAKCbp++qNyTn5LOsV0h9gLKJALBcjg2A14I3804N9UyDhPW2
QKQ2W424u2P1MfKrw/2C+CErGlrADlnco2RQVDAarAIuGdFvBOt5UezqOS7Mk4OS
9MlS7NZnMbc37KuM9UIG5ScJjXR/Z5z9dxeR0I9y3n0Ix6khbV7tOSHobiweI0FI
8LftBS+CQnXr6vbWPcHcW6Z0FHUv7IWhqMWmv9PlZRGe9Y6VzXrRp0PBnZMOnAYf
aMQUwYRswWdm9j9Z1sMdTJ14G+KVmO3Vj6XI6Sm9uIcYhlwG/kORwogJFWlVuP9o
rloFRCjyHJ1d7GZqqnRyHHDDCBms8ed+3YfEYQECAwEAAaNwMG4wLAYDVR0RBCUw
I4IJbG9jYWxob3N0hwR/AAABhxAAAAAAAAAAAAAAAAAAAAABMB0GA1UdDgQWBBQl
J4vxUoCzqqeTwQAiLqE8wYezKzAfBgNVHSMEGDAWgBRXcbxSZ70NsUkAS3Hhy6ir
ugywJDANBgkqhkiG9w0BAQsFAAOCAQEAnXHIBDQ4AHAMV71piTOuI41ShASQed6L
bi7XUMZgZDslLkfU1vnP3BlwpliraBsAytSYQC6kbytOuz1uQ4K7yLb2tAAmUgEO
EdIVt9SXr5tCcIPeLmInF0pysPqjZO8n7vtJyd9gryKqdhm1uzA7WQWq/Az8a9Sk
uW2J6Oc5p6P7Mf3/ixqXzvGRo8rzu0CUJOJ67UTE/HhbJuplQ5dep5CEEOAIsAtH
zn9O4rW92ueBkoBJM++YILS1vQ7jKc2N3RNrnHm7FeootBrtR9mBi0TH97K73ZPZ
2Cdhnym0CsCJggrllFGH32cYo7+K2PO7/4oj5XbBCSWcssicvd8ovg==
-----END CERTIFICATE-----"#;

        // Server private key
        // test-only embedded key -- though it might expire at some point:
        // @lint-ignore PRIVATEKEY insecure-private-key-storage
        const TEST_SERVER_KEY: &str = r#"-----BEGIN PRIVATE KEY-----
MIIEugIBADANBgkqhkiG9w0BAQEFAASCBKQwggSgAgEAAoIBAQCgm6fvqjck5+Sz
rFdIfYCyiQCwXI4NgNeCN/NODfVMg4T1tkCkNluNuLtj9THyq8P9gvghKxpawA5Z
3KNkUFQwGqwCLhnRbwTreVHs6jkuzJODkvTJUuzWZzG3N+yrjPVCBuUnCY10f2ec
/XcXkdCPct59CMepIW1e7Tkh6G4sHiNBSPC37QUvgkJ16+r21j3B3FumdBR1L+yF
oajFpr/T5WURnvWOlc160adDwZ2TDpwGH2jEFMGEbMFnZvY/WdbDHUydeBvilZjt
1Y+lyOkpvbiHGIZcBv5DkcKICRVpVbj/aK5aBUQo8hydXexmaqp0chxwwwgZrPHn
ft2HxGEBAgMBAAECgf8G5qlQov+7ljs9fSpC8yGUik59RXzVF7Qq5DyQHglsQDp2
VF5yr+M/M7DZmq+KvdauDfKbej6np5j2Q4TByrHTX1IExfZWCW8srwnWJDpQyHmO
LcJW5DlI/SYluUFyHZxsOd+ezcpGNzM8i6eSW7GaeFUXCkmJ+uW4LnlF+7bALnnd
D6sak/58EsII+IJyd4lFn+voszlPn3CZGR0jkp21rvpaKgrMIsKVWWQO/sLDU5pr
VbpBThcLU5gRcnQouQX12e2VTCIlFu75WTsJ8V/KnEaOZUVlU/B/Bs+WQF3U+/Jo
eX4N+D6OsEcNQjERAFyWujxsl1WpD4uSsbFMN0ECgYEA2b7AdL+oKPQHku2KcBhr
Zw8K4tMDlr2VPPNwZcBTLo+O71vv/xXjMcXrXmowzkgEQckUmt1VB46riyydhwdP
/n9ciWcz0Va/nwHR6Y9F9unBiyUBP7PRhRyjQyRZZRGDSJvP+Xmc5UJFpRr07VLU
nfgMXDj37vXzKDpfhdEB2nkCgYEAvNMfA8P8w3+6246x5YHflvTkPdw+2oyge+LD
mphB/w7SF8mlyNGloj3+KBZmd9SkvT57wCvO96Y9/n+mBAVisRggc0hK4ymOVYhb
+im/JvqGQMbVeg6iCOHnWdaZf9tL8uVsegQy3kVTN7vAa+CMFgX1dt65cGBX6XkB
44pYmMkCgYALhbiRdQLlB+TOtZs5y1EDpxwgXKI3+9hF3Wv5NnAwapBZwje0++eF
3r9Rw7TJda4j/QwGFehF+hrBxp6fYpetE/hFnRx0225Qb7w368j8A+ql/lNOl6li
rd1F1EqWupKD6RrcTL8sspEU55RGaretlE6zIqCcGI/BdTVQ03qRoQKBgHDC3zWf
d7XD9HGjQGdfbIe4jQjIGxzmd/wjik4q+NZ5IkukVwWa9P/zZ3DHF8Ad05dT1hEH
2FwaAdGWpyyljq9VSiOuG1KXAXHgsZSuE4ISf9P1KYzvaiJFzaPfvOEWs79E9MfU
9A+6dJzG2X1SpjWMr26iSTlrv3QkmFUqzAfJAoGASBkn4wls+oC5rv/Mch43pBv5
UmKru4ltnEHJZdbSi2DJ+AnDLD222JCasb1VT1tm2XgW6DBqrdVRPPP6GOlB0MHU
+3ULtZxAczt7I+ST2bo0/DV2Hse89Cm63w4wLOiVZs7+1wrAzJZLokWF7Q5gesra
u19txmtkiMEH+aNmekk=
-----END PRIVATE KEY-----"#;

        #[async_timed_test(timeout_secs = 30)]
        async fn test_tls_basic() {
            // Ensure ring is installed as the default crypto provider
            // (no-op if already installed, e.g. under Buck with native-tls).
            let _ = rustls::crypto::ring::default_provider().install_default();

            // Set up TLS config using the standard override pattern
            let config = hyperactor_config::global::lock();
            let _guard_cert =
                config.override_key(TLS_CERT, Pem::Value(TEST_SERVER_CERT.as_bytes().to_vec()));
            let _guard_key =
                config.override_key(TLS_KEY, Pem::Value(TEST_SERVER_KEY.as_bytes().to_vec()));
            let _guard_ca =
                config.override_key(TLS_CA, Pem::Value(TEST_CA_CERT.as_bytes().to_vec()));

            // Create a TLS server bound to localhost with dynamic port
            let addr = TlsAddr::new("localhost", 0);

            let (local_addr, mut rx) =
                server::serve::<u64>(ChannelAddr::Tls(addr), None).expect("failed to serve");

            // Dial the server
            let tx: super::NetTx<u64> = super::spawn(
                link(
                    match &local_addr {
                        ChannelAddr::Tls(addr) => addr.clone(),
                        _ => panic!("unexpected address type"),
                    },
                    SessionId::random(),
                    0,
                )
                .expect("failed to create link"),
            );

            // Send a message
            tx.post(42u64);

            // Receive the message
            let received = rx.recv().await.expect("failed to receive");
            assert_eq!(received, 42u64);
        }

        #[async_timed_test(timeout_secs = 30)]
        async fn test_tls_multiple_messages() {
            let _ = rustls::crypto::ring::default_provider().install_default();

            // Set up TLS config using the standard override pattern
            let config = hyperactor_config::global::lock();
            let _guard_cert =
                config.override_key(TLS_CERT, Pem::Value(TEST_SERVER_CERT.as_bytes().to_vec()));
            let _guard_key =
                config.override_key(TLS_KEY, Pem::Value(TEST_SERVER_KEY.as_bytes().to_vec()));
            let _guard_ca =
                config.override_key(TLS_CA, Pem::Value(TEST_CA_CERT.as_bytes().to_vec()));

            let addr = TlsAddr::new("localhost", 0);

            let (local_addr, mut rx) =
                server::serve::<String>(ChannelAddr::Tls(addr), None).expect("failed to serve");
            let tx: super::NetTx<String> = super::spawn(
                link(
                    match &local_addr {
                        ChannelAddr::Tls(addr) => addr.clone(),
                        _ => panic!("unexpected address type"),
                    },
                    SessionId::random(),
                    0,
                )
                .expect("failed to create link"),
            );

            // Send multiple messages
            for i in 0..10 {
                tx.post(format!("message {}", i));
            }

            // Receive all messages
            for i in 0..10 {
                let received = rx.recv().await.expect("failed to receive");
                assert_eq!(received, format!("message {}", i));
            }
        }

        #[test]
        fn test_tls_parse_hostname_port() {
            let addr = parse("localhost:8080").expect("failed to parse");
            assert!(matches!(
                addr,
                ChannelAddr::Tls(TlsAddr { hostname, port })
                    if hostname == "localhost" && port == 8080
            ));
        }

        #[test]
        fn test_tls_parse_socket_addr() {
            let addr = parse("127.0.0.1:8080").expect("failed to parse");
            assert!(matches!(
                addr,
                ChannelAddr::Tls(TlsAddr { hostname, port })
                    if hostname == "127.0.0.1" && port == 8080
            ));
        }

        #[test]
        fn test_tls_certs_parsing() {
            // Verify that the test certificates can be parsed correctly
            let cert_pem = Pem::Value(TEST_SERVER_CERT.as_bytes().to_vec());
            let key_pem = Pem::Value(TEST_SERVER_KEY.as_bytes().to_vec());
            let ca_pem = Pem::Value(TEST_CA_CERT.as_bytes().to_vec());

            let certs = super::load_certs(&cert_pem).expect("failed to load certs");
            assert!(!certs.is_empty(), "expected at least one certificate");

            let _key = super::load_key(&key_pem).expect("failed to load key");

            let root_store = super::build_root_store(&ca_pem).expect("failed to build root store");
            assert!(!root_store.is_empty(), "expected at least one CA cert");
        }

        #[test]
        fn test_tls_acceptor_creation() {
            // Ensure ring is installed as the default crypto provider
            // (no-op if already installed, e.g. under Buck with native-tls).
            let _ = rustls::crypto::ring::default_provider().install_default();

            // Set up TLS config using the standard override pattern
            let config = hyperactor_config::global::lock();
            let _guard_cert =
                config.override_key(TLS_CERT, Pem::Value(TEST_SERVER_CERT.as_bytes().to_vec()));
            let _guard_key =
                config.override_key(TLS_KEY, Pem::Value(TEST_SERVER_KEY.as_bytes().to_vec()));
            let _guard_ca =
                config.override_key(TLS_CA, Pem::Value(TEST_CA_CERT.as_bytes().to_vec()));

            // Verify that we can create a TLS acceptor
            let _acceptor = super::tls_acceptor().expect("failed to create TLS acceptor");
        }

        #[test]
        fn test_tls_connector_creation() {
            // Ensure ring is installed as the default crypto provider
            // (no-op if already installed, e.g. under Buck with native-tls).
            let _ = rustls::crypto::ring::default_provider().install_default();

            // Set up TLS config using the standard override pattern
            let config = hyperactor_config::global::lock();
            let _guard_cert =
                config.override_key(TLS_CERT, Pem::Value(TEST_SERVER_CERT.as_bytes().to_vec()));
            let _guard_key =
                config.override_key(TLS_KEY, Pem::Value(TEST_SERVER_KEY.as_bytes().to_vec()));
            let _guard_ca =
                config.override_key(TLS_CA, Pem::Value(TEST_CA_CERT.as_bytes().to_vec()));

            // Verify that we can create a TLS connector
            let _connector = super::tls_connector().expect("failed to create TLS connector");
        }
    }
}

/// Build the OSS PemBundle from hyperactor_config attributes.
fn oss_pem_bundle() -> crate::config::PemBundle {
    crate::config::PemBundle {
        ca: hyperactor_config::global::get_cloned(crate::config::TLS_CA),
        cert: hyperactor_config::global::get_cloned(crate::config::TLS_CERT),
        key: hyperactor_config::global::get_cloned(crate::config::TLS_KEY),
    }
}

/// Try to find a usable TLS [`PemBundle`](crate::config::PemBundle)
/// by probing the same sources as [`try_tls_acceptor`] /
/// [`try_tls_connector`].
///
/// Returns the first bundle whose CA certificate is readable.
/// Only CA readability is checked — cert and key are returned as-is
/// and may not be valid. Callers that cannot use `tokio_rustls` types
/// directly (e.g. reqwest) can read the raw PEM bytes via
/// [`Pem::reader`](crate::config::Pem::reader).
pub fn try_tls_pem_bundle() -> Option<crate::config::PemBundle> {
    let oss_bundle = oss_pem_bundle();
    if oss_bundle.ca.reader().is_ok() {
        return Some(oss_bundle);
    }
    tracing::debug!("OSS TLS bundle: CA not readable, trying Meta paths");

    let meta_bundle = meta::get_server_pem_bundle();
    if meta_bundle.ca.reader().is_ok() {
        return Some(meta_bundle);
    }
    tracing::debug!("Meta TLS bundle: CA not readable, no TLS available");

    None
}

/// Try to build a [`TlsAcceptor`](tokio_rustls::TlsAcceptor) for an
/// HTTP server by probing for available TLS certificates.
///
/// Detection order:
/// 1. **OSS / explicit config** — `HYPERACTOR_TLS_CERT`,
///    `HYPERACTOR_TLS_KEY`, and `HYPERACTOR_TLS_CA` (read via
///    [`hyperactor_config`]).
/// 2. **Meta default paths** —
///    `/var/facebook/x509_identities/server.pem` and
///    `/var/facebook/rootcanal/ca.pem`. These are present on
///    devservers and in MAST / Tupperware containers.
/// 3. **None** — no usable certificates found; caller should fall
///    back to plain HTTP.
///
/// When `enforce_client_tls` is `true`, the returned acceptor
/// requires clients to present a valid certificate signed by the
/// configured CA (mutual TLS via `WebPkiClientVerifier`). When
/// `false`, the acceptor authenticates itself but does not demand
/// client certificates.
pub fn try_tls_acceptor(enforce_client_tls: bool) -> Option<tokio_rustls::TlsAcceptor> {
    let oss_bundle = oss_pem_bundle();
    if let Ok(acceptor) = tls::tls_acceptor_from_bundle(&oss_bundle, enforce_client_tls) {
        return Some(acceptor);
    }
    tracing::debug!("OSS TLS acceptor failed, trying Meta paths");

    let meta_bundle = meta::get_server_pem_bundle();
    if let Ok(acceptor) = tls::tls_acceptor_from_bundle(&meta_bundle, enforce_client_tls) {
        return Some(acceptor);
    }
    tracing::debug!("Meta TLS acceptor failed, no TLS available");

    None
}

/// Try to build a [`TlsConnector`](tokio_rustls::TlsConnector) for an
/// HTTP client that needs to connect to a TLS-enabled server.
///
/// Detection mirrors [`try_tls_acceptor`]:
/// 1. **OSS** — `HYPERACTOR_TLS_CA` (and optionally
///    `HYPERACTOR_TLS_CERT` + `HYPERACTOR_TLS_KEY` for mutual TLS).
/// 2. **Meta** — root CA at `/var/facebook/rootcanal/ca.pem`,
///    optional client certs from `THRIFT_TLS_CL_CERT_PATH` /
///    `THRIFT_TLS_CL_KEY_PATH`.
/// 3. **None** — no usable CA found; caller should fall back to plain
///    HTTP.
pub fn try_tls_connector() -> Option<tokio_rustls::TlsConnector> {
    let oss_bundle = oss_pem_bundle();
    if let Ok(connector) = tls::tls_connector_from_bundle(&oss_bundle) {
        return Some(connector);
    }
    tracing::debug!("OSS TLS connector failed, trying Meta paths");

    if let Ok(connector) = meta::try_tls_connector() {
        return Some(connector);
    }
    tracing::debug!("Meta TLS connector failed, no TLS available");

    None
}

#[cfg(test)]
mod tests {

    #![expect(
        clippy::await_holding_invalid_type,
        reason = "tracing_test::traced_test macro expansion holds tracing::span::Entered across awaits; can't be fixed in our code"
    )]

    use std::assert_matches;
    use std::collections::VecDeque;
    use std::marker::PhantomData;
    use std::sync::Arc;
    use std::sync::RwLock;
    use std::sync::atomic::AtomicBool;
    use std::sync::atomic::AtomicU64;
    use std::sync::atomic::Ordering;
    use std::time::Duration;
    #[cfg(target_os = "linux")] // uses abstract names
    use std::time::UNIX_EPOCH;

    #[cfg(target_os = "linux")] // uses abstract names
    use anyhow::Result;
    use bytes::Bytes;
    use rand::RngExt as _;
    use rand::SeedableRng as _;
    use rand::distr::Alphanumeric;
    use rand::rngs::SysRng;
    use timed_test::async_timed_test;
    use tokio::io::AsyncWrite;
    use tokio::io::DuplexStream;
    use tokio::io::ReadHalf;
    use tokio::io::WriteHalf;
    use tokio::task::JoinHandle;
    use tokio_util::sync::CancellationToken;

    use super::server;
    use super::*;
    use crate::channel;
    use crate::channel::net::framed::FrameReader;
    use crate::channel::net::framed::FrameWrite;
    use crate::channel::net::server::AcceptorLink;
    use crate::config;
    use crate::metrics;
    use crate::sync::mvar::MVar;

    /// Like the `logs_assert` injected by `#[traced_test]`, but without scope
    /// filtering. Use when asserting on events emitted outside the test's span
    /// (e.g. from spawned tasks or panic hooks).
    fn logs_assert_unscoped(f: impl Fn(&[&str]) -> Result<(), String>) {
        let buf = tracing_test::internal::global_buf().lock().unwrap();
        let logs_str = std::str::from_utf8(&buf).expect("Logs contain invalid UTF8");
        let lines: Vec<&str> = logs_str.lines().collect();
        match f(&lines) {
            Ok(()) => {}
            Err(msg) => panic!("{}", msg),
        }
    }

    #[cfg(target_os = "linux")] // uses abstract names
    #[tracing_test::traced_test]
    #[tokio::test]
    async fn test_unix_basic() -> Result<()> {
        let timestamp = std::time::SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let unique_address = format!("test_unix_basic_{}", timestamp);

        let (addr, mut rx) = server::serve::<u64>(
            ChannelAddr::Unix(unix::SocketAddr::from_abstract_name(&unique_address)?),
            None,
        )
        .unwrap();

        // It is important to keep Tx alive until all expected messages are
        // received. Otherwise, the channel would be closed when Tx is dropped.
        // Although the messages are sent to the server's buffer before the
        // channel was closed, NetRx could still error out before taking them
        // out of the buffer because NetRx could not ack through the closed
        // channel.
        {
            let tx: ChannelTx<u64> = channel::dial::<u64>(addr.clone()).unwrap();
            tx.post(123);
            assert_eq!(rx.recv().await.unwrap(), 123);
        }

        {
            let tx = channel::dial::<u64>(addr.clone()).unwrap();
            tx.post(321);
            tx.post(111);
            tx.post(444);

            assert_eq!(rx.recv().await.unwrap(), 321);
            assert_eq!(rx.recv().await.unwrap(), 111);
            assert_eq!(rx.recv().await.unwrap(), 444);
        }

        {
            let tx = channel::dial::<u64>(addr).unwrap();
            drop(rx);

            let (return_tx, return_rx) = oneshot::channel();
            tx.try_post(123, return_tx);
            assert_matches!(
                return_rx.await,
                Ok(SendError {
                    error: ChannelError::Closed,
                    message: 123,
                    ..
                })
            );
        }

        Ok(())
    }

    #[cfg(target_os = "linux")] // uses abstract names
    #[tracing_test::traced_test]
    #[tokio::test]
    async fn test_unix_basic_client_before_server() -> Result<()> {
        // We run this test on Unix because we can pick our own port names more easily.
        let timestamp = std::time::SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let socket_addr =
            unix::SocketAddr::from_abstract_name(&format!("test_unix_basic_{}", timestamp))
                .unwrap();

        // Dial the channel before we actually serve it.
        let addr = ChannelAddr::Unix(socket_addr.clone());
        let tx = crate::channel::dial::<u64>(addr.clone()).unwrap();
        tx.post(123);

        let (_, mut rx) = server::serve::<u64>(ChannelAddr::Unix(socket_addr), None).unwrap();
        assert_eq!(rx.recv().await.unwrap(), 123);

        tx.post(321);
        tx.post(111);
        tx.post(444);

        assert_eq!(rx.recv().await.unwrap(), 321);
        assert_eq!(rx.recv().await.unwrap(), 111);
        assert_eq!(rx.recv().await.unwrap(), 444);

        Ok(())
    }

    #[tracing_test::traced_test]
    #[async_timed_test(timeout_secs = 60)]
    // TODO: OSS: called `Result::unwrap()` on an `Err` value: Listen(Tcp([::1]:0), Os { code: 99, kind: AddrNotAvailable, message: "Cannot assign requested address" })
    #[cfg_attr(not(fbcode_build), ignore)]
    async fn test_tcp_basic() {
        let (addr, mut rx) =
            server::serve::<u64>(ChannelAddr::Tcp("[::1]:0".parse().unwrap()), None).unwrap();
        {
            let tx = channel::dial::<u64>(addr.clone()).unwrap();
            tx.post(123);
            assert_eq!(rx.recv().await.unwrap(), 123);
        }

        {
            let tx = channel::dial::<u64>(addr.clone()).unwrap();
            tx.post(321);
            tx.post(111);
            tx.post(444);

            assert_eq!(rx.recv().await.unwrap(), 321);
            assert_eq!(rx.recv().await.unwrap(), 111);
            assert_eq!(rx.recv().await.unwrap(), 444);
        }

        {
            let tx = channel::dial::<u64>(addr).unwrap();
            drop(rx);

            let (return_tx, return_rx) = oneshot::channel();
            tx.try_post(123, return_tx);
            assert_matches!(
                return_rx.await,
                Ok(SendError {
                    error: ChannelError::Closed,
                    message: 123,
                    ..
                })
            );
        }
    }

    // The message size is limited by CODEC_MAX_FRAME_LENGTH.
    #[async_timed_test(timeout_secs = 5)]
    // TODO: OSS: called `Result::unwrap()` on an `Err` value: Listen(Tcp([::1]:0), Os { code: 99, kind: AddrNotAvailable, message: "Cannot assign requested address" })
    #[cfg_attr(not(fbcode_build), ignore)]
    async fn test_tcp_message_size() {
        let default_size_in_bytes = 100 * 1024 * 1024;
        // Use temporary config for this test
        let config = hyperactor_config::global::lock();
        let _guard1 = config.override_key(config::MESSAGE_DELIVERY_TIMEOUT, Duration::from_secs(1));
        let _guard2 = config.override_key(config::CODEC_MAX_FRAME_LENGTH, default_size_in_bytes);

        let (addr, mut rx) =
            server::serve::<String>(ChannelAddr::Tcp("[::1]:0".parse().unwrap()), None).unwrap();

        let tx = channel::dial::<String>(addr.clone()).unwrap();
        // Default size is okay
        {
            // Leave some headroom because Tx will wrap the payload in Frame::Message.
            let message = "a".repeat(default_size_in_bytes - 1024);
            tx.post(message.clone());
            assert_eq!(rx.recv().await.unwrap(), message);
        }
        // Bigger than the default size will fail.
        {
            let (return_channel, return_receiver) = oneshot::channel();
            let message = "a".repeat(default_size_in_bytes + 1024);
            tx.try_post(message.clone(), return_channel);
            let returned = return_receiver.await.unwrap();
            assert_eq!(message, returned.message);
        }
    }

    #[async_timed_test(timeout_secs = 30)]
    // TODO: OSS: called `Result::unwrap()` on an `Err` value: Listen(Tcp([::1]:0), Os { code: 99, kind: AddrNotAvailable, message: "Cannot assign requested address" })
    #[cfg_attr(not(fbcode_build), ignore)]
    async fn test_ack_flush() {
        let config = hyperactor_config::global::lock();
        // Set a large value to effectively prevent acks from being sent except
        // during shutdown flush.
        let _guard_message_ack =
            config.override_key(config::MESSAGE_ACK_EVERY_N_MESSAGES, 100000000);
        let _guard_delivery_timeout =
            config.override_key(config::MESSAGE_DELIVERY_TIMEOUT, Duration::from_secs(5));

        let (addr, mut net_rx) =
            server::serve::<u64>(ChannelAddr::Tcp("[::1]:0".parse().unwrap()), None).unwrap();
        let net_tx = channel::dial::<u64>(addr.clone()).unwrap();
        let (tx, rx) = oneshot::channel();
        net_tx.try_post(1, tx);
        assert_eq!(net_rx.recv().await.unwrap(), 1);
        drop(net_rx);
        // Using `is_err` to confirm the message is delivered/acked is confusing,
        // but is correct. See how send is implemented: https://fburl.com/code/ywt8lip2
        assert!(rx.await.is_err());
    }

    #[async_timed_test(timeout_secs = 60)]
    // TODO: OSS: failed to retrieve ipv6 address
    #[cfg_attr(not(fbcode_build), ignore)]
    async fn test_meta_tls_basic() {
        hyperactor_telemetry::initialize_logging_for_test();

        let addr = ChannelAddr::any(ChannelTransport::MetaTls(TlsMode::IpV6));
        let meta_addr = match addr {
            ChannelAddr::MetaTls(meta_addr) => meta_addr,
            _ => panic!("expected MetaTls address"),
        };
        let (local_addr, mut rx) =
            server::serve::<u64>(ChannelAddr::MetaTls(meta_addr), None).unwrap();
        {
            let tx = channel::dial::<u64>(local_addr.clone()).unwrap();
            tx.post(123);
        }
        assert_eq!(rx.recv().await.unwrap(), 123);

        {
            let tx = channel::dial::<u64>(local_addr.clone()).unwrap();
            tx.post(321);
            tx.post(111);
            tx.post(444);
            assert_eq!(rx.recv().await.unwrap(), 321);
            assert_eq!(rx.recv().await.unwrap(), 111);
            assert_eq!(rx.recv().await.unwrap(), 444);
        }

        {
            let tx = channel::dial::<u64>(local_addr).unwrap();
            drop(rx);

            let (return_tx, return_rx) = oneshot::channel();
            tx.try_post(123, return_tx);
            assert_matches!(
                return_rx.await,
                Ok(SendError {
                    error: ChannelError::Closed,
                    message: 123,
                    ..
                })
            );
        }
    }

    #[derive(Clone, Debug, Default)]
    struct NetworkFlakiness {
        // A tuple of:
        //   1. the probability of a network failure when sending a message.
        //   2. the max number of disconnections allowed.
        //   3. the minimum duration between disconnections.
        //
        //   2 and 3 are useful to prevent frequent disconnections leading to
        //   unacked messages being sent repeatedly.
        disconnect_params: Option<(f64, u64, Duration)>,
        // The max possible latency when sending a message. The actual latency
        // is randomly generated between 0 and max_latency.
        latency_range: Option<(Duration, Duration)>,
    }

    impl NetworkFlakiness {
        // Calculate whether to disconnect
        async fn should_disconnect(
            &self,
            rng: &mut impl rand::Rng,
            disconnected_count: u64,
            prev_disconnected_at: &RwLock<Instant>,
        ) -> bool {
            let Some((prob, max_disconnects, duration)) = &self.disconnect_params else {
                return false;
            };

            let disconnected_at = prev_disconnected_at.read().unwrap();
            if disconnected_at.elapsed() > *duration && disconnected_count < *max_disconnects {
                rng.random_bool(*prob)
            } else {
                false
            }
        }
    }

    struct MockLink<M> {
        buffer_size: usize,
        session_id: SessionId,
        receiver_storage: Arc<MVar<DuplexStream>>,
        // If true, `next()` on this link will always return an error.
        fail_connects: Arc<AtomicBool>,
        // Used to break the existing connection, if there is one. It still
        // allows reconnect.
        disconnect_signal: watch::Sender<()>,
        network_flakiness: NetworkFlakiness,
        disconnected_count: Arc<AtomicU64>,
        prev_disconnected_at: Arc<RwLock<Instant>>,
        // If set, print logs every `debug_log_sampling_rate` messages. This
        // is normally set only when debugging a test failure.
        debug_log_sampling_rate: Option<u64>,
        _message_type: PhantomData<M>,
    }

    impl<M> fmt::Debug for MockLink<M> {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("MockLink")
                .field("buffer_size", &self.buffer_size)
                .field("receiver_storage", &"<MVar<DuplexStream>>")
                .field("fail_connects", &self.fail_connects)
                .field("disconnect_signal", &"<watch::Sender>")
                .field("network_flakiness", &self.network_flakiness)
                .field("disconnected_count", &self.disconnected_count)
                .field("prev_disconnected_at", &"<RwLock<Instant>>")
                .field("debug_log_sampling_rate", &self.debug_log_sampling_rate)
                .finish()
        }
    }

    impl<M: RemoteMessage> MockLink<M> {
        fn new() -> Self {
            let (sender, _) = watch::channel(());
            Self {
                buffer_size: 64,
                session_id: SessionId::random(),
                receiver_storage: Arc::new(MVar::empty()),
                fail_connects: Arc::new(AtomicBool::new(false)),
                disconnect_signal: sender,
                network_flakiness: NetworkFlakiness::default(),
                disconnected_count: Arc::new(AtomicU64::new(0)),
                prev_disconnected_at: Arc::new(RwLock::new(tokio::time::Instant::now())),
                debug_log_sampling_rate: None,
                _message_type: PhantomData,
            }
        }

        // If `fail_connects` is true, `next()` on this link will
        // always return an error.
        fn fail_connects() -> Self {
            Self {
                fail_connects: Arc::new(AtomicBool::new(true)),
                ..Self::new()
            }
        }

        fn with_network_flakiness(network_flakiness: NetworkFlakiness) -> Self {
            if let Some((min, max)) = network_flakiness.latency_range {
                assert!(min < max);
            }

            Self {
                network_flakiness,
                ..Self::new()
            }
        }

        fn receiver_storage(&self) -> Arc<MVar<DuplexStream>> {
            self.receiver_storage.clone()
        }

        fn disconnected_count(&self) -> Arc<AtomicU64> {
            self.disconnected_count.clone()
        }

        fn disconnect_signal(&self) -> &watch::Sender<()> {
            &self.disconnect_signal
        }

        fn fail_connects_switch(&self) -> Arc<AtomicBool> {
            self.fail_connects.clone()
        }

        fn set_buffer_size(&mut self, size: usize) {
            self.buffer_size = size;
        }

        fn set_sampling_rate(&mut self, sampling_rate: u64) {
            self.debug_log_sampling_rate = Some(sampling_rate);
        }
    }

    #[async_trait]
    impl<M: RemoteMessage> Link for MockLink<M> {
        type Stream = DuplexStream;

        fn dest(&self) -> ChannelAddr {
            ChannelAddr::Local(u64::MAX)
        }

        fn link_id(&self) -> SessionId {
            self.session_id
        }

        async fn next(&mut self) -> Result<Self::Stream, ClientError> {
            let session_id = self.session_id;
            tracing::debug!("MockLink starts to connect.");
            if self.fail_connects.load(Ordering::Acquire) {
                return Err(ClientError::Connect(
                    self.dest(),
                    std::io::Error::other("intentional error"),
                    "expected failure injected by the mock".to_string(),
                ));
            }

            // Add relays between server and client streams. The
            // relays provides the place to inject network flakiness.
            // The message flow looks like:
            //
            // server <-> server relay <-> injection logic <-> client relay <-> client
            async fn relay_message<M: RemoteMessage>(
                mut disconnect_signal: watch::Receiver<()>,
                network_flakiness: NetworkFlakiness,
                disconnected_count: Arc<AtomicU64>,
                prev_disconnected_at: Arc<RwLock<Instant>>,
                mut reader: FrameReader<ReadHalf<DuplexStream>>,
                mut writer: WriteHalf<DuplexStream>,
                // Used by client and server tokio tasks to coordinate
                // stopping together.
                task_coordination_token: CancellationToken,
                debug_log_sampling_rate: Option<u64>,
                // Whether the relayed message is from client to
                // server.
                is_from_client: bool,
            ) {
                // Used to simulate latency. Briefly, messages are
                // buffered in the queue and wait for the expected
                // latency elapse.
                async fn wait_for_latency_elapse(
                    queue: &VecDeque<(Bytes, Instant)>,
                    network_flakiness: &NetworkFlakiness,
                    rng: &mut impl rand::Rng,
                ) {
                    if let Some((min, max)) = network_flakiness.latency_range {
                        let diff = max.abs_diff(min);
                        let factor = rng.random_range(0.0..=1.0);
                        let latency = min + diff.mul_f64(factor);
                        tokio::time::sleep_until(queue.front().unwrap().1 + latency).await;
                    }
                }

                let mut rng = rand::rngs::SmallRng::try_from_rng(&mut SysRng).unwrap();
                let mut queue: VecDeque<(Bytes, Instant)> = VecDeque::new();
                let mut send_count = 0u64;

                loop {
                    tokio::select! {
                        read_res = reader.next() => {
                            match read_res {
                                Ok(Some((_, data))) => {
                                    queue.push_back((data, tokio::time::Instant::now()));
                                }
                                Ok(None) | Err(_) => {
                                        tracing::debug!("The upstream is closed or dropped. MockLink disconnects");
                                        break;
                                }
                            }
                        }
                        _ = wait_for_latency_elapse(&queue, &network_flakiness, &mut rng), if !queue.is_empty() => {
                            let count = disconnected_count.load(Ordering::Relaxed);
                            if network_flakiness.should_disconnect(&mut rng, count, &prev_disconnected_at).await {
                                tracing::debug!("MockLink disconnects");
                                disconnected_count.fetch_add(1, Ordering::Relaxed);

                                metrics::CHANNEL_RECONNECTIONS.add(
                                    1,
                                    hyperactor_telemetry::kv_pairs!(
                                        "transport" => "mock",
                                        "reason" => "network_flakiness",
                                    ),
                                );

                                let mut w = prev_disconnected_at.write().unwrap();
                                *w = tokio::time::Instant::now();
                                break;
                            }
                            let data = queue.pop_front().unwrap().0;
                            let is_sampled = debug_log_sampling_rate.is_some_and(|sample_rate| send_count % sample_rate == 1);
                            if is_sampled {
                                if is_from_client {
                                    if let Ok((Frame::Message(_seq, _msg), _)) = bincode::serde::decode_from_slice::<Frame<M>, _>(&data, bincode::config::legacy()) {
                                        tracing::debug!("MockLink relays a msg from client. msg type: {}", std::any::type_name::<M>());
                                    }
                                } else {
                                    let result = deserialize_response(data.clone());
                                    if let Ok(NetRxResponse::Ack(seq)) = result {
                                        tracing::debug!("MockLink relays an ack from server. seq: {}", seq);
                                    }
                                }
                            }
                            let mut fw  = FrameWrite::new(writer, data, hyperactor_config::global::get(config::CODEC_MAX_FRAME_LENGTH), 0).unwrap();
                            if fw.send().await.is_err() {
                                break;
                            }
                            writer = fw.complete();
                            send_count += 1;
                        }
                        _ = task_coordination_token.cancelled() => break,

                        changed = disconnect_signal.changed() => {
                            tracing::debug!("MockLink disconnects per disconnect_signal {:?}", changed);
                            break;
                        }
                    }
                }

                task_coordination_token.cancel();
            }

            let (server, mut server_relay) = tokio::io::duplex(self.buffer_size);
            let (client, client_relay) = tokio::io::duplex(self.buffer_size);

            // Write LinkInit on server_relay so it's readable from `server`.
            // This simulates the client sending LinkInit over the wire before
            // the frame-level relay begins.
            write_link_init(&mut server_relay, session_id, 0)
                .await
                .map_err(|err| ClientError::Io(self.dest(), err))?;

            let (server_r, server_writer) = tokio::io::split(server_relay);
            let (client_r, client_writer) = tokio::io::split(client_relay);

            let max_len = hyperactor_config::global::get(config::CODEC_MAX_FRAME_LENGTH);
            let server_reader = FrameReader::new(server_r, max_len);
            let client_reader = FrameReader::new(client_r, max_len);

            let task_coordination_token = CancellationToken::new();
            let _server_relay_task_handle = tokio::spawn(relay_message::<M>(
                self.disconnect_signal.subscribe(),
                self.network_flakiness.clone(),
                self.disconnected_count.clone(),
                self.prev_disconnected_at.clone(),
                server_reader,
                client_writer,
                task_coordination_token.clone(),
                self.debug_log_sampling_rate,
                /*is_from_client*/ false,
            ));
            let _client_relay_task_handle = tokio::spawn(relay_message::<M>(
                self.disconnect_signal.subscribe(),
                self.network_flakiness.clone(),
                self.disconnected_count.clone(),
                self.prev_disconnected_at.clone(),
                client_reader,
                server_writer,
                task_coordination_token,
                self.debug_log_sampling_rate,
                /*is_from_client*/ true,
            ));

            self.receiver_storage.put(server).await;
            Ok(client)
        }
    }

    struct MockLinkListener {
        receiver_storage: Arc<MVar<DuplexStream>>,
        channel_addr: ChannelAddr,
    }

    impl MockLinkListener {
        fn new(receiver_storage: Arc<MVar<DuplexStream>>, channel_addr: ChannelAddr) -> Self {
            Self {
                receiver_storage,
                channel_addr,
            }
        }
    }

    impl fmt::Debug for MockLinkListener {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("MockLinkListener")
                .field("channel_addr", &self.channel_addr)
                .finish()
        }
    }

    #[async_trait]
    impl super::Listener for MockLinkListener {
        type Stream = DuplexStream;

        async fn accept(&mut self) -> Result<(Self::Stream, ChannelAddr), ServerError> {
            let stream = self.receiver_storage.take().await;
            Ok((stream, self.channel_addr.clone()))
        }
    }

    /// Create an AcceptorLink-based server test rig. Returns the
    /// session task handle, the channel sender for dispatching
    /// streams, the message receiver, and a cancellation token.
    fn serve_acceptor_test<M: RemoteMessage>(
        session_id: SessionId,
    ) -> (
        JoinHandle<()>,
        mpsc::UnboundedSender<DuplexStream>,
        mpsc::Receiver<M>,
        CancellationToken,
    ) {
        let (acceptor_tx, acceptor_rx) = mpsc::unbounded_channel::<DuplexStream>();
        let cancel_token = CancellationToken::new();
        let link = AcceptorLink {
            dest: ChannelAddr::Local(u64::MAX),
            session_id,
            stream: acceptor_rx,
            cancel: cancel_token.clone(),
        };
        let (tx, rx) = mpsc::channel::<M>(1024);
        let ct = cancel_token.clone();
        let handle = tokio::spawn(async move {
            let mut session = Session::new(link);
            let mut next = session::Next { seq: 0, ack: 0 };

            loop {
                let connected = match session.connect().await {
                    Ok(s) => s,
                    Err(_) => break,
                };

                let result = {
                    let stream = connected.stream(INITIATOR_TO_ACCEPTOR);
                    tokio::select! {
                        r = session::recv_connected::<M, _, _>(&stream, &tx, &mut next) => r,
                        _ = ct.cancelled() => Err(session::RecvLoopError::Cancelled),
                    }
                };

                // Flush remaining ack if behind.
                if next.ack < next.seq {
                    let ack = serialize_response(NetRxResponse::Ack(next.seq - 1)).unwrap();
                    let stream = connected.stream(INITIATOR_TO_ACCEPTOR);
                    let mut completion = stream.write(ack);
                    match completion.drive().await {
                        Ok(()) => {
                            next.ack = next.seq;
                        }
                        Err(e) => {
                            tracing::debug!(
                                error = %e,
                                "failed to flush acks during cleanup"
                            );
                        }
                    }
                }

                // Send reject or closed response if appropriate.
                let terminal_response = match &result {
                    Err(session::RecvLoopError::SequenceError(reason)) => {
                        Some(NetRxResponse::Reject(reason.clone()))
                    }
                    Err(session::RecvLoopError::Cancelled) => Some(NetRxResponse::Closed),
                    _ => None,
                };
                if let Some(rsp) = terminal_response {
                    let data = serialize_response(rsp).unwrap();
                    let stream = connected.stream(INITIATOR_TO_ACCEPTOR);
                    let mut completion = stream.write(data);
                    let _ = completion.drive().await;
                }

                let recoverable = matches!(&result, Ok(()) | Err(session::RecvLoopError::Io(_)));
                session = connected.release();
                if recoverable {
                    continue;
                }
                break;
            }
        });
        (handle, acceptor_tx, rx, cancel_token)
    }

    async fn write_stream<M, W>(
        mut writer: W,
        _session_id: u64,
        messages: &[(u64, M)],
        _init: bool,
    ) -> W
    where
        M: RemoteMessage + PartialEq + Clone,
        W: AsyncWrite + Unpin,
    {
        for (seq, message) in messages {
            let message =
                serde_multipart::serialize_bincode(&Frame::<M>::Message(*seq, message.clone()))
                    .unwrap();
            let mut fw = FrameWrite::new(
                writer,
                message.framed(),
                hyperactor_config::global::get(config::CODEC_MAX_FRAME_LENGTH),
                0,
            )
            .map_err(|(_w, e)| e)
            .unwrap();
            fw.send().await.unwrap();
            writer = fw.complete();
        }

        writer
    }

    #[async_timed_test(timeout_secs = 60)]
    async fn test_persistent_server_session() {
        let config = hyperactor_config::global::lock();
        let _guard = config.override_key(config::MESSAGE_ACK_EVERY_N_MESSAGES, 1);

        async fn verify_ack(reader: &mut FrameReader<ReadHalf<DuplexStream>>, expected_last: u64) {
            let mut last_acked: i128 = -1;
            loop {
                let (_, bytes) = reader.next().await.unwrap().unwrap();
                let acked = deserialize_response(bytes).unwrap().into_ack().unwrap();
                assert!(
                    acked as i128 > last_acked,
                    "acks should be delivered in ascending order"
                );
                last_acked = acked as i128;
                assert!(acked <= expected_last);
                if acked == expected_last {
                    break;
                }
            }
        }

        let session_id = SessionId(123);
        let (_handle, acceptor_tx, mut rx, cancel_token) = serve_acceptor_test::<u64>(session_id);

        // First connection: send messages, verify delivery and ack.
        {
            let (sender, receiver) = tokio::io::duplex(5000);
            acceptor_tx.send(receiver).unwrap();

            let (r, writer) = tokio::io::split(sender);
            let mut reader = FrameReader::new(
                r,
                hyperactor_config::global::get(config::CODEC_MAX_FRAME_LENGTH),
            );

            let _writer = write_stream(
                writer,
                123,
                &[
                    (0u64, 100u64),
                    (1u64, 101u64),
                    (2u64, 102u64),
                    (3u64, 103u64),
                ],
                true,
            )
            .await;

            assert_eq!(rx.recv().await, Some(100));
            assert_eq!(rx.recv().await, Some(101));
            assert_eq!(rx.recv().await, Some(102));
            assert_eq!(rx.recv().await, Some(103));

            verify_ack(&mut reader, 3).await;
            // Drop reader and writer to close the connection.
        }

        // Second connection (reconnection): retransmitted messages are deduped.
        {
            let (sender2, receiver2) = tokio::io::duplex(5000);
            acceptor_tx.send(receiver2).unwrap();

            let (r2, writer2) = tokio::io::split(sender2);
            let mut reader2 = FrameReader::new(
                r2,
                hyperactor_config::global::get(config::CODEC_MAX_FRAME_LENGTH),
            );

            let _ = write_stream(
                writer2,
                123,
                &[
                    (2u64, 102u64),
                    (3u64, 103u64),
                    (4u64, 104u64),
                    (5u64, 105u64),
                ],
                true,
            )
            .await;

            // 102 and 103 are retransmits; only 104 and 105 are new.
            assert_eq!(rx.recv().await, Some(104));
            assert_eq!(rx.recv().await, Some(105));

            verify_ack(&mut reader2, 5).await;

            cancel_token.cancel();
        }
    }

    #[async_timed_test(timeout_secs = 60)]
    async fn test_ack_from_server_session() {
        let config = hyperactor_config::global::lock();
        let _guard = config.override_key(config::MESSAGE_ACK_EVERY_N_MESSAGES, 1);
        let session_id = SessionId(123);
        let (_handle, acceptor_tx, mut rx, cancel_token) = serve_acceptor_test::<u64>(session_id);

        let (sender, receiver) = tokio::io::duplex(5000);
        acceptor_tx.send(receiver).unwrap();
        let (r, mut writer) = tokio::io::split(sender);
        let mut reader = FrameReader::new(
            r,
            hyperactor_config::global::get(config::CODEC_MAX_FRAME_LENGTH),
        );

        for i in 0u64..100u64 {
            writer = write_stream(writer, 123, &[(i, 100u64 + i)], /*init*/ i == 0u64).await;
            assert_eq!(rx.recv().await, Some(100u64 + i));
            let (_, bytes) = reader.next().await.unwrap().unwrap();
            let acked = deserialize_response(bytes).unwrap().into_ack().unwrap();
            assert_eq!(acked, i);
        }

        // Wait long enough to ensure server processed everything.
        tokio::time::sleep(Duration::from_secs(5)).await;

        cancel_token.cancel();

        // Should send NetRxResponse::Closed before stopping.
        let (_, bytes) = reader.next().await.unwrap().unwrap();
        assert!(deserialize_response(bytes).unwrap().is_closed());
    }

    #[tracing_test::traced_test]
    async fn verify_tx_closed(tx_status: &mut watch::Receiver<TxStatus>, expected_log: &str) {
        match tokio::time::timeout(Duration::from_secs(5), tx_status.changed()).await {
            Ok(Ok(())) => {
                let current_status = tx_status.borrow().clone();
                assert!(current_status.is_closed());
                logs_assert_unscoped(|logs| {
                    if logs.iter().any(|log| log.contains(expected_log)) {
                        Ok(())
                    } else {
                        Err("expected log not found".to_string())
                    }
                });
            }
            Ok(Err(_)) => panic!("watch::Receiver::changed() failed because sender is dropped."),
            Err(_) => panic!("timeout before tx_status changed"),
        }
    }

    #[tracing_test::traced_test]
    #[tokio::test]
    // TODO: OSS: The logs_assert function returned an error: expected log not found
    #[cfg_attr(not(fbcode_build), ignore)]
    async fn test_tcp_tx_delivery_timeout() {
        // This link always fails to connect.
        let link = MockLink::<u64>::fail_connects();
        let tx = spawn::<u64>(link);
        // Override the default (1m) for the purposes of this test.
        let config = hyperactor_config::global::lock();
        let _guard = config.override_key(config::MESSAGE_DELIVERY_TIMEOUT, Duration::from_secs(1));
        let mut tx_receiver = tx.status().clone();
        let (return_channel, _return_receiver) = oneshot::channel();
        tx.try_post(123, return_channel);
        verify_tx_closed(&mut tx_receiver, "failed to deliver message within timeout").await;
    }

    async fn take_receiver(
        receiver_storage: &MVar<DuplexStream>,
    ) -> (FrameReader<ReadHalf<DuplexStream>>, WriteHalf<DuplexStream>) {
        let mut receiver = receiver_storage.take().await;
        // Read and discard the LinkInit header that MockLink::connect() writes.
        let _link_init = read_link_init(&mut receiver).await.expect("read LinkInit");
        let (r, writer) = tokio::io::split(receiver);
        let reader = FrameReader::new(
            r,
            hyperactor_config::global::get(config::CODEC_MAX_FRAME_LENGTH),
        );
        (reader, writer)
    }

    async fn verify_message<M: RemoteMessage + PartialEq + std::fmt::Debug>(
        reader: &mut FrameReader<ReadHalf<DuplexStream>>,
        expect: (u64, M),
        loc: u32,
    ) {
        let expected = Frame::Message(expect.0, expect.1);
        let (_, bytes) = reader.next().await.unwrap().expect("unexpected EOF");
        let message = serde_multipart::Message::from_framed(bytes).unwrap();
        let frame: Frame<M> = serde_multipart::deserialize_bincode(message).unwrap();

        assert_eq!(frame, expected, "from ln={loc}");
    }

    async fn verify_stream<M: RemoteMessage + PartialEq + std::fmt::Debug + Clone>(
        reader: &mut FrameReader<ReadHalf<DuplexStream>>,
        expects: &[(u64, M)],
        _expect_session_id: Option<u64>,
        loc: u32,
    ) {
        for expect in expects {
            verify_message(reader, expect.clone(), loc).await;
        }
    }

    async fn net_tx_send(tx: &NetTx<u64>, msgs: &[u64]) {
        for msg in msgs {
            tx.post(*msg);
        }
    }

    // Happy path: all messages are acked.
    #[async_timed_test(timeout_secs = 30)]
    async fn test_ack_in_net_tx_basic() {
        let link = MockLink::<u64>::new();
        let receiver_storage = link.receiver_storage();
        let tx = spawn::<u64>(link);

        // Send some messages, but not acking any of them.
        net_tx_send(&tx, &[100, 101, 102, 103, 104]).await;
        {
            let (mut reader, mut writer) = take_receiver(&receiver_storage).await;
            verify_stream(
                &mut reader,
                &[
                    (0u64, 100u64),
                    (1u64, 101u64),
                    (2u64, 102u64),
                    (3u64, 103u64),
                    (4u64, 104u64),
                ],
                None,
                line!(),
            )
            .await;

            for i in 0u64..5u64 {
                writer = FrameWrite::write_frame(
                    writer,
                    serialize_response(NetRxResponse::Ack(i)).unwrap(),
                    1024,
                    0,
                )
                .await
                .map_err(|(_, e)| e)
                .unwrap();
            }
            // Wait for the acks to be processed by NetTx.
            tokio::time::sleep(Duration::from_secs(3)).await;
            // Drop both halves to break the in-memory connection (parity with old drop of DuplexStream).
            drop(reader);
            drop(writer);
        };

        // Sent a new message to verify all sent messages will not be resent.
        net_tx_send(&tx, &[105u64]).await;
        {
            let (mut reader, _writer) = take_receiver(&receiver_storage).await;
            verify_stream(&mut reader, &[(5u64, 105u64)], None, line!()).await;
            // Reader/writer dropped here. This breaks the connection.
        };
    }

    // Verify unacked message will be resent after reconnection.
    #[async_timed_test(timeout_secs = 60)]
    async fn test_persistent_net_tx() {
        let link = MockLink::<u64>::new();
        let receiver_storage = link.receiver_storage();

        let tx = spawn::<u64>(link);

        // Send some messages, but not acking any of them.
        net_tx_send(&tx, &[100, 101, 102, 103, 104]).await;

        // How many times to reconnect. Keep this small because the send loop
        // applies exponential backoff between reconnections, and mock connections
        // are too short-lived to trigger the backoff reset.
        let n = 3;

        // Reconnect multiple times. The messages should be resent every time
        // because none of them is acked.
        for i in 0..n {
            {
                let (mut reader, mut writer) = take_receiver(&receiver_storage).await;
                verify_stream(
                    &mut reader,
                    &[
                        (0u64, 100u64),
                        (1u64, 101u64),
                        (2u64, 102u64),
                        (3u64, 103u64),
                        (4u64, 104u64),
                    ],
                    None,
                    line!(),
                )
                .await;

                // In the last iteration, ack part of the messages. This should
                // prune them from future resent.
                if i == n - 1 {
                    writer = FrameWrite::write_frame(
                        writer,
                        serialize_response(NetRxResponse::Ack(1)).unwrap(),
                        1024,
                        0,
                    )
                    .await
                    .map_err(|(_, e)| e)
                    .unwrap();
                    // Wait for the acks to be processed by NetTx.
                    tokio::time::sleep(Duration::from_secs(3)).await;
                }
                // client DuplexStream is dropped here. This breaks the connection.
                drop(reader);
                drop(writer);
            };
        }

        // Verify only unacked are resent.
        for _ in 0..n {
            {
                let (mut reader, mut _writer) = take_receiver(&receiver_storage).await;
                verify_stream(
                    &mut reader,
                    &[(2u64, 102u64), (3u64, 103u64), (4u64, 104u64)],
                    None,
                    line!(),
                )
                .await;
                // drop(reader/_writer) at scope end
            };
        }

        // Now send more messages.
        net_tx_send(&tx, &[105u64, 106u64, 107u64, 108u64, 109u64]).await;
        // Verify the unacked messages from the 1st send will be grouped with
        // the 2nd send.
        for i in 0..n {
            {
                let (mut reader, mut writer) = take_receiver(&receiver_storage).await;
                verify_stream(
                    &mut reader,
                    &[
                        // From the 1st send.
                        (2u64, 102u64),
                        (3u64, 103u64),
                        (4u64, 104u64),
                        // From the 2nd send.
                        (5u64, 105u64),
                        (6u64, 106u64),
                        (7u64, 107u64),
                        (8u64, 108u64),
                        (9u64, 109u64),
                    ],
                    None,
                    line!(),
                )
                .await;

                // In the last iteration, ack part of the messages from the 1st
                // sent.
                if i == n - 1 {
                    // Intentionally ack 1 again to verify it is okay to ack
                    // messages that was already acked.
                    writer = FrameWrite::write_frame(
                        writer,
                        serialize_response(NetRxResponse::Ack(1)).unwrap(),
                        1024,
                        0,
                    )
                    .await
                    .map_err(|(_, e)| e)
                    .unwrap();
                    writer = FrameWrite::write_frame(
                        writer,
                        serialize_response(NetRxResponse::Ack(2)).unwrap(),
                        1024,
                        0,
                    )
                    .await
                    .map_err(|(_, e)| e)
                    .unwrap();
                    writer = FrameWrite::write_frame(
                        writer,
                        serialize_response(NetRxResponse::Ack(3)).unwrap(),
                        1024,
                        0,
                    )
                    .await
                    .map_err(|(_, e)| e)
                    .unwrap();
                    // Wait for the acks to be processed by NetTx.
                    tokio::time::sleep(Duration::from_secs(3)).await;
                }
                // client DuplexStream is dropped here. This breaks the connection.
                drop(reader);
                drop(writer);
            };
        }

        for i in 0..n {
            {
                let (mut reader, mut writer) = take_receiver(&receiver_storage).await;
                verify_stream(
                    &mut reader,
                    &[
                        // From the 1st send.
                        (4u64, 104),
                        // From the 2nd send.
                        (5u64, 105u64),
                        (6u64, 106u64),
                        (7u64, 107u64),
                        (8u64, 108u64),
                        (9u64, 109u64),
                    ],
                    None,
                    line!(),
                )
                .await;

                // In the last iteration, ack part of the messages from the 2nd send.
                if i == n - 1 {
                    writer = FrameWrite::write_frame(
                        writer,
                        serialize_response(NetRxResponse::Ack(7)).unwrap(),
                        1024,
                        0,
                    )
                    .await
                    .map_err(|(_, e)| e)
                    .unwrap();
                    // Wait for the acks to be processed by NetTx.
                    tokio::time::sleep(Duration::from_secs(3)).await;
                }
                // client DuplexStream is dropped here. This breaks the connection.
                drop(reader);
                drop(writer);
            };
        }

        for _ in 0..n {
            {
                let (mut reader, writer) = take_receiver(&receiver_storage).await;
                verify_stream(
                    &mut reader,
                    &[
                        // From the 2nd send.
                        (8u64, 108u64),
                        (9u64, 109u64),
                    ],
                    None,
                    line!(),
                )
                .await;
                // client DuplexStream is dropped here. This breaks the connection.
                drop(reader);
                drop(writer);
            };
        }
    }

    #[async_timed_test(timeout_secs = 15)]
    async fn test_ack_before_redelivery_in_net_tx() {
        let link = MockLink::<u64>::new();
        let receiver_storage = link.receiver_storage();
        let net_tx = spawn::<u64>(link);

        // Verify sent-and-ack a message. This is necessary for the test to
        // trigger a connection.
        let (return_channel_tx, return_channel_rx) = oneshot::channel();
        net_tx.try_post(100, return_channel_tx);
        let (mut reader, mut writer) = take_receiver(&receiver_storage).await;
        verify_stream(&mut reader, &[(0u64, 100u64)], None, line!()).await;
        // ack it
        writer = FrameWrite::write_frame(
            writer,
            serialize_response(NetRxResponse::Ack(0)).unwrap(),
            1024,
            0,
        )
        .await
        .map_err(|(_, e)| e)
        .unwrap();
        // confirm Tx received ack
        //
        // Using `is_err` to confirm the message is delivered/acked is confusing,
        // but is correct. See how send is implemented: https://fburl.com/code/ywt8lip2
        assert!(return_channel_rx.await.is_err());

        // Now fake an unknown delivery for Tx:
        // Although Tx did not actually send seq=1, we still ack it from Rx to
        // pretend Tx already sent it, just it did not know it was sent
        // successfully.
        let _ = FrameWrite::write_frame(
            writer,
            serialize_response(NetRxResponse::Ack(1)).unwrap(),
            1024,
            0,
        )
        .await
        .map_err(|(_, e)| e)
        .unwrap();

        let (return_channel_tx, return_channel_rx) = oneshot::channel();
        net_tx.try_post(101, return_channel_tx);
        // Verify the message is sent to Rx.
        verify_message(&mut reader, (1u64, 101u64), line!()).await;
        // although we did not ack the message after it is sent, since we already
        // acked it previously, Tx will treat it as acked, and considered the
        // message delivered successfully.
        //
        // Using `is_err` to confirm the message is delivered/acked is confusing,
        // but is correct. See how send is implemented: https://fburl.com/code/ywt8lip2
        assert!(return_channel_rx.await.is_err());
    }

    async fn verify_ack_exceeded_limit(disconnect_before_ack: bool) {
        // Use temporary config for this test
        let config = hyperactor_config::global::lock();
        let _guard = config.override_key(config::MESSAGE_DELIVERY_TIMEOUT, Duration::from_secs(2));

        let link: MockLink<u64> = MockLink::<u64>::new();
        let disconnect_signal = link.disconnect_signal().clone();
        let fail_connect_switch = link.fail_connects_switch();
        let receiver_storage = link.receiver_storage();
        let tx = spawn::<u64>(link);
        let mut tx_status = tx.status().clone();
        // send a message
        tx.post(100);
        let (mut reader, writer) = take_receiver(&receiver_storage).await;
        // Confirm message is sent to rx.
        verify_stream(&mut reader, &[(0u64, 100u64)], None, line!()).await;
        // ack it
        let _ = FrameWrite::write_frame(
            writer,
            serialize_response(NetRxResponse::Ack(0)).unwrap(),
            hyperactor_config::global::get(config::CODEC_MAX_FRAME_LENGTH),
            0,
        )
        .await
        .map_err(|(_, e)| e)
        .unwrap();
        tokio::time::sleep(Duration::from_secs(3)).await;
        // Channel should be still alive because ack was sent.
        assert!(!tx_status.has_changed().unwrap());
        assert_eq!(*tx_status.borrow(), TxStatus::Active);

        tx.post(101);
        // Confirm message is sent to rx.
        verify_message(&mut reader, (1u64, 101u64), line!()).await;

        if disconnect_before_ack {
            // Prevent link from reconnect
            fail_connect_switch.store(true, Ordering::Release);
            // Break the existing connection
            disconnect_signal.send(()).unwrap();
        }

        // Verify the channel is closed due to ack timeout based on the log.
        let expected_log: &str = if disconnect_before_ack {
            "failed to receive ack within timeout 2s; link is currently broken"
        } else {
            "failed to receive ack within timeout 2s; link is currently connected"
        };

        verify_tx_closed(&mut tx_status, expected_log).await;
    }

    #[tracing_test::traced_test]
    #[async_timed_test(timeout_secs = 30)]
    // TODO: OSS: The logs_assert function returned an error: expected log not found
    #[cfg_attr(not(fbcode_build), ignore)]
    async fn test_ack_exceeded_limit_with_connected_link() {
        verify_ack_exceeded_limit(false).await;
    }

    #[tracing_test::traced_test]
    #[async_timed_test(timeout_secs = 30)]
    // TODO: OSS: The logs_assert function returned an error: expected log not found
    #[cfg_attr(not(fbcode_build), ignore)]
    async fn test_ack_exceeded_limit_with_broken_link() {
        verify_ack_exceeded_limit(true).await;
    }

    // Verify a large number of messages can be delivered and acked with the
    // presence of flakiness in the network, i.e. random delay and disconnection.
    #[async_timed_test(timeout_secs = 60)]
    async fn test_network_flakiness_in_channel() {
        hyperactor_telemetry::initialize_logging_for_test();

        let sampling_rate = 100;
        let mut link = MockLink::<u64>::with_network_flakiness(NetworkFlakiness {
            disconnect_params: Some((0.001, 15, Duration::from_millis(400))),
            latency_range: Some((Duration::from_millis(100), Duration::from_millis(200))),
        });
        link.set_sampling_rate(sampling_rate);
        // Set a large buffer size to improve throughput.
        link.set_buffer_size(1024000);
        let disconnected_count = link.disconnected_count();
        let receiver_storage = link.receiver_storage();
        let listener = MockLinkListener::new(receiver_storage.clone(), link.dest());
        let local_addr = listener.channel_addr.clone();
        let (_, mut nx): (ChannelAddr, NetRx<u64>) =
            super::server::serve_with_listener(listener, local_addr).unwrap();
        let tx = spawn::<u64>(link);
        let messages: Vec<_> = (0..10001).collect();
        let messages_clone = messages.clone();
        // Put the sender side in a separate task so we can start the receiver
        // side concurrently.
        let send_task_handle = tokio::spawn(async move {
            for message in messages_clone {
                // Add a small delay between messages to give NetRx time to ack.
                // Technically, this test still can pass without this delay. But
                // the test will need a might larger timeout. The reason is
                // fairly convoluted:
                //
                // MockLink uses the number of delivery to calculate the disconnection
                // probability. If NetRx sends messages much faster than NetTx
                // can ack them, there is a higher chance that the messages are
                // not acked before reconnect. Then those message would be redelivered.
                // The repeated redelivery increases the total time of sending
                // these messages.
                tokio::time::sleep(Duration::from_micros(rand::random::<u64>() % 100)).await;
                tx.post(message);
            }
            tracing::debug!("NetTx sent all messages");
            // It is important to return tx instead of dropping it here, because
            // Rx might not receive all messages yet.
            tx
        });

        for message in &messages {
            if message % sampling_rate == 0 {
                tracing::debug!("NetRx received a message: {message}");
            }
            assert_eq!(nx.recv().await.unwrap(), *message);
        }
        tracing::debug!("NetRx received all messages");

        let send_result = send_task_handle.await;
        assert!(send_result.is_ok());

        tracing::debug!(
            "MockLink disconnected {} times.",
            disconnected_count.load(Ordering::SeqCst)
        );
        // TODO(pzhang) after the return_handle work in NetTx is done, add a
        // check here to verify the messages are acked correctly.
    }

    #[async_timed_test(timeout_secs = 60)]
    async fn test_ack_every_n_messages() {
        let config = hyperactor_config::global::lock();
        let _guard_message_ack = config.override_key(config::MESSAGE_ACK_EVERY_N_MESSAGES, 600);
        let _guard_time_interval =
            config.override_key(config::MESSAGE_ACK_TIME_INTERVAL, Duration::from_secs(1000));
        sparse_ack().await;
    }

    #[async_timed_test(timeout_secs = 60)]
    async fn test_ack_every_time_interval() {
        let config = hyperactor_config::global::lock();
        let _guard_message_ack =
            config.override_key(config::MESSAGE_ACK_EVERY_N_MESSAGES, 100000000);
        let _guard_time_interval = config.override_key(
            config::MESSAGE_ACK_TIME_INTERVAL,
            Duration::from_millis(500),
        );
        sparse_ack().await;
    }

    async fn sparse_ack() {
        let mut link = MockLink::<u64>::new();
        // Set a large buffer size to improve throughput.
        link.set_buffer_size(1024000);
        let disconnected_count = link.disconnected_count();
        let receiver_storage = link.receiver_storage();
        let listener = MockLinkListener::new(receiver_storage.clone(), link.dest());
        let local_addr = listener.channel_addr.clone();
        let (_, mut nx): (ChannelAddr, NetRx<u64>) =
            super::server::serve_with_listener(listener, local_addr).unwrap();
        let tx = spawn::<u64>(link);
        let messages: Vec<_> = (0..20001).collect();
        let messages_clone = messages.clone();
        // Put the sender side in a separate task so we can start the receiver
        // side concurrently.
        let send_task_handle = tokio::spawn(async move {
            for message in messages_clone {
                tokio::time::sleep(Duration::from_micros(rand::random::<u64>() % 100)).await;
                tx.post(message);
            }
            tokio::time::sleep(Duration::from_secs(5)).await;
            tracing::debug!("NetTx sent all messages");
            tx
        });

        for message in &messages {
            assert_eq!(nx.recv().await.unwrap(), *message);
        }
        tracing::debug!("NetRx received all messages");

        let send_result = send_task_handle.await;
        assert!(send_result.is_ok());

        tracing::debug!(
            "MockLink disconnected {} times.",
            disconnected_count.load(Ordering::SeqCst)
        );
    }

    #[test]
    fn test_metatls_parsing() {
        // host:port
        let channel: ChannelAddr = "metatls!localhost:1234".parse().unwrap();
        assert_eq!(
            channel,
            ChannelAddr::MetaTls(TlsAddr::new("localhost", 1234))
        );
        // ipv4:port - parsed as hostname with ip normalization
        let channel: ChannelAddr = "metatls!1.2.3.4:1234".parse().unwrap();
        assert_eq!(channel, ChannelAddr::MetaTls(TlsAddr::new("1.2.3.4", 1234)));
        // ipv6:port
        let channel: ChannelAddr = "metatls!2401:db00:33c:6902:face:0:2a2:0:1234"
            .parse()
            .unwrap();
        assert_eq!(
            channel,
            ChannelAddr::MetaTls(TlsAddr::new("2401:db00:33c:6902:face:0:2a2:0", 1234))
        );

        let channel: ChannelAddr = "metatls![::]:1234".parse().unwrap();
        assert_eq!(channel, ChannelAddr::MetaTls(TlsAddr::new("::", 1234)));
    }

    #[async_timed_test(timeout_secs = 300)]
    // TODO: OSS: called `Result::unwrap()` on an `Err` value: Listen(Tcp([::1]:0), Os { code: 99, kind: AddrNotAvailable, message: "Cannot assign requested address" })
    #[cfg_attr(not(fbcode_build), ignore)]
    async fn test_tcp_throughput() {
        let config = hyperactor_config::global::lock();
        let _guard = config.override_key(config::MESSAGE_DELIVERY_TIMEOUT, Duration::from_mins(5));

        let socket_addr: SocketAddr = "[::1]:0".parse().unwrap();
        let (local_addr, mut rx) =
            server::serve::<String>(ChannelAddr::Tcp(socket_addr), None).unwrap();

        // Test with 10 connections (senders), each sends 500K messages, 5M messages in total.
        let total_num_msgs = 500000;

        let receive_handle = tokio::spawn(async move {
            let mut num = 0;
            for _ in 0..10 * total_num_msgs {
                rx.recv().await.unwrap();
                num += 1;

                if num % 100000 == 0 {
                    tracing::info!("total number of received messages: {}", num);
                }
            }
        });

        let mut tx_handles = vec![];
        let mut txs = vec![];
        for _ in 0..10 {
            let server_addr = local_addr.clone();
            let tx = Arc::new(channel::dial::<String>(server_addr).unwrap());
            let tx2 = Arc::clone(&tx);
            txs.push(tx);
            tx_handles.push(tokio::spawn(async move {
                let random_string = rand::rng()
                    .sample_iter(&Alphanumeric)
                    .take(2048)
                    .map(char::from)
                    .collect::<String>();
                for _ in 0..total_num_msgs {
                    tx2.post(random_string.clone());
                }
            }));
        }

        receive_handle.await.unwrap();
        for handle in tx_handles {
            handle.await.unwrap();
        }
    }

    #[tracing_test::traced_test]
    #[async_timed_test(timeout_secs = 60)]
    // TODO: OSS: The logs_assert function returned an error: expected log not found
    #[cfg_attr(not(fbcode_build), ignore)]
    async fn test_net_tx_closed_on_server_reject() {
        let link = MockLink::<u64>::new();
        let receiver_storage = link.receiver_storage();
        let mut tx = spawn::<u64>(link);
        net_tx_send(&tx, &[100]).await;

        {
            let (_reader, writer) = take_receiver(&receiver_storage).await;
            let _ = FrameWrite::write_frame(
                writer,
                serialize_response(NetRxResponse::Reject("testing".to_string())).unwrap(),
                1024,
                0,
            )
            .await
            .map_err(|(_, e)| e);

            // Wait for response to be processed by NetTx before dropping reader/writer. Otherwise
            // the channel will be closed and we will get the wrong error.
            tokio::time::sleep(tokio::time::Duration::from_secs(3)).await;
        }

        verify_tx_closed(&mut tx.status, "server rejected connection").await;
    }

    #[async_timed_test(timeout_secs = 60)]
    async fn test_server_rejects_conn_on_out_of_sequence_message() {
        let config = hyperactor_config::global::lock();
        let _guard = config.override_key(config::MESSAGE_ACK_EVERY_N_MESSAGES, 1);
        let session_id = SessionId(123);
        let (_handle, acceptor_tx, mut rx, _cancel_token) = serve_acceptor_test::<u64>(session_id);

        let (sender, receiver) = tokio::io::duplex(5000);
        acceptor_tx.send(receiver).unwrap();
        let (r, writer) = tokio::io::split(sender);
        let mut reader = FrameReader::new(
            r,
            hyperactor_config::global::get(config::CODEC_MAX_FRAME_LENGTH),
        );

        let _ = write_stream(writer, 123, &[(0, 100u64), (1, 101u64), (3, 103u64)], true).await;
        assert_eq!(rx.recv().await, Some(100u64));
        assert_eq!(rx.recv().await, Some(101u64));
        let (_, bytes) = reader.next().await.unwrap().unwrap();
        let acked = deserialize_response(bytes).unwrap().into_ack().unwrap();
        assert_eq!(acked, 0);
        let (_, bytes) = reader.next().await.unwrap().unwrap();
        let acked = deserialize_response(bytes).unwrap().into_ack().unwrap();
        assert_eq!(acked, 1);
        let (_, bytes) = reader.next().await.unwrap().unwrap();
        assert!(deserialize_response(bytes).unwrap().is_reject());
    }

    #[async_timed_test(timeout_secs = 60)]
    // TODO: OSS: called `Result::unwrap()` on an `Err` value: Listen(Tcp([::1]:0), Os { code: 99, kind: AddrNotAvailable, message: "Cannot assign requested address" })
    #[cfg_attr(not(fbcode_build), ignore)]
    async fn test_stop_net_tx_after_stopping_net_rx() {
        hyperactor_telemetry::initialize_logging_for_test();

        let config = hyperactor_config::global::lock();
        let _guard = config.override_key(config::MESSAGE_DELIVERY_TIMEOUT, Duration::from_mins(5));
        let (addr, mut rx) =
            server::serve::<u64>(ChannelAddr::Tcp("[::1]:0".parse().unwrap()), None).unwrap();
        let socket_addr = match addr {
            ChannelAddr::Tcp(a) => a,
            _ => panic!("unexpected channel type"),
        };
        let tx: NetTx<u64> = spawn(tcp::link(socket_addr, SessionId::random(), 0));
        // NetTx will not establish a connection until it sends the 1st message.
        // Without a live connection, NetTx cannot received the Closed message
        // from NetRx. Therefore, we need to send a message to establish the
        //connection.
        tx.send(100).await.unwrap();
        assert_eq!(rx.recv().await.unwrap(), 100);
        // Drop rx will close the NetRx server.
        rx.2.stop("testing");
        assert!(rx.recv().await.is_err());

        // NetTx will only read from the stream when it needs to send a message
        // or wait for an ack. Therefore we need to send a message to trigger that.
        tx.post(101);
        let mut watcher = tx.status().clone();
        // When NetRx exits, it should notify NetTx to exit as well.
        let _ = watcher.wait_for(|val| val.is_closed()).await;
        // wait_for could return Err due to race between when watch's sender was
        // dropped and when wait_for was called. So we still need to do an
        // equality check.
        assert!(watcher.borrow().is_closed());
    }

    /// Yields pre-built `DuplexStream`s to the accept loop and
    /// blocks once drained. Lets the `rx_join_flushes_pending_ack_*`
    /// tests inspect the wire from the other end.
    struct QueueListener {
        streams: std::collections::VecDeque<DuplexStream>,
        addr: ChannelAddr,
    }

    #[async_trait]
    impl super::Listener for QueueListener {
        type Stream = DuplexStream;

        async fn accept(&mut self) -> Result<(DuplexStream, ChannelAddr), ServerError> {
            match self.streams.pop_front() {
                Some(s) => Ok((s, self.addr.clone())),
                None => std::future::pending().await,
            }
        }
    }

    /// In-memory connection: server end goes into the listener; the
    /// test reads from `client_r`.
    struct PreparedConnection {
        server_side: DuplexStream,
        // Kept alive so the server's recv-loop stays in its `select!`
        // on cancellation rather than exiting on EOF. Tests must
        // exercise the cancel-flush path.
        _client_w: tokio::io::WriteHalf<DuplexStream>,
        client_r: ReadHalf<DuplexStream>,
    }

    /// Write `LinkInit` and the framed `Frame::Message(seq, value)`
    /// payloads on the client side; return both halves.
    async fn prepare_connection(
        session_id: SessionId,
        stream_id: u8,
        messages: &[(u64, u64)],
    ) -> PreparedConnection {
        let (client_side, server_side) = tokio::io::duplex(8192);
        let (client_r, mut client_w) = tokio::io::split(client_side);

        super::write_link_init(&mut client_w, session_id, stream_id)
            .await
            .unwrap();
        let max_len = hyperactor_config::global::get(config::CODEC_MAX_FRAME_LENGTH);
        for (seq, value) in messages {
            let payload =
                serde_multipart::serialize_bincode(&Frame::<u64>::Message(*seq, *value)).unwrap();
            let mut fw = FrameWrite::new(client_w, payload.framed(), max_len, 0)
                .map_err(|(_w, e)| e)
                .unwrap();
            fw.send().await.unwrap();
            client_w = fw.complete();
        }

        PreparedConnection {
            server_side,
            _client_w: client_w,
            client_r,
        }
    }

    /// Test plan for `run_separate_sessions_flush_test`: each entry
    /// describes one connection that lives on its own session.
    struct SeparateSessionPlan {
        session_id: SessionId,
        stream_id: u8,
        messages: Vec<(u64, u64)>,
    }

    /// Drive the rx.join flush test across multiple connections,
    /// each on its own session. Verifies:
    ///
    /// 1. Every message sent reaches the application via `rx.recv()`.
    /// 2. After application delivery completes, *no* ack frames have
    ///    been emitted on any connection's read side — the policy
    ///    thresholds are out of reach.
    /// 3. After `rx.join()` returns, every connection has exactly
    ///    one `NetRxResponse::Ack(highest_seq)` frame on its read
    ///    side, followed by a `NetRxResponse::Closed` terminal frame.
    async fn run_separate_sessions_flush_test(
        plans: Vec<SeparateSessionPlan>,
        stream_id_label: &str,
    ) {
        let config = hyperactor_config::global::lock();
        let _g_msg = config.override_key(config::MESSAGE_ACK_EVERY_N_MESSAGES, 1_000_000);
        let _g_time =
            config.override_key(config::MESSAGE_ACK_TIME_INTERVAL, Duration::from_secs(3600));

        // Build all connections and stage them into a `QueueListener`.
        let mut conns: Vec<PreparedConnection> = Vec::with_capacity(plans.len());
        let mut expected_messages: std::collections::HashSet<u64> =
            std::collections::HashSet::new();
        let mut expected_acks: Vec<u64> = Vec::with_capacity(plans.len());
        for plan in &plans {
            for (_seq, value) in &plan.messages {
                expected_messages.insert(*value);
            }
            expected_acks.push(plan.messages.iter().map(|(s, _)| *s).max().unwrap());
            conns.push(prepare_connection(plan.session_id, plan.stream_id, &plan.messages).await);
        }

        let addr = ChannelAddr::Local(u64::MAX);
        let listener = QueueListener {
            streams: conns
                .iter_mut()
                .map(|c| {
                    // Move server_side out of each PreparedConnection by replacing it with a placeholder.
                    std::mem::replace(&mut c.server_side, tokio::io::duplex(1).0)
                })
                .collect(),
            addr: addr.clone(),
        };
        let (_addr, mut rx) = super::server::serve_with_listener::<u64, _>(listener, addr).unwrap();

        // Drain every expected message off the application channel.
        let mut received: std::collections::HashSet<u64> = std::collections::HashSet::new();
        for _ in 0..expected_messages.len() {
            received.insert(rx.recv().await.unwrap());
        }
        assert_eq!(
            received, expected_messages,
            "{stream_id_label}: every produced message should reach the application"
        );

        // Give any policy-driven ack timer a generous grace period to
        // fire. Because `MESSAGE_ACK_EVERY_N_MESSAGES` and
        // `MESSAGE_ACK_TIME_INTERVAL` are out of reach, no ack should
        // be emitted yet — but if a regression makes them reachable
        // (or adds a new spontaneous emission path), this sleep gives
        // it room to do so before the assertion below.
        tokio::time::sleep(Duration::from_millis(100)).await;

        let max_len = hyperactor_config::global::get(config::CODEC_MAX_FRAME_LENGTH);
        let mut readers: Vec<FrameReader<ReadHalf<DuplexStream>>> = conns
            .into_iter()
            .map(|c| FrameReader::new(c.client_r, max_len))
            .collect();
        for (idx, reader) in readers.iter_mut().enumerate() {
            match tokio::time::timeout(Duration::from_millis(10), reader.next()).await {
                Err(_) => {} // timeout — no frame, expected.
                Ok(Err(e)) => panic!(
                    "{stream_id_label}: connection {idx} frame reader error before rx.join: {e}"
                ),
                Ok(Ok(None)) => {
                    panic!("{stream_id_label}: connection {idx} closed before rx.join()")
                }
                Ok(Ok(Some((_, bytes)))) => {
                    let resp = super::deserialize_response(bytes).unwrap();
                    panic!(
                        "{stream_id_label}: connection {idx} unexpectedly received {resp:?} \
                         before rx.join()"
                    );
                }
            }
        }

        rx.join().await;

        // After rx.join() returns, every connection must have its
        // terminal cleanup frames already written: an `Ack` covering
        // the highest seq it sent, then a `Closed`.
        for (idx, (reader, expected_ack)) in readers.iter_mut().zip(&expected_acks).enumerate() {
            let bytes = tokio::time::timeout(Duration::from_millis(50), reader.next())
                .await
                .unwrap_or_else(|_| {
                    panic!(
                        "{stream_id_label}: connection {idx} produced no Ack frame within 50ms \
                         after rx.join()"
                    )
                })
                .expect("frame reader error")
                .expect("frame reader returned None");
            let acked = super::deserialize_response(bytes.1)
                .unwrap()
                .into_ack()
                .unwrap_or_else(|other| {
                    panic!("{stream_id_label}: connection {idx} expected Ack, got {other:?}")
                });
            assert_eq!(
                acked, *expected_ack,
                "{stream_id_label}: connection {idx} ack mismatch"
            );

            let bytes = tokio::time::timeout(Duration::from_millis(50), reader.next())
                .await
                .unwrap_or_else(|_| {
                    panic!(
                        "{stream_id_label}: connection {idx} produced no Closed frame within 50ms"
                    )
                })
                .expect("frame reader error")
                .expect("frame reader returned None");
            assert!(
                super::deserialize_response(bytes.1).unwrap().is_closed(),
                "{stream_id_label}: connection {idx} expected Closed terminal frame"
            );
        }
    }

    #[async_timed_test(timeout_secs = 30)]
    async fn rx_join_flushes_pending_ack_single_stream() {
        // Three independent single-stream sessions, each with three
        // framed messages. Verifies every recv-loop's terminal flush
        // ran by the time `rx.join()` returns.
        let plans = (1u64..=3)
            .map(|sid| SeparateSessionPlan {
                session_id: SessionId(sid),
                stream_id: 0,
                messages: (0u64..3).map(|seq| (seq, sid * 100 + seq)).collect(),
            })
            .collect();
        run_separate_sessions_flush_test(plans, "single-stream").await;
    }

    #[async_timed_test(timeout_secs = 30)]
    async fn rx_join_flushes_pending_ack_multi_stream() {
        // Three independent multi-stream sessions, each with three
        // framed messages, each on its own session_id with a single
        // stream_id of 1. Exercises `dispatch_multi_stream`'s terminal
        // cleanup once per session.
        let plans = (1u64..=3)
            .map(|sid| SeparateSessionPlan {
                session_id: SessionId(sid),
                stream_id: 1,
                messages: (0u64..3).map(|seq| (seq, sid * 100 + seq)).collect(),
            })
            .collect();
        run_separate_sessions_flush_test(plans, "multi-stream").await;
    }

    /// One session, multiple stream_ids — exercises `AckWatermark`'s
    /// shared-watermark path. Three streams in the same session each
    /// send three messages with disjoint seqs filling the contiguous
    /// range 0..=8. Each stream's cleanup reads `highest_uncommitted`
    /// and emits `Ack(8)` on its own wire so the peer's per-wire
    /// NetTx sees an ack for messages it sent there; the receiver
    /// discards duplicates. Every stream also emits its own `Closed`.
    #[async_timed_test(timeout_secs = 30)]
    async fn rx_join_flushes_pending_ack_shared_multi_stream_session() {
        let config = hyperactor_config::global::lock();
        let _g_msg = config.override_key(config::MESSAGE_ACK_EVERY_N_MESSAGES, 1_000_000);
        let _g_time =
            config.override_key(config::MESSAGE_ACK_TIME_INTERVAL, Duration::from_secs(3600));

        let session_id = SessionId(99);
        let num_streams = 3u8;
        let msgs_per_stream = 3u64;
        let mut conns: Vec<PreparedConnection> = Vec::with_capacity(num_streams as usize);
        let mut expected_messages: std::collections::HashSet<u64> =
            std::collections::HashSet::new();
        for stream_id in 1..=num_streams {
            let messages: Vec<(u64, u64)> = (0u64..msgs_per_stream)
                .map(|i| {
                    let seq = (stream_id as u64 - 1) * msgs_per_stream + i;
                    (seq, 1000 + seq)
                })
                .collect();
            for (_, v) in &messages {
                expected_messages.insert(*v);
            }
            conns.push(prepare_connection(session_id, stream_id, &messages).await);
        }
        let highest_seq = num_streams as u64 * msgs_per_stream - 1;

        let addr = ChannelAddr::Local(u64::MAX);
        let listener = QueueListener {
            streams: conns
                .iter_mut()
                .map(|c| std::mem::replace(&mut c.server_side, tokio::io::duplex(1).0))
                .collect(),
            addr: addr.clone(),
        };
        let (_addr, mut rx) = super::server::serve_with_listener::<u64, _>(listener, addr).unwrap();

        let mut received: std::collections::HashSet<u64> = std::collections::HashSet::new();
        for _ in 0..expected_messages.len() {
            received.insert(rx.recv().await.unwrap());
        }
        assert_eq!(
            received, expected_messages,
            "shared-session multi-stream: every message reaches the application"
        );

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

        let max_len = hyperactor_config::global::get(config::CODEC_MAX_FRAME_LENGTH);
        let mut readers: Vec<FrameReader<ReadHalf<DuplexStream>>> = conns
            .into_iter()
            .map(|c| FrameReader::new(c.client_r, max_len))
            .collect();
        for (idx, reader) in readers.iter_mut().enumerate() {
            match tokio::time::timeout(Duration::from_millis(10), reader.next()).await {
                Err(_) | Ok(Ok(None)) => {} // timeout / EOF — no early ack, good.
                Ok(Err(e)) => {
                    panic!("shared-session multi-stream: stream {idx} frame reader error: {e}")
                }
                Ok(Ok(Some((_, bytes)))) => {
                    let resp = super::deserialize_response(bytes).unwrap();
                    panic!(
                        "shared-session multi-stream: stream {idx} unexpectedly received \
                         {resp:?} before rx.join()"
                    );
                }
            }
        }

        rx.join().await;

        // Drain every frame on every reader. Terminal cleanup may emit
        // `Ack(highest_seq)` on one or more wires before `Closed`, but
        // once one cleanup commits the shared watermark, later streams
        // may emit only `Closed`.
        let mut ack_count = 0;
        let mut closed_count = 0;
        for (idx, reader) in readers.iter_mut().enumerate() {
            loop {
                match tokio::time::timeout(Duration::from_millis(50), reader.next()).await {
                    Err(_) => panic!(
                        "shared-session multi-stream: stream {idx} did not yield expected \
                         frames within 50ms after rx.join()"
                    ),
                    Ok(Err(e)) => panic!("frame reader error: {e}"),
                    Ok(Ok(None)) => break,
                    Ok(Ok(Some((_, bytes)))) => {
                        let resp = super::deserialize_response(bytes).unwrap();
                        match resp {
                            NetRxResponse::Ack(seq) => {
                                assert_eq!(
                                    seq, highest_seq,
                                    "shared-session multi-stream: ack should cover the full \
                                     contiguous range 0..={highest_seq}"
                                );
                                ack_count += 1;
                            }
                            NetRxResponse::Closed => {
                                closed_count += 1;
                                break;
                            }
                            other => panic!(
                                "shared-session multi-stream: stream {idx} unexpected {other:?}"
                            ),
                        }
                    }
                }
            }
        }
        assert!(
            ack_count >= 1,
            "shared-session multi-stream: expected at least one Ack({highest_seq}); \
             got {ack_count}"
        );
        assert!(
            ack_count <= num_streams as usize,
            "shared-session multi-stream: expected at most {num_streams} Ack({highest_seq}) \
             frames; got {ack_count}"
        );
        assert_eq!(
            closed_count, num_streams as usize,
            "shared-session multi-stream: every stream should emit its own Closed frame"
        );
    }

    /// Duplex analog of `rx_join_flushes_pending_ack_single_stream`.
    /// Three independent duplex sessions, each with three framed
    /// messages. Verifies every `dispatch_duplex_stream`'s terminal
    /// flush ran by the time `DuplexServer::join()` returns —
    /// structured concurrency makes the listener task await every
    /// inline recv/send loop before resolving.
    #[async_timed_test(timeout_secs = 30)]
    async fn server_join_flushes_pending_ack_duplex_session() {
        let config = hyperactor_config::global::lock();
        let _g_msg = config.override_key(config::MESSAGE_ACK_EVERY_N_MESSAGES, 1_000_000);
        let _g_time =
            config.override_key(config::MESSAGE_ACK_TIME_INTERVAL, Duration::from_secs(3600));

        let session_count = 3u64;
        let msgs_per_session = 3u64;
        let mut conns: Vec<PreparedConnection> = Vec::with_capacity(session_count as usize);
        let mut expected_messages: std::collections::HashSet<u64> =
            std::collections::HashSet::new();
        let mut expected_acks: Vec<u64> = Vec::with_capacity(session_count as usize);
        for sid in 1..=session_count {
            let messages: Vec<(u64, u64)> = (0u64..msgs_per_session)
                .map(|seq| (seq, sid * 100 + seq))
                .collect();
            for (_, v) in &messages {
                expected_messages.insert(*v);
            }
            expected_acks.push(messages.iter().map(|(s, _)| *s).max().unwrap());
            conns.push(prepare_connection(SessionId(sid), 0, &messages).await);
        }

        let addr = ChannelAddr::Local(u64::MAX);
        let listener = QueueListener {
            streams: conns
                .iter_mut()
                .map(|c| std::mem::replace(&mut c.server_side, tokio::io::duplex(1).0))
                .collect(),
            addr: addr.clone(),
        };
        let mut server = super::duplex::serve_with_listener::<u64, u64, _>(listener, addr).unwrap();

        // Accept one (rx, tx) pair per session. Hold the tx halves
        // alive so the server's send-loop stays parked in `select!`
        // (it never sees an `AppClosed` terminal) — the test must
        // exercise the cancel-driven flush path, not an
        // app-disconnect one.
        let mut all_rx: Vec<super::duplex::DuplexRx<u64>> =
            Vec::with_capacity(session_count as usize);
        let mut all_tx: Vec<super::duplex::DuplexTx<u64>> =
            Vec::with_capacity(session_count as usize);
        for _ in 0..session_count {
            let (rx, tx) = server.accept().await.unwrap();
            all_rx.push(rx);
            all_tx.push(tx);
        }

        // Drain the messages each session sent. Each session's
        // dispatch task delivers exactly its own `msgs_per_session`
        // values to its own `rx`, but the order in which sessions
        // are accepted is non-deterministic — collect into a set.
        let mut received: std::collections::HashSet<u64> = std::collections::HashSet::new();
        for rx in all_rx.iter_mut() {
            for _ in 0..msgs_per_session {
                received.insert(rx.recv().await.unwrap());
            }
        }
        assert_eq!(
            received, expected_messages,
            "duplex: every produced message should reach the application"
        );

        // Policy thresholds are out of reach, so no ack should fire
        // spontaneously. The sleep gives any regression that adds a
        // new emission path room to surface.
        tokio::time::sleep(Duration::from_millis(100)).await;

        let max_len = hyperactor_config::global::get(config::CODEC_MAX_FRAME_LENGTH);
        let mut readers: Vec<FrameReader<ReadHalf<DuplexStream>>> = conns
            .into_iter()
            .map(|c| FrameReader::new(c.client_r, max_len))
            .collect();
        for (idx, reader) in readers.iter_mut().enumerate() {
            match tokio::time::timeout(Duration::from_millis(10), reader.next()).await {
                Err(_) => {} // timeout — no frame, expected.
                Ok(Err(e)) => {
                    panic!("duplex: connection {idx} frame reader error before join: {e}")
                }
                Ok(Ok(None)) => {
                    panic!("duplex: connection {idx} closed before server.join()")
                }
                Ok(Ok(Some((_, bytes)))) => {
                    let resp = super::deserialize_response(bytes).unwrap();
                    panic!(
                        "duplex: connection {idx} unexpectedly received {resp:?} \
                         before server.join()"
                    );
                }
            }
        }

        // Trigger graceful shutdown. Holding `all_tx` / `all_rx`
        // alive keeps the send-loop parked and `inbound_tx` valid,
        // so the dispatch task only exits via the cancel branch of
        // its `select!` — exercising the structured-concurrency
        // flush-on-cancel path.
        server.join().await;

        // After server.join() returns, every connection must have
        // its terminal cleanup frames already on the wire: an
        // `Ack(highest_seq)` covering the messages it sent, then a
        // `Closed`.
        for (idx, (reader, expected_ack)) in readers.iter_mut().zip(&expected_acks).enumerate() {
            let bytes = tokio::time::timeout(Duration::from_millis(50), reader.next())
                .await
                .unwrap_or_else(|_| {
                    panic!(
                        "duplex: connection {idx} produced no Ack frame within 50ms after \
                         server.join()"
                    )
                })
                .expect("frame reader error")
                .expect("frame reader returned None");
            let acked = super::deserialize_response(bytes.1)
                .unwrap()
                .into_ack()
                .unwrap_or_else(|other| {
                    panic!("duplex: connection {idx} expected Ack, got {other:?}")
                });
            assert_eq!(
                acked, *expected_ack,
                "duplex: connection {idx} ack mismatch"
            );

            let bytes = tokio::time::timeout(Duration::from_millis(50), reader.next())
                .await
                .unwrap_or_else(|_| {
                    panic!("duplex: connection {idx} produced no Closed frame within 50ms")
                })
                .expect("frame reader error")
                .expect("frame reader returned None");
            assert!(
                super::deserialize_response(bytes.1).unwrap().is_closed(),
                "duplex: connection {idx} expected Closed terminal frame"
            );
        }
    }

    /// Test-only [`Link`] that yields each pre-built `DuplexStream`
    /// once. Writes `LinkInit` on the stream before returning it so
    /// the test (which holds the other end) sees the same wire
    /// format a real `TcpLink` produces.
    struct DuplexDialMockLink {
        session_id: SessionId,
        streams: std::collections::VecDeque<DuplexStream>,
    }

    impl fmt::Debug for DuplexDialMockLink {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("DuplexDialMockLink")
                .field("session_id", &self.session_id)
                .field("remaining_streams", &self.streams.len())
                .finish()
        }
    }

    #[async_trait]
    impl super::Link for DuplexDialMockLink {
        type Stream = DuplexStream;

        fn dest(&self) -> ChannelAddr {
            ChannelAddr::Local(u64::MAX)
        }

        fn link_id(&self) -> SessionId {
            self.session_id
        }

        async fn next(&mut self) -> Result<DuplexStream, ClientError> {
            match self.streams.pop_front() {
                Some(mut stream) => {
                    super::write_link_init(&mut stream, self.session_id, 0)
                        .await
                        .map_err(|err| ClientError::Io(self.dest(), err))?;
                    Ok(stream)
                }
                None => Err(ClientError::Connect(
                    self.dest(),
                    std::io::Error::other("mock link exhausted"),
                    "no more streams".into(),
                )),
            }
        }
    }

    /// Acceptor-side counterpart of
    /// [`duplex_dial_flushes_pending_ack_on_app_closed`]. The
    /// application drops its `DuplexTx`, the dispatch task's
    /// `send_connected` returns `SendLoopError::AppClosed` —
    /// terminal — so `dispatch_duplex_stream`'s loop exits via the
    /// flush + `Closed` + break path. Verifies the cumulative ack
    /// and the terminal `Closed` are written on
    /// `INITIATOR_TO_ACCEPTOR` before the dispatch task ends.
    #[async_timed_test(timeout_secs = 30)]
    async fn duplex_serve_flushes_pending_ack_on_app_closed() {
        let config = hyperactor_config::global::lock();
        let _g_msg = config.override_key(config::MESSAGE_ACK_EVERY_N_MESSAGES, 1_000_000);
        let _g_time =
            config.override_key(config::MESSAGE_ACK_TIME_INTERVAL, Duration::from_secs(3600));

        let session_id = SessionId(1);
        let messages: Vec<(u64, u64)> = vec![(0, 100), (1, 200), (2, 300)];
        let expected_ack: u64 = 2;
        let conn = prepare_connection(session_id, 0, &messages).await;

        let addr = ChannelAddr::Local(u64::MAX);
        let listener = QueueListener {
            streams: std::collections::VecDeque::from([conn.server_side]),
            addr: addr.clone(),
        };

        let mut server = super::duplex::serve_with_listener::<u64, u64, _>(listener, addr).unwrap();

        let (mut server_rx, server_tx) = server.accept().await.unwrap();

        // Drain the messages the test wrote on `INITIATOR_TO_ACCEPTOR`.
        let mut received: Vec<u64> = Vec::with_capacity(messages.len());
        for _ in &messages {
            received.push(server_rx.recv().await.unwrap());
        }
        let expected_values: Vec<u64> = messages.iter().map(|(_, v)| *v).collect();
        assert_eq!(
            received, expected_values,
            "duplex serve: every message should reach the application"
        );

        // Policy thresholds are unreachable, so no ack should fire
        // spontaneously.
        let max_len = hyperactor_config::global::get(config::CODEC_MAX_FRAME_LENGTH);
        tokio::time::sleep(Duration::from_millis(100)).await;
        let mut reader = FrameReader::new(conn.client_r, max_len);
        match tokio::time::timeout(Duration::from_millis(10), reader.next()).await {
            Err(_) => {} // timeout — expected.
            Ok(Err(e)) => panic!("duplex serve: frame reader error before app close: {e}"),
            Ok(Ok(None)) => panic!("duplex serve: wire closed before app close"),
            Ok(Ok(Some((_, bytes)))) => {
                let resp = super::deserialize_response(bytes).unwrap();
                panic!("duplex serve: unexpectedly received {resp:?} before app close");
            }
        }

        // Drop the application's `DuplexTx` to trigger `AppClosed`
        // in `send_connected` — terminal — so the dispatch task
        // exits. The flush logic must write the cumulative ack on
        // `INITIATOR_TO_ACCEPTOR` and the terminal `Closed` before
        // the loop breaks.
        drop(server_tx);

        let bytes = tokio::time::timeout(Duration::from_millis(100), reader.next())
            .await
            .unwrap_or_else(|_| panic!("duplex serve: produced no Ack frame within 100ms"))
            .expect("frame reader error")
            .expect("frame reader returned None");
        let acked = super::deserialize_response(bytes.1)
            .unwrap()
            .into_ack()
            .unwrap_or_else(|other| panic!("duplex serve: expected Ack, got {other:?}"));
        assert_eq!(
            acked, expected_ack,
            "duplex serve: ack should cover the highest seq received"
        );

        let bytes = tokio::time::timeout(Duration::from_millis(100), reader.next())
            .await
            .unwrap_or_else(|_| panic!("duplex serve: produced no Closed frame within 100ms"))
            .expect("frame reader error")
            .expect("frame reader returned None");
        assert!(
            super::deserialize_response(bytes.1).unwrap().is_closed(),
            "duplex serve: expected Closed terminal frame after Ack"
        );

        drop(server_rx);
        // `server.join()` triggers the listener-task cancel so the
        // outer `accept_loop` returns. The dispatch task already
        // exited via `AppClosed` and was drained from the JoinSet,
        // so this just stops the listener.
        server.join().await;
    }

    /// [`DuplexClient::join`] cancels the recv/send loop's
    /// cancellation token. The `select!`s observe cancel and the
    /// loop exits via the flush + `Closed` + break path. Verifies
    /// the cumulative ack and the terminal `Closed` are written on
    /// `ACCEPTOR_TO_INITIATOR` before the spawned task ends.
    #[async_timed_test(timeout_secs = 30)]
    async fn duplex_client_join_flushes_pending_ack() {
        let config = hyperactor_config::global::lock();
        let _g_msg = config.override_key(config::MESSAGE_ACK_EVERY_N_MESSAGES, 1_000_000);
        let _g_time =
            config.override_key(config::MESSAGE_ACK_TIME_INTERVAL, Duration::from_secs(3600));

        let session_id = SessionId(123);
        let (client_side, server_side) = tokio::io::duplex(8192);
        let (mut test_r, mut test_w) = tokio::io::split(server_side);

        let link = DuplexDialMockLink {
            session_id,
            streams: std::collections::VecDeque::from([client_side]),
        };

        let mut dial_client = super::duplex::spawn::<u64, u64>(link);
        let _dial_tx = dial_client.tx();
        let mut dial_rx = dial_client.take_rx().unwrap();

        // Drain the LinkInit the dial-side wrote on connect.
        super::read_link_init(&mut test_r).await.unwrap();

        // Test acts as the acceptor: write framed `Frame::Message`s
        // on `ACCEPTOR_TO_INITIATOR` so the dial-side recv-loop
        // reads them and forwards to the application via
        // `inbound_tx`.
        let max_len = hyperactor_config::global::get(config::CODEC_MAX_FRAME_LENGTH);
        let messages: Vec<(u64, u64)> = vec![(0, 100), (1, 200), (2, 300)];
        let expected_ack: u64 = 2;
        for (seq, value) in &messages {
            let payload =
                serde_multipart::serialize_bincode(&Frame::<u64>::Message(*seq, *value)).unwrap();
            let mut fw = FrameWrite::new(
                test_w,
                payload.framed(),
                max_len,
                super::ACCEPTOR_TO_INITIATOR,
            )
            .map_err(|(_w, e)| e)
            .unwrap();
            fw.send().await.unwrap();
            test_w = fw.complete();
        }

        // Drain the messages on the dial-side rx so the dial-side's
        // `recv_next.seq` advances past every seq.
        let mut received: Vec<u64> = Vec::with_capacity(messages.len());
        for _ in &messages {
            received.push(dial_rx.recv().await.unwrap());
        }
        let expected_values: Vec<u64> = messages.iter().map(|(_, v)| *v).collect();
        assert_eq!(
            received, expected_values,
            "dial: every message should reach the application"
        );

        // Policy thresholds are unreachable, so no ack should fire
        // spontaneously.
        tokio::time::sleep(Duration::from_millis(100)).await;
        let mut reader = FrameReader::new(test_r, max_len);
        match tokio::time::timeout(Duration::from_millis(10), reader.next()).await {
            Err(_) => {} // timeout — expected.
            Ok(Err(e)) => panic!("dial: frame reader error before join: {e}"),
            Ok(Ok(None)) => panic!("dial: wire closed before join"),
            Ok(Ok(Some((_, bytes)))) => {
                let resp = super::deserialize_response(bytes).unwrap();
                panic!("dial: unexpectedly received {resp:?} before join");
            }
        }

        // Trigger graceful shutdown via the structured-concurrency
        // join handle. Cancel propagates into the spawn loop's
        // `select!`s; the loop exits via Cancelled (terminal) and
        // the flush logic writes the cumulative ack and the
        // `Closed` terminal frame before the task ends.
        dial_client.join().await;

        let bytes = tokio::time::timeout(Duration::from_millis(100), reader.next())
            .await
            .unwrap_or_else(|_| panic!("dial: produced no Ack frame within 100ms after join"))
            .expect("frame reader error")
            .expect("frame reader returned None");
        let acked = super::deserialize_response(bytes.1)
            .unwrap()
            .into_ack()
            .unwrap_or_else(|other| panic!("dial: expected Ack, got {other:?}"));
        assert_eq!(
            acked, expected_ack,
            "dial: ack should cover the highest seq received"
        );

        // After the ack, the dial-side writes the terminal
        // `Closed` response on `ACCEPTOR_TO_INITIATOR` (mirrors
        // `dispatch_duplex_stream`), then releases the connection.
        let bytes = tokio::time::timeout(Duration::from_millis(100), reader.next())
            .await
            .unwrap_or_else(|_| panic!("dial: produced no Closed frame within 100ms after join"))
            .expect("frame reader error")
            .expect("frame reader returned None");
        assert!(
            super::deserialize_response(bytes.1).unwrap().is_closed(),
            "dial: expected Closed terminal frame after Ack"
        );

        drop(dial_rx);
        drop(test_w);
    }

    /// Verifies that an in-progress [`DuplexRx::recv`] on the
    /// receiver returned by [`DuplexClient::take_rx`] resolves with
    /// [`ChannelError::Closed`] when [`DuplexClient::join`] is
    /// called concurrently. Structured concurrency guarantees the
    /// spawned task drops `inbound_tx` before `join` returns, so
    /// the receiver observes the close deterministically.
    #[async_timed_test(timeout_secs = 30)]
    async fn duplex_client_join_terminates_in_progress_recv() {
        let session_id = SessionId(123);
        let (client_side, server_side) = tokio::io::duplex(8192);
        let (mut test_r, _test_w) = tokio::io::split(server_side);

        let link = DuplexDialMockLink {
            session_id,
            streams: std::collections::VecDeque::from([client_side]),
        };

        let mut dial_client = super::duplex::spawn::<u64, u64>(link);
        let mut dial_rx = dial_client.take_rx().unwrap();

        // Drain the LinkInit so the dial-side has finished its
        // setup before we kick off the recv() under test.
        super::read_link_init(&mut test_r).await.unwrap();

        // Park a recv() in a separate task; the dial-side hasn't
        // forwarded any inbound frames so this future will sit in
        // `inbound_rx.recv().await` indefinitely until `join`
        // drops the spawned task's `inbound_tx`.
        let recv_handle: tokio::task::JoinHandle<Result<u64, ChannelError>> =
            tokio::spawn(async move { dial_rx.recv().await });

        // `join` cancels the spawned task; on exit the task drops
        // its `inbound_tx`, which closes `inbound_rx` and resolves
        // the recv() with `ChannelError::Closed`.
        dial_client.join().await;

        let result = tokio::time::timeout(Duration::from_millis(100), recv_handle)
            .await
            .expect("parked recv should resolve within 100ms after join")
            .expect("recv task should not panic");
        assert!(
            matches!(result, Err(ChannelError::Closed)),
            "in-progress recv should resolve with ChannelError::Closed after join, got {result:?}"
        );
    }
}