pea2pea 0.57.4

use std::{
    any::Any,
    collections::{HashMap, hash_map::Entry},
    future::{Future, poll_fn},
    io,
    net::SocketAddr,
    panic::{AssertUnwindSafe, resume_unwind},
    sync::Arc,
    time::Duration,
};

#[cfg(doc)]
use bytes::Bytes;
use futures_util::{FutureExt, SinkExt};
use parking_lot::RwLock;
use tokio::{
    io::AsyncWrite,
    sync::{mpsc, oneshot},
    task::JoinHandle,
    time::timeout,
};
use tokio_util::codec::{Encoder, FramedWrite};
use tracing::*;

#[cfg(doc)]
use crate::{Config, Node, protocols::Handshake};
use crate::{
    Connection, ConnectionSide, Pea2Pea,
    connections::{DisconnectOrigin, create_connection_span},
    node::NodeTask,
    protocols::{
        DisconnectOnDrop, Protocol, ProtocolHandler, ReturnableConnection, catch_setup_panic,
        install_protocol_handler, panic_message, run_setup_handler_loop,
    },
};

/// The per-address message senders; each entry is tagged with its connection's unique
/// instance id (`Connection::id`), so that stale cleanups can recognize address reuse.
type WritingSenders = Arc<RwLock<HashMap<SocketAddr, (u64, Arc<dyn Any + Send + Sync>)>>>;

