aerospike-core 2.2.0

// Copyright 2015-2018 Aerospike, Inc.
//
// Portions may be licensed to Aerospike, Inc. under one or more contributor
// license agreements.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not
// use this file except in compliance with the License. You may obtain a copy of
// the License at http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations under
// the License.

// The timeout arms below match `Err(_)` because `io_with_timeout!` yields `()`
// as its timeout error under rt-tokio and `Elapsed` under rt-async-std; naming
// the unit explicitly, as `clippy::ignored_unit_patterns` asks, would not
// compile for the async-std runtime.
#![allow(clippy::ignored_unit_patterns)]

#[cfg(feature = "tls")]
use std::convert::TryFrom;
#[cfg(feature = "tls")]
use std::sync::Arc;

use crate::commands::admin_command::AdminCommand;
use crate::commands::buffer::{self, Buffer, MAX_BUFFER_SIZE};
use crate::errors::{Error, Result};
use crate::net::Host;
use crate::policy::{AuthMode, ClientPolicy};
#[cfg(feature = "rt-async-std")]
use aerospike_rt::async_std::net::Shutdown;
#[cfg(feature = "rt-tokio")]
use aerospike_rt::io::{AsyncReadExt, AsyncWriteExt};
use aerospike_rt::net::TcpStream;
use aerospike_rt::time::{Duration, Instant};
#[cfg(feature = "rt-async-std")]
use futures::{AsyncReadExt, AsyncWriteExt};
use std::cmp::min;
use std::ops::Add;

#[cfg(feature = "tls")]
use rustls::pki_types::ServerName;
#[cfg(feature = "tls")]
use tokio_rustls::{client::TlsStream, rustls, TlsConnector};

/// State of a connection in the wire protocol.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ConnectionState {
    /// Connection is idle and ready for a command.
    Ready,
    /// Connection is closed.
    Closed,
    /// Writing request data.
    Writing,
    /// Reading response header (payload size in bytes).
    ReadingHeader(usize),
    /// Reading response body.
    ReadingBody(usize),
    /// Reading stream response header.
    ReadingStreamHeader(usize),
    /// Reading stream response body.
    ReadingStreamBody(usize),
}

/// Result of a pool-checkout liveness peek.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Liveness {
    /// Socket open with nothing pending.
    Alive,
    /// Socket open, but bytes are waiting that nobody asked for.
    PendingBytes,
    /// Peer closed the connection, or the socket is broken.
    Closed,
}

/// Underlying socket type for a connection (TCP or TLS).
#[derive(Debug)]
#[cfg_attr(test, allow(dead_code))]
#[allow(clippy::large_enum_variant)]
pub enum Netsocket {
    /// Plain TCP stream.
    Tcp(TcpStream),
    /// TLS-wrapped TCP stream.
    #[cfg(feature = "tls")]
    Tls(TlsStream<TcpStream>),
    /// Test double (tests only).
    #[cfg(test)]
    TestDummy,
}

#[cfg(test)]
thread_local! {
    /// Makes the next [`Connection::new`] fail once, then resets itself.
    ///
    /// The test double never fails, so error paths that only run when
    /// connecting fails — notably the reservation bookkeeping in
    /// [`crate::net::connection_pool::ConnectionPool::make_conn`] — would
    /// otherwise be unreachable in unit tests. `Connection::new` is awaited
    /// inline by its callers, so it runs on the thread that set this.
    pub(crate) static FAIL_NEXT_CONNECT: std::cell::Cell<bool> =
        const { std::cell::Cell::new(false) };
}

#[derive(Debug)]
#[allow(clippy::struct_field_names)]
pub struct Connection {
    pub(crate) addr: String,
    socket_timeout: u32,
    deadline: Option<Instant>,
    timeout_delay: u32,
    // duration after which connection is considered idle
    idle_timeout: Option<Duration>,
    idle_deadline: Option<Instant>,

    // connection object
    pub(crate) conn: Netsocket,

    bytes_read: usize,

    pub buffer: Buffer,

    pub(crate) state: ConnectionState,
    can_recover_connection: bool,

    /// Reusable per-IO timer, reset before each read/write and raced against the
    /// IO future.
    ///
    /// `aerospike_rt::timeout` builds a fresh `Sleep` for every operation, which
    /// allocates and then registers and removes an entry in tokio's timer wheel
    /// — a per-slot lock that every connection on the runtime contends for. One
    /// timer per connection, reset in place, does the same job without touching
    /// the wheel's structure on the hot path.
    #[cfg(feature = "rt-tokio")]
    pub(crate) sleep: std::pin::Pin<Box<aerospike_rt::tokio::time::Sleep>>,
}

/// Races an IO future against the connection's timeout.
///
/// `$holder` owns the `sleep` field: `self` inside [`Connection`], `self.conn`
/// inside [`BufferedConn`]. Yields `Ok(io_result)` or `Err(())` on timeout, so
/// it substitutes for `aerospike_rt::timeout(..).await` at every call site.
///
/// `biased` polls the IO first: when data is already available the timer is not
/// polled at all, and a completed read never loses a race to an expired timer.
macro_rules! io_with_timeout {
    ($holder:expr, $timeout:expr, $io:expr) => {{
        #[cfg(feature = "rt-tokio")]
        {
            $holder
                .sleep
                .as_mut()
                .reset(aerospike_rt::tokio::time::Instant::now() + $timeout);
            let sleep = $holder.sleep.as_mut();
            aerospike_rt::tokio::select! {
                biased;
                r = $io => Ok::<_, ()>(r),
                _ = sleep => Err(()),
            }
        }
        #[cfg(feature = "rt-async-std")]
        {
            aerospike_rt::timeout($timeout, $io).await
        }
    }};
}

impl Connection {
    #[cfg_attr(test, allow(dead_code))]
    #[cfg(feature = "tls")]
    async fn get_netsocket(
        stream: TcpStream,
        host: &Host,
        policy: &ClientPolicy,
    ) -> Result<Netsocket> {
        if let Some(tls_config) = policy.tls_config.clone() {
            let connector = TlsConnector::from(Arc::new(tls_config));
            let server_name = host
                .tls_name
                .clone()
                .unwrap_or_else(|| policy.cluster_name.clone().unwrap_or_default());
            let domain = ServerName::try_from(server_name.as_str())
                .map_err(|e| Error::ClientError(e.to_string()))?
                .to_owned();
            Ok(Netsocket::Tls(connector.connect(domain, stream).await?))
        } else {
            Ok(Netsocket::Tcp(stream))
        }
    }

