hyperactor/

channel.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.
 */

//! One-way, multi-process, typed communication channels. These are used
//! to send messages between mailboxes residing in different processes.

use core::net::SocketAddr;
use std::fmt;
use std::future::Future;
use std::net::IpAddr;
use std::net::Ipv6Addr;
#[cfg(target_os = "linux")]
use std::os::linux::net::SocketAddrExt;
use std::os::unix::io::FromRawFd;
use std::os::unix::io::RawFd;
use std::panic::Location;
use std::pin::Pin;
use std::str::FromStr;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicU8;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::task::Context;
use std::task::Poll;

use async_trait::async_trait;
use enum_as_inner::EnumAsInner;
use futures::task::AtomicWaker;
use hyperactor_config::attrs::AttrValue;
use serde::Deserialize;
use serde::Serialize;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio::time::Instant;
use tokio_util::sync::CancellationToken;

use crate as hyperactor;
use crate::RemoteMessage;
pub(crate) mod local;
pub(crate) mod net;

// Public TLS API for HTTP services (mesh admin, TUI, etc.). The
// implementation lives in `net` but we re-export here to keep `net`'s
// internal types out of the public API surface.
pub use net::ServerError;
pub use net::try_tls_acceptor_with_pem_bundle;
pub use net::try_tls_connector;
pub use net::try_tls_pem_bundle;

/// Duplex channel API: a single connection carries messages in both directions.
pub mod duplex {
    pub use super::net::duplex::DuplexClient;
    pub use super::net::duplex::DuplexRx;
    pub use super::net::duplex::DuplexServer;
    pub use super::net::duplex::DuplexTx;
    pub use super::net::duplex::dial;
    pub use super::net::duplex::serve;
}

/// The type of error that can occur on channel operations.
#[derive(thiserror::Error, Debug)]
pub enum ChannelError {
    /// An operation was attempted on a closed channel.
    #[error("channel closed")]
    Closed,

    /// An error occurred during send.
    #[error("send: {0}")]
    Send(#[source] anyhow::Error),

    /// A network client error.
    #[error(transparent)]
    Client(#[from] net::ClientError),

    /// The address was not valid.
    #[error("invalid address {0:?}")]
    InvalidAddress(String),

    /// A serving error was encountered.
    #[error(transparent)]
    Server(#[from] net::ServerError),

    /// A bincode encoding error occurred.
    #[error(transparent)]
    BincodeEncode(#[from] bincode::error::EncodeError),

    /// A bincode decoding error occurred.
    #[error(transparent)]
    BincodeDecode(#[from] bincode::error::DecodeError),

    /// Data encoding errors.
    #[error(transparent)]
    Data(#[from] wirevalue::Error),

    /// Some other error.
    #[error(transparent)]
    Other(#[from] anyhow::Error),

    /// An operation timeout occurred.
    #[error("operation timed out after {0:?}")]
    Timeout(std::time::Duration),
}

/// Structured context for a send error.
#[derive(thiserror::Error, Debug, Clone, PartialEq, Eq)]
pub enum SendErrorReason {
    /// The serialized frame exceeded the configured channel frame limit.
    #[error(
        "rejecting oversize frame: len={len} > max={max}. \
        ack will not arrive before timeout; increase CODEC_MAX_FRAME_LENGTH to allow."
    )]
    OversizedFrame {
        /// The serialized frame length.
        len: usize,

        /// The configured frame limit.
        max: usize,
    },

    /// Other human-readable context.
    #[error("{0}")]
    Other(String),
}

/// An error that occurred during send. Returns the message that failed to send.
#[derive(thiserror::Error, Debug)]
#[error("{error} for reason {reason:?}")]
pub struct SendError<M: RemoteMessage> {
    /// Inner channel error
    #[source]
    pub error: ChannelError,
    /// Message that couldn't be sent
    pub message: M,
    /// Reason that message couldn't be sent, if any.
    pub reason: Option<SendErrorReason>,
}

#[repr(u8)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum CompletionStatus {
    Pending = 0,
    Accepted = 1,
    Rejected = 2,
}

impl CompletionStatus {
    fn from_u8(value: u8) -> Self {
        match value {
            value if value == Self::Pending as u8 => Self::Pending,
            value if value == Self::Accepted as u8 => Self::Accepted,
            value if value == Self::Rejected as u8 => Self::Rejected,
            _ => panic!("invalid completion state"),
        }
    }
}

struct CompletionState<M: RemoteMessage> {
    state: AtomicU8,
    waker: AtomicWaker,
    rejected: Mutex<Option<Box<SendError<M>>>>,
}

pub(crate) struct CompletionSender<M: RemoteMessage> {
    inner: Arc<CompletionState<M>>,
}

/// Future that resolves when a posted message is accepted or rejected.
pub struct CompletionReceipt<M: RemoteMessage> {
    inner: Arc<CompletionState<M>>,
}

impl<M: RemoteMessage> CompletionSender<M> {
    fn pair() -> (Self, CompletionReceipt<M>) {
        let inner = Arc::new(CompletionState {
            state: AtomicU8::new(CompletionStatus::Pending as u8),
            waker: AtomicWaker::new(),
            rejected: Mutex::new(None),
        });

        (
            Self {
                inner: Arc::clone(&inner),
            },
            CompletionReceipt { inner },
        )
    }

    fn accept(self) {
        self.inner
            .state
            .store(CompletionStatus::Accepted as u8, Ordering::Release);
        self.inner.waker.wake();
    }

    fn reject(self, error: SendError<M>) {
        *self.inner.rejected.lock().unwrap() = Some(Box::new(error));
        self.inner
            .state
            .store(CompletionStatus::Rejected as u8, Ordering::Release);
        self.inner.waker.wake();
    }
}

impl<M: RemoteMessage> Drop for CompletionSender<M> {
    fn drop(&mut self) {
        if CompletionStatus::from_u8(self.inner.state.load(Ordering::Acquire))
            == CompletionStatus::Pending
        {
            self.inner
                .state
                .store(CompletionStatus::Accepted as u8, Ordering::Release);
            self.inner.waker.wake();
        }
    }
}

impl<M: RemoteMessage> Future for CompletionReceipt<M> {
    type Output = Result<(), SendError<M>>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        match self.poll_ready() {
            Poll::Ready(result) => Poll::Ready(result),
            Poll::Pending => {
                self.inner.waker.register(cx.waker());
                self.poll_ready()
            }
        }
    }
}

impl<M: RemoteMessage> CompletionReceipt<M> {
    fn poll_ready(&self) -> Poll<Result<(), SendError<M>>> {
        match CompletionStatus::from_u8(self.inner.state.load(Ordering::Acquire)) {
            CompletionStatus::Accepted => Poll::Ready(Ok(())),
            CompletionStatus::Rejected => {
                let error = self
                    .inner
                    .rejected
                    .lock()
                    .unwrap()
                    .take()
                    .expect("rejected completion should store send error");
                Poll::Ready(Err(*error))
            }
            CompletionStatus::Pending => Poll::Pending,
        }
    }
}

/// Shared completion counter for sinks that need to wake flush waiters.
pub(crate) struct CompletionTracker {
    completed: Arc<AtomicUsize>,
    completed_notify: Arc<tokio::sync::Notify>,
}

impl CompletionTracker {
    pub(crate) fn new(
        completed: Arc<AtomicUsize>,
        completed_notify: Arc<tokio::sync::Notify>,
    ) -> Self {
        Self {
            completed,
            completed_notify,
        }
    }

    fn complete(&self) {
        self.completed.fetch_add(1, Ordering::Relaxed);
        self.completed_notify.notify_waiters();
    }
}

enum CompletionSinkInner<M: RemoteMessage> {
    Ignore,
    Receipt(CompletionSender<M>),
    OnReject(Box<dyn FnOnce(SendError<M>) + Send + Sync>),
    Tracked {
        tracker: CompletionTracker,
        on_reject: Box<dyn FnOnce(SendError<M>) + Send + Sync>,
    },
}

/// Sink for the terminal outcome of a posted message.
pub struct CompletionSink<M: RemoteMessage>(CompletionSinkInner<M>);

impl<M: RemoteMessage> CompletionSink<M> {
    /// Ignore the message completion.
    pub fn ignore() -> Self {
        Self(CompletionSinkInner::Ignore)
    }

    /// Invoke `f` only when the channel rejects the message.
    pub fn on_reject(f: impl FnOnce(SendError<M>) + Send + Sync + 'static) -> Self {
        Self(CompletionSinkInner::OnReject(Box::new(f)))
    }

    /// Return a completion sink and a receipt that observes its terminal outcome.
    pub fn receipt() -> (Self, CompletionReceipt<M>) {
        let (sender, receipt) = CompletionSender::pair();
        (Self(CompletionSinkInner::Receipt(sender)), receipt)
    }

    /// Track every completion and invoke `on_reject` only for rejected messages.
    pub(crate) fn tracked(
        tracker: CompletionTracker,
        on_reject: impl FnOnce(SendError<M>) + Send + Sync + 'static,
    ) -> Self {
        Self(CompletionSinkInner::Tracked {
            tracker,
            on_reject: Box::new(on_reject),
        })
    }

    /// Adapt rejected send errors for a wrapped message type.
    pub fn contramap_rejected<N: RemoteMessage>(
        self,
        f: impl FnOnce(SendError<N>) -> Option<SendError<M>> + Send + Sync + 'static,
    ) -> CompletionSink<N> {
        match self.0 {
            CompletionSinkInner::Ignore => CompletionSink::ignore(),
            CompletionSinkInner::Receipt(sender) => CompletionSink::on_reject(move |error| {
                if let Some(error) = f(error) {
                    sender.reject(error);
                } else {
                    sender.accept();
                }
            }),
            CompletionSinkInner::OnReject(on_reject) => CompletionSink::on_reject(move |error| {
                if let Some(error) = f(error) {
                    on_reject(error);
                }
            }),
            CompletionSinkInner::Tracked { tracker, on_reject } => {
                CompletionSink::tracked(tracker, move |error| {
                    if let Some(error) = f(error) {
                        on_reject(error);
                    }
                })
            }
        }
    }

    /// Report that the channel accepted the message.
    pub fn accept(self) {
        match self.0 {
            CompletionSinkInner::Ignore => {}
            CompletionSinkInner::Receipt(sender) => sender.accept(),
            CompletionSinkInner::OnReject(_) => {}
            CompletionSinkInner::Tracked { tracker, .. } => tracker.complete(),
        }
    }

    /// Report that the channel rejected the message.
    pub fn reject(self, error: SendError<M>) {
        match self.0 {
            CompletionSinkInner::Ignore => {}
            CompletionSinkInner::Receipt(sender) => sender.reject(error),
            CompletionSinkInner::OnReject(on_reject) => on_reject(error),
            CompletionSinkInner::Tracked { tracker, on_reject } => {
                on_reject(error);
                tracker.complete();
            }
        }
    }
}

impl<M: RemoteMessage> From<SendError<M>> for ChannelError {
    fn from(error: SendError<M>) -> Self {
        error.error
    }
}

/// Reason a [`TxStatus`] transitioned to `Closed`. Callers should branch on
/// the typed variants for cases they care about (e.g. cache-eviction logic in
/// `DialMailboxRouter` keys on `SequenceMismatch`); everything else falls
/// into `Other` and is for display/logging only.
#[derive(Debug, Clone, PartialEq)]
pub enum CloseReason {
    /// The peer rejected our session because our sequence number did not
    /// match what the peer's dispatcher expected — the K8s "out-of-sequence
    /// message, expected seq 0, got N" case where the peer GC'd the
    /// `SessionId` while we still hold an `Outbox.next_seq` past 0.
    /// Re-dialing produces a fresh session that the peer accepts.
    SequenceMismatch(String),
    /// The peer rejected a frame whose length exceeded
    /// `config::CODEC_MAX_FRAME_LENGTH`. Re-dialing will not help — the
    /// message itself is the problem.
    OversizedFrame {
        /// Actual frame length in bytes.
        size: usize,
        /// `CODEC_MAX_FRAME_LENGTH` at the time of rejection.
        max: usize,
    },
    /// Any close reason the transport hasn't classified further. The string
    /// is for display/logging only — do not parse it. If a caller needs to
    /// branch on a sub-case, lift it into its own variant on this enum.
    Other(String),
}

impl fmt::Display for CloseReason {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::SequenceMismatch(s) => write!(f, "stale session: {}", s),
            Self::OversizedFrame { size, max } => {
                write!(f, "oversized frame: len={size} > max={max}")
            }
            Self::Other(s) => f.write_str(s),
        }
    }
}

/// The possible states of a `Tx`.
#[derive(Debug, Clone, PartialEq, EnumAsInner)]
pub enum TxStatus {
    /// The tx is good.
    Active,
    /// The tx cannot be used for message delivery.
    Closed(CloseReason),
}

/// The transmit end of an M-typed channel.
#[async_trait]
pub trait Tx<M: RemoteMessage> {
    /// Post a message and report its terminal outcome to `completion`.
    ///
    /// Users should use the `try_post`, and `post` variants directly.
    fn do_post(&self, message: M, completion: CompletionSink<M>);

