cdrs-tokio 9.0.2

//!This module contains a declaration of `CdrsTransport` trait which should be implemented
//!for particular transport to be able using it as transport of CDRS client.
//!
//!Currently, CDRS provides concrete transport which implements `CdrsTransport` trait.
//!There are:
//!
//! * [`TransportTcp`] is default TCP transport which is usually used to establish
//!connection and exchange frames.
//!
//! * [`TransportRustls`] is a transport which is used to establish SSL encrypted connection
//!with Apache Cassandra server. **Note:** this option is available if and only if CDRS is imported
//!with the `rust-tls` feature.
use cassandra_protocol::compression::Compression;
use cassandra_protocol::frame::frame_decoder::FrameDecoder;
use cassandra_protocol::frame::frame_encoder::FrameEncoder;
use cassandra_protocol::frame::message_result::ResultKind;
use cassandra_protocol::frame::{Envelope, StreamId, MAX_FRAME_SIZE};
use cassandra_protocol::frame::{FromBytes, Opcode, EVENT_STREAM_ID};
use cassandra_protocol::types::INT_LEN;
use derive_more::Constructor;
use futures::FutureExt;
use fxhash::FxHashMap;
use itertools::Itertools;
use std::io;
use std::net::SocketAddr;
use std::sync::atomic::{AtomicBool, AtomicI16, Ordering};
use std::sync::{Arc, Mutex};
use tokio::io::{
    split, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader, BufWriter, ReadHalf,
    WriteHalf,
};
use tokio::net::TcpStream;
use tokio::sync::{mpsc, oneshot};
use tokio::task::JoinHandle;
#[cfg(feature = "rust-tls")]
use tokio_rustls::rustls::{pki_types::ServerName, ClientConfig};
#[cfg(feature = "rust-tls")]
use tokio_rustls::TlsConnector as RustlsConnector;
use tracing::*;

#[cfg(test)]
use mockall::*;

use crate::cluster::KeyspaceHolder;
use crate::envelope_parser::{convert_envelope_into_result, parse_envelope};
use crate::future::BoxFuture;
use crate::Error;
use crate::Result;

const INITIAL_STREAM_ID: i16 = 1;

/// General CDRS transport trait.
pub trait CdrsTransport: Send + Sync {
    /// Schedules data envelope for writing and waits for a response. Handshake envelopes need to
    /// be marked as such, since their wire representation is different.
    fn write_envelope<'a>(
        &'a self,
        envelope: &'a Envelope,
        handshake: bool,
    ) -> BoxFuture<'a, Result<Envelope>>;

    /// Checks if the connection is broken (e.g. after read or write errors).
    fn is_broken(&self) -> bool;

    /// Returns associated node address.
    fn address(&self) -> SocketAddr;
}

#[cfg(test)]
mock! {
    pub CdrsTransport {
    }

    impl CdrsTransport for CdrsTransport {
        fn write_envelope(
            &self,
            envelope: &Envelope,
            handshake: bool,
        ) -> BoxFuture<'static, Result<Envelope>>;

        fn is_broken(&self) -> bool;

        fn address(&self) -> SocketAddr;
    }
}

/// Default Tcp transport.
#[derive(Debug)]
pub struct TransportTcp {
    inner: AsyncTransport,
}

impl TransportTcp {
    #[allow(clippy::too_many_arguments)]
    pub async fn new(
        addr: SocketAddr,
        keyspace_holder: Arc<KeyspaceHolder>,
        event_handler: Option<mpsc::Sender<Envelope>>,
        error_handler: Option<mpsc::Sender<Error>>,
        compression: Compression,
        frame_encoder: Box<dyn FrameEncoder + Send + Sync>,
        frame_decoder: Box<dyn FrameDecoder + Send + Sync>,
        buffer_size: usize,
        tcp_nodelay: bool,
    ) -> io::Result<TransportTcp> {
        TcpStream::connect(addr).await.and_then(move |socket| {
            socket.set_nodelay(tcp_nodelay)?;
            Self::with_stream(
                socket,
                addr,
                keyspace_holder,
                event_handler,
                error_handler,
                compression,
                frame_encoder,
                frame_decoder,
                buffer_size,
            )
        })
    }