    #[cfg(not(feature = "tls"))]
    async fn get_netsocket(
        stream: TcpStream,
        _host: &Host,
        _policy: &ClientPolicy,
    ) -> Result<Netsocket> {
        Ok(Netsocket::Tcp(stream))
    }

    #[cfg(not(test))]
    pub async fn new(
        host: &Host,
        policy: &ClientPolicy,
        hashed_pass: Option<&String>,
    ) -> Result<Self> {
        let addr = host.address();
        let stream =
            aerospike_rt::timeout(policy.timeout(), TcpStream::connect(addr.clone())).await;
        if stream.is_err() {
            return Err(Error::Connection(
                "Could not open network connection".to_string(),
            ));
        }

        let stream = stream.unwrap()?;
        let stream = Self::get_netsocket(stream, host, policy).await?;

        let idle_timeout = if policy.idle_timeout > 0 {
            Some(Duration::from_millis(u64::from(policy.idle_timeout)))
        } else {
            None
        };

        let mut conn = Connection {
            addr,
            buffer: Buffer::new(policy.buffer_reclaim_threshold),
            bytes_read: 0,
            conn: stream,
            socket_timeout: policy.timeout().as_millis() as u32,
            timeout_delay: 0,
            deadline: None,
            idle_timeout,
            idle_deadline: idle_timeout.map(|timeout| Instant::now() + timeout),
            state: ConnectionState::Ready,
            can_recover_connection: false,
            // Far-future deadline; every IO resets it before use, so it never
            // fires first on its own.
            #[cfg(feature = "rt-tokio")]
            sleep: Box::pin(aerospike_rt::tokio::time::sleep(Duration::from_secs(3600))),
        };
        conn.authenticate(&policy.auth_mode, hashed_pass).await?;
        conn.refresh();
        Ok(conn)
    }

    #[cfg(test)]
    pub async fn new(
        host: &Host,
        policy: &ClientPolicy,
        _hashed_pass: Option<&String>,
    ) -> Result<Self> {
        if FAIL_NEXT_CONNECT.with(std::cell::Cell::take) {
            return Err(crate::Error::Connection(
                "forced connection failure (test)".to_string(),
            ));
        }

        let addr = host.address();
        let stream = Netsocket::TestDummy;

        let idle_timeout = if policy.idle_timeout > 0 {
            Some(Duration::from_millis(policy.idle_timeout as u64))
        } else {
            None
        };

        let mut conn = Connection {
            addr: addr.into(),
            buffer: Buffer::new(policy.buffer_reclaim_threshold),
            bytes_read: 0,
            conn: stream,
            socket_timeout: policy.timeout().as_millis() as u32,
            timeout_delay: 0,
            deadline: None,
            idle_timeout: idle_timeout,
            idle_deadline: idle_timeout.map(|timeout| Instant::now() + timeout),
            state: ConnectionState::Ready,
            can_recover_connection: false,
            // Far-future deadline; every IO resets it before use, so it never
            // fires first on its own.
            #[cfg(feature = "rt-tokio")]
            sleep: Box::pin(aerospike_rt::tokio::time::sleep(Duration::from_secs(3600))),
        };
        conn.refresh();
        Ok(conn)
    }

    pub fn close(&mut self) {
        self.state = ConnectionState::Closed;
        #[allow(clippy::let_underscore_future)]
        let () = match self.conn {
            Netsocket::Tcp(ref mut conn) => {
                #[cfg(feature = "rt-tokio")]
                let _ = conn.shutdown();
                #[cfg(feature = "rt-async-std")]
                let _ = conn.shutdown(Shutdown::Both);
            }
            #[cfg(feature = "tls")]
            Netsocket::Tls(ref mut conn) => {
                #[cfg(feature = "rt-tokio")]
                let _ = conn.shutdown();
                #[cfg(feature = "rt-async-std")]
                let _ = conn.shutdown(Shutdown::Both);
            }
            #[cfg(test)]
            _ => (),
        };
    }

    pub async fn flush(&mut self) -> Result<()> {
        self.state = ConnectionState::Writing;
        let timeout = self.deadline();

        let data = &self.buffer.data_buffer;
        let res = match self.conn {
            Netsocket::Tcp(ref mut conn) => {
                io_with_timeout!(self, timeout, async {
                    conn.write_all(data).await?;
                    conn.flush().await
                })
            }
            #[cfg(feature = "tls")]
            Netsocket::Tls(ref mut conn) => {
                io_with_timeout!(self, timeout, async {
                    conn.write_all(data).await?;
                    conn.flush().await
                })
            }
            #[cfg(test)]
            _ => unreachable!(),
        };

        match res {
            Ok(Ok(())) => (),
            // classify socket I/O errors as Connection err and hence command retries.
            Ok(Err(e)) => return Err(Error::Connection(format!("flush: {e}"))),
            Err(_) => {
                return Err(Error::Timeout(
                    "Timeout writing to network connection".to_string(),
                ));
            }
        }

        self.refresh();
        Ok(())
    }

    pub(crate) const fn set_state(&mut self, state: ConnectionState) {
        self.state = state;
        self.bytes_read = 0;
    }

    pub(crate) const fn reset_state(&mut self) {
        self.state = ConnectionState::Ready;
        self.bytes_read = 0;
    }

    /// Sets the timeout delay for the connection.
    pub(crate) const fn set_timeout_delay(
        &mut self,
        can_recover_connection: bool,
        timeout_delay: u32,
    ) {
        self.can_recover_connection = can_recover_connection;
        self.timeout_delay = timeout_delay;
    }

    /// Sets the timeout for the connection.
    pub const fn set_socket_timeout(&mut self, deadline: Option<Instant>, socket_timeout: u32) {
        self.deadline = deadline;
        if socket_timeout > 0 {
            self.socket_timeout = socket_timeout;
        } else {
            self.socket_timeout = 30_000; // 30 secs
        }
    }

    /// Reads the socket deadline for the connection.
    pub fn deadline(&self) -> Duration {
        let now = Instant::now();
        let socket_deadline = now + self.socket_timeout();

        let deadline = self
            .deadline
            .map_or(socket_deadline, |deadline| min(deadline, socket_deadline));

        deadline - now
    }

