dope-session 0.2.3

#![cfg(all(target_os = "linux", target_arch = "aarch64"))]

use std::io::{Read, Write};
use std::net::{SocketAddr, TcpListener};
use std::ptr::NonNull;
use std::time::{Duration, Instant};

use dambi::{Bytes, SlotId};
use dope_core::driver::{Backend, DriverConfig};
use dope_core::profile::Tail;
use dope_runtime::runtime::Manifold;
use dope_session::Plain;
use dope_session::pool::Pool;
use dope_session::pool::client::connect_source::StaticUpstreams;
use dope_session::pool::client::{Client, TcpClientConfig};
use dope_session::protocol::client::ClientProtocol;
use dope_session::session::Egress;
use dope_transport::Tcp;
use dope_transport::config::{SocketToggle, StreamSocketConfig};
use dope_uring::{CurrentSlot, UringBackend, set_current};

const MAX_CONN: usize = 4;
const SLOT_BASE: usize = 0;
const TICK_BUDGET: Duration = Duration::from_secs(5);
const ECHO_LEN: usize = 11;

#[derive(Default)]
struct EchoState;

#[derive(Default)]
struct EchoProto {
    received: Vec<Vec<u8>>,
    connected: bool,
}

impl ClientProtocol for EchoProto {
    type Head = Bytes;
    type ConnState = EchoState;
    type CodecLayer = Plain;

    fn parse(&self, _state: &mut EchoState, buf: &Bytes) -> Option<(Bytes, usize)> {
        if buf.len() < ECHO_LEN {
            None
        } else {
            Some((buf.slice(0..ECHO_LEN), ECHO_LEN))
        }
    }

    fn on_connect(
        &mut self,
        _conn_id: SlotId,
        _conn_state: &mut EchoState,
        _out: Egress<'_, Plain>,
    ) {
        self.connected = true;
    }

    fn on_response(
        &mut self,
        _conn_id: SlotId,
        _conn_state: &mut EchoState,
        head: Bytes,
        _out: Egress<'_, Plain>,
    ) {
        self.received.push(head.as_ref().to_vec());
    }
}

struct EchoServer {
    addr: SocketAddr,
    join: Option<std::thread::JoinHandle<()>>,
}

impl EchoServer {
    fn spawn() -> Self {
        let listener = TcpListener::bind("127.0.0.1:0").expect("bind echo server");
        let addr = listener.local_addr().expect("local_addr");
        let join = std::thread::spawn(move || {
            let (mut stream, _) = match listener.accept() {
                Ok(v) => v,
                Err(_) => return,
            };
            stream.set_read_timeout(Some(Duration::from_secs(2))).ok();
            let mut buf = [0u8; 4096];
            match stream.read(&mut buf) {
                Ok(0) => {}
                Ok(n) => {
                    let _ = stream.write_all(&buf[..n]);
                }
                Err(_) => {}
            }
        });
        Self {
            addr,
            join: Some(join),
        }
    }
}

impl Drop for EchoServer {
    fn drop(&mut self) {
        if let Some(j) = self.join.take() {
            let _ = j.join();
        }
    }
}

#[test]
fn pool_client_echo_roundtrip() {
    let server = EchoServer::spawn();

    let driver_cfg = <UringBackend as Backend>::Config::for_tcp_profile(256, MAX_CONN, 4096);
    let mut driver = <UringBackend as Backend>::new_driver(driver_cfg).expect("driver");

    let proto = EchoProto::default();
    let upstreams = StaticUpstreams::<Tcp>::new(vec![server.addr], Duration::from_millis(50));
    let mut socket_opts = StreamSocketConfig::default();
    socket_opts.nodelay = SocketToggle::Enabled;
    let cfg = TcpClientConfig {
        max_conn: MAX_CONN,
        socket_opts,
        shutdown: None,
        cpu_id: None,
    };
    let mut pool: Pool<Client<EchoProto, Tcp, StaticUpstreams<Tcp>>, Plain, Tail, UringBackend> =
        Pool::build_raw(proto, upstreams, cfg, &mut driver, SLOT_BASE).expect("build_raw");

    let driver_ptr = NonNull::from(&mut driver);
    set_current(Some(CurrentSlot::new(driver_ptr, std::ptr::null())));

    let deadline = Instant::now() + TICK_BUDGET;
    let mut sent = false;
    loop {
        pool.drive(&mut driver);

        if !sent {
            let conn_opt = pool.live_conns().next();
            if let Some(conn_id) = conn_opt {
                if let Some(session) = pool.session_mut(conn_id) {
                    session.enqueue(Bytes::copy_from_slice(b"hello-pool!"));
                    let _ = pool.poke_outbound(&mut driver, conn_id);
                    sent = true;
                }
            }
        }

        let done = pool.direction_mut().protocol().received.len() >= 1;
        if done {
            break;
        }
        if Instant::now() >= deadline {
            let proto = pool.direction_mut().protocol();
            panic!(
                "echo response timeout: connected={}, received_count={}",
                proto.connected,
                proto.received.len()
            );
        }
        let _ = driver.park_for(Duration::from_millis(5));
    }

    let proto = pool.direction_mut().protocol();
    assert!(proto.connected, "on_connect should fire");
    assert_eq!(proto.received.len(), 1, "expected exactly one response");
    assert_eq!(proto.received[0], b"hello-pool!", "echo body mismatch");
}