pea2pea 0.57.4

use std::{
    collections::HashMap,
    fmt,
    io::{self, ErrorKind},
    net::SocketAddr,
    ops::Deref,
    sync::{
        Arc,
        atomic::{AtomicBool, AtomicUsize, Ordering::*},
    },
    time::{Duration, Instant},
};

use futures_util::stream::{FuturesUnordered, StreamExt};
use parking_lot::Mutex;
use tokio::{
    net::{TcpListener, TcpSocket, TcpStream},
    sync::{Notify, RwLock, Semaphore, oneshot, watch},
    task::JoinHandle,
    time::{sleep, timeout},
};
use tracing::*;

#[cfg(doc)]
use crate::protocols::{Handshake, OnConnect, OnDisconnect, Reading, Writing};
use crate::{
    Config, Heuristics, Stats,
    connections::{
        Connection, ConnectionGuard, ConnectionInfo, ConnectionSide, Connections, DisconnectOrigin,
        create_connection_span,
    },
    protocols::{Protocol, Protocols},
};

// Starts the selected protocol handler for a new connection
macro_rules! enable_protocol {
    ($handler_type:ident, $node:expr, $conn:expr) => {
        if let Some(handler) = $node.protocols.$handler_type.get() {
            let (conn_returner, conn_retriever) = oneshot::channel();

            handler.trigger(($conn, conn_returner)).await;

            match crate::protocols::await_handler_response(conn_retriever, handler.closed()).await {
                Some(Ok(conn)) => conn,
                // the request was dropped by the handler's shutdown drain, the handler is
                // gone, or the message was stranded in its closed channel after an abort -
                // all of which mean the node is shutting down
                None => return Err(shutting_down_error()),
                Some(e) => return e,
            }
        } else {
            $conn
        }
    };
}

/// Waits until `drained` holds, (re)registering with `notify` before every check so that
/// a notification between the check and the await cannot be missed.
async fn wait_for_drain(notify: &Notify, drained: impl Fn() -> bool) {
    loop {
        let notified = notify.notified();
        tokio::pin!(notified);
        notified.as_mut().enable(); // register the waiter
        if drained() {
            break;
        }
        notified.await;
    }
}

/// The payload carried by every [`io::Error`] the library returns because the node is shutting
/// down - from [`Node::connect`], [`Node::toggle_listener`], or any other operation that can no
/// longer be admitted.
///
/// The [`io::ErrorKind`] of those errors is [`io::ErrorKind::Other`], since none of the standard
/// kinds means "this component is going away permanently". This type is therefore the contract:
/// match on it rather than on the message, which is not stable.
///
/// ```no_run
/// # use pea2pea::{Node, ShuttingDown};
/// # async fn f(node: &Node, addr: std::net::SocketAddr) {
/// if let Err(e) = node.connect(addr).await {
///     if ShuttingDown::caused(&e) {
///         // the node is gone for good; don't retry
///     }
/// }
/// # }
/// ```
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ShuttingDown;

impl fmt::Display for ShuttingDown {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("shutting down")
    }
}

impl std::error::Error for ShuttingDown {}

impl ShuttingDown {
    /// Whether `err` was produced because the node is shutting down.
    ///
    /// note: This inspects `err` itself, not a chain of sources - it recognizes the errors the
    /// library returns, not ones an application has wrapped further.
    pub fn caused(err: &io::Error) -> bool {
        err.get_ref().is_some_and(|inner| inner.is::<Self>())
    }
}

/// The error returned by every path that rejects work due to node shutdown.
pub(crate) fn shutting_down_error() -> io::Error {
    io::Error::other(ShuttingDown)
}

/// A sequential numeric identifier assigned to `Node`s that were not provided with a name.
static SEQUENTIAL_NODE_ID: AtomicUsize = AtomicUsize::new(0);

/// The types of long-running tasks supported by the Node. The declaration order is
/// immaterial (storage is keyed); the graceful wind-down order is the array in `shut_down`.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub(crate) enum NodeTask {
    Listener,
    OnDisconnect,
    Handshake,
    OnConnect,
    Reading,
    Writing,
}

/// The central object responsible for handling connections.
///
/// note: Due to the architecture of protocol handlers capturing the node, a reference cycle exists
/// that prevents the Node from being dropped automatically. You must call [`Node::shut_down`] when
/// you are finished with a node to ensure all background tasks are aborted and sockets are closed.
#[derive(Clone)]
pub struct Node(Arc<InnerNode>);

impl Deref for Node {
    type Target = Arc<InnerNode>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

/// The actual node object that gets wrapped in an Arc in the Node.
#[doc(hidden)]
pub struct InnerNode {
    /// The tracing span.
    span: Span,
    /// The node's configuration.
    config: Config,
    /// The node's current listening address.
    listening_addr: RwLock<Option<SocketAddr>>,
    /// Contains objects used by the protocols implemented by the node.
    pub(crate) protocols: Protocols,
    /// Contains objects related to the node's active connections.
    pub(crate) connections: Connections,
    /// Bounds the number of concurrent connection *setups* to [`Config::max_connecting`],
    /// covering inbound accepts **and** outbound connects from a single budget. Sharing the
    /// budget means heavy outbound dialing applies real backpressure to the inbound accept
    /// loop (all but one of the surplus connections wait in the kernel accept queue; a single
    /// already-accepted one waits for a permit in the accept loop).
    connecting_permits: Arc<Semaphore>,
    /// Collects statistics related to the node itself.
    stats: Stats,
    /// Node-level operational signals not attributable to any single connection.
    heuristics: Heuristics,
    /// The node's tasks.
    pub(crate) tasks: Mutex<HashMap<NodeTask, JoinHandle<()>>>,
    /// The node's shutdown lifecycle state.
    pub(crate) shutdown: ShutdownState,
}

/// The node's shutdown lifecycle state. The shutdown is two-phase: `begin` (a
/// synchronous flag, checked inside lock-held sections, which stops new work
/// from being admitted) fires first; `signal_handlers` (an asynchronous watch,
/// which makes the protocol handler tasks drain their queues and exit) fires
/// later, once `shut_down` has drained the active connections and the
/// in-flight `OnConnect` scheduling.
pub(crate) struct ShutdownState {
    /// Indicates whether the shutdown sequence has begun.
    flag: AtomicBool,
    /// Tells the protocol handler tasks to wind down gracefully.
    handler_signal: watch::Sender<bool>,
    /// The number of in-flight `OnConnect` scheduling tasks; `shut_down`
    /// waits for it to reach zero before signaling the handlers, so that a
    /// scheduling task delayed by executor lag cannot fire its trigger into
    /// an already-closed channel and silently skip the hook.
    sched_in_flight: AtomicUsize,
    /// Notifies `shut_down` when `sched_in_flight` reaches zero.
    sched_drained: Notify,
}

impl Default for ShutdownState {
    fn default() -> Self {
        Self {
            flag: Default::default(),
            handler_signal: watch::Sender::new(false),
            sched_in_flight: Default::default(),
            sched_drained: Notify::new(),
        }
    }
}

impl ShutdownState {
    /// Indicates whether the shutdown sequence has begun; no new work may be
    /// admitted once it returns `true`.
    pub(crate) fn is_underway(&self) -> bool {
        self.flag.load(Acquire)
    }