    /// Enqueue a `message` on the local end of the channel and return a receipt
    /// that resolves when the message is accepted or rejected.
    fn try_post(&self, message: M) -> CompletionReceipt<M> {
        let (completion, receipt) = CompletionSink::receipt();
        self.do_post(message, completion);
        receipt
    }

    /// Enqueue a message to be sent on the channel.
    #[hyperactor::instrument_infallible]
    fn post(&self, message: M) {
        self.do_post(message, CompletionSink::ignore());
    }

    /// Send a message synchronously, returning when the message has
    /// been delivered to the remote end of the channel.
    async fn send(&self, message: M) -> Result<(), SendError<M>> {
        self.try_post(message).await
    }

    /// The channel address to which this Tx is sending.
    fn addr(&self) -> ChannelAddr;

    /// A means to monitor the health of a `Tx`.
    fn status(&self) -> &watch::Receiver<TxStatus>;
}

/// The receive end of an M-typed channel.
#[async_trait]
pub trait Rx<M: RemoteMessage> {
    /// Receive the next message from the channel. If the channel returns
    /// an error it is considered broken and should be discarded.
    async fn recv(&mut self) -> Result<M, ChannelError>;

    /// The channel address from which this Rx is receiving.
    fn addr(&self) -> ChannelAddr;

    /// Gracefully shut down the channel receiver, flushing any pending
    /// acks before returning. Implementations must ensure all pending
    /// acks are sent before this method returns.
    async fn join(self)
    where
        Self: Sized;
}

/// The hostname to use for TLS connections.
#[derive(
    Clone,
    Debug,
    PartialEq,
    Eq,
    Hash,
    Serialize,
    Deserialize,
    strum::EnumIter,
    strum::Display,
    strum::EnumString
)]
pub enum TcpMode {
    /// Use localhost/loopback for the connection.
    Localhost,
    /// Use a non-loopback address resolved from the host domain name.
    Hostname,
}

/// The hostname to use for TLS connections.
#[derive(
    Clone,
    Debug,
    PartialEq,
    Eq,
    Hash,
    Serialize,
    Deserialize,
    strum::EnumIter,
    strum::Display,
    strum::EnumString
)]
pub enum TlsMode {
    /// Use IpV6 address for TLS connections.
    IpV6,
    /// Use host domain name for TLS connections.
    Hostname,
    // TODO: consider adding IpV4 support.
}

/// Address format for TLS channels.
#[derive(
    Clone,
    Debug,
    PartialEq,
    Eq,
    Hash,
    Serialize,
    Deserialize,
    Ord,
    PartialOrd
)]
pub struct TlsAddr {
    /// The hostname to connect to.
    pub hostname: Hostname,
    /// The port to connect to.
    pub port: Port,
}

impl TlsAddr {
    /// Creates a new TLS address with a normalized hostname.
    pub fn new(hostname: impl Into<Hostname>, port: Port) -> Self {
        Self {
            hostname: normalize_host(&hostname.into()),
            port,
        }
    }

    /// Returns the port number for this address.
    pub fn port(&self) -> Port {
        self.port
    }

    /// Returns the hostname for this address.
    pub fn hostname(&self) -> &str {
        &self.hostname
    }
}

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

    fn from_str(addr: &str) -> Result<Self, Self::Err> {
        let (hostname, port_str) = addr
            .rsplit_once(':')
            .ok_or_else(|| anyhow::anyhow!("invalid TLS address: {}", addr))?;
        let port = port_str
            .parse()
            .map_err(|_| anyhow::anyhow!("invalid TLS address port: {}", port_str))?;
        Ok(Self::new(hostname, port))
    }
}

impl fmt::Display for TlsAddr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}:{}", self.hostname, self.port)
    }
}

/// Types of channel transports.
#[derive(
    Clone,
    Debug,
    PartialEq,
    Eq,
    Hash,
    Serialize,
    Deserialize,
    typeuri::Named
)]
pub enum ChannelTransport {
    /// Transport over a TCP connection.
    Tcp(TcpMode),

    /// Transport over a TCP connection with TLS support within Meta
    MetaTls(TlsMode),

    /// Transport over a TCP connection with configurable TLS support
    Tls,

    /// Transport over a QUIC connection with configurable TLS support.
    Quic,

    /// Transport over a QUIC connection with TLS support within Meta.
    MetaQuic(TlsMode),

    /// Local transports use a process-local registry and private Unix socket
    /// pairs.
    Local,

    /// Transport over unix domain socket.
    Unix,
}

impl fmt::Display for ChannelTransport {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Tcp(mode) => write!(f, "tcp({:?})", mode),
            Self::MetaTls(mode) => write!(f, "metatls({:?})", mode),
            Self::Tls => write!(f, "tls"),
            Self::Quic => write!(f, "quic"),
            Self::MetaQuic(mode) => write!(f, "metaquic({:?})", mode),
            Self::Local => write!(f, "local"),
            Self::Unix => write!(f, "unix"),
        }
    }
}

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

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            // Default to TcpMode::Hostname, if the mode isn't set
            "tcp" => Ok(ChannelTransport::Tcp(TcpMode::Hostname)),
            s if s.starts_with("tcp(") => {
                let inner = &s["tcp(".len()..s.len() - 1];
                let mode = inner.parse()?;
                Ok(ChannelTransport::Tcp(mode))
            }
            "local" => Ok(ChannelTransport::Local),
            "unix" => Ok(ChannelTransport::Unix),
            "tls" => Ok(ChannelTransport::Tls),
            "quic" => Ok(ChannelTransport::Quic),
            s if s.starts_with("metatls(") && s.ends_with(")") => {
                let inner = &s["metatls(".len()..s.len() - 1];
                let mode = inner.parse()?;
                Ok(ChannelTransport::MetaTls(mode))
            }
            s if s.starts_with("metaquic(") && s.ends_with(")") => {
                let inner = &s["metaquic(".len()..s.len() - 1];
                let mode = inner.parse()?;
                Ok(ChannelTransport::MetaQuic(mode))
            }
            unknown => Err(anyhow::anyhow!("unknown channel transport: {}", unknown)),
        }
    }
}

impl ChannelTransport {
    /// All known channel transports.
    pub fn all() -> [ChannelTransport; 3] {
        [
            // TODO: @rusch add back once figuring out unspecified override for OSS CI
            // ChannelTransport::Tcp(TcpMode::Localhost),
            ChannelTransport::Tcp(TcpMode::Hostname),
            ChannelTransport::Local,
            ChannelTransport::Unix,
            // Tls requires certificate configuration, tested separately in tls::tests
            // TODO add MetaTls (T208303369)
        ]
    }

    /// Return an "any" address for this transport.
    pub fn any(&self) -> ChannelAddr {
        ChannelAddr::any(self.clone())
    }

    /// Returns true if this transport type represents a remote channel.
    pub fn is_remote(&self) -> bool {
        match self {
            ChannelTransport::Tcp(_) => true,
            ChannelTransport::MetaTls(_) => true,
            ChannelTransport::Tls => true,
            ChannelTransport::Quic => true,
            ChannelTransport::MetaQuic(_) => true,
            ChannelTransport::Local => false,
            ChannelTransport::Unix => false,
        }
    }

    /// Returns true if this transport is served by the `net` module
    /// (i.e., a kernel-level socket: TCP, Unix, or a TLS variant
    /// thereof). The only non-net transport is the in-process
    /// [`Local`](ChannelTransport::Local) channel.
    pub fn is_net(&self) -> bool {
        match self {
            ChannelTransport::Tcp(_) => true,
            ChannelTransport::MetaTls(_) => true,
            ChannelTransport::Tls => true,
            ChannelTransport::Quic => false,
            ChannelTransport::MetaQuic(_) => false,
            ChannelTransport::Unix => true,
            ChannelTransport::Local => false,
        }
    }

    /// Returns true if this transport uses TLS encryption.
    pub fn is_tls(&self) -> bool {
        matches!(self, ChannelTransport::Tls | ChannelTransport::MetaTls(_))
    }

