142 lines
4.1 KiB
Rust
142 lines
4.1 KiB
Rust
//! 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 server = match TcpStream::connect(&upstream) {
|
|
Ok(s) => s,
|
|
Err(_) => return,
|
|
};
|
|
let c2s = {
|
|
let _guard = OpenGuard::new(open);
|
|
match forward(&client, &server) {
|
|
Ok(c2s) => c2s,
|
|
Err(_) => return,
|
|
}
|
|
};
|
|
let _ = client.shutdown(Shutdown::Both);
|
|
let _ = server.shutdown(Shutdown::Both);
|
|
let _ = c2s.join();
|
|
}
|
|
|
|
fn forward(client: &UnixStream, server: &TcpStream) -> io::Result<thread::JoinHandle<()>> {
|
|
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()?;
|
|
let c2s = thread::spawn(move || {
|
|
let _ = copy(&mut c2s_in, &mut c2s_out);
|
|
let _ = c2s_out.shutdown(Shutdown::Write);
|
|
});
|
|
let _ = thread::spawn(move || {
|
|
let _ = copy(&mut s2c_in, &mut s2c_out);
|
|
})
|
|
.join();
|
|
Ok(c2s)
|
|
}
|