fanout_requests/fanout_requests.rs
1//! The pool's canonical job — pooling **upstream connections** — shown entirely
2//! on one compio runtime, with the concurrency coming from compio *tasks* rather
3//! than threads.
4//!
5//! Everything runs in a single [`Runtime`] on one thread:
6//!
7//! 1. an in-process echo **backend** on an ephemeral loopback port, spawned as a
8//! task (so the example is self-contained — no second process to start);
9//! 2. a [`Pool`] of upstream [`TcpStream`]s to that backend, capped at
10//! `CAPACITY`;
11//! 3. `CLIENTS` **client tasks**, all [`compio::runtime::spawn`]ed on this same
12//! runtime, sharing one [`LocalPool`] and racing to drain a shared work
13//! counter of `REQUESTS` round trips.
14//!
15//! The lesson is backpressure: however many client tasks you spawn, the pool only
16//! ever opens up to `CAPACITY` upstreams, and [`LocalPool::get`] makes a task
17//! *wait* for a free one when they are all busy. The program proves it by printing
18//! how many upstreams were ever opened and the peak number in use at once — both
19//! bounded by `CAPACITY`, no matter how large `CLIENTS` is.
20//!
21//! Because it is all one thread, the shared counters are plain [`Rc`]`<`[`Cell`]`>`
22//! — no atomics, no locks — which is exactly the `!Send`, single-thread world the
23//! pool is built for.
24//!
25//! ```text
26//! cargo run --release --example fanout_requests -- [REQUESTS] [CAPACITY] [CLIENTS] [SIZE]
27//! REQUESTS total round trips default 50000
28//! CAPACITY pooled upstreams (max) default 8
29//! CLIENTS concurrent client tasks default 128
30//! SIZE bytes per request default 256
31//! ```
32
33use std::{
34 cell::Cell,
35 io,
36 net::SocketAddr,
37 rc::Rc,
38 sync::atomic::{AtomicU64, Ordering},
39 time::Instant,
40};
41
42use compio::{
43 BufResult,
44 io::{AsyncRead, AsyncReadExt, AsyncWriteExt},
45 net::{TcpListener, TcpStream},
46 runtime::Runtime,
47};
48use compio_pool::{LocalPool, ManageConnection, Pool};
49
50/// How many upstream connections the pool has ever opened. A real pool's whole
51/// value is keeping this far below the number of concurrent callers.
52static OPENS: AtomicU64 = AtomicU64::new(0);
53
54/// Opens pooled upstream connections to the in-process backend. This is the real
55/// job of a pool — a Redis, Postgres or HTTP backend would connect here the same
56/// way; here the backend is just a loopback echo so the example needs nothing
57/// external.
58struct Upstreams {
59 addr: SocketAddr,
60}
61
62impl ManageConnection for Upstreams {
63 type Connection = TcpStream;
64 type Error = io::Error;
65
66 async fn connect(&self) -> io::Result<TcpStream> {
67 let stream = TcpStream::connect(self.addr).await?;
68 stream.set_nodelay(true)?;
69 OPENS.fetch_add(1, Ordering::Relaxed);
70 Ok(stream)
71 }
72
73 async fn is_valid(&self, _conn: &mut TcpStream) -> io::Result<()> {
74 Ok(())
75 }
76
77 fn has_broken(&self, _conn: &mut TcpStream) -> bool {
78 false
79 }
80}
81
82fn main() -> io::Result<()> {
83 let requests: usize = arg(1).unwrap_or(50_000);
84 let capacity: u32 = arg(2).unwrap_or(8);
85 let clients: usize = arg(3).unwrap_or(128);
86 let size: usize = arg(4).unwrap_or(256);
87
88 let runtime = Runtime::builder().build().expect("build compio runtime");
89 runtime.block_on(async move {
90 // 1. An in-process echo backend on an ephemeral port, run as a task.
91 let listen: SocketAddr = "127.0.0.1:0".parse().unwrap();
92 let backend = TcpListener::bind(listen).await?;
93 let backend_addr = backend.local_addr()?;
94 compio::runtime::spawn(backend_loop(backend)).detach();
95
96 // 2. One pool of upstreams to the backend, and this thread's LocalPool.
97 let pool = Pool::builder()
98 .max_size(capacity)
99 .build(Upstreams { addr: backend_addr });
100 let local = pool.local();
101
102 println!(
103 "fanout_requests -> {backend_addr}: {clients} client tasks over {capacity} pooled \
104 upstreams, {requests} requests x {size} B"
105 );
106
107 // Shared, single-thread state: plain Rc<Cell>, no atomics needed.
108 let next = Rc::new(Cell::new(0usize)); // next request index to claim
109 let bytes = Rc::new(Cell::new(0u64)); // total bytes echoed back
110 let inflight = Rc::new(Cell::new(0u32)); // upstreams in use right now
111 let peak = Rc::new(Cell::new(0u32)); // high-water mark of `inflight`
112
113 // 3. Fan out: spawn the client tasks and await them all.
114 let started = Instant::now();
115 let mut tasks = Vec::with_capacity(clients);
116 for _ in 0..clients {
117 tasks.push(compio::runtime::spawn(client(
118 local.clone(),
119 next.clone(),
120 requests,
121 size,
122 bytes.clone(),
123 inflight.clone(),
124 peak.clone(),
125 )));
126 }
127 for task in tasks {
128 let _ = task.await;
129 }
130 let elapsed = started.elapsed();
131
132 println!(
133 "done: {requests} requests in {:.3}s ({:.0} req/s), {} B echoed",
134 elapsed.as_secs_f64(),
135 requests as f64 / elapsed.as_secs_f64(),
136 bytes.get(),
137 );
138 println!(
139 "pool: opened {} upstreams total, peak {} in use at once (cap {capacity}); final {:?}",
140 OPENS.load(Ordering::Relaxed),
141 peak.get(),
142 local.state(),
143 );
144 Ok(())
145 })
146}
147
148/// One client task: claim request indices from the shared counter until they run
149/// out, leasing a pooled upstream for each round trip and handing it straight
150/// back. Many of these share `CAPACITY` upstreams.
151async fn client(
152 local: LocalPool<Upstreams>,
153 next: Rc<Cell<usize>>,
154 total: usize,
155 size: usize,
156 bytes: Rc<Cell<u64>>,
157 inflight: Rc<Cell<u32>>,
158 peak: Rc<Cell<u32>>,
159) {
160 loop {
161 // Claim the next index. There is no `.await` between the read and the
162 // write, so on this single thread no other task can observe or take the
163 // same index — a lock-free counter that needs no lock.
164 let i = next.get();
165 if i >= total {
166 break;
167 }
168 next.set(i + 1);
169
170 // Wait for a free upstream. `get` only returns once one is available, so
171 // `inflight` can never climb past `CAPACITY` — that is the backpressure.
172 let mut up = match local.get().await {
173 Ok(up) => up,
174 Err(_) => return,
175 };
176 let now = inflight.get() + 1;
177 inflight.set(now);
178 if now > peak.get() {
179 peak.set(now);
180 }
181
182 // One round trip: send `size` bytes, read exactly that many back.
183 let BufResult(written, _) = up.write_all(vec![b'x'; size]).await;
184 if written.is_ok() {
185 let BufResult(read, reply) = up.read_exact(Vec::with_capacity(size)).await;
186 if read.is_ok() {
187 bytes.set(bytes.get() + reply.len() as u64);
188 }
189 }
190
191 inflight.set(inflight.get() - 1);
192 // `up` drops here, returning the upstream to the pool for the next task.
193 }
194}
195
196/// The in-process echo backend: accept forever, echo each connection on its own
197/// task. Stands in for whatever real service the pool would front.
198async fn backend_loop(listener: TcpListener) {
199 loop {
200 let Ok((mut stream, _peer)) = listener.accept().await else {
201 continue;
202 };
203 let _ = stream.set_nodelay(true);
204 compio::runtime::spawn(async move {
205 let mut buf = Vec::with_capacity(16 * 1024);
206 loop {
207 buf.clear();
208 let BufResult(read, b) = stream.read(buf).await;
209 buf = b;
210 match read {
211 Ok(0) | Err(_) => return,
212 Ok(_) => {}
213 }
214 let BufResult(written, b) = stream.write_all(buf).await;
215 buf = b;
216 if written.is_err() {
217 return;
218 }
219 }
220 })
221 .detach();
222 }
223}
224
225/// Parse the nth CLI argument, if present and valid.
226fn arg<T: std::str::FromStr>(n: usize) -> Option<T> {
227 std::env::args().nth(n).and_then(|s| s.parse().ok())
228}