ntex-io 4.3.0

//! In-memory I/O transport and helpers for tests.
#![allow(clippy::missing_panics_doc)]
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker};
use std::{any, cell::RefCell, cmp, fmt, future::poll_fn, io, mem, net, rc::Rc};

use ntex_bytes::{BufMut, BytePages, Bytes, BytesMut};
use ntex_util::time::{Millis, sleep};

use crate::{Handle, IoContext, IoStream, IoTaskStatus, Readiness, types};

#[derive(Default)]
struct AtomicWaker(Arc<Mutex<RefCell<Option<Waker>>>>);

impl AtomicWaker {
    fn wake(&self) -> bool {
        if let Some(waker) = self.0.lock().unwrap().borrow_mut().take() {
            waker.wake();
            true
        } else {
            false
        }
    }
}

impl fmt::Debug for AtomicWaker {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "AtomicWaker")
    }
}

/// One endpoint of an in-memory asynchronous byte stream.
///
/// Use [`IoTest::create`] to construct connected client and server endpoints.
/// Bytes written by one endpoint become readable from the other.
#[derive(Debug)]
pub struct IoTest {
    tp: Type,
    peer_addr: Option<net::SocketAddr>,
    state: Arc<Mutex<RefCell<State>>>,
    local: Arc<Mutex<RefCell<Channel>>>,
    remote: Arc<Mutex<RefCell<Channel>>>,
}

bitflags::bitflags! {
    #[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
    struct IoTestFlags: u8 {
        const FLUSHED = 0b0000_0001;
        const CLOSED  = 0b0000_0010;
    }
}

#[derive(Copy, Clone, PartialEq, Eq, Debug)]
enum Type {
    Client,
    Server,
    ClientClone,
    ServerClone,
}

#[derive(Copy, Clone, Default, Debug)]
struct State {
    client_dropped: bool,
    server_dropped: bool,
}

#[derive(Default, Debug)]
struct Channel {
    buf: BytesMut,
    buf_cap: usize,
    flags: IoTestFlags,
    waker: AtomicWaker,
    read: IoTestState,
    write: IoTestState,
}

unsafe impl Sync for Channel {}
unsafe impl Send for Channel {}

impl Channel {
    fn is_closed(&self) -> bool {
        self.flags.contains(IoTestFlags::CLOSED)
    }
}

impl Default for IoTestFlags {
    fn default() -> Self {
        IoTestFlags::empty()
    }
}

#[derive(Debug, Default)]
enum IoTestState {
    #[default]
    Ok,
    Pending,
    Close,
    Err(io::Error),
}

impl IoTest {
    /// Creates connected client and server endpoints.
    pub fn create() -> (IoTest, IoTest) {
        let local = Arc::new(Mutex::new(RefCell::new(Channel::default())));
        let remote = Arc::new(Mutex::new(RefCell::new(Channel::default())));
        let state = Arc::new(Mutex::new(RefCell::new(State::default())));

        (
            IoTest {
                tp: Type::Client,
                peer_addr: None,
                local: local.clone(),
                remote: remote.clone(),
                state: state.clone(),
            },
            IoTest {
                state,
                peer_addr: None,
                tp: Type::Server,
                local: remote,
                remote: local,
            },
        )
    }

    /// Returns `true` after the client endpoint has been dropped.
    pub fn is_client_dropped(&self) -> bool {
        self.state.lock().unwrap().borrow().client_dropped
    }

    /// Returns `true` after the server endpoint has been dropped.
    pub fn is_server_dropped(&self) -> bool {
        self.state.lock().unwrap().borrow().server_dropped
    }

    /// Returns `true` after the peer has closed its write side.
    pub fn is_closed(&self) -> bool {
        self.remote.lock().unwrap().borrow().is_closed()
    }

    /// Sets the socket address returned by transport queries.
    #[must_use]
    pub fn set_peer_addr(mut self, addr: net::SocketAddr) -> Self {
        self.peer_addr = Some(addr);
        self
    }

    /// Forces subsequent reads from the peer endpoint to remain pending.
    pub fn read_pending(&self) {
        self.remote.lock().unwrap().borrow_mut().read = IoTestState::Pending;
    }