    /// Marks the beginning of the shutdown sequence (phase one).
    fn begin(&self) {
        self.flag.store(true, Release);
    }

    /// Tells the protocol handler tasks to drain and exit (phase two).
    fn signal_handlers(&self) {
        let _ = self.handler_signal.send(true);
    }

    /// Subscribes to the phase-two signal.
    pub(crate) fn handler_signal(&self) -> watch::Receiver<bool> {
        self.handler_signal.subscribe()
    }

    fn sched_started(&self) {
        self.sched_in_flight.fetch_add(1, AcqRel);
    }

    fn sched_finished(&self) {
        if self.sched_in_flight.fetch_sub(1, AcqRel) == 1 {
            self.sched_drained.notify_waiters();
        }
    }

    async fn wait_sched_drained(&self) {
        wait_for_drain(&self.sched_drained, || {
            self.sched_in_flight.load(Acquire) == 0
        })
        .await;
    }
}

impl Node {
    /// Creates a new [`Node`] using the given [`Config`].
    ///
    /// # Panics
    ///
    /// Panics if any of [`Config::max_connections`], [`Config::max_connections_per_ip`], or
    /// [`Config::max_connecting`] is `0`; such values would render the node inoperable.
    pub fn new(mut config: Config) -> Self {
        // reject limit values that would render the node inoperable; a zero for any of them
        // would otherwise wedge the accept loop or break protocol setup in ways that are much
        // harder to diagnose than an upfront panic
        assert!(
            config.max_connections != 0,
            "Config::max_connections must not be 0"
        );
        assert!(
            config.max_connections_per_ip != 0,
            "Config::max_connections_per_ip must not be 0"
        );
        assert!(
            config.max_connecting != 0,
            "Config::max_connecting must not be 0"
        );

        // if there is no pre-configured name, assign a sequential numeric identifier
        if config.name.is_none() {
            config.name = Some(SEQUENTIAL_NODE_ID.fetch_add(1, Relaxed).to_string());
        }

        // create a tracing span containing the node's name
        let span = create_span(config.name.as_deref().unwrap());

        // pending connections count towards `max_connections`, so any surplus of the
        // connection-setup budget could never be used; clamp it, which also keeps the
        // default `max_connecting` valid when only `max_connections` is lowered
        if config.max_connecting > config.max_connections {
            debug!(
                parent: &span,
                "clamping max_connecting ({}) to max_connections ({})",
                config.max_connecting, config.max_connections,
            );
            config.max_connecting = config.max_connections;
        }

        // shared budget for concurrent connection setups
        let connecting_permits = Arc::new(Semaphore::new(config.max_connecting as usize));

        let node = Node(Arc::new(InnerNode {
            span,
            config,
            listening_addr: Default::default(),
            shutdown: Default::default(),
            protocols: Default::default(),
            connections: Default::default(),
            connecting_permits,
            stats: Default::default(),
            heuristics: Default::default(),
            tasks: Default::default(),
        }));

        debug!(parent: node.span(), "the node is ready");

        node
    }

    /// Enables or disables listening for inbound connections. Returns `Ok(Some(_))` with the
    /// actual bound address when the listener was just enabled - it will differ from the one in
    /// [`Config::listener_addr`] if that one's port was unspecified (i.e. `0`) - and `Ok(None)`
    /// when it was just disabled.
    ///
    /// note: Once the node is shutting down this fails with a [`ShuttingDown`] payload; see that
    /// type for recognizing it. Enabling the listener when [`Config::listener_addr`] is `None`
    /// fails with [`ErrorKind::AddrNotAvailable`].
    ///
    /// note: Disabling the listener aborts the accept loop, so no *new* inbound connections are
    /// admitted after this returns. It does **not** abort inbound connections already accepted and
    /// mid-setup (handshake/protocol wiring): those proceed to completion and may appear as active
    /// connections shortly after the listener is reported disabled (an accepted connection still
    /// awaiting a connection-setup slot is dropped instead). To reject them, gate acceptance
    /// in [`Handshake`], or follow up with [`Node::disconnect`] once they register.
    ///
    /// # Cancel safety
    ///
    /// If this future is dropped before completion, the node's listening state is left as it was:
    /// a listener that was in the middle of being enabled is aborted and its socket closed.
    pub async fn toggle_listener(&self) -> io::Result<Option<SocketAddr>> {
        // we deliberately maintain the write guard for the entirety of this method
        let mut listening_addr = self.listening_addr.write().await;

        if let Some(old_listening_addr) = *listening_addr {
            // validate before mutating: if the listener task is gone, the node is shutting
            // down, and the address is `shut_down`'s to clear
            let Some(listener_task) = self.tasks.lock().remove(&NodeTask::Listener) else {
                return Err(shutting_down_error());
            };
            listener_task.abort();
            trace!(parent: self.span(), "aborted the listening task");
            debug!(parent: self.span(), "no longer listening on {old_listening_addr}");
            *listening_addr = None;

            Ok(None)
        } else {
            let listener_addr = self.config().listener_addr.ok_or_else(|| {
                error!(parent: self.span(), "the listener was toggled on, but Config::listener_addr is not set");
                ErrorKind::AddrNotAvailable
            })?;
            trace!(parent: self.span(), "attempting to listen on {listener_addr}");
            let socket = match listener_addr {
                SocketAddr::V4(_) => TcpSocket::new_v4()?,
                SocketAddr::V6(_) => TcpSocket::new_v6()?,
            };

            socket.set_reuseaddr(true)?;

            #[cfg(all(unix, not(any(target_os = "solaris", target_os = "illumos"))))]
            if self.config().reuse_listener_port {
                socket.set_reuseport(true)?;
            }
            #[cfg(not(all(unix, not(any(target_os = "solaris", target_os = "illumos")))))]
            if self.config().reuse_listener_port {
                return Err(io::Error::new(
                    ErrorKind::Unsupported,
                    "Config::reuse_listener_port is set, but SO_REUSEPORT is not supported on this platform",
                ));
            }

            socket.bind(listener_addr)?;
            let listener = socket.listen(self.config().listener_backlog)?; // capped by somaxconn
            let port = listener.local_addr()?.port(); // discover the port if it was unspecified
            let new_listening_addr = (listener_addr.ip(), port).into();

            // start listening
            self.start_listening(listener).await?;
            debug!(parent: self.span(), "listening on {new_listening_addr}");

            // update the node's listening address
            *listening_addr = Some(new_listening_addr);

            Ok(Some(new_listening_addr))
        }
    }