    /// Reads the socket timeout for the connection.
    /// If the timeout is zero, it will return the default (30 000 ms)
    pub fn socket_timeout(&self) -> Duration {
        if self.socket_timeout > 0 {
            Duration::from_millis(u64::from(self.socket_timeout))
        } else {
            Duration::from_secs(30) // 30 secs
        }
    }

    // This function validates the message header.
    pub(crate) fn validate_header(&self, header: u64) -> Result<()> {
        let msg_version = (header & 0xFF00_0000_0000_0000) >> 56;
        if msg_version != 2 {
            return Err(Error::ClientError(format!(
                "Invalid Message Header: Expected version to be 2, but got {msg_version}"
            )));
        }

        let msg_type = (header & 0x00FF_0000_0000_0000) >> 49;
        if !(msg_type == 1 || msg_type == 3 || msg_type == 4) {
            return Err(Error::ClientError(format!(
                "Invalid Message Header: Expected type to be 1, 3 or 4, but got {msg_type}"
            )));
        }

        Ok(())
    }

    // This function reads a standard header, setting the state correctly.
    pub(crate) async fn read_header(&mut self) -> Result<usize> {
        let header_size = buffer::MSG_TOTAL_HEADER_SIZE as usize;
        self.set_state(ConnectionState::ReadingHeader(header_size));
        let res = self.read_buffer(header_size).await?;
        self.set_state(ConnectionState::Ready);

        let proto = self.buffer.read_u64(Some(0));
        self.validate_header(proto)?;

        Ok(res)
    }

    // This function reads a standard header, setting the state correctly.
    pub(crate) async fn read_body(&mut self, receive_size: usize) -> Result<usize> {
        self.set_state(ConnectionState::ReadingBody(receive_size));
        let res = self.read_buffer(receive_size).await?;
        self.set_state(ConnectionState::Ready);
        Ok(res)
    }

    pub(crate) async fn read_buffer(&mut self, size: usize) -> Result<usize> {
        self.read_buffer_at(0, size).await
    }

    pub(crate) async fn read_buffer_at(&mut self, pos: usize, size: usize) -> Result<usize> {
        self.buffer.resize_buffer(size + pos)?;

        let timeout = self.deadline();
        let read_result = match self.conn {
            Netsocket::Tcp(ref mut conn) => {
                io_with_timeout!(
                    self,
                    timeout,
                    conn.read_exact(&mut self.buffer.data_buffer[pos..])
                )
            }

            #[cfg(feature = "tls")]
            Netsocket::Tls(ref mut conn) => {
                io_with_timeout!(
                    self,
                    timeout,
                    conn.read_exact(&mut self.buffer.data_buffer[pos..])
                )
            }
            #[cfg(test)]
            _ => unreachable!(),
        };

        match read_result {
            Ok(Ok(_)) => self.bytes_read += size,
            Ok(Err(e)) => return Err(Error::Connection(format!("read: {e}"))),
            Err(_) => {
                return Err(Error::Timeout(
                    "Timeout reading from the network connection".into(),
                ))
            }
        }

        self.buffer.reset_offset();
        self.refresh();
        Ok(size)
    }

    /// Writes to the connection until done or timeout has been reached.
    pub async fn write_all(&mut self, buf: &[u8]) -> Result<()> {
        self.state = ConnectionState::Writing;

        let timeout = self.deadline();
        // Flush after the write, exactly as `flush()` does (CLIENT-5268): on a
        // rustls stream `write_all` can leave the encrypted records in the
        // session buffer, and this is the info/tend path, which reads the reply
        // straight after writing — so an unflushed request would wait on an
        // answer the server never received. A no-op on plain TCP.
        let res = match self.conn {
            Netsocket::Tcp(ref mut conn) => {
                io_with_timeout!(self, timeout, async {
                    conn.write_all(buf).await?;
                    conn.flush().await
                })
            }
            #[cfg(feature = "tls")]
            Netsocket::Tls(ref mut conn) => {
                io_with_timeout!(self, timeout, async {
                    conn.write_all(buf).await?;
                    conn.flush().await
                })
            }
            #[cfg(test)]
            _ => unreachable!(),
        };

        match res {
            Ok(Ok(())) => (),
            Ok(Err(e)) => {
                return Err(Error::Connection(format!("write: {e}")));
            }
            // The timer carries no detail worth printing, and the two runtimes
            // report it as different types.
            Err(_) => {
                return Err(Error::Timeout(
                    "Timeout writing to the network connection".to_string(),
                ));
            }
        }

        self.refresh();
        Ok(())
    }

    /// Reads from the connection until the buffer is full or timeout has been reached.
    pub async fn read_all(&mut self, buf: &mut [u8]) -> Result<()> {
        self.state = ConnectionState::ReadingBody(buf.len());

        let timeout = self.deadline();
        let res = match self.conn {
            Netsocket::Tcp(ref mut conn) => {
                io_with_timeout!(self, timeout, conn.read_exact(buf))
            }
            #[cfg(feature = "tls")]
            Netsocket::Tls(ref mut conn) => {
                io_with_timeout!(self, timeout, conn.read_exact(buf))
            }
            #[cfg(test)]
            _ => unreachable!(),
        };

        match res {
            Ok(Ok(_)) => (),
            Ok(Err(e)) => return Err(Error::Connection(format!("read_all: {e}"))),
            Err(_) => {
                return Err(Error::Timeout(
                    "Timeout reading from the network connection".to_string(),
                ))
            }
        }

        self.bytes_read += buf.len();
        self.refresh();
        Ok(())
    }

    pub fn is_idle(&self) -> bool {
        self.idle_deadline
            .is_some_and(|idle_dl| Instant::now() >= idle_dl)
    }

    /// What a one-byte peek says about a socket.
    fn peek_liveness(sock: &socket2::SockRef<'_>) -> Liveness {
        // MSG_PEEK, so nothing is consumed: a byte seen here is still there for
        // the command that follows. Both runtimes keep the fd non-blocking, so
        // this never waits.
        let mut probe = [std::mem::MaybeUninit::<u8>::uninit(); 1];
        match sock.peek(&mut probe) {
            // Nothing to read on an open socket: the healthy idle case.
            Err(ref e)
                if matches!(
                    e.kind(),
                    std::io::ErrorKind::WouldBlock | std::io::ErrorKind::Interrupted
                ) =>
            {
                Liveness::Alive
            }
            // A peer that sent FIN reads as end-of-stream; an RST, EBADF or
            // anything else is equally unusable.
            Ok(0) | Err(_) => Liveness::Closed,
            Ok(_) => Liveness::PendingBytes,
        }
    }