    /// Makes the peer endpoint's next read fail with `err`.
    pub fn read_error(&self, err: io::Error) {
        let channel = self.remote.lock().unwrap();
        channel.borrow_mut().read = IoTestState::Err(err);
        channel.borrow().waker.wake();
    }

    /// Makes this endpoint's next transport write fail with `err`.
    pub fn write_error(&self, err: io::Error) {
        self.local.lock().unwrap().borrow_mut().write = IoTestState::Err(err);
        self.remote.lock().unwrap().borrow().waker.wake();
    }

    /// Makes this endpoint's next transport write return zero bytes.
    pub fn write_zero(&self) {
        self.local.lock().unwrap().borrow_mut().write = IoTestState::Close;
        self.remote.lock().unwrap().borrow().waker.wake();
    }

    /// Provides mutable access to bytes readable by this endpoint.
    pub fn local_buffer<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&mut BytesMut) -> R,
    {
        let guard = self.local.lock().unwrap();
        let mut ch = guard.borrow_mut();
        f(&mut ch.buf)
    }

    /// Provides mutable access to bytes readable by the peer endpoint.
    pub fn remote_buffer<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&mut BytesMut) -> R,
    {
        let guard = self.remote.lock().unwrap();
        let mut ch = guard.borrow_mut();
        f(&mut ch.buf)
    }

    /// Simulates the peer closing its write side.
    ///
    /// This wakes a pending reader and yields once so the close can be
    /// observed by asynchronous test code.
    pub async fn close(&self) {
        {
            let guard = self.remote.lock().unwrap();
            let mut remote = guard.borrow_mut();
            remote.read = IoTestState::Close;
            let res = remote.waker.wake();
            log::debug!("close remote socket, waker:{res}");
        }
        sleep(Millis(35)).await;
    }

    /// Writes bytes for the peer endpoint to read and wakes its reader.
    pub fn write<T: AsRef<[u8]>>(&self, data: T) {
        let guard = self.remote.lock().unwrap();
        let mut write = guard.borrow_mut();
        write.buf.extend_from_slice(data.as_ref());
        write.waker.wake();
    }

    /// Sets how many bytes the peer may write before becoming blocked.
    ///
    /// Increasing the capacity wakes the peer's pending writer.
    pub fn remote_buffer_cap(&self, cap: usize) {
        // change cap
        self.local.lock().unwrap().borrow_mut().buf_cap = cap;
        // wake remote
        self.remote.lock().unwrap().borrow().waker.wake();
    }

    /// Takes all bytes currently readable by this endpoint without waiting.
    pub fn read_any(&self) -> Bytes {
        self.local.lock().unwrap().borrow_mut().buf.take()
    }

    /// Waits for readable data or peer closure, then takes all available bytes.
    pub async fn read(&self) -> Result<Bytes, io::Error> {
        if self.local.lock().unwrap().borrow().buf.is_empty() {
            poll_fn(|cx| {
                let guard = self.local.lock().unwrap();
                let read = guard.borrow_mut();
                if read.buf.is_empty() {
                    let closed = match self.tp {
                        Type::Client | Type::ClientClone => {
                            self.is_server_dropped() || read.is_closed()
                        }
                        Type::Server | Type::ServerClone => self.is_client_dropped(),
                    };
                    if closed {
                        Poll::Ready(())
                    } else {
                        *read.waker.0.lock().unwrap().borrow_mut() = Some(cx.waker().clone());
                        drop(read);
                        drop(guard);
                        Poll::Pending
                    }
                } else {
                    Poll::Ready(())
                }
            })
            .await;
        }
        Ok(self.local.lock().unwrap().borrow_mut().buf.take())
    }

    /// Polls a transport read into `buf`.
    ///
    /// Returns the number of bytes copied, zero on simulated peer closure, or
    /// `Pending` when no input is available.
    pub fn poll_read_buf(
        &self,
        cx: &mut Context<'_>,
        buf: &mut BytesMut,
    ) -> Poll<io::Result<usize>> {
        let guard = self.local.lock().unwrap();
        let mut ch = guard.borrow_mut();
        *ch.waker.0.lock().unwrap().borrow_mut() = Some(cx.waker().clone());

        if !ch.buf.is_empty() {
            let size = std::cmp::min(ch.buf.len(), buf.remaining_mut());
            assert!(size > 0, "Supplied buffer is zero sized");
            let b = ch.buf.split_to(size);
            buf.put_slice(&b);
            return Poll::Ready(Ok(size));
        }

        match mem::take(&mut ch.read) {
            IoTestState::Ok | IoTestState::Pending => Poll::Pending,
            IoTestState::Close => {
                ch.read = IoTestState::Close;
                Poll::Ready(Ok(0))
            }
            IoTestState::Err(e) => Poll::Ready(Err(e)),
        }
    }

    /// Polls a transport write from `buf`.
    ///
    /// The write is limited by the capacity configured with
    /// [`remote_buffer_cap`](Self::remote_buffer_cap).
    pub fn poll_write_buf(&self, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        let guard = self.remote.lock().unwrap();
        let mut ch = guard.borrow_mut();

        match mem::take(&mut ch.write) {
            IoTestState::Ok => {
                let cap = cmp::min(buf.len(), ch.buf_cap);
                if cap > 0 {
                    ch.buf.extend(&buf[..cap]);
                    ch.buf_cap -= cap;
                    ch.flags.remove(IoTestFlags::FLUSHED);
                    ch.waker.wake();
                    Poll::Ready(Ok(cap))
                } else {
                    *self
                        .local
                        .lock()
                        .unwrap()
                        .borrow_mut()
                        .waker
                        .0
                        .lock()
                        .unwrap()
                        .borrow_mut() = Some(cx.waker().clone());
                    Poll::Pending
                }
            }
            IoTestState::Close => Poll::Ready(Ok(0)),
            IoTestState::Pending => {
                *self
                    .local
                    .lock()
                    .unwrap()
                    .borrow_mut()
                    .waker
                    .0
                    .lock()
                    .unwrap()
                    .borrow_mut() = Some(cx.waker().clone());
                Poll::Pending
            }
            IoTestState::Err(e) => Poll::Ready(Err(e)),
        }
    }
}