    /// Spawn a task responsible for listening for inbound connections.
    async fn start_listening(&self, listener: TcpListener) -> io::Result<()> {
        // use a channel to know when the listening task is ready
        let (tx, rx) = oneshot::channel();

        let node = self.clone();
        // the task belongs to the guard until the node adopts it below, so that a `toggle_listener`
        // future dropped at the readiness await doesn't leave a listener nothing can abort
        let listening_task = UnregisteredTask::new(tokio::spawn(async move {
            trace!(parent: node.span(), "spawned the listening task");
            if tx.send(()).is_err() {
                error!(parent: node.span(), "listener setup interrupted; shutting down the listening task");
                return;
            }

            // inbound accepts draw from the node-wide connection-setup budget that outbound
            // connects also consume; the hard connection limits are still enforced inside
            // `handle_connection_request` via `check_and_reserve`, but gating on this shared
            // budget keeps surplus inbound connections in the kernel accept queue (cheap
            // backpressure) instead of accepting, spawning, and immediately rejecting them
            // while outbound dialing holds the `connecting` slots. The permit is only acquired
            // *after* `accept` returns: a listener idling in `accept` must not park a permit
            // that outbound connects could otherwise use
            let inbound_permits = node.connecting_permits.clone();

            loop {
                match listener.accept().await {
                    Ok((stream, addr)) => {
                        // wait for setup capacity before spawning the handler; while the budget
                        // is exhausted, no further connections are accepted
                        let permit = match inbound_permits.clone().acquire_owned().await {
                            Ok(p) => p,
                            Err(_) => {
                                // semaphore is never closed in practice; bail defensively
                                error!(parent: node.span(), "inbound permit semaphore closed unexpectedly");
                                debug_assert!(
                                    false,
                                    "acquiring an owned listening semaphore failed"
                                );
                                return;
                            }
                        };

                        // handle connection requests asynchronously
                        let node = node.clone();
                        tokio::spawn(async move {
                            // the permit is released when the connection is fully processed
                            let _permit = permit;
                            node.handle_connection_request(stream, addr).await.inspect_err(|e|
                                match e.kind() {
                                    ErrorKind::QuotaExceeded | ErrorKind::AlreadyExists => {
                                        debug!(parent: node.span(), "rejecting connection from {addr}: {e}");
                                    }
                                    // a shutdown racing the accept loop's own abort is expected,
                                    // and so is anything else failing once one is under way
                                    _ if node.shutdown.is_underway() => {
                                        debug!(parent: node.span(), "dropping the connection from {addr}: {e}");
                                    }
                                    _ => {
                                        error!(parent: node.span(), "couldn't accept a connection from {addr}: {e}");
                                    }
                                }
                            )
                        });
                    }
                    Err(e) => {
                        match e.kind() {
                            // a peer aborted/reset before accept completed; no backoff - the listener is healthy
                            ErrorKind::ConnectionAborted | ErrorKind::ConnectionReset => {
                                debug!(parent: node.span(), "transient accept error: {e}");
                            }
                            // otherwise, assume fd / memory exhaustion (EMFILE, ENFILE, ENOBUFS, ...)
                            // and back off so we don't spin at 100% CPU waiting for a slot to free
                            _ => {
                                node.heuristics.register_accept_error();
                                error!(parent: node.span(), "couldn't accept a connection: {e}");
                                sleep(Duration::from_millis(500)).await;
                            }
                        }
                    }
                }
            }
        }));
        let _ = rx.await;
        listening_task.register(self, NodeTask::Listener)?;

        Ok(())
    }

    /// Processes a single inbound connection request. Only used in [`Node::start_listening`].
    async fn handle_connection_request(
        &self,
        stream: TcpStream,
        addr: SocketAddr,
    ) -> io::Result<()> {
        // immediately reject if shutting down
        if self.shutdown.is_underway() {
            return Err(shutting_down_error());
        }

        // check connection limits and set up a connection guard; a rejection here means admission
        // control turned the peer away before any protocol ran - invisible to user hooks, so record
        // it as a node-level signal of inbound pressure
        let guard = self.check_and_reserve(addr).inspect_err(|_| {
            self.heuristics.register_inbound_rejection();
        })?;

        // finalize the connection
        self.adapt_stream(stream, addr, ConnectionSide::Responder, guard)
            .await
    }

    /// Returns the name assigned to the node.
    #[inline]
    pub fn name(&self) -> &str {
        // safe; can be set as None in Config, but receives a default value on Node creation
        self.config.name.as_deref().unwrap()
    }

    /// Returns a reference to the node's config.
    #[inline]
    pub fn config(&self) -> &Config {
        &self.config
    }

    /// Returns a reference to the node's stats.
    #[inline]
    pub fn stats(&self) -> &Stats {
        &self.stats
    }

    /// Returns a reference to the node's operational [`Heuristics`] - node-level health signals (e.g.
    /// connection-setup budget exhaustion) that the library observes internally.
    #[inline]
    pub fn heuristics(&self) -> &Heuristics {
        &self.heuristics
    }

    /// Returns the tracing [`Span`] associated with the node.
    #[inline]
    pub fn span(&self) -> &Span {
        &self.span
    }

    /// Returns the node's current listening address; returns an error if the node was configured
    /// to not listen for inbound connections or if the listener is currently disabled.
    ///
    /// note: This returns the address the listener is *bound* to, which is not necessarily an
    /// address peers can reach. If [`Config::listener_addr`] specified a wildcard IP (e.g.
    /// `0.0.0.0` or `::`), this returns that same wildcard IP with the resolved port - e.g.
    /// `0.0.0.0:34567` - which is **not routable** and must not be handed to peers as a dial
    /// target. Resolving a wildcard bind to a concrete, advertisable address requires enumerating
    /// local interfaces (and, behind NAT, discovering the external address), which the library
    /// deliberately leaves to the application. The port, however, is always accurate.
    pub async fn listening_addr(&self) -> io::Result<SocketAddr> {
        self.listening_addr
            .read()
            .await
            .as_ref()
            .copied()
            .ok_or_else(|| ErrorKind::AddrNotAvailable.into())
    }