/// Can be used to specify and enable writing, i.e. sending outbound messages. If the [`Handshake`]
/// protocol is enabled too, it goes into force only after the handshake has been concluded.
pub trait Writing: Pea2Pea
where
    Self: Clone + Send + Sync + 'static,
{
    /// The depth of per-connection queues used to send outbound messages; the greater it is, the more outbound
    /// messages the node can enqueue. Setting it to a large value is not recommended, as doing it might
    /// obscure potential issues with your implementation (like slow serialization) or network.
    ///
    /// note: Must not be `0`.
    const MESSAGE_QUEUE_DEPTH: usize = 64;

    /// The initial size of a per-connection buffer for writing outbound messages. Can be set to the maximum expected size
    /// of the outbound message in order to only allocate it once.
    ///
    /// note: It doubles as the point at which a batch of outbound messages is flushed to the
    /// socket mid-encoding, so a greater value means fewer, larger writes; values below the
    /// underlying 8KiB default don't lower that threshold.
    const INITIAL_BUFFER_SIZE: usize = 64 * 1024;

    /// The maximum time (in milliseconds) allowed for a batch of outbound messages (up to
    /// [`Writing::MESSAGE_QUEUE_DEPTH`] of them) to be written and flushed to the underlying
    /// stream before the connection is considered dead. Under sustained load, the effective
    /// per-message allowance is therefore this value divided by the batch size.
    ///
    /// note: `0` does not disable this timeout - it fails every write (almost) instantly.
    const TIMEOUT_MS: u64 = 10_000;

    /// The type of the outbound messages; unless their serialization is expensive and the message
    /// is broadcasted (in which case it would get serialized multiple times), serialization should
    /// be done in the implementation of [`Writing::Codec`].
    type Message: Send;

    /// The user-supplied [`Encoder`] used to write outbound messages to the target stream.
    type Codec: Encoder<Self::Message, Error = io::Error> + Send;

    /// Prepares the node to send messages.
    ///
    /// note: A no-op if the node is already shutting down; the protocol is not enabled.
    ///
    /// # Panics
    ///
    /// Panics if called more than once on the same [`Node`], or if
    /// [`Writing::MESSAGE_QUEUE_DEPTH`] is `0`.
    fn enable_writing(&self) -> impl Future<Output = ()> {
        async {
            // a zero depth would panic (obscurely) at the first connection instead
            assert!(
                Self::MESSAGE_QUEUE_DEPTH != 0,
                "Writing::MESSAGE_QUEUE_DEPTH must not be 0"
            );

            let (conn_sender, conn_receiver) =
                mpsc::channel(self.node().config().max_connecting as usize);

            // the conn_senders are used to send messages from the Node to individual connections
            let conn_senders: WritingSenders = Default::default();
            // procure a clone to create the WritingHandler with
            let senders = conn_senders.clone();

            // the task spawning tasks sending messages to all the streams
            let self_clone = self.clone();
            let handler_loop = async move {
                let node = self_clone.node().clone();
                run_setup_handler_loop(node, "Writing", conn_receiver, |conn, setup_tasks| {
                    let self_clone = self_clone.clone();
                    let senders = conn_senders.clone();
                    setup_tasks.spawn(async move {
                        self_clone.handle_new_connection(conn, &senders).await;
                    });
                })
                .await;
            };

            install_protocol_handler(
                self.node(),
                NodeTask::Writing,
                "Writing",
                |protocols| &protocols.writing,
                WritingHandler {
                    handler: ProtocolHandler(conn_sender),
                    senders,
                },
                handler_loop,
            )
            .await;
        }
    }

    /// Creates an [`Encoder`] used to write the outbound messages to the target stream.
    /// The `side` param indicates the connection side **from the node's perspective**.
    fn codec(&self, addr: SocketAddr, side: ConnectionSide) -> Self::Codec;

    /// Sends the provided message to the specified [`SocketAddr`]. Returns as soon as the message
    /// is queued to be sent; the caller is provided with a [`oneshot::Receiver`] which resolves
    /// once the message's fate is known:
    ///
    /// - `Ok(Ok(()))`: the message was encoded and flushed into the OS socket send buffer. This
    ///   is **not** a peer-side acknowledgement - whether the peer received or processed the
    ///   message can only be determined via an application-level reply.
    /// - `Ok(Err(e))`: the write batch containing the message failed, and the connection is
    ///   going down. Since writes are batched, this does not guarantee that the message did not
    ///   reach the socket - preceding messages of the failed batch may have been flushed along
    ///   the way.
    /// - `Err(RecvError)`: the connection was torn down before the write concluded - either the
    ///   message was still queued (never written), or the writer task was aborted mid-batch, in
    ///   which case the message *may* have reached the socket. Like the `Ok(Err(_))` case, this
    ///   is not a safe "not sent" signal for deduplication purposes.
    ///
    /// # Errors
    ///
    /// The following errors can be returned:
    /// - [`io::ErrorKind::BrokenPipe`] if the outbound message channel is down
    /// - [`io::ErrorKind::NotConnected`] if the node is not connected to the provided address
    /// - [`io::ErrorKind::QuotaExceeded`] if the outbound message queue for this address is full
    /// - [`io::ErrorKind::Unsupported`] if [`Writing::enable_writing`] hadn't been called yet, or
    ///   if a duplicative owner of the underlying node has a different `Writing` impl (API misuse)
    fn unicast(
        &self,
        addr: SocketAddr,
        message: Self::Message,
    ) -> io::Result<oneshot::Receiver<io::Result<()>>> {
        self.queue_message(addr, message, true)
            .map(|delivery| delivery.unwrap()) // always requested, so infallible
    }

    /// Sends the provided message to the specified [`SocketAddr`], and returns as soon as the
    /// message is queued to be sent, providing no feedback on the write (as opposed to
    /// [`Writing::unicast`], which does).
    ///
    /// # Errors
    ///
    /// See the error section for [`Writing::unicast`].
    fn unicast_fast(&self, addr: SocketAddr, message: Self::Message) -> io::Result<()> {
        self.queue_message(addr, message, false).map(|_| ())
    }

    /// Broadcasts the provided message to all connected peers. Returns as soon as the message
    /// is queued to be sent to all the peers, without waiting for the actual delivery. For any
    /// peer whose queue is full or whose channel has been closed, the message is silently
    /// dropped and the failure is logged at `ERROR` level - the call still returns `Ok(())` for
    /// the broadcast as a whole and continues delivering to the remaining peers.
    ///
    /// If you need to know which peers received (or even queued) the message, call
    /// [`Writing::unicast`] for each address returned by [`Node::connected_addrs`] and inspect
    /// the returned [`oneshot::Receiver`]s.
    ///
    /// note: This method clones the message for every connected peer, and serialization via
    /// [`Writing::Codec`] happens individually for each connection. If your serialization is
    /// expensive (e.g., large JSON/Bincode structs), manually serialize your message into
    /// [`Bytes`] *before* calling broadcast. This ensures serialization happens only once,
    /// rather than N times.
    ///
    /// # Errors
    ///
    /// Returns [`io::ErrorKind::Unsupported`] if [`Writing::enable_writing`] hadn't been called
    /// yet, or if a duplicative owner of the underlying node has a different `Writing` impl
    /// (API misuse).
    #[deprecated(
        since = "0.57.0",
        note = "use Node::connected_addrs + unicast(_fast) depending on your specific use case"
    )]
    fn broadcast(&self, message: Self::Message) -> io::Result<()>
    where
        Self::Message: Clone,
    {
        // access the protocol handler
        let Some(handler) = self.node().protocols.writing.get() else {
            return Err(io::ErrorKind::Unsupported.into());
        };

        // collect only the addresses, so the senders map isn't cloned wholesale
        let addrs: Vec<_> = handler.senders.read().keys().copied().collect();
        for addr in addrs {
            if let Err(e) = self.queue_message(addr, message.clone(), false) {
                // API misuse concerns the whole broadcast; per-peer failures are
                // logged by `queue_message` and skipped, as documented
                if e.kind() == io::ErrorKind::Unsupported {
                    return Err(e);
                }
            }
        }

        Ok(())
    }
}