impl Clone for IoTest {
    fn clone(&self) -> Self {
        let tp = match self.tp {
            Type::Server => Type::ServerClone,
            Type::Client => Type::ClientClone,
            val => val,
        };

        IoTest {
            tp,
            local: self.local.clone(),
            remote: self.remote.clone(),
            state: self.state.clone(),
            peer_addr: self.peer_addr,
        }
    }
}

impl Drop for IoTest {
    fn drop(&mut self) {
        let mut state = *self.state.lock().unwrap().borrow();
        match self.tp {
            Type::Server => state.server_dropped = true,
            Type::Client => state.client_dropped = true,
            _ => (),
        }
        *self.state.lock().unwrap().borrow_mut() = state;

        let guard = self.remote.lock().unwrap();
        let mut remote = guard.borrow_mut();
        remote.read = IoTestState::Close;
        remote.waker.wake();
        log::debug!("drop remote socket");
    }
}

impl IoStream for IoTest {
    fn start(self, ctx: IoContext) -> Box<dyn Handle> {
        let io = Rc::new(self);
        ntex_util::spawn(run(io.clone(), ctx));
        Box::new(io)
    }
}

impl Handle for Rc<IoTest> {
    fn query(&self, id: any::TypeId) -> Option<Box<dyn any::Any>> {
        if id == any::TypeId::of::<types::PeerAddr>()
            && let Some(addr) = self.peer_addr
        {
            return Some(Box::new(types::PeerAddr(addr)));
        }
        None
    }
}

async fn run(io: Rc<IoTest>, ctx: IoContext) {
    poll_fn(|cx| turn(&io, &ctx, cx)).await;

    log::debug!("{}: Shuting down io", ctx.tag());

    // shutdown WRITE side
    io.local
        .lock()
        .unwrap()
        .borrow_mut()
        .flags
        .insert(IoTestFlags::CLOSED);

    log::debug!("{}: Shutdown complete", ctx.tag());
    ctx.stopped(None);
}