    /// Enable the applicable protocols for a new connection.
    async fn enable_protocols(&self, conn: Connection) -> io::Result<Connection> {
        let mut conn = enable_protocol!(handshake, self, conn);

        // split the stream after the handshake (if not done before)
        if let Some(stream) = conn.stream.take() {
            let (reader, writer) = stream.into_split();
            conn.reader = Some(Box::new(reader));
            conn.writer = Some(Box::new(writer));
        }

        let conn = enable_protocol!(reading, self, conn);
        let conn = enable_protocol!(writing, self, conn);

        Ok(conn)
    }

    /// Prepares the freshly acquired connection to handle the protocols the Node implements.
    async fn adapt_stream(
        &self,
        stream: TcpStream,
        peer_addr: SocketAddr,
        own_side: ConnectionSide,
        guard: ConnectionGuard<'_>,
    ) -> io::Result<()> {
        let conn_span = create_connection_span(peer_addr, self.span());
        debug!(parent: &conn_span, "establishing connection as the {own_side:?}");

        // register the port seen by the peer; `local_addr` is a syscall, so it is only
        // worth making when the log it feeds is enabled
        if own_side == ConnectionSide::Initiator && enabled!(Level::TRACE) {
            if let Ok(addr) = stream.local_addr() {
                trace!(parent: &conn_span, "the peer is connected on port {}", addr.port());
            } else {
                warn!(parent: &conn_span, "couldn't determine the peer-side port");
            }
        }

        let connection = Connection::new(peer_addr, stream, !own_side, conn_span.clone());

        // enact the enabled protocols
        let mut connection = self.enable_protocols(connection).await?;

        // if Reading is enabled, we'll notify the related task when the connection is fully ready
        let conn_ready_tx = connection.readiness_notifier.take();

        // capture the connection's instance id before it's moved into the active set, so OnConnect's
        // cleanup guard can be scoped to this exact connection (see DisconnectOnDrop)
        let conn_id = connection.id;

        // account for the upcoming OnConnect scheduling (if the protocol is enabled at
        // all; it can never be enabled retroactively) before the connection becomes
        // visible in `active`, so that `shut_down` - which drains `active` first - can
        // never miss it
        let sched_guard = self
            .protocols
            .on_connect
            .get()
            .is_some()
            .then(|| SchedulingGuard::new(self.clone()));
        self.connections.add(connection, guard, &self.shutdown)?;

        // send the aforementioned notification so that reading from the socket can commence
        if let Some(tx) = conn_ready_tx {
            let _ = tx.send(());
        }

        debug!(parent: &conn_span, "fully connected");

        // if enabled, enact OnConnect; keyed off the guard rather than a fresh protocol check,
        // so that a concurrent `enable_on_connect` can't produce a scheduling that the sched
        // accounting (and thus `shut_down`) doesn't know about
        if sched_guard.is_some() {
            trace!(parent: &conn_span, "executing OnConnect logic...");

            // the scheduling runs in a dedicated task, so that it is carried through even if the
            // future driving this method (e.g. a user-side `Node::connect` wrapped in a timeout)
            // is dropped at this point; the connection is already live, so the hook must neither
            // be skipped nor - if `OnConnect::ABORTABLE` - have its task escape cleanup
            let node = self.clone();
            let scheduling_task = tokio::spawn(async move {
                let _sched_guard = sched_guard;

                let Some(handler) = node.protocols.on_connect.get() else {
                    return; // unreachable: checked above, and protocols are never disabled
                };
                let (sender, receiver) = oneshot::channel();
                handler.trigger(((peer_addr, conn_id), sender)).await;

                // receive the handle for the running task
                if let Some((handle, abortable)) =
                    crate::protocols::await_handler_response(receiver, handler.closed()).await
                {
                    if !abortable {
                        // leak the OnConnect task handle in order to ensure that its logic gets
                        // executed in full
                        drop(handle);
                    } else if let Some(conn) = node
                        .connections
                        .active
                        .write()
                        .get_mut(&peer_addr)
                        .filter(|conn| conn.id == conn_id)
                    {
                        // add the task to *this* connection so it gets aborted on its disconnect;
                        // the id check ensures we don't attach it to a different connection that
                        // has since reused the address after a rapid reconnect (mirrors
                        // DisconnectOnDrop)
                        conn.tasks.push(handle);
                    } else {
                        // the connection was terminated (or a newer one already replaced it at
                        // this address); abort the OnConnect work rather than attach it to a
                        // stranger
                        handle.abort();
                    }
                }
            });

            // preserve the completion order: the connection setup doesn't conclude until the
            // hook has been scheduled; cancelling this await doesn't cancel the task above
            let _ = scheduling_task.await;
        }

        Ok(())
    }

    // A helper method to facilitate a common potential disconnect at the callsite.
    async fn create_stream(
        &self,
        addr: SocketAddr,
        socket: Option<TcpSocket>,
    ) -> io::Result<TcpStream> {
        match timeout(
            Duration::from_millis(self.config().connection_timeout_ms.into()),
            self.create_stream_inner(addr, socket),
        )
        .await
        {
            Ok(Ok(stream)) => Ok(stream),
            Ok(err) => err,
            Err(err) => Err(io::Error::new(ErrorKind::TimedOut, err)),
        }
    }

    /// A wrapper method for greater readability.
    async fn create_stream_inner(
        &self,
        addr: SocketAddr,
        socket: Option<TcpSocket>,
    ) -> io::Result<TcpStream> {
        if let Some(socket) = socket {
            socket.connect(addr).await
        } else {
            TcpStream::connect(addr).await
        }
    }