    #[allow(clippy::too_many_arguments)]
    pub fn with_stream<T: AsyncRead + AsyncWrite + Send + Sync + 'static>(
        stream: T,
        addr: SocketAddr,
        keyspace_holder: Arc<KeyspaceHolder>,
        event_handler: Option<mpsc::Sender<Envelope>>,
        error_handler: Option<mpsc::Sender<Error>>,
        compression: Compression,
        frame_encoder: Box<dyn FrameEncoder + Send + Sync>,
        frame_decoder: Box<dyn FrameDecoder + Send + Sync>,
        buffer_size: usize,
    ) -> io::Result<TransportTcp> {
        let (read_half, write_half) = split(stream);
        Ok(TransportTcp {
            inner: AsyncTransport::new(
                addr,
                compression,
                frame_encoder,
                frame_decoder,
                buffer_size,
                read_half,
                write_half,
                event_handler,
                error_handler,
                keyspace_holder,
            ),
        })
    }
}

impl CdrsTransport for TransportTcp {
    //noinspection DuplicatedCode
    #[inline]
    fn write_envelope<'a>(
        &'a self,
        envelope: &'a Envelope,
        handshake: bool,
    ) -> BoxFuture<'a, Result<Envelope>> {
        self.inner.write_envelope(envelope, handshake).boxed()
    }

    #[inline]
    fn is_broken(&self) -> bool {
        self.inner.is_broken()
    }

    #[inline]
    fn address(&self) -> SocketAddr {
        self.inner.addr()
    }
}

#[cfg(feature = "rust-tls")]
#[derive(Debug)]
pub struct TransportRustls {
    inner: AsyncTransport,
}

#[cfg(feature = "rust-tls")]
impl TransportRustls {
    #[allow(clippy::too_many_arguments)]
    pub async fn new(
        addr: SocketAddr,
        dns_name: ServerName<'static>,
        config: Arc<ClientConfig>,
        keyspace_holder: Arc<KeyspaceHolder>,
        event_handler: Option<mpsc::Sender<Envelope>>,
        error_handler: Option<mpsc::Sender<Error>>,
        compression: Compression,
        frame_encoder: Box<dyn FrameEncoder + Send + Sync>,
        frame_decoder: Box<dyn FrameDecoder + Send + Sync>,
        buffer_size: usize,
        tcp_nodelay: bool,
    ) -> io::Result<Self> {
        let stream = TcpStream::connect(addr).await?;
        stream.set_nodelay(tcp_nodelay)?;

        Self::with_stream(
            stream,
            addr,
            dns_name,
            config,
            keyspace_holder,
            event_handler,
            error_handler,
            compression,
            frame_encoder,
            frame_decoder,
            buffer_size,
        )
        .await
    }

    #[allow(clippy::too_many_arguments)]
    pub async fn with_stream<T: AsyncRead + AsyncWrite + Send + Sync + Unpin + 'static>(
        stream: T,
        addr: SocketAddr,
        dns_name: ServerName<'static>,
        config: Arc<ClientConfig>,
        keyspace_holder: Arc<KeyspaceHolder>,
        event_handler: Option<mpsc::Sender<Envelope>>,
        error_handler: Option<mpsc::Sender<Error>>,
        compression: Compression,
        frame_encoder: Box<dyn FrameEncoder + Send + Sync>,
        frame_decoder: Box<dyn FrameDecoder + Send + Sync>,
        buffer_size: usize,
    ) -> io::Result<Self> {
        let connector = RustlsConnector::from(config.clone());
        let stream = connector.connect(dns_name, stream).await?;
        let (read_half, write_half) = split(stream);

        Ok(Self {
            inner: AsyncTransport::new(
                addr,
                compression,
                frame_encoder,
                frame_decoder,
                buffer_size,
                read_half,
                write_half,
                event_handler,
                error_handler,
                keyspace_holder,
            ),
        })
    }
}