    /// Non-blocking one-byte peek for pool checkout, so a socket the peer closed
    /// while it sat in the pool is discarded instead of handed to a command that
    /// would fail on its first read.
    ///
    /// This is the *only* mechanism that sheds dead pooled sockets — it is
    /// deliberately independent of `idle_timeout`, because a socket can die long
    /// before its idle deadline (a server restart kills sockets that were in use
    /// a millisecond earlier).
    pub(crate) fn is_alive(&self) -> bool {
        match self.conn {
            Netsocket::Tcp(ref s) => {
                // Unsolicited bytes on a plain connection mean the stream is out
                // of step with the protocol, which is not recoverable here.
                !matches!(
                    Self::peek_liveness(&socket2::SockRef::from(s)),
                    Liveness::Closed | Liveness::PendingBytes
                )
            }
            #[cfg(feature = "tls")]
            Netsocket::Tls(ref s) => {
                // Peek the TCP socket underneath the TLS session.
                //
                // Pending bytes are treated as ALIVE here, unlike the plain
                // arm: post-handshake TLS control records (a TLS 1.3
                // NewSessionTicket, a KeyUpdate) legitimately arrive while a
                // connection sits idle in the pool, and they are not
                // application data. Calling those dead would evict a healthy
                // connection, reconnect, receive a fresh ticket and evict
                // again — churn caused by the probe itself. Only a closed or
                // broken socket is fatal.
                let (tcp, _session) = s.get_ref();
                !matches!(
                    Self::peek_liveness(&socket2::SockRef::from(tcp)),
                    Liveness::Closed
                )
            }
            #[cfg(test)]
            _ => true,
        }
    }

    fn refresh(&mut self) {
        self.idle_deadline = None;
        self.deadline = None;
        if let Some(idle_to) = self.idle_timeout {
            self.idle_deadline = Some(Instant::now().add(idle_to));
        }
    }

    #[cfg_attr(test, allow(dead_code))]
    async fn authenticate(
        &mut self,
        auth_mode: &AuthMode,
        hashed_pass: Option<&String>,
    ) -> Result<()> {
        self.state = ConnectionState::Writing;
        return match AdminCommand::authenticate(self, auth_mode, hashed_pass).await {
            Ok(()) => {
                // Restore Ready so PooledConnection::Drop puts the conn back
                // in the pool instead of taking the non-recoverable close arm.
                self.set_state(ConnectionState::Ready);
                Ok(())
            }
            Err(err) => {
                self.close();
                Err(err)
            }
        };
    }

    pub const fn bookmark(&mut self) {
        self.bytes_read = 0;
    }

    pub const fn bytes_read(&self) -> usize {
        self.bytes_read
    }

    pub(crate) const fn should_attempt_recovery(&self) -> bool {
        self.can_recover_connection && self.timeout_delay > 0
    }

    // reads the rest of the message to empty the connection buffer
    // before returning the connection back to the pool.
    async fn drain(&mut self, mut limit: usize, timeout: Duration) -> Result<()> {
        while limit > 0 {
            let count = match self.conn {
                Netsocket::Tcp(ref mut conn) => aerospike_rt::timeout(
                    timeout,
                    aerospike_rt::io::copy(
                        &mut conn.take(limit as u64),
                        &mut aerospike_rt::io::sink(),
                    ),
                )
                .await
                .map_err(|e| Error::Timeout(format!("Timeout draining the connection {e}")))?,

                #[cfg(feature = "tls")]
                Netsocket::Tls(ref mut conn) => aerospike_rt::timeout(
                    timeout,
                    aerospike_rt::io::copy(
                        &mut conn.take(limit as u64),
                        &mut aerospike_rt::io::sink(),
                    ),
                )
                .await
                .map_err(|e| Error::Timeout(format!("Timeout draining the connection {e}")))?,
                #[cfg(test)]
                _ => unreachable!(),
            }?;

            limit -= count as usize;
            self.bytes_read += count as usize;
        }

        Ok(())
    }
}

/***********************************************************************************/
/*  Buffered Connection                                                            */
/***********************************************************************************/

// Holds data buffer for the command
#[derive(Debug)]
pub struct BufferedConn<'a> {
    pub(crate) conn: &'a mut Connection,

    cache: Vec<u8>,
    pos: usize,

    pub(crate) limit: usize,
    bytes_read: usize,
}

impl<'a> BufferedConn<'a> {
    pub fn new(conn: &'a mut Connection) -> Self {
        BufferedConn {
            conn,
            cache: Vec::with_capacity(4 * 1024),
            limit: 0,
            pos: 0,
            bytes_read: 0,
        }
    }

    pub(crate) const fn bookmark(&mut self) {
        self.bytes_read = 0;
        self.conn.bookmark();
    }

    #[inline]
    pub(crate) const fn buffer(&mut self) -> &mut Buffer {
        &mut self.conn.buffer
    }

    #[inline]
    pub(crate) const fn bytes_read(&self) -> usize {
        self.bytes_read
    }

    pub(crate) fn set_limit_header(&mut self, size: usize) -> Result<()> {
        self.conn
            .set_state(ConnectionState::ReadingStreamHeader(size));
        self.set_limit(size)
    }

    pub(crate) fn set_limit_body(&mut self, size: usize) -> Result<()> {
        self.conn
            .set_state(ConnectionState::ReadingStreamBody(size));
        self.set_limit(size)
    }

    fn set_limit(&mut self, size: usize) -> Result<()> {
        self.limit = size;
        self.pos = 0;
        self.bytes_read = 0;
        self.resize_cache(0)
    }

    fn resize_cache(&mut self, size: usize) -> Result<()> {
        // Corrupted data streams can result in a huge length.
        // Do a sanity check here.
        if size > MAX_BUFFER_SIZE {
            return Err(Error::InvalidArgument(format!(
                "Invalid size for buffer: {size}"
            )));
        }

        self.cache.resize(size, 0);

        Ok(())
    }