/// This trait is used to restrict access to methods that would otherwise be public in [`Writing`].
trait WritingInternal: Writing {
    /// Queues a message to be sent to the given address, returning the delivery notification
    /// receiver if one was requested. The error mapping is the contract documented on
    /// [`Writing::unicast`].
    fn queue_message(
        &self,
        addr: SocketAddr,
        message: Self::Message,
        confirm_delivery: bool,
    ) -> io::Result<Option<oneshot::Receiver<io::Result<()>>>>;

    /// Feeds a whole batch of messages into the stream's buffer, flushes once,
    /// and returns the number of written messages and their cumulative size in bytes.
    async fn write_batch<W: AsyncWrite + Unpin + Send>(
        &self,
        messages: &mut Vec<WrappedMessage<Self::Message>>,
        writer: &mut FramedWrite<W, Self::Codec>,
    ) -> io::Result<(usize, usize)>;

    /// Applies the [`Writing`] protocol to a single connection.
    async fn handle_new_connection(
        &self,
        conn_with_returner: ReturnableConnection,
        conn_senders: &WritingSenders,
    );
}

impl<W: Writing> WritingInternal for W {
    fn queue_message(
        &self,
        addr: SocketAddr,
        message: Self::Message,
        confirm_delivery: bool,
    ) -> io::Result<Option<oneshot::Receiver<io::Result<()>>>> {
        // access the protocol handler
        let Some(handler) = self.node().protocols.writing.get() else {
            return Err(io::ErrorKind::Unsupported.into());
        };

        // find the message sender for the given address; the guard may be held across the
        // send, since `try_send` never blocks
        let senders = handler.senders.read();
        let Some((_, erased)) = senders.get(&addr) else {
            return Err(io::ErrorKind::NotConnected.into());
        };
        let sender = erased
            .downcast_ref::<mpsc::Sender<WrappedMessage<Self::Message>>>()
            // under valid API use, same Self => same Message => TypeId always matches
            .ok_or(io::ErrorKind::Unsupported)?;

        let (msg, delivery) = WrappedMessage::new(message, confirm_delivery);
        let result = sender.try_send(msg);
        drop(senders);

        result.map(|_| delivery).map_err(|e| {
            let conn_span = create_connection_span(addr, self.node().span());
            error!(parent: conn_span, "can't send a message: {e}");
            match e {
                mpsc::error::TrySendError::Full(_) => io::ErrorKind::QuotaExceeded.into(),
                mpsc::error::TrySendError::Closed(_) => io::ErrorKind::BrokenPipe.into(),
            }
        })
    }

