task_echo/
task_echo.rs

1//! The thread-free counterpart of `echo.rs`: a TCP echo server on **one** compio
2//! runtime, where concurrency is compio *tasks*, not OS threads.
3//!
4//! `echo.rs` spawns a pinned OS thread per core and gives each its own runtime,
5//! listener and [`LocalPool`]. This example does none of that: it builds a single
6//! [`Runtime`], binds one ordinary [`TcpListener`], and for each accepted
7//! connection calls [`compio::runtime::spawn`] — a task on this same thread and
8//! this same `io_uring` ring. There is no [`std::thread::spawn`], no core pinning
9//! and no `SO_REUSEPORT`; thousands of connections are served concurrently by
10//! thousands of cheap tasks cooperatively scheduled on one ring.
11//!
12//! The [`Pool`] is used exactly as everywhere else — one [`LocalPool`] for the
13//! thread, a buffer leased per connection — which is the point: the pool's model
14//! is *per runtime*, so a single-runtime server is just the one-thread case of
15//! it, with nothing extra to set up.
16//!
17//! ```text
18//! cargo run --release --example task_echo -- [ADDR] [CAPACITY]
19//!   ADDR       bind address               default 127.0.0.1:7000
20//!   CAPACITY   max connections (buffers)  default 1024
21//!
22//! # talk to it:
23//! echo -n hello | nc -w1 127.0.0.1 7000
24//! ```
25
26use std::{io, net::SocketAddr};
27
28use compio::{
29    BufResult,
30    io::{AsyncRead, AsyncWriteExt},
31    net::{TcpListener, TcpStream},
32    runtime::Runtime,
33};
34use compio_pool::{LocalPool, ManageConnection, Pool};
35
36/// Hands out reusable 16 KiB read buffers, one per in-flight connection.
37struct Buffers;
38
39impl ManageConnection for Buffers {
40    type Connection = Vec<u8>;
41    type Error = io::Error;
42
43    async fn connect(&self) -> io::Result<Vec<u8>> {
44        Ok(Vec::with_capacity(16 * 1024))
45    }
46
47    async fn is_valid(&self, _buf: &mut Vec<u8>) -> io::Result<()> {
48        Ok(())
49    }
50
51    fn has_broken(&self, _buf: &mut Vec<u8>) -> bool {
52        false
53    }
54}
55
56fn main() -> io::Result<()> {
57    let addr: SocketAddr = std::env::args()
58        .nth(1)
59        .unwrap_or_else(|| "127.0.0.1:7000".into())
60        .parse()
61        .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, format!("bad address: {e}")))?;
62    let capacity: u32 = std::env::args()
63        .nth(2)
64        .and_then(|s| s.parse().ok())
65        .unwrap_or(1024);
66
67    let pool = Pool::builder().max_size(capacity).build(Buffers);
68
69    // One runtime, one ring. Everything below runs on this single thread.
70    let runtime = Runtime::builder().build().expect("build compio runtime");
71    runtime.block_on(async move {
72        // This runtime's pool. `!Send`, so it and its buffers stay on this thread.
73        let local = pool.local();
74
75        let listener = TcpListener::bind(addr).await?;
76        println!("task_echo on {addr}: 1 runtime, capacity {capacity}");
77
78        // Accept forever; hand each connection to its own task on this same ring.
79        loop {
80            let Ok((stream, _peer)) = listener.accept().await else {
81                continue;
82            };
83            let _ = stream.set_nodelay(true);
84            compio::runtime::spawn(serve(local.clone(), stream)).detach();
85        }
86    })
87}
88
89/// Borrow a buffer from this runtime's pool for the life of the connection and
90/// echo until the peer closes. Dropping the lease returns the buffer.
91async fn serve(local: LocalPool<Buffers>, mut stream: TcpStream) {
92    let mut lease = match local.get().await {
93        Ok(lease) => lease,
94        Err(_) => return,
95    };
96    loop {
97        let mut buf = std::mem::take(&mut *lease);
98        buf.clear();
99        let BufResult(read, buf) = stream.read(buf).await;
100        match read {
101            Ok(0) | Err(_) => {
102                *lease = buf;
103                return;
104            }
105            Ok(_) => {}
106        }
107        let BufResult(written, buf) = stream.write_all(buf).await;
108        *lease = buf;
109        if written.is_err() {
110            return;
111        }
112    }
113}
task_echo/task_echo.rs

task_echo/
task_echo.rs
