pea2pea 0.57.0

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

use futures_util::stream::{FuturesUnordered, StreamExt};
use parking_lot::Mutex;
use tokio::{
    io::split,
    net::{TcpListener, TcpSocket, TcpStream},
    sync::{RwLock, Semaphore, oneshot},
    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 conn_retriever.await {
                Ok(Ok(conn)) => conn,
                // the handler (and the channel with the connection's returner) is gone, which
                // only happens when its task was aborted, i.e. the node is shutting down; match
                // the error used by the other shutdown-race paths
                Err(_) => return Err(io::Error::other("shutting down")),
                Ok(e) => return e,
            }
        } else {
            $conn
        }
    };
}

/// 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.
#[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<()>>>,
    /// Indicates whether the shutdown sequence has begun.
    shutting_down: Arc<AtomicBool>,
}

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());

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

        // create the shutdown flag
        let shutting_down = Arc::new(AtomicBool::default());

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

        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: 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.
    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(io::Error::other("shutting down"));
            };
            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();
        let listening_task = 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}");
                                    }
                                    _ => {
                                        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;
        self.register_task(NodeTask::Listener, listening_task)?;

        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.shutting_down.load(Acquire) {
            return Err(io::Error::other("shutting down"));
        }

        // 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) = split(stream);
            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
        if own_side == ConnectionSide::Initiator {
            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;

        // connecting -> connected
        self.connections.add(connection, guard)?;

        // 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
        if self.protocols.on_connect.get().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 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 Ok((handle, abortable)) = receiver.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`.
    ///
    /// note: `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, you must implement a tie-breaking mechanism in your application logic
    /// in the [`Handshake`] protocol.
    ///
    /// note: A best-effort self-connect check is performed against the node's listening address
    /// and the loopback variant of its port. It does **not** enumerate local network interfaces:
    /// if the node listens on a wildcard address (e.g. `0.0.0.0`) and `connect` is called with
    /// one of the host's own non-loopback addresses (e.g. its LAN IP, or its public IP), the
    /// connection will succeed and the node will end up talking to itself over a real TCP loop.
    /// The library cannot detect this without OS-specific interface enumeration, which it
    /// deliberately avoids. If your application must reject such connections, do it in
    /// [`Handshake`] - typically by exchanging a unique node identifier and refusing matches,
    /// which also handles the simultaneous-connection case above.
    ///
    /// note: A disconnect is not instantaneous. From the moment a disconnect 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.
    ///
    /// note: A return of [`ErrorKind::QuotaExceeded`] covers several distinct limits - the per-IP
    /// cap ([`Config::max_connections_per_ip`]), the global connection cap ([`Config::max_connections`]),
    /// and exhaustion of the shared connection-setup budget ([`Config::max_connecting`]). The intended
    /// usage is to check your own outbound conditions (e.g. [`Node::num_connecting`] /
    /// [`Node::num_connected`] against the configured caps) *before* dialing, so a `QuotaExceeded`
    /// here is not expected during normal operation. The budget case is the notable exception: because
    /// inbound accepts and outbound connects share that budget, a hostile inbound flood can exhaust it
    /// and make *your own* outbound dials fail here even when you are below your own caps. For
    /// monitoring that specifically, [`Heuristics::connect_budget_rejections`] counts these rejections
    /// for rate-based detection; react by shedding inbound load (e.g. [`Node::toggle_listener`]) or
    /// filtering peers in [`Handshake`].
    ///
    /// note: If this future is dropped mid-call (e.g. due to an enclosing timeout) 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. After a cancellation, 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.shutting_down.load(Acquire) {
            return Err(io::Error::other("shutting down"));
        }

        // a simple self-connect attempt check
        if let Ok(listening_addr) = self.listening_addr().await
            && (addr == listening_addr
                || 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.
    ///
    /// note: If this future is dropped before completion (e.g. due to an enclosing timeout), 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
        if let Some(conn) = self.connections.active.read().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 false;
            }
            if conn.disconnecting.swap(true, AcqRel) {
                // valid connection, but someone else is already disconnecting it
                return false;
            }
        } else {
            // not connected
            return false;
        };

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

        // the claim is taken, so from this point on the teardown must run even if this future is
        // dropped at one of the awaits below (e.g. due to a timeout around `Node::disconnect`);
        // no other path can reclaim the connection, so a skipped teardown would leave it in the
        // active set forever and hang the drain loop in `Node::shut_down`
        let finalizer = DisconnectFinalizer { node: self, addr };

        // 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 Ok((handle, waiter)) = receiver.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
    }

    /// 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();

        // reject duplicates regardless of side (although is very niche on responder side)
        if self.connections.is_connected(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
        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 = self.connections.num_connected();
        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.
    ///
    /// note: If this future is dropped before completion (e.g. due to an enclosing timeout), 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.shutting_down.store(true, Release);

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

        // move the task handles into a guard that aborts them on drop; they are detached 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 - and,
        // as each holds a `Node` clone, keep the node alive for the lifetime of the runtime
        let mut tasks = NodeTaskAborter(std::mem::take(&mut *self.tasks.lock()));

        // abort the listening task first (if it exists)
        if let Some(listening_task) = tasks.0.remove(&NodeTask::Listener) {
            listening_task.abort();
        }

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

        // wait for any concurrent disconnects to finish their cleanup; a concurrent
        // disconnect (e.g. a user-initiated may have won the `disconnecting` swap
        // before this `shut_down` began; in that case our spawned disconnect task
        // above returned `false` without firing `on_disconnect`, and the concurrent
        // disconnect is still 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
        loop {
            let notified = self.connections.drain_notify.notified();
            tokio::pin!(notified);
            notified.as_mut().enable(); // register the waiter
            if self.connections.active.read().is_empty() {
                break;
            }
            notified.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
    /// the Release store + `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.shutting_down.load(Acquire) {
            handle.abort();
            return Err(io::Error::other("shutting down"));
        }
        // 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(())
    }
}

/// Aborts the node's long-running tasks on drop. [`Node::shut_down`] moves the task handles out
/// of the node before its first await, so if its future is dropped mid-execution (e.g. due to an
/// enclosing timeout), this guard is what prevents them from being silently detached; the tasks
/// must nonetheless outlive the disconnect fan-out in `shut_down`, which relies on the
/// `OnDisconnect` handler, hence a guard rather than an upfront abort.
struct NodeTaskAborter(HashMap<NodeTask, JoinHandle<()>>);

impl Drop for NodeTaskAborter {
    fn drop(&mut self) {
        for handle in self.0.values() {
            handle.abort();
        }
    }
}

/// 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()
        });
    }
}

#[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);
    }
}