    /// Connects to the provided `SocketAddr`.
    ///
    /// # Notes
    ///
    /// - `pea2pea` identifies connections by their socket address (IP + port). If Node A connects
    ///   to Node B, and Node B simultaneously connects to Node A, the library considers these two
    ///   distinct connections (one outgoing, one incoming). To ensure a single logical connection
    ///   per peer, implement a tie-breaking mechanism in the [`Handshake`] protocol.
    /// - The self-connect check is best-effort. It compares against the node's listening address
    ///   and - for a wildcard-bound listener - the loopback variant of its port, but does not
    ///   enumerate local interfaces, so dialing one of the host's own non-loopback addresses (its
    ///   LAN or public IP) will connect the node to itself over a real TCP loop. Rejecting that is
    ///   the same job as the tie-breaking above: exchange a node identifier in [`Handshake`] and
    ///   refuse matches. A self-connect the check *does* catch is refused with
    ///   [`ErrorKind::AddrInUse`].
    /// - A dial that hasn't produced a TCP connection within [`Config::connection_timeout_ms`]
    ///   fails with [`ErrorKind::TimedOut`]; the timeout covers only that stage, with
    ///   [`Handshake::TIMEOUT_MS`] bounding the next one.
    /// - A disconnect is not instantaneous. From the moment one is initiated (by
    ///   [`Node::disconnect`], a read/write error, or peer-side close) until the connection is
    ///   fully removed, the address remains registered and `connect` to it returns
    ///   [`io::ErrorKind::AlreadyExists`]. If [`OnDisconnect`] is enabled this window spans the
    ///   hook's execution, bounded by [`OnDisconnect::TIMEOUT_MS`]. Reconnection logic that races a
    ///   disconnect should treat `AlreadyExists` as retriable and back off, rather than as a
    ///   permanent failure.
    /// - [`ErrorKind::QuotaExceeded`] covers three distinct limits - the per-IP cap
    ///   ([`Config::max_connections_per_ip`]), the global cap ([`Config::max_connections`]), and
    ///   exhaustion of the shared connection-setup budget ([`Config::max_connecting`]). Checking
    ///   [`Node::num_connecting`] / [`Node::num_connected`] against your caps before dialing keeps
    ///   the first two from firing. The budget is the exception: it is shared with inbound accepts,
    ///   so a flood can exhaust it and fail your dials while you are below your own caps.
    /// - Once the node is shutting down every dial is refused, carrying a [`ShuttingDown`] payload
    ///   ([`ShuttingDown::caused`] recognizes it). Unlike the cases above, that one is permanent -
    ///   there is no point in retrying.
    ///
    /// # Cancel safety
    ///
    /// If this future is dropped before the connection is finalized, the attempt is cleanly rolled
    /// back; if it is dropped after finalization, the connection remains established and
    /// [`OnConnect`] (if enabled) still runs for it. Use [`Node::is_connected`] to tell the two
    /// outcomes apart.
    pub async fn connect(&self, addr: SocketAddr) -> io::Result<()> {
        self.connect_inner(addr, None)
            .await
            .inspect_err(|e| error!(parent: self.span(), "couldn't connect to {addr}: {e}"))
    }

    /// Connects to a `SocketAddr` using the provided `TcpSocket`.
    pub async fn connect_using_socket(
        &self,
        addr: SocketAddr,
        socket: TcpSocket,
    ) -> io::Result<()> {
        self.connect_inner(addr, Some(socket))
            .await
            .inspect_err(|e| error!(parent: self.span(), "couldn't connect to {addr}: {e}"))
    }

    /// Connects to the provided `SocketAddr` using an optional `TcpSocket`.
    async fn connect_inner(&self, addr: SocketAddr, socket: Option<TcpSocket>) -> io::Result<()> {
        // immediately abort if shutting down
        if self.shutdown.is_underway() {
            return Err(shutting_down_error());
        }

        // a simple self-connect attempt check; the loopback heuristic only applies to a
        // wildcard-bound listener, the only case where a loopback address with the listening
        // port actually reaches the node itself
        if let Ok(listening_addr) = self.listening_addr().await
            && (addr == listening_addr
                || listening_addr.ip().is_unspecified()
                    && addr.ip().is_loopback()
                    && addr.port() == listening_addr.port())
        {
            return Err(io::Error::new(
                ErrorKind::AddrInUse,
                format!("can't connect to node's own listening address ({addr})"),
            ));
        }

        // the address may already be occupied by a live connection, e.g. one still inside its
        // teardown window; check before drawing from the connection-setup budget, so that the
        // documented `AlreadyExists` is returned even if the budget happens to be exhausted at
        // the same time (the check is repeated atomically in `check_and_reserve`)
        if self.connections.is_connected(addr) {
            return Err(io::Error::new(
                ErrorKind::AlreadyExists,
                "already connected",
            ));
        }

        // take a slot in the shared connection-setup budget before reserving, so an
        // in-flight outbound connect counts against the same `max_connecting` ceiling the
        // inbound accept loop honors; the permit is held for the whole setup and released
        // when this function returns
        let _permit = self
            .connecting_permits
            .clone()
            .try_acquire_owned()
            .map_err(|_| {
                self.heuristics.register_connect_budget_rejection();
                io::Error::new(
                    ErrorKind::QuotaExceeded,
                    format!(
                        "maximum number ({}) of pending connections reached",
                        self.config.max_connecting
                    ),
                )
            })?;

        // attempt to reserve a connection slot atomically
        let guard = self.check_and_reserve(addr)?;

        // attempt to physically connect to the specified address
        let stream = self.create_stream(addr, socket).await?;

        // attempt to finalize the connection
        self.adapt_stream(stream, addr, ConnectionSide::Initiator, guard)
            .await
    }

    /// Disconnects from the provided `SocketAddr`; returns `true` if an actual disconnect took place.
    ///
    /// If [`OnDisconnect`] is enabled, its hook runs to completion (subject to
    /// [`OnDisconnect::TIMEOUT_MS`]) before the connection is removed, and is the appropriate place
    /// to send any final messages. Messages queued via [`Writing`] but not awaited to delivery
    /// confirmation may be dropped once this function returns.
    ///
    /// note: The address is not immediately reusable. See [`Node::connect`] for the reconnection
    /// contract during the teardown window.
    ///
    /// # Cancel safety
    ///
    /// If this future is dropped before completion, the connection is still fully torn down and
    /// the accounting remains consistent; only the [`OnDisconnect`] hook's run-to-completion
    /// guarantee is lost - it may be skipped, aborted mid-flight, or overlap the connection's
    /// removal.
    pub async fn disconnect(&self, addr: SocketAddr) -> bool {
        self.disconnect_w_origin(addr, DisconnectOrigin::User, None)
            .await
    }