    async fn fill_buffer(&mut self) -> Result<usize> {
        // fill_buffer fills the buffer from the beginning.
        // The buffer should have been completely consumed before calling this function
        if self.pos != self.cache.len() || self.limit == 0 {
            return Ok(0);
        }

        let size = min(self.cache.capacity(), self.limit);
        self.resize_cache(size)?;

        let deadline = self.conn.deadline();
        let read_result = match self.conn.conn {
            Netsocket::Tcp(ref mut conn) => {
                io_with_timeout!(self.conn, deadline, conn.read_exact(&mut self.cache))
            }

            #[cfg(feature = "tls")]
            Netsocket::Tls(ref mut conn) => {
                io_with_timeout!(self.conn, deadline, conn.read_exact(&mut self.cache))
            }
            #[cfg(test)]
            _ => unreachable!(),
        };

        match read_result {
            Ok(Ok(_)) => {
                self.limit -= self.cache.len();
                self.conn.bytes_read += self.cache.len();
            }
            Ok(Err(e)) => return Err(Error::Connection(format!("buffered_read: {e}"))),
            Err(_) => {
                return Err(Error::Timeout(
                    "Timeout reading from the network connection".into(),
                ))
            }
        }

        self.pos = 0;
        Ok(size)
    }

    pub(crate) async fn drain(&mut self, timeout: Duration) -> Result<()> {
        while self.limit > 0 {
            let count = match self.conn.conn {
                Netsocket::Tcp(ref mut conn) => aerospike_rt::timeout(
                    timeout,
                    aerospike_rt::io::copy(
                        &mut conn.take(self.limit as u64),
                        &mut aerospike_rt::io::sink(),
                    ),
                )
                .await
                .map_err(|e| Error::Timeout(format!("Timeout draining the connection {e}")))?,
                #[cfg(feature = "tls")]
                Netsocket::Tls(ref mut conn) => aerospike_rt::timeout(
                    timeout,
                    aerospike_rt::io::copy(
                        &mut conn.take(self.limit as u64),
                        &mut aerospike_rt::io::sink(),
                    ),
                )
                .await
                .map_err(|e| Error::Timeout(format!("Timeout draining the connection {e}")))?,
                #[cfg(test)]
                _ => unreachable!(),
            }?;

            self.limit -= count as usize;
            self.bytes_read += count as usize;
            self.conn.bytes_read += count as usize;
        }

        let _ = self.resize_cache(0);
        self.pos = 0;
        assert!(self.exhausted());

        self.conn.state = ConnectionState::Ready;

        Ok(())
    }

    #[inline]
    pub(crate) const fn exhausted(&self) -> bool {
        self.limit == 0 && self.empty()
    }

    #[inline]
    const fn len(&self) -> usize {
        self.cache.len() - self.pos
    }

    #[inline]
    const fn empty(&self) -> bool {
        self.len() == 0
    }

    fn cached_read_rest(&mut self) -> usize {
        if !self.empty() {
            return self.cached_read(0, self.len());
        }
        0
    }

    fn cached_read(&mut self, pos: usize, size: usize) -> usize {
        self.conn.buffer.data_buffer[pos..pos + size]
            .copy_from_slice(&self.cache[self.pos..self.pos + size]);

        self.pos += size;
        size
    }

    pub async fn read_buffer(&mut self, size: usize) -> Result<usize> {
        self.conn.buffer.resize_buffer(size)?;

        if self.limit > 0 && self.empty() {
            self.fill_buffer().await?;
        }

        if size <= self.len() {
            self.cached_read(0, size);
        } else if size > self.len() {
            // we have data left in the buffer, but we need more
            let cached = self.cached_read_rest();
            let remaining = size - cached;
            if remaining > self.cache.capacity() / 2 {
                // read directly
                self.conn.read_buffer_at(cached, remaining).await?;
                self.limit -= remaining;
            } else {
                // fill the buffer and read the rest of requested bytes
                self.fill_buffer().await?;
                self.cached_read(cached, remaining);
            }
        }

        self.bytes_read += size;

        self.conn.buffer.reset_offset();
        self.conn.refresh();

        Ok(size)
    }
}

impl Drop for Connection {
    fn drop(&mut self) {
        self.close();
    }
}

pub struct ConnectionRecovery<'a> {
    conn: &'a mut Connection,
}

impl<'a> ConnectionRecovery<'a> {
    pub const fn new(conn: &'a mut Connection) -> Self {
        Self { conn }
    }

    pub async fn recover(&mut self) {
        if !self.conn.can_recover_connection || self.conn.timeout_delay == 0 {
            return;
        }

        self.conn.set_socket_timeout(None, self.conn.timeout_delay);
        match self.conn.state {
            ConnectionState::Ready | ConnectionState::Closed | ConnectionState::Writing => (),
            ConnectionState::ReadingHeader(total_size) => {
                let Ok(receive_size) = self.read_header(total_size).await else {
                    return;
                };

                self.conn
                    .set_state(ConnectionState::ReadingBody(receive_size));

                if self.read_body(receive_size).await.is_ok() {
                    self.conn.reset_state();
                }
            }

            ConnectionState::ReadingBody(total_size) => {
                if self.read_body(total_size).await.is_ok() {
                    self.conn.reset_state();
                }
            }

            ConnectionState::ReadingStreamHeader(total_size) => {
                let Ok(mut receive_size) = self.read_stream_header(total_size).await else {
                    return;
                };

                while receive_size > 0 {
                    self.conn
                        .set_state(ConnectionState::ReadingStreamBody(receive_size));
                    match self.read_stream_body(receive_size).await {
                        Ok(true) => {
                            self.conn.reset_state();
                            return;
                        }
                        Err(_) => return,
                        _ => (),
                    }

                    self.conn
                        .set_state(ConnectionState::ReadingStreamHeader(receive_size));
                    receive_size = match self.read_stream_header(total_size).await {
                        Ok(v) => v,
                        Err(_) => return,
                    };
                }
            }

            ConnectionState::ReadingStreamBody(mut receive_size) => {
                while receive_size > 0 {
                    match self.read_stream_body(receive_size).await {
                        Ok(true) => {
                            self.conn.reset_state();
                            return;
                        }
                        Err(_) => return,
                        _ => (),
                    }

                    self.conn.set_state(ConnectionState::ReadingStreamHeader(8));
                    receive_size = match self.read_stream_header(8).await {
                        Ok(v) => v,
                        Err(_) => return,
                    };

                    self.conn
                        .set_state(ConnectionState::ReadingStreamBody(receive_size));
                }
            }
        }
    }