#[cfg(feature = "rust-tls")]
impl CdrsTransport for TransportRustls {
    //noinspection DuplicatedCode
    #[inline]
    fn write_envelope<'a>(
        &'a self,
        envelope: &'a Envelope,
        handshake: bool,
    ) -> BoxFuture<'a, Result<Envelope>> {
        self.inner.write_envelope(envelope, handshake).boxed()
    }

    #[inline]
    fn is_broken(&self) -> bool {
        self.inner.is_broken()
    }

    #[inline]
    fn address(&self) -> SocketAddr {
        self.inner.addr()
    }
}

#[derive(Debug)]
struct AsyncTransport {
    addr: SocketAddr,
    compression: Compression,
    write_sender: mpsc::Sender<Request>,
    is_broken: Arc<AtomicBool>,
    processing_handle: JoinHandle<()>,
}

impl Drop for AsyncTransport {
    fn drop(&mut self) {
        self.processing_handle.abort();
    }
}

impl AsyncTransport {
    #[allow(clippy::too_many_arguments)]
    fn new<T: AsyncRead + AsyncWrite + Send + 'static>(
        addr: SocketAddr,
        compression: Compression,
        frame_encoder: Box<dyn FrameEncoder + Send + Sync>,
        frame_decoder: Box<dyn FrameDecoder + Send + Sync>,
        buffer_size: usize,
        read_half: ReadHalf<T>,
        write_half: WriteHalf<T>,
        event_handler: Option<mpsc::Sender<Envelope>>,
        error_handler: Option<mpsc::Sender<Error>>,
        keyspace_holder: Arc<KeyspaceHolder>,
    ) -> Self {
        let (write_sender, write_receiver) = mpsc::channel(buffer_size);
        let is_broken = Arc::new(AtomicBool::new(false));

        let processing_handle = tokio::spawn(Self::start_processing(
            write_receiver,
            event_handler,
            error_handler,
            read_half,
            write_half,
            keyspace_holder,
            is_broken.clone(),
            compression,
            addr,
            frame_encoder,
            frame_decoder,
        ));

        AsyncTransport {
            addr,
            compression,
            write_sender,
            is_broken,
            processing_handle,
        }
    }

    #[inline]
    fn is_broken(&self) -> bool {
        self.is_broken.load(Ordering::Relaxed)
    }

    #[inline]
    fn addr(&self) -> SocketAddr {
        self.addr
    }

    async fn write_envelope(&self, envelope: &Envelope, handshake: bool) -> Result<Envelope> {
        let (sender, receiver) = oneshot::channel();

        // leave stream id empty for now and generate it later

        // handshake messages are never compressed
        let data = if handshake {
            envelope.encode_with(Compression::None)?
        } else {
            envelope.encode_with(self.compression)?
        };

        self.write_sender
            .send(Request::new(data, sender, handshake))
            .await
            .map_err(|_| Error::General("Connection closed when writing data!".into()))?;

        receiver
            .await
            .map_err(|_| Error::General("Connection closed while waiting for response!".into()))?
    }

    #[allow(clippy::too_many_arguments)]
    async fn start_processing<T: AsyncRead + AsyncWrite>(
        write_receiver: mpsc::Receiver<Request>,
        event_handler: Option<mpsc::Sender<Envelope>>,
        error_handler: Option<mpsc::Sender<Error>>,
        read_half: ReadHalf<T>,
        write_half: WriteHalf<T>,
        keyspace_holder: Arc<KeyspaceHolder>,
        is_broken: Arc<AtomicBool>,
        compression: Compression,
        addr: SocketAddr,
        frame_encoder: Box<dyn FrameEncoder + Send + Sync>,
        frame_decoder: Box<dyn FrameDecoder + Send + Sync>,
    ) {
        let response_handler_map = ResponseHandlerMap::new();

        let writer = Self::start_writing(
            write_receiver,
            BufWriter::new(write_half),
            &response_handler_map,
            frame_encoder,
        );

        let reader = Self::start_reading_handshake_frames(
            BufReader::with_capacity(MAX_FRAME_SIZE, read_half),
            event_handler,
            compression,
            addr,
            keyspace_holder,
            &response_handler_map,
            frame_decoder,
        );

        let result = tokio::try_join!(writer, reader);
        if let Err(error) = result {
            error!(%error, "Transport error!");

            is_broken.store(true, Ordering::Relaxed);
            response_handler_map.signal_general_error(&error);

            if let Some(error_handler) = error_handler {
                match error_handler.try_send(error) {
                    Ok(_) => debug!("Error handler notified!"),
                    Err(e) => warn!("Error handler failed to notify: {e}"),
                }
            }
        }
    }

    async fn start_reading_handshake_frames(
        mut read_half: impl AsyncRead + Unpin,
        event_handler: Option<mpsc::Sender<Envelope>>,
        compression: Compression,
        addr: SocketAddr,
        keyspace_holder: Arc<KeyspaceHolder>,
        response_handler_map: &ResponseHandlerMap,
        frame_decoder: Box<dyn FrameDecoder + Send + Sync>,
    ) -> Result<()> {
        // before Authenticate or Ready, envelopes are unframed
        loop {
            let result = parse_envelope(&mut read_half, compression, addr).await;
            match result {
                Ok(envelope) => {
                    if envelope.stream_id >= 0 {
                        let opcode = envelope.opcode;
                        response_handler_map.send_response(envelope.stream_id, Ok(envelope))?;

                        if opcode == Opcode::Authenticate || opcode == Opcode::Ready {
                            // all frames should now be encoded
                            return Self::start_reading_normal_frames(
                                read_half,
                                event_handler,
                                compression,
                                addr,
                                keyspace_holder,
                                response_handler_map,
                                frame_decoder,
                            )
                            .await;
                        }
                    } else if envelope.stream_id == EVENT_STREAM_ID {
                        // server event
                        if let Some(event_handler) = &event_handler {
                            let _ = event_handler.send(envelope).await;
                        }
                    }
                }
                Err(error) => return Err(error),
            }
        }
    }

    async fn start_reading_normal_frames(
        mut read_half: impl AsyncRead + Unpin,
        event_handler: Option<mpsc::Sender<Envelope>>,
        compression: Compression,
        addr: SocketAddr,
        keyspace_holder: Arc<KeyspaceHolder>,
        response_handler_map: &ResponseHandlerMap,
        mut frame_decoder: Box<dyn FrameDecoder + Send + Sync>,
    ) -> Result<()> {
        let mut buffer = Vec::with_capacity(MAX_FRAME_SIZE);
        loop {
            let num_read = read_half.read_buf(&mut buffer).await?;
            if num_read == 0 {
                break Err(Error::Io(io::Error::new(
                    io::ErrorKind::UnexpectedEof,
                    "EOF",
                )));
            }

            let envelopes = frame_decoder.consume(&mut buffer, compression)?;
            for envelope in envelopes {
                if envelope.stream_id >= 0 {
                    // in case we get a SetKeyspace result, we need to store current keyspace
                    // checks are done manually for speed
                    if envelope.opcode == Opcode::Result {
                        let result_kind = ResultKind::from_bytes(&envelope.body[..INT_LEN])?;
                        if result_kind == ResultKind::SetKeyspace {
                            let response_body = envelope.response_body()?;
                            let set_keyspace =
                                response_body.into_set_keyspace().ok_or_else(|| {
                                    Error::General(
                                        "SetKeyspace not found with SetKeyspace opcode!".into(),
                                    )
                                })?;

                            keyspace_holder.update_current_keyspace(set_keyspace.body);
                        }
                    }

                    // normal response to query
                    response_handler_map.send_response(
                        envelope.stream_id,
                        convert_envelope_into_result(envelope, addr),
                    )?;
                } else if envelope.stream_id == EVENT_STREAM_ID {
                    // server event
                    if let Some(event_handler) = &event_handler {
                        let _ = event_handler.send(envelope).await;
                    }
                }
            }
        }
    }

    async fn start_writing(
        mut write_receiver: mpsc::Receiver<Request>,
        mut write_half: impl AsyncWrite + Unpin,
        response_handler_map: &ResponseHandlerMap,
        mut frame_encoder: Box<dyn FrameEncoder + Send + Sync>,
    ) -> Result<()> {
        let mut frame_stream_ids = Vec::with_capacity(1);

        while let Some(mut request) = write_receiver.recv().await {
            frame_stream_ids.clear();

            loop {
                let stream_id = response_handler_map.next_stream_id();
                frame_stream_ids.push(stream_id);

                request.set_stream_id(stream_id);
                response_handler_map.add_handler(stream_id, request.handler);

                if request.handshake {
                    // handshake messages are not framed, so let's write them directly
                    if let Err(error) = write_half.write_all(&request.data).await {
                        response_handler_map.send_response(stream_id, Err(error.into()))?;
                        return Err(Error::General("Write channel failure!".into()));
                    }
                } else {
                    // post-handshake messages can be aggregated in frames by the encoder
                    loop {
                        if frame_encoder.can_fit(request.data.len()) {
                            frame_encoder.add_envelope(request.data);
                            break;
                        }

                        // flush the previous frame or create a non-self-contained one
                        if frame_encoder.has_envelopes() {
                            // we have some envelopes => flush current frame
                            Self::write_self_contained_frame(
                                &mut write_half,
                                response_handler_map,
                                &mut frame_stream_ids,
                                frame_encoder.as_mut(),
                            )
                            .await?;
                        } else {
                            // non-self-contained
                            let data_len = request.data.len();
                            let mut data_start = 0;

                            while data_start < data_len {
                                let (data_start_offset, frame) = frame_encoder
                                    .finalize_non_self_contained(&request.data[data_start..]);

                                data_start += data_start_offset;

                                Self::write_frame(
                                    &mut write_half,
                                    response_handler_map,
                                    &mut frame_stream_ids,
                                    frame,
                                )
                                .await?;

                                frame_encoder.reset();
                            }

                            break;
                        }
                    }
                }

                request = match write_receiver.try_recv() {
                    Ok(request) => request,
                    Err(_) => {
                        if frame_encoder.has_envelopes() {
                            Self::write_self_contained_frame(
                                &mut write_half,
                                response_handler_map,
                                &mut frame_stream_ids,
                                frame_encoder.as_mut(),
                            )
                            .await?;
                        }

                        if let Err(error) = write_half.flush().await {
                            Self::notify_error_handlers(
                                response_handler_map,
                                &mut frame_stream_ids,
                                error.into(),
                            )?;
                            return Err(Error::General("Write channel failure!".into()));
                        }

                        break;
                    }
                }
            }
        }

        Ok(())
    }

    async fn write_self_contained_frame(
        write_half: &mut (impl AsyncWrite + Unpin),
        response_handler_map: &ResponseHandlerMap,
        frame_stream_ids: &mut Vec<StreamId>,
        frame_encoder: &mut (dyn FrameEncoder + Send + Sync),
    ) -> Result<()> {
        Self::write_frame(
            write_half,
            response_handler_map,
            frame_stream_ids,
            frame_encoder.finalize_self_contained(),
        )
        .await?;

        frame_encoder.reset();

        // Drop the stream ids that just left over the wire. Their handlers
        // are now waiting for actual server responses (or the connection
        // tearing down), and must NOT be notified of write failures from a
        // *subsequent* frame in the same batch. Without this clear we'd
        // accumulate stream ids across every flushed frame and a single late
        // write error would falsely fail every previously-sent envelope.
        frame_stream_ids.clear();
        Ok(())
    }

    async fn write_frame(
        write_half: &mut (impl AsyncWrite + Unpin),
        response_handler_map: &ResponseHandlerMap,
        frame_stream_ids: &mut Vec<StreamId>,
        frame: &[u8],
    ) -> Result<()> {
        // If the underlying socket write fails, fan the error out to every
        // handler whose envelope is in this frame, and then propagate the
        // failure so the writer task can shut down cleanly. The previous
        // implementation returned Ok(()) here, swallowing the original write
        // error, so the writer would happily keep looping on a dead socket.
        if let Err(error) = write_half.write_all(frame).await {
            let propagated: Error = error.into();
            Self::notify_error_handlers(
                response_handler_map,
                frame_stream_ids,
                propagated.clone(),
            )?;
            return Err(propagated);
        }

        Ok(())
    }

    fn notify_error_handlers(
        response_handler_map: &ResponseHandlerMap,
        frame_stream_ids: &mut Vec<StreamId>,
        error: Error,
    ) -> Result<()> {
        frame_stream_ids
            .drain(..)
            .map(|stream_id| response_handler_map.send_response(stream_id, Err(error.clone())))
            .try_collect()
    }
}