    /// `conn_id`, when provided, restricts the disconnect to that exact connection instance; if
    /// the address has since been reused by a newer connection, the call is a no-op. Internal
    /// callers acting on behalf of a specific connection (e.g. its cleanup guards) must pass it,
    /// as it is checked atomically with the claim below - an earlier id check outside this
    /// function is only an optimization and cannot rule out an address reuse in between.
    pub(crate) async fn disconnect_w_origin(
        &self,
        addr: SocketAddr,
        origin: DisconnectOrigin,
        conn_id: Option<u64>,
    ) -> bool {
        // claim the disconnect to avoid duplicate executions, or return early if already claimed;
        // the claim comes with the finalizer that completes the teardown, so from this point on it
        // runs even if this future is dropped at one of the awaits below
        let Some(finalizer) = self.claim_disconnect(addr, conn_id) else {
            return false;
        };

        let conn_span = create_connection_span(addr, self.span());
        debug!(parent: &conn_span, "disconnecting (origin: {origin:?})...");

        // if the OnDisconnect protocol is enabled, trigger it
        if let Some(handler) = self.protocols.on_disconnect.get() {
            trace!(parent: &conn_span, "executing OnDisconnect logic...");
            let (sender, receiver) = oneshot::channel();
            handler.trigger(((addr, origin), sender)).await;
            if let Some((handle, waiter)) =
                crate::protocols::await_handler_response(receiver, handler.closed()).await
            {
                // register the associated task with the connection, in case
                // it gets terminated before its completion
                if let Some(conn) = self.connections.active.write().get_mut(&addr) {
                    conn.tasks.push(handle);
                } else {
                    // can't really happen, since disconnects are exclusive and atomic
                    debug_assert!(
                        false,
                        "disconnect of {addr} claimed, yet the connection vanished before OnDisconnect registration"
                    );
                    handle.abort();
                }
                // wait for the OnDisconnect protocol to perform its specified actions
                // time out, or even panic - we're already disconnecting, so ignore the
                // result
                let _ = waiter.await;
            }
        }

        // the OnDisconnect hook (above) is the appropriate place to send any final
        // messages; by the time we get here, the user has already had their chance
        // to flush
        drop(finalizer);

        debug!(parent: &conn_span, "fully disconnected");

        true
    }

    /// Claims the disconnect of `addr`, returning the [`DisconnectFinalizer`] that completes the
    /// teardown on drop; `None` if the address isn't connected, belongs to a newer connection
    /// than `conn_id`, or is already being disconnected by someone else.
    ///
    /// Taking the claim and arming the finalizer must not be separated by a yield point: no other
    /// path can reclaim the connection, so a disconnect cancelled in between would leave it in the
    /// active set forever and hang the drain loop in [`Node::shut_down`]. This being a plain `fn`
    /// rather than inline `async` code is what enforces that.
    fn claim_disconnect(
        &self,
        addr: SocketAddr,
        conn_id: Option<u64>,
    ) -> Option<DisconnectFinalizer<'_>> {
        let active = self.connections.active.read();
        let conn = active.get(&addr)?;
        if conn_id.is_some_and(|id| conn.id != id) {
            // the address now belongs to a newer connection; this call was aimed at its
            // defunct predecessor
            return None;
        }
        if conn.disconnecting.swap(true, AcqRel) {
            // valid connection, but someone else is already disconnecting it
            return None;
        }

        Some(DisconnectFinalizer { node: self, addr })
    }

    /// Returns a list containing addresses of active connections.
    pub fn connected_addrs(&self) -> Vec<SocketAddr> {
        self.connections.addrs()
    }

    /// Checks whether the provided address is connected.
    pub fn is_connected(&self, addr: SocketAddr) -> bool {
        self.connections.is_connected(addr)
    }

    /// Checks if the node is currently setting up a connection with the provided address.
    pub fn is_connecting(&self, addr: SocketAddr) -> bool {
        self.connections.limits.lock().connecting.contains(&addr)
    }

    /// Returns the number of active connections.
    pub fn num_connected(&self) -> usize {
        self.connections.num_connected()
    }

    /// Returns the number of connections that are currently being set up.
    pub fn num_connecting(&self) -> usize {
        self.connections.limits.lock().connecting.len()
    }

    /// Returns basic information related to a connection.
    pub fn connection_info(&self, addr: SocketAddr) -> Option<ConnectionInfo> {
        self.connections.get_info(addr)
    }

    /// Returns a list of all active connections and their basic information.
    pub fn connection_infos(&self) -> HashMap<SocketAddr, ConnectionInfo> {
        self.connections.infos()
    }

    /// Atomically checks connection limits and reserves a slot if available.
    fn check_and_reserve(&self, addr: SocketAddr) -> io::Result<ConnectionGuard<'_>> {
        // this lock is held for the duration of the check to prevent races
        let mut limits = self.connections.limits.lock();
        // a single read guard serves both the duplicate check and the connection count;
        // the `limits` -> `active` acquisition order is preserved
        let active = self.connections.active.read();

        // reject duplicates regardless of side (although is very niche on responder side)
        if active.contains_key(&addr) {
            return Err(io::Error::new(
                ErrorKind::AlreadyExists,
                "already connected",
            ));
        }

        // check the per-IP limit first
        let num_ip_conns = limits.ip_count(addr);
        let per_ip_limit = self.config.max_connections_per_ip as usize;
        if num_ip_conns >= per_ip_limit {
            return Err(io::Error::new(
                ErrorKind::QuotaExceeded,
                format!("maximum number ({per_ip_limit}) of per-IP connections reached"),
            ));
        }

        // check the global connecting count limit; in practice this never rejects: both callers
        // hold a permit from `connecting_permits` (capacity `max_connecting`) for longer than the
        // reservation this counts - `ConnectionGuard` clears `connecting` inside
        // `Connections::add`, while the permit lives until the whole setup returns - so the
        // semaphore is always the tighter bound and caps `connecting.len()` at
        // `max_connecting - 1` by the time we get here
        let num_connecting = limits.connecting.len();
        let connecting_limit = self.config.max_connecting as usize;
        if num_connecting >= connecting_limit {
            return Err(io::Error::new(
                ErrorKind::QuotaExceeded,
                format!("maximum number ({connecting_limit}) of pending connections reached"),
            ));
        }

        // check the global connection count limit
        let num_connected = active.len();
        let connection_limit = self.config.max_connections as usize;
        if num_connected + num_connecting >= connection_limit {
            return Err(io::Error::new(
                ErrorKind::QuotaExceeded,
                format!("maximum number ({connection_limit}) of connections reached"),
            ));
        }

        // check if already connecting (duplicate connection attempt from same node)
        if limits.connecting.contains(&addr) {
            return Err(io::Error::new(
                ErrorKind::AlreadyExists,
                "already connecting",
            ));
        }

        // reserve a connecting slot
        limits.reserve(addr);