    /// Returns true if this transport can carry the duplex byte-stream
    /// protocol (see [`crate::channel::net::duplex`]). This is a
    /// distinct predicate from [`is_net`](Self::is_net): the in-process
    /// [`Local`](ChannelTransport::Local) transport is not a kernel
    /// socket (so `is_net` is false) yet is still served over the net
    /// stack and carries duplex.
    pub fn supports_duplex(&self) -> bool {
        match self {
            ChannelTransport::Tcp(_) => true,
            ChannelTransport::MetaTls(_) => true,
            ChannelTransport::Tls => true,
            // Quic actually supports duplex byte streams, but they are not yet tested.
            ChannelTransport::Quic => false,
            ChannelTransport::MetaQuic(_) => false,
            ChannelTransport::Unix => true,
            ChannelTransport::Local => true,
        }
    }
}

impl AttrValue for ChannelTransport {
    fn display(&self) -> String {
        self.to_string()
    }

    fn parse(s: &str) -> Result<Self, anyhow::Error> {
        s.parse()
    }
}

/// Specifies how to bind a channel server.
#[derive(
    Clone,
    Debug,
    PartialEq,
    Eq,
    Hash,
    Serialize,
    Deserialize,
    typeuri::Named
)]
pub enum BindSpec {
    /// Bind to any available address for the given transport.
    Any(ChannelTransport),

    /// Bind to a specific channel address.
    Addr(ChannelAddr),
}

impl BindSpec {
    /// Return an "any" address for this bind spec.
    pub fn binding_addr(&self) -> ChannelAddr {
        match self {
            BindSpec::Any(transport) => ChannelAddr::any(transport.clone()),
            BindSpec::Addr(addr) => addr.clone(),
        }
    }
}

impl From<ChannelTransport> for BindSpec {
    fn from(transport: ChannelTransport) -> Self {
        BindSpec::Any(transport)
    }
}

impl From<ChannelAddr> for BindSpec {
    fn from(addr: ChannelAddr) -> Self {
        BindSpec::Addr(addr)
    }
}

impl fmt::Display for BindSpec {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Any(transport) => write!(f, "{}", transport),
            Self::Addr(addr) => write!(f, "{}", addr),
        }
    }
}

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

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if let Ok(transport) = ChannelTransport::from_str(s) {
            Ok(BindSpec::Any(transport))
        } else if let Ok(addr) = ChannelAddr::from_zmq_url(s) {
            Ok(BindSpec::Addr(addr))
        } else if let Ok(addr) = ChannelAddr::from_str(s) {
            Ok(BindSpec::Addr(addr))
        } else {
            Err(anyhow::anyhow!("invalid bind spec: {}", s))
        }
    }
}

impl AttrValue for BindSpec {
    fn display(&self) -> String {
        self.to_string()
    }

    fn parse(s: &str) -> Result<Self, anyhow::Error> {
        Self::from_str(s)
    }
}

/// The type of (TCP) hostnames.
pub type Hostname = String;

/// The type of (TCP) ports.
pub type Port = u16;

/// The type of a channel address, used to multiplex different underlying
/// channel implementations. ChannelAddrs also have a concrete syntax:
/// the address type (e.g., "tcp" or "local"), followed by ":", and an address
/// parseable to that type. For example:
///
/// - `tcp:127.0.0.1:1234` - localhost port 1234 over TCP
/// - `tcp:192.168.0.1:1111` - 192.168.0.1 port 1111 over TCP
/// - `quic:example.com:1234` - example.com port 1234 over QUIC
/// - `local:123` - the (in-process) local port 123
/// - `unix:/some/path` - the Unix socket at `/some/path`
///
/// Both local and TCP ports 0 are reserved to indicate "any available
/// port" when serving.
///
/// ```
/// # use hyperactor::channel::ChannelAddr;
/// let addr: ChannelAddr = "tcp:127.0.0.1:1234".parse().unwrap();
/// let ChannelAddr::Tcp(socket_addr) = addr else {
///     panic!()
/// };
/// assert_eq!(socket_addr.port(), 1234);
/// assert_eq!(socket_addr.is_ipv4(), true);
/// ```
#[derive(
    Clone,
    Debug,
    PartialEq,
    Eq,
    Ord,
    PartialOrd,
    Serialize,
    Deserialize,
    Hash,
    typeuri::Named
)]
pub enum ChannelAddr {
    /// A socket address used to establish TCP channels. Supports
    /// both  IPv4 and IPv6 address / port pairs.
    Tcp(SocketAddr),

    /// An address to establish TCP channels with TLS support within Meta.
    /// Uses TlsAddr with hostname and port.
    MetaTls(TlsAddr),

    /// An address to establish TCP channels with configurable TLS support.
    /// Uses TlsAddr with hostname and port.
    Tls(TlsAddr),

    /// An address to establish QUIC channels with configurable TLS support.
    /// Uses TlsAddr with hostname and port.
    Quic(TlsAddr),

    /// An address to establish QUIC channels with TLS support within Meta.
    /// Uses TlsAddr with hostname and port.
    MetaQuic(TlsAddr),

    /// Local addresses are registered in-process and given an integral
    /// index.
    Local(u64),

    /// A unix domain socket address. Supports both absolute path names as
    ///  well as "abstract" names per https://manpages.debian.org/unstable/manpages/unix.7.en.html#Abstract_sockets
    Unix(net::unix::SocketAddr),

    /// A pair of addresses, one for the client and one for the server:
    ///   - The client should dial to the `dial_to` address.
    ///   - The server should bind to the `bind_to` address.
    ///
    /// The user is responsible for ensuring the traffic to the `dial_to` address
    /// is routed to the `bind_to` address.
    ///
    /// This is useful for scenarios where the network is configured in a way,
    /// that the bound address is not directly accessible from the client.
    ///
    /// For example, in AWS, the client could be provided with the public IP
    /// address, yet the server is bound to a private IP address or simply
    /// INADDR_ANY. Traffic to the public IP address is mapped to the private
    /// IP address through network address translation (NAT).
    ///
    /// `Alias` is serve-side syntax. [`serve`] consumes it by binding to
    /// `bind_to` and advertising `dial_to`; identity-bearing values such as
    /// proc addresses, actor addresses, host references, and routing keys
    /// should store only `dial_to`. Dial helpers canonicalize aliases the same
    /// way.
    Alias {
        /// The address to which the client should dial to.
        dial_to: Box<ChannelAddr>,
        /// The address to which the server should bind to.
        bind_to: Box<ChannelAddr>,
    },
}

impl From<SocketAddr> for ChannelAddr {
    fn from(value: SocketAddr) -> Self {
        Self::Tcp(value)
    }
}

impl From<net::unix::SocketAddr> for ChannelAddr {
    fn from(value: net::unix::SocketAddr) -> Self {
        Self::Unix(value)
    }
}

impl From<std::os::unix::net::SocketAddr> for ChannelAddr {
    fn from(value: std::os::unix::net::SocketAddr) -> Self {
        Self::Unix(net::unix::SocketAddr::new(value))
    }
}

impl From<tokio::net::unix::SocketAddr> for ChannelAddr {
    fn from(value: tokio::net::unix::SocketAddr) -> Self {
        std::os::unix::net::SocketAddr::from(value).into()
    }
}

fn is_routable_address(addr: &IpAddr) -> bool {
    !addr.is_loopback()
        && match addr {
            IpAddr::V6(v6) => !v6.is_unicast_link_local(),
            IpAddr::V4(v4) => !v4.is_link_local(),
        }
}

/// Return the first non-loopback, non-link-local address from a list.
fn find_routable_address(addresses: &[IpAddr]) -> Option<IpAddr> {
    addresses
        .iter()
        .find(|addr| is_routable_address(addr))
        .cloned()
}

impl ChannelAddr {
    /// Returns the "any" address for the given transport.
    ///
    /// For [`TcpMode::Hostname`], this resolves the host name to a
    /// non-loopback, non-link-local address with port zero.
    ///
    /// # Panics
    ///
    /// Panics if host name resolution fails or produces no routable address.
    pub fn any(transport: ChannelTransport) -> Self {
        match transport {
            ChannelTransport::Tcp(mode) => {
                let ip = match mode {
                    TcpMode::Localhost => IpAddr::V6(Ipv6Addr::LOCALHOST),
                    TcpMode::Hostname => {
                        let hostname = hostname::get().expect("failed to retrieve hostname");
                        let hostname = hostname.to_str().expect("hostname is not valid UTF-8");
                        // TODO: Avoid using DNS directly once we figure out a good extensibility story here
                        let addresses =
                            dns_lookup::lookup_host(hostname).expect("failed to resolve hostname");
                        find_routable_address(&addresses).expect(
                            "hostname resolved to no non-loopback, non-link-local address; configure an explicit transport address",
                        )
                    }
                };
                Self::Tcp(SocketAddr::new(ip, 0))
            }
            ChannelTransport::MetaTls(mode) => {
                let host_address = match mode {
                    TlsMode::Hostname => hostname::get()
                        .ok()
                        .and_then(|hostname| hostname.to_str().map(|s| s.to_string()))
                        .unwrap_or("unknown_host".to_string()),
                    TlsMode::IpV6 => {
                        get_host_ipv6_address().expect("failed to retrieve ipv6 address")
                    }
                };
                Self::MetaTls(TlsAddr::new(host_address, 0))
            }
            ChannelTransport::MetaQuic(mode) => {
                let host_address = match mode {
                    TlsMode::Hostname => hostname::get()
                        .ok()
                        .and_then(|hostname| hostname.to_str().map(|s| s.to_string()))
                        .unwrap_or("unknown_host".to_string()),
                    TlsMode::IpV6 => {
                        get_host_ipv6_address().expect("failed to retrieve ipv6 address")
                    }
                };
                Self::MetaQuic(TlsAddr::new(host_address, 0))
            }
            ChannelTransport::Local => Self::Local(0),
            ChannelTransport::Tls => {
                let host_address = hostname::get()
                    .ok()
                    .and_then(|hostname| hostname.to_str().map(|s| s.to_string()))
                    .unwrap_or("localhost".to_string());
                Self::Tls(TlsAddr::new(host_address, 0))
            }
            ChannelTransport::Quic => {
                let host_address = hostname::get()
                    .ok()
                    .and_then(|hostname| hostname.to_str().map(|s| s.to_string()))
                    .unwrap_or("localhost".to_string());
                Self::Quic(TlsAddr::new(host_address, 0))
            }
            // This works because the file will be deleted but we know we have a unique file by this point.
            ChannelTransport::Unix => Self::Unix(net::unix::SocketAddr::from_str("").unwrap()),
        }
    }