type ResponseHandler = oneshot::Sender<Result<Envelope>>;

struct ResponseHandlerMap {
    stream_handlers: Mutex<FxHashMap<StreamId, ResponseHandler>>,
    available_stream_id: AtomicI16,
}

impl ResponseHandlerMap {
    #[inline]
    pub fn new() -> Self {
        ResponseHandlerMap {
            stream_handlers: Default::default(),
            available_stream_id: AtomicI16::new(INITIAL_STREAM_ID),
        }
    }

    #[inline]
    pub fn add_handler(&self, stream_id: StreamId, handler: ResponseHandler) {
        self.stream_handlers
            .lock()
            .unwrap()
            .insert(stream_id, handler);
    }

    pub fn send_response(&self, stream_id: StreamId, response: Result<Envelope>) -> Result<()> {
        match self.stream_handlers.lock().unwrap().remove(&stream_id) {
            Some(handler) => {
                let _ = handler.send(response);
                Ok(())
            }
            // unmatched stream - probably a bug somewhere
            None => Err(Error::General(format!("Unmatched stream id: {stream_id}"))),
        }
    }

    pub fn signal_general_error(&self, error: &Error) {
        for (_, handler) in self.stream_handlers.lock().unwrap().drain() {
            let _ = handler.send(Err(error.clone()));
        }
    }

