Skip to main content

unix_echo/
unix_echo.rs

1//! An **AF_UNIX** (Unix domain socket) echo server — the local-IPC cousin of the
2//! TCP `echo.rs`.
3//!
4//! A filesystem socket path can only be bound once, so there is no `SO_REUSEPORT`
5//! fan-out here: this example runs a single pinned worker with one
6//! [`UnixListener`], and still borrows a per-connection buffer from the crate's
7//! [`Pool`] exactly as the thread-per-core examples do. It is the example to copy
8//! when the thing you are pooling talks to a backend over a Unix socket — a local
9//! Postgres, Redis or a sidecar — rather than TCP.
10//!
11//! ```text
12//! cargo run --release --example unix_echo -- [PATH] [CAPACITY]
13//!   PATH       socket path                 default /tmp/compio-pool-echo.sock
14//!   CAPACITY   buffers (connections)       default 1024
15//!
16//! # talk to it:
17//! echo -n hello | socat - UNIX-CONNECT:/tmp/compio-pool-echo.sock
18//! # ...or interactively:  nc -U /tmp/compio-pool-echo.sock
19//! ```
20
21use std::io;
22
23use compio::{
24    BufResult,
25    io::{AsyncRead, AsyncWriteExt},
26    net::{UnixListener, UnixStream},
27    runtime::Runtime,
28};
29use compio_pool::{LocalPool, ManageConnection, Pool, cpu};
30
31/// Hands out reusable 16 KiB read buffers, one per connection.
32struct Buffers;
33
34impl ManageConnection for Buffers {
35    type Connection = Vec<u8>;
36    type Error = io::Error;
37
38    async fn connect(&self) -> io::Result<Vec<u8>> {
39        Ok(Vec::with_capacity(16 * 1024))
40    }
41
42    async fn is_valid(&self, _buf: &mut Vec<u8>) -> io::Result<()> {
43        Ok(())
44    }
45
46    fn has_broken(&self, _buf: &mut Vec<u8>) -> bool {
47        false
48    }
49}
50
51fn main() -> io::Result<()> {
52    let path = std::env::args()
53        .nth(1)
54        .unwrap_or_else(|| "/tmp/compio-pool-echo.sock".into());
55    let capacity: u32 = std::env::args()
56        .nth(2)
57        .and_then(|s| s.parse().ok())
58        .unwrap_or(1024);
59
60    // A stale socket file from a previous run would make `bind` fail with
61    // EADDRINUSE, so clear it first.
62    let _ = std::fs::remove_file(&path);
63
64    let pool = Pool::builder().max_size(capacity).build(Buffers);
65
66    // One listener, so one worker. Pin it to the first permitted core so the ring
67    // and the socket share a CPU, the same discipline the fan-out examples apply
68    // per thread.
69    if let Some(core) = cpu::cores().into_iter().next() {
70        cpu::pin_current(core);
71    }
72    println!("unix_echo on {path}: 1 worker, capacity {capacity}");
73
74    let runtime = Runtime::builder().build().expect("build compio runtime");
75    runtime.block_on(async move {
76        let local = pool.local();
77        let listener = UnixListener::bind(&path).await?;
78
79        loop {
80            let Ok((stream, _peer)) = listener.accept().await else {
81                continue;
82            };
83            compio::runtime::spawn(serve(local.clone(), stream)).detach();
84        }
85    })
86}
87
88/// Borrow a buffer for the life of the connection and echo until the peer closes.
89async fn serve(local: LocalPool<Buffers>, mut stream: UnixStream) {
90    let mut lease = match local.get().await {
91        Ok(lease) => lease,
92        Err(_) => return,
93    };
94    loop {
95        let mut buf = std::mem::take(&mut *lease);
96        buf.clear();
97        let BufResult(read, buf) = stream.read(buf).await;
98        match read {
99            Ok(0) | Err(_) => {
100                *lease = buf;
101                return;
102            }
103            Ok(_) => {}
104        }
105        let BufResult(written, buf) = stream.write_all(buf).await;
106        *lease = buf;
107        if written.is_err() {
108            return;
109        }
110    }
111}