fn turn(io: &IoTest, ctx: &IoContext, cx: &mut Context<'_>) -> Poll<()> {
    let read = match ctx.poll_read_ready(cx) {
        Poll::Ready(Readiness::Ready) => read(io, ctx, cx),
        Poll::Ready(Readiness::Close | Readiness::Terminate) => Poll::Ready(()),
        Poll::Pending => Poll::Pending,
    };

    let write = match ctx.poll_write_ready(cx) {
        Poll::Ready(Readiness::Ready) => write(io, ctx, cx),
        Poll::Ready(Readiness::Close | Readiness::Terminate) => Poll::Ready(()),
        Poll::Pending => Poll::Pending,
    };

    if read.is_pending() && write.is_pending() {
        Poll::Pending
    } else {
        Poll::Ready(())
    }
}

fn write(io: &IoTest, ctx: &IoContext, cx: &mut Context<'_>) -> Poll<()> {
    let result = ctx.with_write_dst(|buf| write_io(io, buf, cx, ctx));
    if ctx.update_write_status(result) == IoTaskStatus::Stop {
        Poll::Ready(())
    } else {
        Poll::Pending
    }
}

fn read(io: &IoTest, ctx: &IoContext, cx: &mut Context<'_>) -> Poll<()> {
    loop {
        let mut pending = false;
        let result = ctx.with_read_buf(|buf| {
            let result = io.poll_read_buf(cx, buf);
            pending = result.is_pending();
            result
        });
        return match result {
            IoTaskStatus::Io => {
                if pending {
                    Poll::Pending
                } else {
                    continue;
                }
            }
            IoTaskStatus::Stop => Poll::Ready(()),
            IoTaskStatus::Pause => Poll::Pending,
        };
    }
}

/// Flush write buffer to underlying I/O stream.
pub(super) fn write_io(
    io: &IoTest,
    buf: &mut BytePages,
    cx: &mut Context<'_>,
    ctx: &IoContext,
) -> io::Result<usize> {
    let tag = ctx.tag();
    let mut written = 0;

    while let Some(mut page) = buf.take() {
        log::debug!("{tag}: flushing framed transport: {}", page.len());

        let result = io.poll_write_buf(cx, &page)?;
        match result {
            Poll::Ready(0) => {
                log::trace!("{tag}: disconnected during flush, written {written}");
                buf.prepend(page);
                return Err(io::Error::new(
                    io::ErrorKind::WriteZero,
                    "failed to write frame to transport",
                ));
            }
            Poll::Ready(n) => {
                written += n;
                page.advance_to(n);
                buf.prepend(page);
            }
            Poll::Pending => {
                buf.prepend(page);
                break;
            }
        }
    }

    log::debug!("{tag}: flushed {written} bytes, remaining: {}", buf.len());
    Ok(written)
}

#[cfg(test)]
#[allow(clippy::redundant_clone)]
mod tests {
    use super::*;
    use ntex_util::future::lazy;

    #[ntex::test]
    async fn basic() {
        let (client, server) = IoTest::create();
        assert_eq!(client.tp, Type::Client);
        assert_eq!(client.clone().tp, Type::ClientClone);
        assert_eq!(server.tp, Type::Server);
        assert_eq!(server.clone().tp, Type::ServerClone);
        assert!(format!("{server:?}").contains("IoTest"));
        assert!(format!("{:?}", AtomicWaker::default()).contains("AtomicWaker"));

        server.read_pending();
        let mut buf = BytesMut::new();
        let res = lazy(|cx| client.poll_read_buf(cx, &mut buf)).await;
        assert!(res.is_pending());

        server.read_pending();
        let res = lazy(|cx| server.poll_write_buf(cx, b"123")).await;
        assert!(res.is_pending());

        assert!(!server.is_client_dropped());
        drop(client);
        assert!(server.is_client_dropped());

        let server2 = server.clone();
        assert!(!server2.is_server_dropped());
        drop(server);
        assert!(server2.is_server_dropped());

        let res = lazy(|cx| server2.poll_write_buf(cx, b"123")).await;
        assert!(res.is_pending());

        let (client, _) = IoTest::create();
        let addr: net::SocketAddr = "127.0.0.1:8080".parse().unwrap();
        let client = crate::Io::from(client.set_peer_addr(addr));
        let item = client.query::<crate::types::PeerAddr>();
        assert!(format!("{item:?}").contains("QueryItem(127.0.0.1:8080)"));
    }
}