    pub fn next_stream_id(&self) -> StreamId {
        // We allocate stream ids in [INITIAL_STREAM_ID, i16::MAX] inclusive,
        // wrapping back to INITIAL_STREAM_ID once the maximum has been used.
        //
        // The previous implementation called `fetch_add` and then tried to
        // compare-and-swap the *pre-increment* value back to INITIAL when it
        // saw a negative result. By construction the post-increment value is
        // already on the counter at that point, so the CAS expected the wrong
        // value and almost never succeeded under contention - the counter
        // would walk through ~32K negative values before yielding usable ids
        // again, and could yield stream id 0 (which we never want).
        //
        // Using a CAS loop that observes the current value and atomically
        // computes both the value to return AND the next counter state keeps
        // every returned id strictly inside the allowed range and skips no
        // ids on wrap.
        //
        // Memory ordering: the caller publishes the freshly-allocated
        // stream id by writing it into the request body and then registering
        // a handler in the Mutex-protected stream_handlers map. The Mutex
        // already provides the necessary happens-before for the reader, so
        // Relaxed would be sufficient today; we use Release on the success
        // leg as defence in depth so any future change to a lock-free
        // handler map keeps the publish ordering intact.
        loop {
            let current = self.available_stream_id.load(Ordering::Relaxed);
            let (return_value, new_value) = if current < INITIAL_STREAM_ID {
                // defensive: counter somehow ended up below INITIAL (e.g. a
                // future caller stored a bad value). Recover by snapping back
                // to the start of the range.
                (INITIAL_STREAM_ID, INITIAL_STREAM_ID + 1)
            } else if current == i16::MAX {
                // last id in the range - return it and wrap the counter
                // straight back to INITIAL_STREAM_ID for the next caller.
                (i16::MAX, INITIAL_STREAM_ID)
            } else {
                (current, current + 1)
            };

            if self
                .available_stream_id
                .compare_exchange_weak(current, new_value, Ordering::Release, Ordering::Relaxed)
                .is_ok()
            {
                return return_value;
            }
            // CAS lost the race; another thread updated the counter. Reload
            // and try again - we'll get the next id in sequence.
        }
    }
}