    async fn read_header(&mut self, total_size: usize) -> Result<usize> {
        if total_size > self.conn.bytes_read {
            // read the rest of the header
            if self
                .conn
                .read_buffer_at(self.conn.bytes_read, total_size - self.conn.bytes_read)
                .await
                .is_err()
            {
                // return early and don't update the connection state
                return Err(Error::StreamTerminatedError());
            };
        }

        self.conn.buffer.reset_offset();
        let sz = self.conn.buffer.read_u64(Some(0));
        let header_length = self.conn.buffer.read_u8(Some(8));

        let receive_size = ((sz & 0xFFFF_FFFF_FFFF) - u64::from(header_length)) as usize;
        Ok(receive_size)
    }

    async fn read_body(&mut self, total_size: usize) -> Result<()> {
        if total_size > self.conn.bytes_read {
            // read the rest of the body
            if self
                .conn
                .drain(
                    total_size - self.conn.bytes_read,
                    Duration::from_millis(u64::from(self.conn.timeout_delay)),
                )
                .await
                .is_err()
            {
                // return early and don't update the connection state
                return Err(Error::StreamTerminatedError());
            }
        }

        assert!(self.conn.bytes_read == total_size);
        Ok(())
    }

    async fn read_stream_header(&mut self, total_size: usize) -> Result<usize> {
        assert_eq!(total_size, 8);
        if total_size > self.conn.bytes_read {
            // read the rest of the header
            if self
                .conn
                .read_buffer_at(self.conn.bytes_read, total_size - self.conn.bytes_read)
                .await
                .is_err()
            {
                // return early and don't update the connection state
                return Err(Error::StreamTerminatedError());
            };
        }

        let receive_size = self.conn.buffer.read_msg_size(Some(0));
        Ok(receive_size)
    }

    async fn read_stream_body(&mut self, total_size: usize) -> Result<bool> {
        // The message has been bigger than a header only last part. Drain it straight away.
        if self.conn.bytes_read > usize::from(crate::commands::buffer::MSG_TOTAL_HEADER_SIZE) {
            // we are past the header portion, clearly not the last message.
            // We can safely drain the connection
            if total_size > self.conn.bytes_read
                && self
                    .conn
                    .drain(
                        total_size - self.conn.bytes_read,
                        Duration::from_millis(u64::from(self.conn.timeout_delay)),
                    )
                    .await
                    .is_err()
            {
                // return early and don't update the connection state
                return Err(Error::StreamTerminatedError());
            }

            assert!(self.conn.bytes_read == total_size);
            return Ok(false);
        }

        // Still the header portion, so we need to read the rest of it and
        // figure out if this is the last message in the stream.
        if usize::from(crate::commands::buffer::MSG_TOTAL_HEADER_SIZE) > self.conn.bytes_read {
            let remaining = min(
                total_size,
                usize::from(crate::commands::buffer::MSG_TOTAL_HEADER_SIZE) - self.conn.bytes_read,
            );
            if self
                .conn
                .read_buffer_at(self.conn.bytes_read, remaining)
                .await
                .is_err()
            {
                // return early and don't update the connection state
                return Err(Error::StreamTerminatedError());
            }
        }

        let info3 = self.conn.buffer.read_u8(Some(3));
        let last_record =
            info3 & crate::commands::buffer::INFO3_LAST == crate::commands::buffer::INFO3_LAST;

        // read the rest of the body
        if total_size > self.conn.bytes_read
            && self
                .conn
                .drain(
                    total_size - self.conn.bytes_read,
                    Duration::from_millis(u64::from(self.conn.timeout_delay)),
                )
                .await
                .is_err()
        {
            // return early and don't update the connection state
            return Err(Error::StreamTerminatedError());
        }

        assert!(self.conn.bytes_read == total_size);
        Ok(last_record)
    }
}

/// The pool-checkout liveness probe, on whichever runtime is compiled.
///
/// The larger `tests_eof_loopback` module below is tokio-only; these cases are
/// runtime-agnostic on purpose, because the async-std arm of [`Connection::is_alive`]
/// had no coverage at all while it was a hardcoded `true`.
#[cfg(test)]
mod liveness_probe_tests {
    use super::*;
    use aerospike_rt::net::{TcpListener, TcpStream};

    /// Half-close the accepted socket, so the client side sees FIN.
    async fn spawn_finning_peer() -> String {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap().to_string();
        aerospike_rt::spawn(async move {
            if let Ok((mut sock, _)) = listener.accept().await {
                #[cfg(feature = "rt-tokio")]
                {
                    use aerospike_rt::io::AsyncWriteExt;
                    let _ = sock.shutdown().await;
                }
                #[cfg(feature = "rt-async-std")]
                {
                    let _ = sock.shutdown(aerospike_rt::async_std::net::Shutdown::Both);
                }
                drop(sock);
            }
        });
        addr
    }

    /// Accept and hold the socket open, saying nothing.
    async fn spawn_quiet_peer() -> String {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap().to_string();
        aerospike_rt::spawn(async move {
            if let Ok((sock, _)) = listener.accept().await {
                aerospike_rt::sleep(Duration::from_secs(30)).await;
                drop(sock);
            }
        });
        addr
    }

    /// Build a `Connection` around a real socket, bypassing the handshake.
    fn conn_over(stream: TcpStream) -> Connection {
        let mut conn = Connection {
            addr: "127.0.0.1:0".into(),
            buffer: Buffer::new(0),
            bytes_read: 0,
            conn: Netsocket::Tcp(stream),
            socket_timeout: 5_000,
            timeout_delay: 0,
            deadline: None,
            idle_timeout: None,
            idle_deadline: None,
            state: ConnectionState::Ready,
            can_recover_connection: false,
            #[cfg(feature = "rt-tokio")]
            sleep: Box::pin(aerospike_rt::tokio::time::sleep(Duration::from_secs(3600))),
        };
        conn.refresh();
        conn
    }

