Add the inferproxy byte forwarder with its two limits

Implemented-By: Laguna S 2.1 (OpenCode)

crates/inferproxy/src/lib.rs

@@ -1 +1,135 @@
 //! Forwards bytes between `infer.sock` and the llama-server router. It logs nothing.
+
+use std::io::{self, copy};
+use std::net::{Shutdown, TcpStream};
+use std::os::unix::net::{UnixListener, UnixStream};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::thread;
+use std::time::Instant;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct Limits {
+    pub max_connections: usize,
+    pub burst: u32,
+    pub per_second: u32,
+}
+
+impl Default for Limits {
+    fn default() -> Self {
+        Limits {
+            max_connections: 8,
+            burst: 10,
+            per_second: 2,
+        }
+    }
+}
+
+/// Holds at most `burst` tokens and gains `per_second` tokens each second.
+pub struct TokenBucket {
+    tokens: u64,
+    burst: u64,
+    per_second: u32,
+    last: Instant,
+}
+
+impl TokenBucket {
+    pub fn new(burst: u32, per_second: u32, now: Instant) -> Self {
+        let burst = u64::from(burst) * 1000;
+        TokenBucket {
+            tokens: burst,
+            burst,
+            per_second,
+            last: now,
+        }
+    }
+
+    pub fn take(&mut self, now: Instant) -> bool {
+        let elapsed = now.saturating_duration_since(self.last);
+        let earned = refill_thousandths(self.per_second, elapsed);
+        let total = u128::from(self.tokens).saturating_add(earned);
+        let capped = total.min(u128::from(self.burst));
+        self.tokens = u64::try_from(capped).unwrap_or(self.burst);
+        self.last = now;
+        if self.tokens >= 1000 {
+            self.tokens -= 1000;
+            true
+        } else {
+            false
+        }
+    }
+}
+
+/// Thousandths of a token earned over `elapsed` at `per_second` tokens per second.
+/// per_second * elapsed_nanoseconds / 1_000_000 keeps partial refills additive without
+/// floating point, so checks 100 ms apart at 2/s accumulate correctly.
+fn refill_thousandths(per_second: u32, elapsed: std::time::Duration) -> u128 {
+    let nanos = u128::from(per_second).saturating_mul(elapsed.as_nanos());
+    nanos.div_euclid(1_000_000)
+}
+
+struct OpenGuard {
+    open: Arc<AtomicUsize>,
+}
+
+impl OpenGuard {
+    fn new(open: Arc<AtomicUsize>) -> Self {
+        open.fetch_add(1, Ordering::SeqCst);
+        OpenGuard { open }
+    }
+}
+
+impl Drop for OpenGuard {
+    fn drop(&mut self) {
+        self.open.fetch_sub(1, Ordering::SeqCst);
+    }
+}
+
+/// Accepts connections forever. Returns only if `accept` itself fails.
+pub fn serve(listener: UnixListener, upstream: String, limits: Limits) -> io::Result<()> {
+    let open = Arc::new(AtomicUsize::new(0));
+    let mut bucket = TokenBucket::new(limits.burst, limits.per_second, Instant::now());
+    loop {
+        let (client, _) = listener.accept()?;
+        if !bucket.take(Instant::now()) {
+            eprintln!("inferproxy: refused (rate limit)");
+            drop(client);
+            continue;
+        }
+        if open.load(Ordering::SeqCst) >= limits.max_connections {
+            eprintln!("inferproxy: refused (connection limit)");
+            drop(client);
+            continue;
+        }
+        let open = Arc::clone(&open);
+        let up = upstream.clone();
+        thread::spawn(move || handle(client, up, open));
+    }
+}
+
+fn handle(client: UnixStream, upstream: String, open: Arc<AtomicUsize>) {
+    let _guard = OpenGuard::new(open);
+    let server = match TcpStream::connect(&upstream) {
+        Ok(s) => s,
+        Err(_) => return,
+    };
+    let _ = forward(&client, &server);
+    drop(_guard);
+}
+
+fn forward(client: &UnixStream, server: &TcpStream) -> io::Result<()> {
+    let mut c2s_in = client.try_clone()?;
+    let mut s2c_out = client.try_clone()?;
+    let mut s2c_in = server.try_clone()?;
+    let mut c2s_out = server.try_clone()?;
+    thread::scope(|s| {
+        s.spawn(move || {
+            let _ = copy(&mut c2s_in, &mut c2s_out);
+            let _ = c2s_out.shutdown(Shutdown::Write);
+        });
+        s.spawn(move || {
+            let _ = copy(&mut s2c_in, &mut s2c_out);
+        });
+    });
+    Ok(())
+}