#[derive(Constructor)]
struct Request {
    data: Vec<u8>,
    handler: ResponseHandler,
    handshake: bool,
}

impl Request {
    #[inline]
    fn set_stream_id(&mut self, stream_d: StreamId) {
        self.data[2..4].copy_from_slice(&stream_d.to_be_bytes());
    }
}

#[cfg(test)]
mod write_buffer_tests {
    use super::*;
    use cassandra_protocol::frame::frame_encoder::UncompressedFrameEncoder;
    use tokio::io::sink;

    // Verifies that once a frame has been successfully written, the stream
    // ids that were aggregated into that frame are dropped from the
    // bookkeeping vector. Otherwise a later write failure would notify those
    // handlers - whose envelopes have already been sent and may yet receive
    // a real server response - with a spurious "Write channel failure".
    #[tokio::test]
    async fn write_self_contained_frame_clears_stream_ids_on_success() {
        let map = ResponseHandlerMap::new();
        let mut sink = sink();
        let mut encoder: Box<dyn FrameEncoder + Send + Sync> =
            Box::new(UncompressedFrameEncoder::default());

        // Pretend two envelopes' worth of work has been added to the encoder.
        // We don't actually need to register handlers on the map for this
        // test - we're checking that the frame_stream_ids buffer is emptied
        // once the frame goes out, regardless of map state.
        encoder.add_envelope(vec![0; 16]);
        let mut frame_stream_ids: Vec<StreamId> = vec![1, 2];

        AsyncTransport::write_self_contained_frame(
            &mut sink,
            &map,
            &mut frame_stream_ids,
            encoder.as_mut(),
        )
        .await
        .expect("sink writes always succeed");

        // After a successful frame write the vector must be empty so the
        // *next* frame's failure can't accidentally notify handlers from
        // already-sent frames.
        assert!(
            frame_stream_ids.is_empty(),
            "frame_stream_ids should be empty after successful write, was {:?}",
            frame_stream_ids
        );
    }
}

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

    #[test]
    fn next_stream_id_starts_at_initial_value() {
        let map = ResponseHandlerMap::new();
        assert_eq!(map.next_stream_id(), INITIAL_STREAM_ID);
        assert_eq!(map.next_stream_id(), INITIAL_STREAM_ID + 1);
    }

    #[test]
    fn next_stream_id_wraps_to_initial_after_overflow() {
        let map = ResponseHandlerMap::new();
        // arrange the counter so the very next fetch_add overflows i16
        map.available_stream_id.store(i16::MAX, Ordering::Relaxed);

        // the call right at the boundary still returns the last positive id
        assert_eq!(map.next_stream_id(), i16::MAX);

        // after the wrap, all returned ids must remain valid (positive) and
        // the sequence must restart at INITIAL_STREAM_ID. Without the fix the
        // function would either burn ~32K negative ids before yielding 0, or
        // return 0 itself, neither of which is correct for our protocol.
        assert_eq!(map.next_stream_id(), INITIAL_STREAM_ID);
        assert_eq!(map.next_stream_id(), INITIAL_STREAM_ID + 1);
        assert_eq!(map.next_stream_id(), INITIAL_STREAM_ID + 2);
    }
}