    async fn write_batch<A: AsyncWrite + Unpin + Send>(
        &self,
        messages: &mut Vec<WrappedMessage<Self::Message>>,
        writer: &mut FramedWrite<A, Self::Codec>,
    ) -> io::Result<(usize, usize)> {
        // Both the `feed`s and the final `flush` can block on a full socket send
        // buffer - `feed` flushes internally at the backpressure boundary - so
        // the entire sequence is guarded.
        let write = async move {
            let msgs = messages.len();
            let mut bytes = 0;

            for wrapped in messages {
                let msg = wrapped.msg.take().unwrap(); // guaranteed to be present here
                // `feed` is `poll_ready` followed by `start_send`, and `poll_ready` flushes the
                // *entire* write buffer once it reaches the backpressure boundary; the two steps
                // are therefore driven separately, so that the buffer is measured after any such
                // flush and the encoded size is exact. Diffing the buffer across the whole `feed`
                // can't tell a flush apart from a no-op, and undercounts by the flushed amount
                // whenever the message is at least as large as it - which, for messages at or
                // above the boundary, means every one of them but the first is counted as 0B.
                poll_fn(|cx| writer.poll_ready_unpin(cx)).await?;
                let prev = writer.write_buffer().len();
                writer.start_send_unpin(msg)?;
                bytes += writer.write_buffer().len().saturating_sub(prev);
            }
            writer.flush().await?;
            Ok::<_, io::Error>((msgs, bytes))
        };

        match timeout(Duration::from_millis(Self::TIMEOUT_MS), write).await {
            Ok(Ok(stats)) => Ok(stats),
            Ok(Err(e)) => Err(e),
            Err(_) => {
                self.node().heuristics().register_write_timeout();
                Err(io::Error::new(io::ErrorKind::TimedOut, "write timed out"))
            }
        }
    }

    async fn handle_new_connection(
        &self,
        (mut conn, conn_returner): ReturnableConnection,
        conn_senders: &WritingSenders,
    ) {
        let addr = conn.addr();
        let conn_id = conn.id;
        let codec = match catch_setup_panic(conn.span(), "Writing::codec", || {
            self.codec(addr, !conn.side())
        }) {
            Ok(codec) => codec,
            Err(err) => {
                let _ = conn_returner.send(Err(err));
                return;
            }
        };
        let Some(writer) = conn.writer.take() else {
            let err = io::Error::other("the stream was not returned during the handshake");
            error!(parent: conn.span(), "{err}");
            let _ = conn_returner.send(Err(err));
            return;
        };
        let mut framed = FramedWrite::new(writer, codec);

        if Self::INITIAL_BUFFER_SIZE != 0 {
            framed.write_buffer_mut().reserve(Self::INITIAL_BUFFER_SIZE);
            // the write buffer is flushed as soon as it reaches the backpressure boundary,
            // which defaults to 8KiB regardless of the buffer reserved above; raise it - but
            // never lower it - so that a batch can coalesce up to the requested size
            let boundary = framed
                .backpressure_boundary()
                .max(Self::INITIAL_BUFFER_SIZE);
            framed.set_backpressure_boundary(boundary);
        }

        let (outbound_message_sender, mut outbound_message_receiver) =
            mpsc::channel(Self::MESSAGE_QUEUE_DEPTH);

        // this will automatically drop the sender upon a disconnect
        let sender_cleanup = SenderCleanup {
            addr,
            conn_id,
            senders: Arc::clone(conn_senders),
        };

        // use a channel to know when the writer task is ready
        let (tx_writer, rx_writer) = oneshot::channel();

        // the task for writing outbound messages
        let self_clone = self.clone();
        let conn_stats = conn.stats().clone();
        let conn_span = conn.span().clone();
        let writer_task = tokio::spawn(Box::pin(async move {
            let node = self_clone.node();
            trace!(parent: &conn_span, "spawned a task for writing messages");
            if tx_writer.send(()).is_err() {
                error!(parent: &conn_span, "Writing was interrupted; shutting down its task");
                return;
            }

            // move the sender cleanup into this task
            let _sender_cleanup = sender_cleanup;

            // disconnect automatically regardless of how this task concludes
            let _conn_cleanup =
                DisconnectOnDrop::new(node.clone(), addr, conn_id, DisconnectOrigin::Writing);

            let writing = AssertUnwindSafe(async {
                let mut batch: Vec<WrappedMessage<Self::Message>> = Vec::new();
                // recv_many blocks for the first message, then takes whatever else is
                // already queued (capped by the queue depth) - so we coalesce without
                // ever waiting to fill a batch. Returns 0 only once the channel closes.
                while outbound_message_receiver
                    .recv_many(&mut batch, Self::MESSAGE_QUEUE_DEPTH)
                    .await
                    > 0
                {
                    match self_clone.write_batch(&mut batch, &mut framed).await {
                        Ok((msgs, bytes)) => {
                            for tx in batch.drain(..).filter_map(|w| w.delivery_notification) {
                                let _ = tx.send(Ok(()));
                            }
                            conn_stats.register_sent_messages(msgs, bytes);
                            node.stats().register_sent_messages(msgs, bytes);
                            trace!(parent: &conn_span, "wrote {bytes}B ({msgs} messages)");
                        }
                        Err(e) => {
                            error!(parent: &conn_span, "couldn't write a batch of {} message(s): {e}", batch.len());
                            // the connection is going down; fail every queued delivery
                            for tx in batch.drain(..).filter_map(|w| w.delivery_notification) {
                                let _ = tx.send(Err(io::Error::new(e.kind(), e.to_string())));
                            }
                            break;
                        }
                    }
                }
            });
            if let Err(payload) = writing.catch_unwind().await {
                error!(parent: &conn_span, "Writing::Codec panicked while encoding: {}", panic_message(&*payload));
                resume_unwind(payload);
            }
        }));
        let _ = rx_writer.await;

        // only now that the writer task is confirmed to be running: registering the sender
        // before the await above would mean this setup task getting aborted there (e.g. due to
        // node shutdown) leaks a running, unowned writer task, kept alive indefinitely by the
        // registered sender
        adopt_writer(
            &mut conn,
            conn_senders,
            conn_id,
            Arc::new(outbound_message_sender),
            writer_task,
        );

        // return the Connection to the Node, resuming Node::adapt_stream
        let conn_span = conn.span().clone();
        if conn_returner.send(Ok(conn)).is_err() {
            error!(parent: &conn_span, "couldn't return a Connection from the Writing handler");
        }
    }
}