    #[aerospike_macro::test]
    async fn probe_says_alive_for_an_open_idle_socket() {
        let addr = spawn_quiet_peer().await;
        let stream = TcpStream::connect(&*addr).await.unwrap();
        aerospike_rt::sleep(Duration::from_millis(20)).await;
        assert!(
            conn_over(stream).is_alive(),
            "an open socket with nothing pending must probe alive"
        );
    }

    #[aerospike_macro::test]
    async fn probe_says_dead_after_peer_fin() {
        let addr = spawn_finning_peer().await;
        let stream = TcpStream::connect(&*addr).await.unwrap();
        // Let the FIN land in our kernel before probing.
        aerospike_rt::sleep(Duration::from_millis(50)).await;
        assert!(
            !conn_over(stream).is_alive(),
            "a socket the peer closed must probe dead"
        );
    }

    /// The peek must not consume: probing twice has to give the same answer, and
    /// a command that follows still sees the pending byte.
    #[aerospike_macro::test]
    async fn probe_does_not_consume_pending_bytes() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap().to_string();
        aerospike_rt::spawn(async move {
            if let Ok((mut sock, _)) = listener.accept().await {
                #[cfg(feature = "rt-tokio")]
                {
                    use aerospike_rt::io::AsyncWriteExt;
                    let _ = sock.write_all(b"XY").await;
                }
                #[cfg(feature = "rt-async-std")]
                {
                    use futures::AsyncWriteExt;
                    let _ = sock.write_all(b"XY").await;
                }
                aerospike_rt::sleep(Duration::from_secs(30)).await;
            }
        });

        let stream = TcpStream::connect(&*addr).await.unwrap();
        aerospike_rt::sleep(Duration::from_millis(50)).await;
        let mut conn = conn_over(stream);

        // Unsolicited bytes on a plain connection: not usable, twice over.
        assert!(!conn.is_alive());
        assert!(!conn.is_alive(), "the verdict must be stable, not consumed");

        // And the bytes are still on the socket for whoever reads next.
        let mut buf = [0_u8; 2];
        conn.read_all(&mut buf).await.expect("bytes still readable");
        assert_eq!(&buf, b"XY", "MSG_PEEK must leave the data in place");
    }
}

///  socket-level liveness probe and the `Error::Connection` classification of socket I/O failures.
#[cfg(all(test, feature = "rt-tokio"))]
mod tests_eof_loopback {
    use super::*;
    use crate::commands::is_network_error;
    use aerospike_rt::net::{TcpListener, TcpStream};
    use std::net::SocketAddr;

    /// Build a `Connection` over a real TCP stream. Local to this module —
    /// inaccessible from other code, no API surface added.
    fn conn_from_stream(stream: TcpStream) -> Connection {
        let mut conn = Connection {
            addr: "127.0.0.1:0".into(),
            buffer: Buffer::new(0),
            bytes_read: 0,
            conn: Netsocket::Tcp(stream),
            socket_timeout: 5_000,
            timeout_delay: 0,
            deadline: None,
            idle_timeout: None,
            idle_deadline: None,
            state: ConnectionState::Ready,
            can_recover_connection: false,
            // Far-future deadline; every IO resets it before use, so it never
            // fires first on its own.
            #[cfg(feature = "rt-tokio")]
            sleep: Box::pin(aerospike_rt::tokio::time::sleep(Duration::from_secs(3600))),
        };
        conn.refresh();
        conn
    }

    /// Spawn a one-shot peer that accepts and immediately half-closes the
    /// socket — simulates a server-side FIN like `asd` exiting.
    async fn spawn_fin_peer() -> SocketAddr {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        aerospike_rt::spawn(async move {
            if let Ok((mut sock, _)) = listener.accept().await {
                use tokio::io::AsyncWriteExt;
                let _ = sock.shutdown().await;
                drop(sock);
            }
        });
        addr
    }

    /// Spawn a peer that accepts and holds the socket open silently —
    /// the live-and-idle case the liveness probe must accept.
    async fn spawn_idle_peer() -> SocketAddr {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        aerospike_rt::spawn(async move {
            if let Ok((sock, _)) = listener.accept().await {
                aerospike_rt::sleep(std::time::Duration::from_secs(60)).await;
                drop(sock);
            }
        });
        addr
    }

    /// Spawn a peer that pushes some bytes at the client without being
    /// asked — exercises the "stray bytes pending" branch of the probe.
    async fn spawn_chatty_peer() -> SocketAddr {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        aerospike_rt::spawn(async move {
            if let Ok((mut sock, _)) = listener.accept().await {
                use tokio::io::AsyncWriteExt;
                let _ = sock.write_all(b"unsolicited").await;
                aerospike_rt::sleep(std::time::Duration::from_secs(60)).await;
            }
        });
        addr
    }

    // ─── Reusable per-IO timer ────────────────────────────────────────────

    /// A read against a peer that never answers must still time out. The timer
    /// now lives on the connection and is reset per IO rather than built per
    /// call, so this is the check that resetting actually arms it.
    #[tokio::test(flavor = "current_thread")]
    async fn read_times_out_against_a_silent_peer() {
        let addr = spawn_idle_peer().await;
        let stream = TcpStream::connect(addr).await.unwrap();
        let mut conn = conn_from_stream(stream);
        conn.set_socket_timeout(None, 60);

        let mut buf = [0_u8; 8];
        let err = conn
            .read_all(&mut buf)
            .await
            .expect_err("a silent peer must produce a timeout");
        assert!(
            matches!(err, Error::Timeout(_)),
            "expected Timeout, got: {0:?}",
            err
        );
    }

    /// The timer is reused, so it must re-arm: an expired one has to be reset
    /// for the *next* IO instead of firing immediately. Without a working
    /// reset, the read after a timeout would fail even though data is waiting.
    #[tokio::test(flavor = "current_thread")]
    async fn timer_rearms_after_an_earlier_timeout() {
        // Peer stays quiet long enough for the first read to expire, then
        // sends exactly what the second read wants.
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        aerospike_rt::spawn(async move {
            if let Ok((mut sock, _)) = listener.accept().await {
                use tokio::io::AsyncWriteExt;
                aerospike_rt::sleep(std::time::Duration::from_millis(200)).await;
                let _ = sock.write_all(b"abcd").await;
                aerospike_rt::sleep(std::time::Duration::from_secs(60)).await;
            }
        });

        let stream = TcpStream::connect(addr).await.unwrap();
        let mut conn = conn_from_stream(stream);

        conn.set_socket_timeout(None, 50);
        let mut buf = [0_u8; 4];
        assert!(
            conn.read_all(&mut buf).await.is_err(),
            "first read must time out while the peer is quiet"
        );

        // Same connection, same timer: a generous timeout must now succeed.
        conn.set_socket_timeout(None, 5_000);
        conn.read_all(&mut buf)
            .await
            .expect("second read must succeed on the reused timer");
        assert_eq!(&buf, b"abcd");
    }