        Ok(ConnectionGuard {
            addr,
            connections: &self.connections,
            completed: false,
        })
    }

    /// Gracefully shuts the node down. This is permanent - the node cannot be restarted,
    /// you need to create a new one.
    ///
    /// note: Do not `await` this from inside a per-connection hook ([`Reading::process_message`],
    /// [`OnConnect::on_connect`], [`OnDisconnect::on_disconnect`]). `shut_down` tears down the very
    /// connection - and aborts the very task - the hook runs on; instead, signal shutdown to a
    /// separate task and call it there.
    ///
    /// # Cancel safety
    ///
    /// If this future is dropped before completion, the node's long-running tasks are aborted
    /// right away - active [`OnDisconnect`] hooks may then be cut short - and the shutdown can be
    /// concluded with another `shut_down` call.
    pub async fn shut_down(&self) {
        // immediately mark the node as shutting down
        self.shutdown.begin();

        debug!(parent: self.span(), "shutting down");

        // move the task handles into guards that abort them on drop; they are unregistered from
        // the node from this point on, so if this future is dropped at one of the awaits below,
        // they would otherwise keep running - unreachable even to a repeated `shut_down` call
        let mut tasks: HashMap<_, _> = std::mem::take(&mut *self.tasks.lock())
            .into_iter()
            .map(|(kind, handle)| (kind, UnregisteredTask::new(handle)))
            .collect();

        // abort the listening task first (if it exists); dropping its guard is what aborts it
        drop(tasks.remove(&NodeTask::Listener));

        // disconnect from all the peers
        let mut disconnects: FuturesUnordered<_> = self
            .connected_addrs()
            .into_iter()
            .map(|addr| self.disconnect_w_origin(addr, DisconnectOrigin::Shutdown, None))
            .collect();
        while disconnects.next().await.is_some() {}

        // wait for any concurrent disconnects to finish their cleanup; a concurrent
        // disconnect (e.g. a user-initiated one) may have won the `disconnecting` swap
        // before this `shut_down` began - the fan-out above then returned `false` for
        // that address without firing `on_disconnect`, and the concurrent disconnect
        // remains responsible for firing it
        //
        // the wait is bounded: `shutting_down` now blocks new entries from
        // being added (see `Connections::add`), and every remaining entry has
        // a single owner committed to removing it
        wait_for_drain(&self.connections.drain_notify, || {
            self.connections.active.read().is_empty()
        })
        .await;

        // wait out any in-flight OnConnect scheduling tasks; every connection
        // that reached `active` was preceded by its guard, and new ones can
        // no longer be added, so the count can only fall from here. Without
        // this, a scheduling task delayed by executor lag could fire its
        // trigger into the closed channel below and silently skip the hook
        self.shutdown.wait_sched_drained().await;

        // signal the protocol handler tasks to wind down gracefully: each one
        // stops accepting new work and then drains its queue - the setup
        // handlers (Handshake/Reading/Writing) by failing the queued requests,
        // the hook handlers (OnConnect/OnDisconnect) by still running the
        // queued hooks, preserving the pairing guarantees. Draining via `recv`
        // (as opposed to dropping the receiver) also waits out any message
        // that is still mid-send, so nothing gets stranded in the channels.
        self.shutdown.signal_handlers();
        // all five handlers share a single drain deadline (the drains are
        // near-instant; the deadline only matters on pathological paths)
        const DRAIN_DEADLINE: Duration = Duration::from_secs(3);
        let deadline = Instant::now() + DRAIN_DEADLINE;
        for kind in [
            NodeTask::Handshake,
            NodeTask::Reading,
            NodeTask::Writing,
            NodeTask::OnConnect,
            NodeTask::OnDisconnect,
        ] {
            if let Some(task) = tasks.get_mut(&kind) {
                // the handles stay inside their guards, so if this future is
                // dropped mid-await, they still get aborted; aborting an
                // already-completed handle is a no-op
                let remaining = deadline.saturating_duration_since(Instant::now());
                let _ = timeout(remaining, task.handle_mut()).await;
            }
        }

        // abort the remaining tasks, which should now be inert
        drop(tasks);

        // erase the listening address in case toggle_listener is called afterwards
        *self.listening_addr.write().await = None;
    }

    /// Registers a long-running node task, unless the node is shutting down, in which case the
    /// task is aborted and an error is returned. The flag check under the `tasks` lock pairs with
    /// `ShutdownState::begin` + the `mem::take` in `shut_down`.
    pub(crate) fn register_task(&self, kind: NodeTask, handle: JoinHandle<()>) -> io::Result<()> {
        let mut tasks = self.tasks.lock();
        if self.shutdown.is_underway() {
            handle.abort();
            return Err(shutting_down_error());
        }
        // a duplicate registration could only result from a concurrent double-enable of a
        // protocol - API misuse that also panics in the enable itself; surface it in debug
        // builds, but don't abort the displaced handle, as depending on the order of the racing
        // registrations it may belong to the handler that ends up wired into the node
        let prev = tasks.insert(kind, handle);
        debug_assert!(prev.is_none(), "a Node task was registered more than once");
        Ok(())
    }
}

/// A node task the node does not own, aborted on drop unless something claims it first. Both
/// ends of a task's life need this, and for the same reason - a task nothing owns keeps running
/// unreachably, holding a `Node` clone that keeps the node alive for the lifetime of the runtime.
pub(crate) struct UnregisteredTask(Option<JoinHandle<()>>);

impl UnregisteredTask {
    pub(crate) fn new(handle: JoinHandle<()>) -> Self {
        Self(Some(handle))
    }

    /// Hands the task over to the node; see [`Node::register_task`].
    pub(crate) fn register(mut self, node: &Node, kind: NodeTask) -> io::Result<()> {
        node.register_task(kind, self.0.take().unwrap()) // guaranteed to be present here
    }

    /// The underlying handle, for awaiting the task's completion; the guard keeps its claim on
    /// the task, so a wait cut short still ends in an abort.
    fn handle_mut(&mut self) -> &mut JoinHandle<()> {
        self.0.as_mut().unwrap() // only ever taken by `register`, which consumes `self`
    }
}

impl Drop for UnregisteredTask {
    fn drop(&mut self) {
        if let Some(handle) = &self.0 {
            handle.abort();
        }
    }
}

/// Tracks one in-flight `OnConnect` scheduling task (or the decision not to
/// spawn one); the count is incremented at creation - crucially *before* the
/// connection enters `active` - and decremented on drop, waking `shut_down`
/// once no scheduling work remains.
struct SchedulingGuard(Node);

impl SchedulingGuard {
    fn new(node: Node) -> Self {
        node.shutdown.sched_started();
        Self(node)
    }
}

impl Drop for SchedulingGuard {
    fn drop(&mut self) {
        self.0.shutdown.sched_finished();
    }
}

/// Completes a claimed disconnect. Created in `Node::disconnect_w_origin` right after the
/// `disconnecting` claim is taken; its `Drop` performs the synchronous teardown tail, so the
/// connection is removed and its accounting released even if the disconnect future is dropped
/// at an await point (e.g. by a user-side timeout around [`Node::disconnect`]). Without this,
/// a cancelled disconnect would leak the claim, making the connection permanently stuck (every
/// other teardown path honors the claim) and hanging the drain loop in [`Node::shut_down`].
struct DisconnectFinalizer<'a> {
    node: &'a Node,
    addr: SocketAddr,
}