/// Publishes the connection's outbound message sender and hands the writer task over to the
/// connection, which from then on owns it.
///
/// These two steps must not be separated by a yield point: if the calling setup task were
/// aborted between them, the published sender would keep a running, unowned writer task alive
/// for the lifetime of the node. This being a plain `fn` rather than inline `async` code is what
/// enforces that - an `.await` cannot be introduced here without changing the signature.
///
/// The publication still precedes `Connections::add` in `Node::adapt_stream`, so [`Writing`] is
/// usable slightly before the connection is fully finalized. That is deliberate: writing is
/// user-initiated and safe to accept early, and gating it would only add complexity.
fn adopt_writer(
    conn: &mut Connection,
    conn_senders: &WritingSenders,
    conn_id: u64,
    sender: Arc<dyn Any + Send + Sync>,
    writer_task: JoinHandle<()>,
) {
    conn_senders.write().insert(conn.addr(), (conn_id, sender));
    conn.tasks.push(writer_task);
}

/// Used to queue messages for delivery and return its confirmation.
pub(crate) struct WrappedMessage<T> {
    msg: Option<T>,
    delivery_notification: Option<oneshot::Sender<io::Result<()>>>,
}

impl<T> WrappedMessage<T> {
    fn new(msg: T, confirmation: bool) -> (Self, Option<oneshot::Receiver<io::Result<()>>>) {
        let (tx, rx) = if confirmation {
            let (tx, rx) = oneshot::channel();
            (Some(tx), Some(rx))
        } else {
            (None, None)
        };

        let wrapped_msg = Self {
            msg: Some(msg),
            delivery_notification: tx,
        };

        (wrapped_msg, rx)
    }
}

/// The handler object dedicated to the [`Writing`] protocol.
pub(crate) struct WritingHandler {
    handler: ProtocolHandler<Connection, io::Result<Connection>>,
    pub(crate) senders: WritingSenders,
}

impl WritingHandler {
    /// See [`ProtocolHandler::closed`].
    pub(crate) async fn closed(&self) {
        self.handler.closed().await;
    }
}

impl Protocol<Connection, io::Result<Connection>> for WritingHandler {
    async fn trigger(&self, item: ReturnableConnection) {
        self.handler.trigger(item).await;
    }
}

struct SenderCleanup {
    addr: SocketAddr,
    senders: WritingSenders,
    /// The owning connection's instance id (`Connection::id`).
    conn_id: u64,
}

impl Drop for SenderCleanup {
    fn drop(&mut self) {
        let mut map = self.senders.write();
        if let Entry::Occupied(e) = map.entry(self.addr) {
            // only remove if this is still *our* sender; otherwise a newer
            // connection has reused the addr and we must not touch it
            if e.get().0 == self.conn_id {
                e.remove();
            }
        }
    }
}