    /// A write on a live socket must not be cut short by a leftover timer
    /// state, and the connection must remain usable afterwards.
    #[tokio::test(flavor = "current_thread")]
    async fn repeated_writes_reuse_the_timer() {
        let addr = spawn_idle_peer().await;
        let stream = TcpStream::connect(addr).await.unwrap();
        let mut conn = conn_from_stream(stream);
        conn.set_socket_timeout(None, 5_000);

        for round in 0..5 {
            conn.write_all(b"ping")
                .await
                .unwrap_or_else(|e| panic!("round {0} write failed: {1:?}", round, e));
        }
    }

    // ─── Bug 2: liveness probe ────────────────────────────────────────────

    #[tokio::test(flavor = "current_thread")]
    async fn is_alive_returns_true_for_idle_socket() {
        let addr = spawn_idle_peer().await;
        let stream = TcpStream::connect(addr).await.unwrap();
        aerospike_rt::sleep(std::time::Duration::from_millis(20)).await;
        let conn = conn_from_stream(stream);
        assert!(conn.is_alive(), "idle live socket must probe alive");
    }

    #[tokio::test(flavor = "current_thread")]
    async fn is_alive_returns_false_after_peer_fin() {
        let addr = spawn_fin_peer().await;
        let stream = TcpStream::connect(addr).await.unwrap();
        // Give the peer's shutdown a moment to arrive at our kernel.
        aerospike_rt::sleep(std::time::Duration::from_millis(50)).await;
        let conn = conn_from_stream(stream);
        assert!(!conn.is_alive(), "FIN'd socket must probe dead");
    }

    #[tokio::test(flavor = "current_thread")]
    async fn is_alive_returns_false_when_stray_bytes_pending() {
        let addr = spawn_chatty_peer().await;
        let stream = TcpStream::connect(addr).await.unwrap();
        aerospike_rt::sleep(std::time::Duration::from_millis(50)).await;
        let conn = conn_from_stream(stream);
        assert!(
            !conn.is_alive(),
            "socket with unread bytes (protocol desync) must probe dead"
        );
    }

    // ─── Bug 1: socket I/O errors classified as Error::Connection ─────────

    #[tokio::test(flavor = "current_thread")]
    async fn read_header_after_peer_fin_yields_error_connection() {
        let addr = spawn_fin_peer().await;
        let stream = TcpStream::connect(addr).await.unwrap();
        let mut conn = conn_from_stream(stream);

        let err = conn
            .read_header()
            .await
            .expect_err("read on FIN'd socket must fail");

        assert!(
            matches!(err, Error::Connection(_)),
            "expected Error::Connection on peer FIN, got: {:?}",
            err
        );
        assert!(
            is_network_error(&err),
            "is_network_error must accept this so the retry gate engages; err = {:?}",
            err
        );
    }

    #[tokio::test(flavor = "current_thread")]
    async fn write_all_after_peer_fin_yields_error_connection() {
        let addr = spawn_fin_peer().await;
        let stream = TcpStream::connect(addr).await.unwrap();
        let mut conn = conn_from_stream(stream);

        // Wait for FIN to propagate so the next write surfaces ECONNRESET
        // before the kernel send-buffer can absorb a small write.
        aerospike_rt::sleep(std::time::Duration::from_millis(50)).await;

        // Force the failure with a write large enough that the kernel can't
        // hide it inside the send buffer.
        let mut last_err: Option<Error> = None;
        for _ in 0..10 {
            let big = vec![0u8; 256 * 1024];
            if let Err(e) = conn.write_all(&big).await {
                last_err = Some(e);
                break;
            }
        }
        let err = last_err.expect("a write must eventually fail after peer FIN");

        assert!(
            matches!(err, Error::Connection(_)),
            "expected Error::Connection on peer-closed write, got: {:?}",
            err
        );
        assert!(
            is_network_error(&err),
            "is_network_error must accept this; err = {:?}",
            err
        );
    }

    // ─── Bug 2 end-to-end: Queue::get evicts dead, returns live ───────────
    //
    // Lives here (rather than in connection_pool.rs) because the
    // `conn_from_stream` helper needs access to `Connection`'s private
    // fields, which are only visible from within `connection.rs`.

    #[tokio::test(flavor = "current_thread")]
    async fn queue_get_evicts_peer_finned_socket() {
        use crate::net::connection_pool::Queue;
        use crate::net::Host;
        use crate::policy::ClientPolicy;

        let host = Host::new("127.0.0.1", 0);
        let policy = ClientPolicy::default();
        let q = Queue::with_capacity(1, host, policy);

        let addr = spawn_fin_peer().await;
        let stream = TcpStream::connect(addr).await.unwrap();
        aerospike_rt::sleep(std::time::Duration::from_millis(50)).await;

        let conn = conn_from_stream(stream);
        assert!(q.reserve_capacity());
        q.put_back(conn);

        let result = q.get();
        assert!(
            result.is_err(),
            "Queue::get() must not return a peer-FIN'd socket"
        );
    }

    #[tokio::test(flavor = "current_thread")]
    async fn queue_get_returns_live_socket() {
        use crate::net::connection_pool::Queue;
        use crate::net::Host;
        use crate::policy::ClientPolicy;

        let host = Host::new("127.0.0.1", 0);
        let policy = ClientPolicy::default();
        let q = Queue::with_capacity(1, host, policy);

        let addr = spawn_idle_peer().await;
        let stream = TcpStream::connect(addr).await.unwrap();
        aerospike_rt::sleep(std::time::Duration::from_millis(50)).await;

        let conn = conn_from_stream(stream);
        assert!(q.reserve_capacity());
        q.put_back(conn);

        let result = q.get();
        assert!(
            result.is_ok(),
            "Queue::get() must return a live socket; got Err({:?})",
            result.err()
        );
    }
}