    /// The transport used by this address.
    pub fn transport(&self) -> ChannelTransport {
        match self {
            Self::Tcp(addr) => {
                if addr.ip().is_loopback() {
                    ChannelTransport::Tcp(TcpMode::Localhost)
                } else {
                    ChannelTransport::Tcp(TcpMode::Hostname)
                }
            }
            Self::MetaTls(addr) => match addr.hostname.parse::<IpAddr>() {
                Ok(IpAddr::V6(_)) => ChannelTransport::MetaTls(TlsMode::IpV6),
                Ok(IpAddr::V4(_)) => ChannelTransport::MetaTls(TlsMode::Hostname),
                Err(_) => ChannelTransport::MetaTls(TlsMode::Hostname),
            },
            Self::Tls(_) => ChannelTransport::Tls,
            Self::Quic(_) => ChannelTransport::Quic,
            Self::MetaQuic(addr) => match addr.hostname.parse::<IpAddr>() {
                Ok(IpAddr::V6(_)) => ChannelTransport::MetaQuic(TlsMode::IpV6),
                Ok(IpAddr::V4(_)) => ChannelTransport::MetaQuic(TlsMode::Hostname),
                Err(_) => ChannelTransport::MetaQuic(TlsMode::Hostname),
            },
            Self::Local(_) => ChannelTransport::Local,
            Self::Unix(_) => ChannelTransport::Unix,
            // bind_to's transport is what is actually used in communication.
            // Therefore we use its transport to represent the Alias.
            Self::Alias { bind_to, .. } => bind_to.transport(),
        }
    }
}

#[cfg(fbcode_build)]
fn get_host_ipv6_address() -> anyhow::Result<String> {
    crate::meta::host_ip::host_ipv6_address()
}

#[cfg(not(fbcode_build))]
fn get_host_ipv6_address() -> anyhow::Result<String> {
    Ok(local_ip_address::local_ipv6()?.to_string())
}

