echo/
echo.rs

1//! The reference server for the bb8-style pool. Unlike the old `Server`, the
2//! crate owns none of the runtime here: *this example* spawns one pinned thread
3//! per core, builds a `compio` runtime on each, binds its own `SO_REUSEPORT`
4//! listener, and borrows a per-connection buffer from the crate's [`Pool`].
5//!
6//! That is the whole point of the redesign — the pool is a component you pull
7//! into your own thread-per-core loop, not a framework that runs the loop for
8//! you.
9//!
10//! ```text
11//! cargo run --release --example echo -- [ADDR] [CAPACITY]
12//!   ADDR       bind address               default 0.0.0.0:7000
13//!   CAPACITY   max connections per worker default 1024
14//! ```
15//!
16//! Drive it with `examples/load.rs`.
17
18use std::{io, net::SocketAddr, thread};
19
20use compio::{
21    BufResult,
22    io::{AsyncRead, AsyncWriteExt},
23    net::{TcpListener, TcpStream},
24    runtime::Runtime,
25};
26use compio_pool::{LocalPool, ManageConnection, Pool, bind_reuseport, cpu};
27
28/// Hands out reusable 16 KiB read buffers, one per in-flight connection. A real
29/// manager would open a backend connection (Redis, Postgres, an upstream socket)
30/// in `connect` instead.
31struct Buffers;
32
33impl ManageConnection for Buffers {
34    type Connection = Vec<u8>;
35    type Error = io::Error;
36
37    async fn connect(&self) -> io::Result<Vec<u8>> {
38        Ok(Vec::with_capacity(16 * 1024))
39    }
40
41    async fn is_valid(&self, _buf: &mut Vec<u8>) -> io::Result<()> {
42        Ok(())
43    }
44
45    fn has_broken(&self, _buf: &mut Vec<u8>) -> bool {
46        false
47    }
48}
49
50fn main() -> io::Result<()> {
51    let addr: SocketAddr = std::env::args()
52        .nth(1)
53        .unwrap_or_else(|| "0.0.0.0:7000".into())
54        .parse()
55        .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, format!("bad address: {e}")))?;
56    let capacity: u32 = std::env::args()
57        .nth(2)
58        .and_then(|s| s.parse().ok())
59        .unwrap_or(1024);
60
61    // Build the pool once. `Pool` is `Send + Clone`; every worker gets a clone.
62    let pool = Pool::builder().max_size(capacity).build(Buffers);
63
64    let cores = cpu::cores();
65    if cores.is_empty() {
66        return Err(io::Error::new(io::ErrorKind::Unsupported, "no cores to run on"));
67    }
68    println!(
69        "echo on {addr}: {} workers, capacity {capacity}/worker",
70        cores.len()
71    );
72
73    let mut handles = Vec::new();
74    for (index, core) in cores.into_iter().enumerate() {
75        let pool = pool.clone();
76        let handle = thread::Builder::new()
77            .name(format!("worker/{index}"))
78            .spawn(move || worker(core, addr, pool))?;
79        handles.push(handle);
80    }
81    for handle in handles {
82        let _ = handle.join();
83    }
84    Ok(())
85}
86
87/// One pinned thread: its own `compio` runtime, its own `SO_REUSEPORT` listener,
88/// and its own [`LocalPool`] carved from the shared [`Pool`].
89fn worker(core: cpu::CoreId, addr: SocketAddr, pool: Pool<Buffers>) {
90    cpu::pin_current(core);
91
92    let runtime = Runtime::builder().build().expect("build compio runtime");
93    runtime.block_on(async move {
94        // This thread's pool. `!Send`, so it — and its ring-bound buffers —
95        // cannot escape the thread.
96        let local = pool.local();
97
98        let std_listener = bind_reuseport(addr, 1024, None).expect("bind SO_REUSEPORT");
99        let listener = TcpListener::from_std(std_listener).expect("wrap listener in ring");
100
101        loop {
102            let Ok((stream, _peer)) = listener.accept().await else {
103                continue;
104            };
105            compio::runtime::spawn(serve(local.clone(), stream)).detach();
106        }
107    });
108}
109
110/// Borrow a buffer from this thread's pool for the life of the connection and
111/// echo until the peer closes. Dropping the lease returns the buffer.
112async fn serve(local: LocalPool<Buffers>, mut stream: TcpStream) {
113    let mut lease = match local.get().await {
114        Ok(lease) => lease,
115        Err(_) => return,
116    };
117    loop {
118        let mut buf = std::mem::take(&mut *lease);
119        buf.clear();
120        let BufResult(read, buf) = stream.read(buf).await;
121        match read {
122            Ok(0) | Err(_) => {
123                *lease = buf;
124                return;
125            }
126            Ok(_) => {}
127        }
128        let BufResult(written, buf) = stream.write_all(buf).await;
129        *lease = buf;
130        if written.is_err() {
131            return;
132        }
133    }
134}
echo/echo.rs

echo/
echo.rs