impl Drop for DisconnectFinalizer<'_> {
    fn drop(&mut self) {
        // closing the sender channel signals the writer task to wind down; the writer task's
        // own `SenderCleanup` guard will attempt the same removal on drop - that's a harmless
        // no-op second removal and serves as a safety net for non-disconnect exits (e.g. write
        // errors)
        if let Some(writing) = self.node.protocols.writing.get() {
            writing.senders.write().remove(&self.addr);
        }

        // drop the connection from the active set under the limits lock, so any concurrent
        // `check_and_reserve` has a consistent view of connection counts; the object itself
        // (whose destruction aborts the connection's tasks) is destroyed only after the lock
        // is released, so that no destructor reachable from `Connection` can deadlock against it
        let conn = {
            let mut limits = self.node.connections.limits.lock();
            let conn = self.node.connections.remove(self.addr);
            limits.release_ip(self.addr);
            conn
        };
        drop(conn);
    }
}

/// Creates the node's tracing span based on its name.
fn create_span(node_name: &str) -> Span {
    macro_rules! try_span {
        ($lvl:expr) => {
            let s = span!($lvl, "node", name = node_name);
            if !s.is_disabled() {
                return s;
            }
        };
    }

    try_span!(Level::TRACE);
    try_span!(Level::DEBUG);
    try_span!(Level::INFO);
    try_span!(Level::WARN);

    error_span!("node", name = node_name)
}

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

    // A zero value for any of the connection limits would render the node inoperable (and break
    // internal channel/semaphore setup), so it must be rejected upfront.
    #[test]
    #[should_panic(expected = "Config::max_connections must not be 0")]
    fn zero_max_connections_is_rejected() {
        let _ = Node::new(Config {
            max_connections: 0,
            ..Default::default()
        });
    }

    #[test]
    #[should_panic(expected = "Config::max_connections_per_ip must not be 0")]
    fn zero_max_connections_per_ip_is_rejected() {
        let _ = Node::new(Config {
            max_connections_per_ip: 0,
            ..Default::default()
        });
    }

    #[test]
    #[should_panic(expected = "Config::max_connecting must not be 0")]
    fn zero_max_connecting_is_rejected() {
        let _ = Node::new(Config {
            max_connecting: 0,
            ..Default::default()
        });
    }

    // Pending connections count towards `max_connections`, so a setup budget above it could
    // never be used.
    #[test]
    fn excessive_max_connecting_is_clamped() {
        let node = Node::new(Config {
            max_connections: 10,
            max_connecting: 11,
            ..Default::default()
        });
        assert_eq!(node.config().max_connecting, 10);
    }
}

#[cfg(test)]
mod shutdown_tests {
    use std::time::Instant;

    use super::*;
    use crate::{Pea2Pea, protocols::OnDisconnect};

    // A `shut_down` future dropped mid-execution (e.g. due to an enclosing timeout) must not
    // detach the node's long-running tasks: each of them holds a `Node` clone, so a leak would
    // keep the node alive for the lifetime of the runtime, and a repeated `shut_down` call
    // would no longer be able to abort them.
    #[tokio::test]
    async fn cancelled_shut_down_aborts_node_tasks() {
        #[derive(Clone)]
        struct SlowDisconnect(Node);
        impl Pea2Pea for SlowDisconnect {
            fn node(&self) -> &Node {
                &self.0
            }
        }
        impl OnDisconnect for SlowDisconnect {
            async fn on_disconnect(&self, _: SocketAddr, _: DisconnectOrigin) {
                sleep(Duration::from_millis(300)).await;
            }
        }

        let slow = SlowDisconnect(Node::new(Config {
            listener_addr: Some("127.0.0.1:0".parse().unwrap()),
            ..Default::default()
        }));
        slow.enable_on_disconnect().await;
        let slow_addr = slow.0.toggle_listener().await.unwrap().unwrap();

        let peer = Node::new(Default::default());
        peer.connect(slow_addr).await.unwrap();
        let deadline = Instant::now() + Duration::from_secs(2);
        while slow.0.num_connected() != 1 {
            assert!(
                Instant::now() < deadline,
                "the test connection was never registered",
            );
            sleep(Duration::from_millis(10)).await;
        }

        // cancel the shutdown while its disconnect fan-out is waiting on the slow hook
        assert!(
            timeout(Duration::from_millis(50), slow.0.shut_down())
                .await
                .is_err()
        );

        // the shutdown can be concluded with another call
        timeout(Duration::from_secs(2), slow.0.shut_down())
            .await
            .expect("a repeated shut_down hung");

        // the cancellation must not have detached the node's tasks: once the runtime reaps
        // the aborted ones, no clone of the node may remain
        let node = slow.0.clone();
        drop(slow);
        let deadline = Instant::now() + Duration::from_secs(2);
        while Arc::strong_count(&node.0) > 1 {
            assert!(
                Instant::now() < deadline,
                "the node is still referenced, most likely by leaked tasks",
            );
            sleep(Duration::from_millis(10)).await;
        }

        peer.shut_down().await;
    }
}

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

    // A saturated connection-setup budget must reject an outbound dial with a `QuotaExceeded` error
    // and bump the node-level counter. Drains the private semaphore directly so the assertion is
    // deterministic (no network, no timing).
    #[tokio::test]
    async fn exhausted_connecting_budget_is_signalled() {
        let config = Config {
            max_connecting: 1,
            listener_addr: None, // no listener needed; the self-connect check is skipped
            ..Default::default()
        };
        let node = Node::new(config);

        // hold the sole permit, simulating a fully saturated budget
        let _held = node.connecting_permits.clone().try_acquire_owned().unwrap();

        let err = node
            .connect("127.0.0.1:9".parse().unwrap())
            .await
            .unwrap_err();

        assert_eq!(err.kind(), ErrorKind::QuotaExceeded);
        assert_eq!(node.heuristics().connect_budget_rejections(), 1);
    }

    // A listener idling in `accept` must not hold a permit from the shared connection-setup
    // budget: with `max_connecting = 1`, a listening node must still be able to dial out.
    #[tokio::test]
    async fn idle_listener_does_not_hold_the_connect_budget() {
        let node = Node::new(Config {
            max_connecting: 1,
            listener_addr: Some("127.0.0.1:0".parse().unwrap()),
            ..Default::default()
        });
        node.toggle_listener().await.unwrap();

        let peer = Node::new(Config {
            listener_addr: Some("127.0.0.1:0".parse().unwrap()),
            ..Default::default()
        });
        let peer_addr = peer.toggle_listener().await.unwrap().unwrap();

        // make sure the listening task is up and parked in `accept` before dialing
        sleep(Duration::from_millis(50)).await;

        node.connect(peer_addr).await.unwrap();
        assert_eq!(node.heuristics().connect_budget_rejections(), 0);
    }
}