impl fmt::Display for ChannelAddr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Tcp(addr) => write!(f, "tcp:{}", addr),
            Self::MetaTls(addr) => write!(f, "metatls:{}", addr),
            Self::Tls(addr) => write!(f, "tls:{}", addr),
            Self::Quic(addr) => write!(f, "quic:{}", addr),
            Self::MetaQuic(addr) => write!(f, "metaquic:{}", addr),
            Self::Local(index) => write!(f, "local:{}", index),
            Self::Unix(addr) => write!(f, "unix:{}", addr),
            Self::Alias { dial_to, bind_to } => {
                write!(f, "alias:dial_to={};bind_to={}", dial_to, bind_to)
            }
        }
    }
}

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

    fn from_str(addr: &str) -> Result<Self, Self::Err> {
        match addr.split_once('!').or_else(|| addr.split_once(':')) {
            Some(("local", rest)) => rest
                .parse::<u64>()
                .map(Self::Local)
                .map_err(anyhow::Error::from),
            Some(("tcp", rest)) => rest
                .parse::<SocketAddr>()
                .map(Self::Tcp)
                .map_err(anyhow::Error::from),
            Some(("metatls", rest)) => net::meta::parse(rest).map_err(|e| e.into()),
            Some(("tls", rest)) => net::tls::parse(rest).map_err(|e| e.into()),
            Some(("quic", rest)) => TlsAddr::from_str(rest).map(Self::Quic),
            Some(("metaquic", rest)) => TlsAddr::from_str(rest).map(Self::MetaQuic),
            Some(("unix", rest)) => Ok(Self::Unix(net::unix::SocketAddr::from_str(rest)?)),
            Some(("alias", _)) => Err(anyhow::anyhow!(
                "detect possible alias address, but we currently do not support \
                parsing alias' string representation since we only want to \
                support parsing its zmq url format."
            )),
            Some((r#type, _)) => Err(anyhow::anyhow!("no such channel type: {type}")),
            None => Err(anyhow::anyhow!("no channel type specified")),
        }
    }
}

/// Normalize a host string. If the host is an IP address, parse and
/// re-format it to produce a canonical string representation.
pub(crate) fn normalize_host(host: &str) -> String {
    // Strip URI-style brackets (e.g., "[::1]") because IpAddr::from_str
    // rejects them — it only accepts bare addresses.
    let host_clean = host
        .strip_prefix('[')
        .and_then(|h| h.strip_suffix(']'))
        .unwrap_or(host);

    if let Ok(ip_addr) = host_clean.parse::<IpAddr>() {
        ip_addr.to_string()
    } else {
        host.to_string()
    }
}

impl ChannelAddr {
    /// Return the canonical address that remote peers should dial.
    ///
    /// For regular addresses this is the address itself. For aliases, this
    /// recursively consumes the alias and returns its `dial_to` address. Use
    /// this before storing an address in identity-bearing state; aliases are
    /// intended as input to [`serve`].
    pub fn into_dial_addr(self) -> Self {
        match self {
            Self::Alias { dial_to, .. } => (*dial_to).into_dial_addr(),
            addr => addr,
        }
    }

    /// Parse ZMQ-style URL format: scheme://address
    /// Supports:
    /// - tcp://hostname:port or tcp://*:port (wildcard binding)
    /// - inproc://endpoint-name (equivalent to local)
    /// - ipc://path (equivalent to unix)
    /// - metatls://hostname:port or metatls://*:port
    /// - quic://hostname:port or quic://*:port
    /// - metaquic://hostname:port or metaquic://*:port
    /// - Alias format: dial_to_url@bind_to_url (e.g., tcp://host:port@tcp://host:port)
    ///   Note: Alias format is currently only supported for TCP addresses
    ///
    /// Alias format is meant for serving. Callers that will dial or store the
    /// result as an identity should canonicalize it with
    /// [`ChannelAddr::into_dial_addr`].
    pub fn from_zmq_url(address: &str) -> Result<Self, anyhow::Error> {
        let (addr, _listener) = Self::from_zmq_url_with_listener(address)?;
        Ok(addr)
    }

    /// Parse ZMQ-style URL format, with support for pre-opened file descriptors.
    ///
    /// When the port portion of a URL is `fdNNN` (e.g. `tcp://myhost:fd5`),
    /// the file descriptor is adopted as a pre-bound `TcpListener`. The
    /// returned `ChannelAddr` will contain the real port that the fd is bound
    /// to, and the `Option<TcpListener>` will be `Some`.
    ///
    /// # Safety
    /// When using the `fd` syntax, the caller must ensure the file descriptor
    /// is a valid, bound TCP socket that is not used elsewhere. The socket
    /// does not need to be in a listening state — `listen()` will be called
    /// automatically.
    pub fn from_zmq_url_with_listener(
        address: &str,
    ) -> Result<(Self, Option<std::net::TcpListener>), anyhow::Error> {
        // Check for Alias format: dial_to_url@bind_to_url
        // The @ character separates two valid ZMQ URLs.
        if let Some(at_pos) = address
            .find('@')
            .filter(|&pos| address[..pos].starts_with("tcp://"))
        {
            let dial_to_str = &address[..at_pos];
            let bind_to_str = &address[at_pos + 1..];

            // Validate that both addresses use TCP scheme
            if !dial_to_str.starts_with("tcp://") {
                return Err(anyhow::anyhow!(
                    "alias format is only supported for TCP addresses, got dial_to: {}",
                    dial_to_str
                ));
            }
            if !bind_to_str.starts_with("tcp://") {
                return Err(anyhow::anyhow!(
                    "alias format is only supported for TCP addresses, got bind_to: {}",
                    bind_to_str
                ));
            }

            let dial_to = Self::from_zmq_url(dial_to_str)?;
            let bind_to = Self::from_zmq_url(bind_to_str)?;

            return Ok((
                Self::Alias {
                    dial_to: Box::new(dial_to),
                    bind_to: Box::new(bind_to),
                },
                None,
            ));
        }

        // Try ZMQ-style URL format first (scheme://...)
        let (scheme, address) = address.split_once("://").ok_or_else(|| {
            anyhow::anyhow!("address must be in url form scheme://endppoint {}", address)
        })?;

        match scheme {
            "tcp" => {
                let (host, port, listener) = Self::parse_host_port_or_fd(address)?;
                let socket_addr = if host == "*" {
                    SocketAddr::new("::".parse().unwrap(), port)
                } else {
                    Self::resolve_hostname_to_socket_addr(host, port)?
                };
                Ok((Self::Tcp(socket_addr), listener))
            }
            "inproc" => {
                let port = address.parse::<u64>().map_err(|_| {
                    anyhow::anyhow!("inproc endpoint must be a valid port number: {}", address)
                })?;
                Ok((Self::Local(port), None))
            }
            "ipc" => Ok((Self::Unix(net::unix::SocketAddr::from_str(address)?), None)),
            "metatls" | "tls" | "quic" | "metaquic" => {
                let (host, port, listener) = Self::parse_host_port_or_fd(address)?;
                let hostname = if host == "*" {
                    std::net::Ipv6Addr::UNSPECIFIED.to_string()
                } else {
                    host.to_string()
                };
                let addr = match scheme {
                    "metatls" => Self::MetaTls(TlsAddr::new(hostname, port)),
                    "metaquic" => Self::MetaQuic(TlsAddr::new(hostname, port)),
                    "quic" => Self::Quic(TlsAddr::new(hostname, port)),
                    _ => Self::Tls(TlsAddr::new(hostname, port)),
                };
                Ok((addr, listener))
            }
            scheme => Err(anyhow::anyhow!("unsupported ZMQ scheme: {}", scheme)),
        }
    }

    /// Parse host:port where the port may be either a numeric port or `fdNNN`
    /// referencing a pre-opened file descriptor. Returns (host, resolved_port, optional_listener).
    fn parse_host_port_or_fd(
        address: &str,
    ) -> Result<(&str, u16, Option<std::net::TcpListener>), anyhow::Error> {
        let (host, port_str) = address
            .rsplit_once(':')
            .ok_or_else(|| anyhow::anyhow!("invalid address format: {}", address))?;

        if let Some(fd_str) = port_str.strip_prefix("fd") {
            let fd_num: RawFd = fd_str
                .parse()
                .map_err(|_| anyhow::anyhow!("invalid file descriptor number: {}", port_str))?;
            // Ensure the socket is in listening state. This is a no-op if
            // listen() was already called, and required if only bind() was done.
            // Safety: fd_num is valid and we are about to take ownership of it.
            let borrowed = unsafe { std::os::unix::io::BorrowedFd::borrow_raw(fd_num) };
            nix::sys::socket::listen(&borrowed, nix::sys::socket::Backlog::new(128)?)?;
            // Safety: caller guarantees the fd is a valid bound TCP socket.
            let std_listener = unsafe { std::net::TcpListener::from_raw_fd(fd_num) };
            let local_addr = std_listener.local_addr()?;
            Ok((host, local_addr.port(), Some(std_listener)))
        } else {
            let port: u16 = port_str
                .parse()
                .map_err(|_| anyhow::anyhow!("invalid port: {}", port_str))?;
            Ok((host, port, None))
        }
    }

    /// Render as a ZMQ-style URL, the inverse of [`from_zmq_url`](Self::from_zmq_url).
    pub fn to_zmq_url(&self) -> String {
        match self {
            Self::Tcp(addr) => format!("tcp://{}", addr),
            Self::MetaTls(addr) => format!("metatls://{}:{}", addr.hostname, addr.port),
            Self::Tls(addr) => format!("tls://{}:{}", addr.hostname, addr.port),
            Self::Quic(addr) => format!("quic://{}:{}", addr.hostname, addr.port),
            Self::MetaQuic(addr) => format!("metaquic://{}:{}", addr.hostname, addr.port),
            Self::Local(index) => format!("inproc://{}", index),
            Self::Unix(addr) => format!("ipc://{}", addr),
            Self::Alias { dial_to, bind_to } => {
                format!("{}@{}", dial_to.to_zmq_url(), bind_to.to_zmq_url())
            }
        }
    }

    /// Resolve hostname to SocketAddr, handling both IP addresses and hostnames
    fn resolve_hostname_to_socket_addr(host: &str, port: u16) -> Result<SocketAddr, anyhow::Error> {
        // Handle IPv6 addresses in brackets by stripping the brackets
        let host_clean = if host.starts_with('[') && host.ends_with(']') {
            &host[1..host.len() - 1]
        } else {
            host
        };

        // First try to parse as an IP address directly
        if let Ok(ip_addr) = host_clean.parse::<IpAddr>() {
            return Ok(SocketAddr::new(ip_addr, port));
        }

        // If not an IP, try hostname resolution
        use std::net::ToSocketAddrs;
        let mut addrs = (host_clean, port)
            .to_socket_addrs()
            .map_err(|e| anyhow::anyhow!("failed to resolve hostname '{}': {}", host_clean, e))?;

        addrs
            .next()
            .ok_or_else(|| anyhow::anyhow!("no addresses found for hostname '{}'", host_clean))
    }
}

/// Universal channel transmitter. Manages the link state, reconnections,
/// etc. on top of a [`net::Link`].
pub struct ChannelTx<M: RemoteMessage> {
    sender: mpsc::UnboundedSender<(M, CompletionSink<M>, Instant)>,
    dest: ChannelAddr,
    status: watch::Receiver<TxStatus>,
}

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

#[async_trait]
impl<M: RemoteMessage> Tx<M> for ChannelTx<M> {
    // HOT PATH: Be mindful of performance when making changes here.
    // To test how performance is affected by a change, run the RPC benchmarks in
    // `monarch/python/benches/`.
    fn do_post(&self, message: M, completion: CompletionSink<M>) {
        tracing::trace!(
            name = "post",
            dest = %self.dest,
            "sending message"
        );

        if let Err(mpsc::error::SendError((message, completion, _))) =
            self.sender.send((message, completion, Instant::now()))
        {
            let reason = self
                .status
                .borrow()
                .as_closed()
                .map(|r| SendErrorReason::Other(r.to_string()));
            completion.reject(SendError {
                error: ChannelError::Closed,
                message,
                reason,
            });
        }
    }

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

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

/// Universal channel receiver.
pub struct ChannelRx<M: RemoteMessage> {
    receiver: mpsc::Receiver<M>,
    dest: ChannelAddr,
    server: net::ServerHandle,
}

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

impl<M: RemoteMessage> ChannelRx<M> {
    /// Stop the channel server, tagging the log with what triggered shutdown.
    fn stop(&self, trigger: &str) {
        self.server.stop(&format!(
            "ChannelRx {trigger}; channel address: {}",
            self.dest
        ));
    }
}

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

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

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

impl<M: RemoteMessage> Drop for ChannelRx<M> {
    fn drop(&mut self) {
        self.stop("dropped");
    }
}

/// Dial the provided address, returning the corresponding Tx, or error
/// if the channel cannot be established. The underlying connection is
/// dropped whenever the returned Tx is dropped.
#[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `ChannelError`.
#[track_caller]
pub fn dial<M: RemoteMessage>(addr: ChannelAddr) -> Result<ChannelTx<M>, ChannelError> {
    let addr = addr.into_dial_addr();
    tracing::debug!(name = "dial", caller = %Location::caller(), %addr, "dialing channel {}", addr);
    Ok(net::spawn::<M>(net::link(
        addr,
        net::SessionId::random(),
        0,
        net::ProtocolKind::Simplex,
    )?))
}

/// Channels that may deliver messages out of send order.
pub mod unordered {
    use super::*;

    /// Dial with out-of-order delivery and N parallel streams.
    ///
    /// Opens N links sharing a single `SessionId` (distinct
    /// `stream_id` in `1..=num_streams` so the server routes them
    /// through the multi-stream receive path). Frames are
    /// load-balanced across streams via a shared MPMC work queue —
    /// idle writers pull next.
    ///
    /// # Semantics (how this differs from [`super::dial`])
    ///
    /// Multi-stream trades several of [`super::dial`]'s delivery
    /// guarantees for aggregate bandwidth. Use [`super::dial`] when
    /// any of these matter:
    ///
    /// - **Ordering.** Messages are delivered to the receiver in
    ///   *arrival order across streams*, not send order. Two messages
    ///   posted back to back on the sender may reach the receiver out
    ///   of order when they're carried on different streams.
    ///   [`super::dial`] is strictly in-order.
    ///
    /// - **Retransmission on reconnect.** If a writer's connection
    ///   drops after the bytes of a message have been written but
    ///   before the peer acks, that message is **not retransmitted**
    ///   on the new connection. It may be silently lost.
    ///   [`super::dial`] re-sends all unacked messages on reconnect.
    ///
    /// - **Delivery timeouts.** No delivery timeout is enforced. On
    ///   sustained peer outage, writers reconnect indefinitely
    ///   (backoff capped at 5s); senders block in `send().await` with
    ///   no bound. [`super::dial`] fails unacked sends after
    ///   `MESSAGE_DELIVERY_TIMEOUT`.
    ///
    /// - **`Tx::send` return semantics.** On session shutdown,
    ///   messages still in the unacked buffer have their
    ///   `return_channel` dropped. Per the [`Tx::send`] contract,
    ///   this makes `send().await` return `Ok(())` for messages that
    ///   may never have been delivered — i.e. a "success" return does
    ///   not actually confirm delivery here. [`super::dial`] delivers
    ///   a structured `SendError` instead.
    ///
    /// # Ack semantics
    ///
    /// Receivers ack a cumulative watermark: the highest `N` such
    /// that all of `0..=N` have been observed across all streams.
    /// Acks may stall behind a single missing seq on a slow stream.
    #[track_caller]
    pub fn dial<M: RemoteMessage>(
        addr: ChannelAddr,
        num_streams: usize,
    ) -> Result<ChannelTx<M>, ChannelError> {
        assert!(num_streams > 0);
        let addr = addr.into_dial_addr();
        let session_id = net::SessionId::random();
        let links: Vec<net::NetLink> = (1..=num_streams)
            .map(|i| {
                net::link(
                    addr.clone(),
                    session_id,
                    i as u8,
                    net::ProtocolKind::Simplex,
                )
            })
            .collect::<Result<_, _>>()?;
        Ok(net::spawn_unordered::<M>(links))
    }

    /// Serve a receiver that accepts unordered senders.
    ///
    /// The network server already routes multi-stream sessions by
    /// `SessionId`, so unordered serving uses the same listener as the
    /// ordered channel API.
    #[track_caller]
    pub fn serve<M: RemoteMessage>(
        addr: ChannelAddr,
    ) -> Result<(ChannelAddr, ChannelRx<M>), ChannelError> {
        super::serve(addr)
    }

    /// Serve with an optional pre-opened listener.
    ///
    /// Pre-opened listeners are only supported for TCP-based transports,
    /// matching [`super::serve_with_listener`].
    #[track_caller]
    pub fn serve_with_listener<M: RemoteMessage>(
        addr: ChannelAddr,
        listener: Option<std::net::TcpListener>,
    ) -> Result<(ChannelAddr, ChannelRx<M>), ChannelError> {
        super::serve_with_listener(addr, listener)
    }
}

/// Serve on the provided channel address. The server is turned down
/// when the returned Rx is dropped.
#[track_caller]
pub fn serve<M: RemoteMessage>(
    addr: ChannelAddr,
) -> Result<(ChannelAddr, ChannelRx<M>), ChannelError> {
    serve_with_listener(addr, None)
}

/// Serve on the provided channel address, optionally using a pre-opened TCP listener.
/// When `listener` is `Some`, the provided listener is used instead of binding a new socket.
/// The server is turned down when the returned Rx is dropped.
#[track_caller]
pub fn serve_with_listener<M: RemoteMessage>(
    addr: ChannelAddr,
    listener: Option<std::net::TcpListener>,
) -> Result<(ChannelAddr, ChannelRx<M>), ChannelError> {
    let caller = Location::caller();
    serve_inner(addr, listener).map(|(addr, rx)| {
        tracing::debug!(
            name = "serve",
            %addr,
            %caller,
        );
        (addr, rx)
    })
}

/// Serve a muxed listener on `addr`. Simplex clients (dialed via
/// [`channel::dial`](dial)) deliver into the bundled [`ChannelRx<M>`];
/// duplex clients (dialed via [`channel::duplex::dial`](net::duplex::dial))
/// populate the bundled [`DuplexServer<In, Out>`]. Only net transports
/// are supported; the caller picks transports that implement both
/// protocol styles.
///
/// The returned [`MuxServer`] owns both halves and a shared
/// [`MuxShutdown`]. Dropping or stopping any of these tears down the
/// listener and the other half together — see [`MuxServer`] for the
/// full lifecycle contract.
#[track_caller]
pub fn serve_mux<M: RemoteMessage, In: RemoteMessage, Out: RemoteMessage>(
    addr: ChannelAddr,
    prebound_listener: Option<std::net::TcpListener>,
) -> Result<MuxServer<M, In, Out>, ChannelError> {
    if !addr.transport().is_net() {
        return Err(ChannelError::InvalidAddress(format!(
            "serve_mux requires a net transport; got {}",
            addr
        )));
    }
    let parts = net::mux::serve::<M, In, Out>(addr, prebound_listener)?;
    Ok(MuxServer {
        addr: parts.addr,
        simplex: parts.simplex,
        duplex: parts.duplex,
        shutdown: MuxShutdown {
            join_handle: parts.join_handle,
            cancel: parts.cancel,
        },
    })
}

/// A muxed server bundling a simplex receiver, a duplex accept
/// server, and the shared shutdown signal that ties them together.
///
/// All three components share one underlying listener. Lifecycle:
///
/// - Dropping the [`MuxServer`] cancels the shared shutdown and
///   tears down the listener and any in-flight sessions.
/// - [`MuxServer::stop`] does the same explicitly.
/// - [`MuxServer::split`] hands out the address, simplex half,
///   duplex half, and a [`MuxShutdown`] separately. After splitting,
///   dropping the simplex half, the duplex half, or the
///   [`MuxShutdown`] guard cancels the shared shutdown. The address
///   is a plain value and does not own any resources.
///
/// The simplex half is a [`ChannelRx<M>`] you `recv()` on; the duplex
/// half is a [`DuplexServer<In, Out>`] you `accept()` on. Neither is
/// `join()`-able — there is no separate per-half task to await.
pub struct MuxServer<M: RemoteMessage, In: RemoteMessage, Out: RemoteMessage> {
    addr: ChannelAddr,
    simplex: ChannelRx<M>,
    duplex: net::duplex::DuplexServer<In, Out>,
    shutdown: MuxShutdown,
}

impl<M: RemoteMessage, In: RemoteMessage, Out: RemoteMessage> MuxServer<M, In, Out> {
    /// The address the muxed listener is bound to.
    pub fn addr(&self) -> &ChannelAddr {
        &self.addr
    }

    /// Borrow the simplex receiver.
    pub fn simplex_mut(&mut self) -> &mut ChannelRx<M> {
        &mut self.simplex
    }

    /// Borrow the duplex accept server.
    pub fn duplex_mut(&mut self) -> &mut net::duplex::DuplexServer<In, Out> {
        &mut self.duplex
    }

    /// Cancel the shared shutdown and tear down both halves.
    pub fn stop(&self, reason: &str) {
        self.shutdown.stop(reason);
    }

    /// Move the bound address, simplex half, duplex half, and a
    /// [`MuxShutdown`] guard out of this wrapper. After splitting,
    /// dropping the simplex half, the duplex half, or the
    /// [`MuxShutdown`] guard cancels the shared shutdown and tears
    /// the rest down. The address is a plain value and does not own
    /// any resources.
    pub fn split(
        self,
    ) -> (
        ChannelAddr,
        ChannelRx<M>,
        net::duplex::DuplexServer<In, Out>,
        MuxShutdown,
    ) {
        (self.addr, self.simplex, self.duplex, self.shutdown)
    }

    /// Wire up handlers for both halves, spawn a background task that
    /// owns the orderly shutdown (duplex drain → simplex pump → listener),
    /// and return a [`MailboxServerHandle`](crate::mailbox::MailboxServerHandle).
    /// Calling `.stop()` on the handle drives the drain.
    ///
    /// `simplex_handler` consumes the simplex receiver and produces
    /// the simplex pump's `MailboxServerHandle` (typically by calling
    /// [`MailboxServer::serve`](crate::mailbox::MailboxServer::serve)
    /// on a forwarder). `duplex_handler` receives the
    /// [`DuplexServer`](net::duplex::DuplexServer) and a stop signal,
    /// and returns the future driving the duplex pump.
    ///
    /// **Shutdown ordering.** On stop, the coordinator awaits
    /// `duplex_task` first so the duplex handler can drive its own
    /// internal drain (e.g., the host's per-connection forwarder
    /// pumps stop and drop their `AttachSender`s, which lets
    /// `send_connected` flush queued outbound naturally as
    /// `SendLoopError::AppClosed`, before the handler's terminal
    /// `duplex_server.stop` signals listener shutdown and `join`
    /// waits for it). Stopping the simplex pump and listener happens
    /// after the duplex drain to avoid cascading cancellation through
    /// `dispatch_duplex_stream`'s `select!` while the duplex side
    /// still has queued sends.
    ///
    /// Mirrors [`MailboxServer::serve`](crate::mailbox::MailboxServer::serve)'s
    /// shape: take the work to do, return a `MailboxServerHandle`.
    pub fn serve<SH, DH, DF>(
        self,
        simplex_handler: SH,
        duplex_handler: DH,
    ) -> crate::mailbox::MailboxServerHandle
    where
        SH: FnOnce(ChannelRx<M>) -> crate::mailbox::MailboxServerHandle,
        DH: FnOnce(net::duplex::DuplexServer<In, Out>, tokio::sync::watch::Receiver<bool>) -> DF,
        DF: std::future::Future<Output = ()> + Send + 'static,
    {
        let (stopped_tx, mut stopped_rx) = tokio::sync::watch::channel(false);
        let duplex_stop = stopped_rx.clone();
        let simplex_handle = simplex_handler(self.simplex);
        let duplex_task = tokio::spawn(duplex_handler(self.duplex, duplex_stop));
        let shutdown = self.shutdown;
        let join_handle = tokio::spawn(async move {
            // Pend forever if `stopped_tx` is silently dropped (caller
            // discarded the handle without `stop()`). Mirrors the
            // existing `MailboxServer::serve` behavior of holding the
            // server open absent an explicit stop signal — otherwise
            // we'd tear down the mux as soon as the handle drops.
            let ok = stopped_rx.wait_for(|stopped| *stopped).await.is_ok();
            if !ok {
                std::future::pending::<()>().await;
            }
            const REASON: &str = "MuxServer shutdown";
            // 1. Wait for the duplex handler to complete its own drain.
            //    The handler already saw `duplex_stop` (a clone of our
            //    `stopped_rx`) at the same instant we did, so it is
            //    already winding down. Awaiting before any cancel fires
            //    preserves the natural app-closed path through
            //    `send_connected` so queued outbound is not abandoned.
            let _ = duplex_task.await;
            // 2. The duplex handler's drop fires `cancel_token`, which
            //    cascades to simplex per-session cancels and closes
            //    `simplex_rx`; the simplex pump then exits naturally
            //    via its `rx.recv()` returning `Closed`. We only need
            //    to await it. Calling `simplex_handle.stop` here would
            //    race the natural exit and panic on the watch send if
            //    the pump's receiver has already dropped.
            let _ = simplex_handle.await;
            // 3. Backstop: cancel the listener explicitly (idempotent
            //    if already cancelled), then await the accept-loop's
            //    full drain.
            shutdown.stop(REASON);
            let _ = shutdown.await;
            Ok::<(), crate::mailbox::MailboxServerError>(())
        });
        crate::mailbox::MailboxServerHandle::from_parts(join_handle, stopped_tx)
    }
}

/// Awaitable shutdown handle for a [`MuxServer`]. Wraps the muxed
/// listener's accept-loop [`JoinHandle`](tokio::task::JoinHandle); the
/// caller signals teardown with [`stop`](Self::stop) and then `.await`s
/// the handle to confirm the listener task has fully exited. Drop
/// cancels the shared signal as a backstop, so a forgotten handle does
/// not leak the listener.
///
/// Mirrors the [`MailboxServerHandle`](crate::mailbox::MailboxServerHandle)
/// shape: signal stop, then await the handle.
pub struct MuxShutdown {
    join_handle: tokio::task::JoinHandle<Result<(), net::ServerError>>,
    cancel: CancellationToken,
}

impl MuxShutdown {
    /// Signal the muxed listener to stop accepting new connections and
    /// tear down. The caller should subsequently `.await` the handle
    /// to confirm shutdown.
    pub fn stop(&self, reason: &str) {
        tracing::info!(
            name = "MuxServerStatus",
            status = "Stop::Sent",
            reason,
            "muxed frontend stop signalled",
        );
        self.cancel.cancel();
    }

    /// Resolve when the shared shutdown has been cancelled (without
    /// awaiting the listener task itself). Useful for outer tasks that
    /// drive per-half pumps and need to wake on teardown.
    pub async fn cancelled(&self) {
        self.cancel.cancelled().await;
    }
}

impl std::future::Future for MuxShutdown {
    type Output =
        <tokio::task::JoinHandle<Result<(), net::ServerError>> as std::future::Future>::Output;

    fn poll(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Self::Output> {
        // `JoinHandle` is `Unpin`, so we can re-pin a mutable borrow of
        // it without unsafe pin projection. `MuxShutdown` is `Unpin` by
        // virtue of its `Unpin` fields, which lets us reach through the
        // outer `Pin`.
        std::pin::Pin::new(&mut self.join_handle).poll(cx)
    }
}

impl Drop for MuxShutdown {
    fn drop(&mut self) {
        self.cancel.cancel();
    }
}

fn serve_inner<M: RemoteMessage>(
    addr: ChannelAddr,
    listener: Option<std::net::TcpListener>,
) -> Result<(ChannelAddr, ChannelRx<M>), ChannelError> {
    match addr {
        ChannelAddr::Unix(_) => {
            assert!(
                listener.is_none(),
                "pre-opened listener not supported for Unix transport"
            );
            let (addr, rx) = net::server::serve::<M>(addr, listener)?;
            Ok((addr, rx))
        }
        ChannelAddr::Tcp(_)
        | ChannelAddr::Local(_)
        | ChannelAddr::Tls(_)
        | ChannelAddr::MetaTls(_)
        | ChannelAddr::Quic(_)
        | ChannelAddr::MetaQuic(_)
        // The `Alias` variant binds on its `bind_to` address but advertises
        // `dial_to`; `listen_with_prebound` resolves this, so it routes through
        // the same net serve path as the other TCP-based transports.
        | ChannelAddr::Alias { .. } => {
            let (addr, rx) = net::server::serve::<M>(addr, listener)?;
            Ok((addr, rx))
        }
    }
}

/// Serve on the local address. The server is turned down
/// when the returned Rx is dropped.
pub fn serve_local<M: RemoteMessage>() -> (ChannelAddr, ChannelRx<M>) {
    serve::<M>(ChannelAddr::Local(0)).expect("fresh local stream port must bind")
}

/// Reserve a local channel address that can be served later.
///
/// Local channels are backed by a process-local port registry, so reserving a
/// concrete address is a synchronous allocation that does not bind an OS
/// listener. Gateways use this to have a stable advertised local location
/// immediately, including when the process-wide gateway is initialized from a
/// [`std::sync::OnceLock`]. Serving is a separate step that binds the reserved
/// port to a receiver.
///
/// Network transports do not have an equivalent reservation API here: their
/// concrete addresses come from binding sockets and starting the corresponding
/// channel server.
pub fn reserve_local_addr() -> ChannelAddr {
    ChannelAddr::Local(local::reserve())
}

#[cfg(test)]
mod tests {
    use std::assert_matches;
    use std::collections::HashSet;
    use std::net::IpAddr;
    use std::net::Ipv4Addr;
    use std::net::Ipv6Addr;
    use std::time::Duration;

    use rand::RngExt as _;
    use rand::distr::Uniform;
    use tokio::task::JoinSet;

    use super::net::*;
    use super::*;
    #[test]
    fn test_channel_addr() {
        let cases_ok = vec![
            (
                "tcp<DELIM>[::1]:1234",
                ChannelAddr::Tcp(SocketAddr::new(
                    IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)),
                    1234,
                )),
            ),
            (
                "tcp<DELIM>127.0.0.1:8080",
                ChannelAddr::Tcp(SocketAddr::new(
                    IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
                    8080,
                )),
            ),
            (
                "quic<DELIM>example.com:443",
                ChannelAddr::Quic(TlsAddr::new("example.com", 443)),
            ),
            (
                "metaquic<DELIM>example.com:443",
                ChannelAddr::MetaQuic(TlsAddr::new("example.com", 443)),
            ),
            #[cfg(target_os = "linux")]
            ("local<DELIM>123", ChannelAddr::Local(123)),
            (
                "unix<DELIM>@yolo",
                ChannelAddr::Unix(
                    unix::SocketAddr::from_abstract_name("yolo")
                        .expect("can't make socket from abstract name"),
                ),
            ),
            (
                "unix<DELIM>/cool/socket-path",
                ChannelAddr::Unix(
                    unix::SocketAddr::from_pathname("/cool/socket-path")
                        .expect("can't make socket from path"),
                ),
            ),
        ];

        for (raw, parsed) in cases_ok {
            for delim in ["!", ":"] {
                let raw = raw.replace("<DELIM>", delim);
                assert_eq!(raw.parse::<ChannelAddr>().unwrap(), parsed);
            }
        }

        let cases_err = vec![
            ("tcp:abcdef..123124", "invalid socket address syntax"),
            ("xxx:foo", "no such channel type: xxx"),
            ("127.0.0.1", "no channel type specified"),
            ("local:abc", "invalid digit found in string"),
        ];

        for (raw, error) in cases_err {
            let Err(err) = raw.parse::<ChannelAddr>() else {
                panic!("expected error parsing: {}", raw)
            };
            assert_eq!(format!("{}", err), error);
        }
    }

    #[test]
    fn test_zmq_style_channel_addr() {
        // Test TCP addresses
        assert_eq!(
            ChannelAddr::from_zmq_url("tcp://127.0.0.1:8080").unwrap(),
            ChannelAddr::Tcp("127.0.0.1:8080".parse().unwrap())
        );

        // Test TCP wildcard binding
        assert_eq!(
            ChannelAddr::from_zmq_url("tcp://*:5555").unwrap(),
            ChannelAddr::Tcp("[::]:5555".parse().unwrap())
        );

        // Test inproc (maps to local with numeric endpoint)
        assert_eq!(
            ChannelAddr::from_zmq_url("inproc://12345").unwrap(),
            ChannelAddr::Local(12345)
        );

        // Test ipc (maps to unix)
        assert_eq!(
            ChannelAddr::from_zmq_url("ipc:///tmp/my-socket").unwrap(),
            ChannelAddr::Unix(unix::SocketAddr::from_pathname("/tmp/my-socket").unwrap())
        );

        // Test metatls with hostname
        assert_eq!(
            ChannelAddr::from_zmq_url("metatls://example.com:443").unwrap(),
            ChannelAddr::MetaTls(TlsAddr::new("example.com", 443))
        );

        // Test metatls with IP address (should be normalized)
        assert_eq!(
            ChannelAddr::from_zmq_url("metatls://192.168.1.1:443").unwrap(),
            ChannelAddr::MetaTls(TlsAddr::new("192.168.1.1", 443))
        );

        // Test quic with hostname
        assert_eq!(
            ChannelAddr::from_zmq_url("quic://example.com:443").unwrap(),
            ChannelAddr::Quic(TlsAddr::new("example.com", 443))
        );

        // Test quic wildcard binding
        assert_eq!(
            ChannelAddr::from_zmq_url("quic://*:8443").unwrap(),
            ChannelAddr::Quic(TlsAddr::new("::", 8443))
        );

        // Test metaquic with hostname
        assert_eq!(
            ChannelAddr::from_zmq_url("metaquic://example.com:443").unwrap(),
            ChannelAddr::MetaQuic(TlsAddr::new("example.com", 443))
        );

        // Test metaquic wildcard binding
        assert_eq!(
            ChannelAddr::from_zmq_url("metaquic://*:8443").unwrap(),
            ChannelAddr::MetaQuic(TlsAddr::new("::", 8443))
        );

        // Test metatls with wildcard (should use IPv6 unspecified address)
        assert_eq!(
            ChannelAddr::from_zmq_url("metatls://*:8443").unwrap(),
            ChannelAddr::MetaTls(TlsAddr::new("::", 8443))
        );

        // Test TCP hostname resolution (should resolve hostname to IP)
        // Note: This test may fail in environments without proper DNS resolution
        // We test that it at least doesn't fail to parse
        let tcp_hostname_result = ChannelAddr::from_zmq_url("tcp://localhost:8080");
        assert!(tcp_hostname_result.is_ok());

        // Test IPv6 address
        assert_eq!(
            ChannelAddr::from_zmq_url("tcp://[::1]:1234").unwrap(),
            ChannelAddr::Tcp("[::1]:1234".parse().unwrap())
        );

        // Test error cases
        assert!(ChannelAddr::from_zmq_url("invalid://scheme").is_err());
        assert!(ChannelAddr::from_zmq_url("tcp://invalid-port").is_err());
        assert!(ChannelAddr::from_zmq_url("metatls://no-port").is_err());
        assert!(ChannelAddr::from_zmq_url("inproc://not-a-number").is_err());

        // IPv6 normalization: leading zeros are stripped
        assert_eq!(
            ChannelAddr::from_zmq_url("metatls://2a03:83e4:5000:c000:56d7:00cf:75ce:144a:443")
                .unwrap(),
            ChannelAddr::MetaTls(TlsAddr::new("2a03:83e4:5000:c000:56d7:cf:75ce:144a", 443))
        );

        // Short and long forms of the same IPv6 produce equal ChannelAddr values
        assert_eq!(
            ChannelAddr::from_zmq_url("metatls://2a03:83e4:5000:c000:56d7:00cf:75ce:144a:443")
                .unwrap(),
            ChannelAddr::from_zmq_url("metatls://2a03:83e4:5000:c000:56d7:cf:75ce:144a:443")
                .unwrap(),
        );

        // Bracketed IPv6 is normalized
        assert_eq!(
            ChannelAddr::from_zmq_url("metatls://[::1]:443").unwrap(),
            ChannelAddr::MetaTls(TlsAddr::new("::1", 443))
        );

        // Same tests for tls://
        assert_eq!(
            ChannelAddr::from_zmq_url("tls://2a03:83e4:5000:c000:56d7:00cf:75ce:144a:443").unwrap(),
            ChannelAddr::Tls(TlsAddr::new("2a03:83e4:5000:c000:56d7:cf:75ce:144a", 443))
        );
        assert_eq!(
            ChannelAddr::from_zmq_url("tls://2a03:83e4:5000:c000:56d7:00cf:75ce:144a:443").unwrap(),
            ChannelAddr::from_zmq_url("tls://2a03:83e4:5000:c000:56d7:cf:75ce:144a:443").unwrap(),
        );
        assert_eq!(
            ChannelAddr::from_zmq_url("tls://[::1]:443").unwrap(),
            ChannelAddr::Tls(TlsAddr::new("::1", 443))
        );
    }

    #[tokio::test]
    async fn test_reserved_local_addr_can_be_served() {
        let addr = reserve_local_addr();
        assert!(dial::<u64>(addr.clone()).is_err());

        let (bound_addr, mut rx) = serve::<u64>(addr.clone()).unwrap();
        assert_eq!(bound_addr, addr);

        let tx = dial::<u64>(addr.clone()).unwrap();
        tx.post(123);
        assert_eq!(rx.recv().await.unwrap(), 123);
        rx.join().await;

        let (rebound_addr, _rx) = serve::<u64>(addr.clone()).unwrap();
        assert_eq!(rebound_addr, addr);
    }

    #[test]
    fn test_normalize_host() {
        // Plain IPv4 passes through
        assert_eq!(normalize_host("192.168.1.1"), "192.168.1.1");

        // Plain hostname passes through
        assert_eq!(normalize_host("example.com"), "example.com");

        // IPv6 with leading zeros gets normalized
        assert_eq!(
            normalize_host("2a03:83e4:5000:c000:56d7:00cf:75ce:144a"),
            "2a03:83e4:5000:c000:56d7:cf:75ce:144a"
        );

        // Bracketed IPv6 is stripped and normalized
        assert_eq!(normalize_host("[::1]"), "::1");

        // Without bracket stripping, IpAddr::from_str rejects bracketed
        // addresses. This demonstrates that the bracket stripping in
        // normalize_host is necessary.
        assert!("[::1]".parse::<IpAddr>().is_err());
    }

    #[test]
    fn test_find_routable_address_skips_loopback() {
        let addresses = [
            "127.0.1.1".parse::<IpAddr>().unwrap(),
            "169.254.1.1".parse::<IpAddr>().unwrap(),
            "::1".parse::<IpAddr>().unwrap(),
            "fe80::1".parse::<IpAddr>().unwrap(),
            "192.0.2.10".parse::<IpAddr>().unwrap(),
        ];

        assert_eq!(
            find_routable_address(&addresses),
            Some("192.0.2.10".parse::<IpAddr>().unwrap())
        );
    }

    #[test]
    fn test_find_routable_address_rejects_unroutable_addresses() {
        let addresses = [
            "127.0.1.1".parse::<IpAddr>().unwrap(),
            "169.254.1.1".parse::<IpAddr>().unwrap(),
            "::1".parse::<IpAddr>().unwrap(),
            "fe80::1".parse::<IpAddr>().unwrap(),
        ];

        assert_eq!(find_routable_address(&addresses), None);
    }

    #[test]
    fn test_zmq_style_alias_channel_addr() {
        // Test Alias format: dial_to_url@bind_to_url
        // The format is: dial_to_url@bind_to_url where both are valid ZMQ URLs
        // Note: Alias format is only supported for TCP addresses

        // Test Alias with tcp on both sides
        let alias_addr = ChannelAddr::from_zmq_url("tcp://127.0.0.1:9000@tcp://[::]:8800").unwrap();
        match alias_addr {
            ChannelAddr::Alias { dial_to, bind_to } => {
                assert_eq!(
                    *dial_to,
                    ChannelAddr::Tcp("127.0.0.1:9000".parse().unwrap())
                );
                assert_eq!(*bind_to, ChannelAddr::Tcp("[::]:8800".parse().unwrap()));
            }
            _ => panic!("Expected Alias"),
        }

        // Non-tcp left side: alias branch is skipped, parsed as regular address.
        // metatls:// with garbage host is not an alias.
        let non_alias = ChannelAddr::from_zmq_url("metatls://example.com:443@tcp://127.0.0.1:8080");
        assert!(
            !matches!(non_alias, Ok(ChannelAddr::Alias { .. })),
            "non-tcp left side must not produce Alias"
        );

        // Test error: alias with non-tcp bind_to (not supported)
        assert!(
            ChannelAddr::from_zmq_url("tcp://127.0.0.1:8080@metatls://example.com:443").is_err()
        );

        // Test error: invalid scheme falls through to scheme parsing, errors there
        assert!(ChannelAddr::from_zmq_url("invalid://scheme@tcp://127.0.0.1:8080").is_err());

        // Test error: invalid bind_to URL in Alias
        assert!(ChannelAddr::from_zmq_url("tcp://127.0.0.1:8080@invalid://scheme").is_err());

        // Test error: missing port in dial_to
        assert!(ChannelAddr::from_zmq_url("tcp://host@tcp://127.0.0.1:8080").is_err());

        // Test error: missing port in bind_to
        assert!(ChannelAddr::from_zmq_url("tcp://127.0.0.1:8080@tcp://example.com").is_err());
    }

    #[tokio::test]
    async fn test_multiple_connections() {
        for addr in ChannelTransport::all().map(ChannelAddr::any) {
            let (listen_addr, mut rx) = crate::channel::serve::<u64>(addr).unwrap();

            let mut sends: JoinSet<()> = JoinSet::new();
            for message in 0u64..100u64 {
                let addr = listen_addr.clone();
                sends.spawn(async move {
                    let tx = dial::<u64>(addr).unwrap();
                    tx.post(message);
                });
            }

            let mut received: HashSet<u64> = HashSet::new();
            while received.len() < 100 {
                received.insert(rx.recv().await.unwrap());
            }

            for message in 0u64..100u64 {
                assert!(received.contains(&message));
            }

            loop {
                match sends.join_next().await {
                    Some(Ok(())) => (),
                    Some(Err(err)) => panic!("{}", err),
                    None => break,
                }
            }
        }
    }

    #[tokio::test]
    async fn test_server_close() {
        for addr in ChannelTransport::all().map(ChannelAddr::any) {
            if net::is_net_addr(&addr) {
                // Net has store-and-forward semantics. We don't expect failures
                // on closure.
                continue;
            }

            let (listen_addr, rx) = crate::channel::serve::<u64>(addr).unwrap();

            let tx = dial::<u64>(listen_addr).unwrap();
            tx.post(123);
            drop(rx);

            // New transmits should fail... but there is buffering, etc.,
            // which can cause the failure to be delayed. We give it
            // a deadline, but it can still technically fail -- the test
            // should be considered a kind of integration test.
            let start = tokio::time::Instant::now();

            let result = loop {
                let result = tx.try_post(123).await;

                if result.is_err() || start.elapsed() > Duration::from_secs(10) {
                    break result;
                }
            };
            assert_matches!(
                result,
                Err(SendError {
                    error: ChannelError::Closed,
                    message: 123,
                    reason: None
                })
            );
        }
    }

    fn addrs() -> Vec<ChannelAddr> {
        let rng = rand::rng();
        let uniform = Uniform::new_inclusive('a', 'z').unwrap();
        vec![
            "tcp:[::1]:0".parse().unwrap(),
            "local:0".parse().unwrap(),
            #[cfg(target_os = "linux")]
            "unix:".parse().unwrap(),
            #[cfg(target_os = "linux")]
            format!(
                "unix:@{}",
                rng.sample_iter(uniform).take(10).collect::<String>()
            )
            .parse()
            .unwrap(),
        ]
    }

    #[test]
    fn test_bind_spec_from_str() {
        // Test parsing ChannelTransport strings -> BindSpec::Any
        assert_eq!(
            BindSpec::from_str("tcp").unwrap(),
            BindSpec::Any(ChannelTransport::Tcp(TcpMode::Hostname))
        );
        assert_eq!(
            BindSpec::from_str("metatls(Hostname)").unwrap(),
            BindSpec::Any(ChannelTransport::MetaTls(TlsMode::Hostname))
        );

        // Test parsing ChannelAddr strings -> BindSpec::Addr
        assert_eq!(
            BindSpec::from_str("tcp:127.0.0.1:8080").unwrap(),
            BindSpec::Addr(ChannelAddr::Tcp("127.0.0.1:8080".parse().unwrap()))
        );

        // Test parsing ZMQ URL format -> BindSpec::Addr
        assert_eq!(
            BindSpec::from_str("tcp://127.0.0.1:9000").unwrap(),
            BindSpec::Addr(ChannelAddr::Tcp("127.0.0.1:9000".parse().unwrap()))
        );
        assert_eq!(
            BindSpec::from_str("tcp://127.0.0.1:9000@tcp://[::1]:7200").unwrap(),
            BindSpec::Addr(
                ChannelAddr::from_zmq_url("tcp://127.0.0.1:9000@tcp://[::1]:7200").unwrap()
            )
        );

        // Test error cases
        assert!(BindSpec::from_str("invalid_spec").is_err());
        assert!(BindSpec::from_str("unknown://scheme").is_err());
        assert!(BindSpec::from_str("").is_err());
    }

    #[tokio::test]
    // TODO: OSS: called `Result::unwrap()` on an `Err` value: Server(Listen(Tcp([::1]:0), Os { code: 99, kind: AddrNotAvailable, message: "Cannot assign requested address" }))
    #[cfg_attr(not(fbcode_build), ignore)]
    async fn test_dial_serve() {
        for addr in addrs() {
            let (listen_addr, mut rx) = crate::channel::serve::<i32>(addr).unwrap();
            let tx = crate::channel::dial(listen_addr).unwrap();
            tx.post(123);
            assert_eq!(rx.recv().await.unwrap(), 123);
        }
    }

    #[tokio::test]
    // TODO: OSS: called `Result::unwrap()` on an `Err` value: Server(Listen(Tcp([::1]:0), Os { code: 99, kind: AddrNotAvailable, message: "Cannot assign requested address" }))
    #[cfg_attr(not(fbcode_build), ignore)]
    async fn test_serve_alias_advertises_dial_to() {
        // Reserve an ephemeral port, then release it so the alias can bind to
        // it via `bind_to`.
        let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
        let port = probe.local_addr().unwrap().port();
        drop(probe);

        // `dial_to` advertises a reachable loopback address; `bind_to` listens
        // on the wildcard interface (the case that matters where the dial
        // address cannot be bound directly, e.g. behind NAT).
        let alias =
            ChannelAddr::from_zmq_url(&format!("tcp://127.0.0.1:{port}@tcp://0.0.0.0:{port}"))
                .unwrap();
        assert_matches!(alias, ChannelAddr::Alias { .. });

        let (listen_addr, mut rx) = crate::channel::serve::<i32>(alias).unwrap();

        // Serving an alias consumes it: the advertised address is the plain
        // `dial_to`. Remote peers independently construct this same `Tcp`
        // address from the dial string, so the proc namespace derived from it
        // must match. If serving left the address an `Alias`, that match would
        // fail and messages would not route to the server.
        assert_eq!(
            listen_addr,
            ChannelAddr::Tcp(format!("127.0.0.1:{port}").parse().unwrap()),
            "serving an alias must advertise dial_to, not the alias itself"
        );

        let tx = crate::channel::dial(listen_addr).unwrap();
        tx.post(123);
        assert_eq!(rx.recv().await.unwrap(), 123);
    }

    #[tokio::test]
    // TODO: OSS: called `Result::unwrap()` on an `Err` value: Server(Listen(Tcp([::1]:0), Os { code: 99, kind: AddrNotAvailable, message: "Cannot assign requested address" }))
    #[cfg_attr(not(fbcode_build), ignore)]
    async fn test_send() {
        let config = hyperactor_config::global::lock();

        // Use temporary config for this test
        let _guard1 = config.override_key(
            crate::config::MESSAGE_DELIVERY_TIMEOUT,
            Duration::from_secs(1),
        );
        let _guard2 = config.override_key(crate::config::MESSAGE_ACK_EVERY_N_MESSAGES, 1);
        for addr in addrs() {
            let (listen_addr, mut rx) = crate::channel::serve::<i32>(addr).unwrap();
            let tx = crate::channel::dial(listen_addr).unwrap();
            tx.send(123).await.unwrap();
            assert_eq!(rx.recv().await.unwrap(), 123);

            drop(rx);
            assert_matches!(
                tx.send(123).await.unwrap_err(),
                SendError {
                    error: ChannelError::Closed,
                    message: 123,
                    ..
                }
            );
        }
    }

    #[test]
    fn test_find_routable_address_skips_link_local_ipv6() {
        let link_local_v6: IpAddr = "fe80::1".parse().unwrap();
        let routable_v6: IpAddr = "2001:db8::1".parse().unwrap();
        let addrs = vec![link_local_v6, routable_v6];
        assert_eq!(find_routable_address(&addrs), Some(routable_v6));
    }

    #[test]
    fn test_find_routable_address_skips_link_local_ipv4() {
        let link_local_v4: IpAddr = "169.254.1.1".parse().unwrap();
        let routable_v4: IpAddr = "192.168.1.1".parse().unwrap();
        let addrs = vec![link_local_v4, routable_v4];
        assert_eq!(find_routable_address(&addrs), Some(routable_v4));
    }

    #[test]
    fn test_find_routable_address_returns_none_when_all_link_local() {
        let link_local_v6: IpAddr = "fe80::1".parse().unwrap();
        let link_local_v4: IpAddr = "169.254.1.1".parse().unwrap();
        let addrs = vec![link_local_v6, link_local_v4];
        assert_eq!(find_routable_address(&addrs), None);
    }

    #[test]
    fn test_find_routable_address_mixed() {
        let link_local_v6: IpAddr = "fe80::1".parse().unwrap();
        let link_local_v4: IpAddr = "169.254.0.1".parse().unwrap();
        let routable_v4: IpAddr = "10.0.0.1".parse().unwrap();
        let routable_v6: IpAddr = "2001:db8::2".parse().unwrap();

        // First routable address in list order should be returned.
        let addrs = vec![link_local_v6, link_local_v4, routable_v4, routable_v6];
        assert_eq!(find_routable_address(&addrs), Some(routable_v4));
    }
}