Bound local HTTP resource use

src/server.rs

@@ -34,13 +34,16 @@ use std::fs::File;
 use std::io::{Read, Write};
 use std::net::{TcpListener, TcpStream};
 use std::path::PathBuf;
-use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex, RwLock};
 use std::thread::{self, JoinHandle};
 use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
 
 const MAX_HEADER_BYTES: usize = 64 * 1024;
 const MAX_BODY_BYTES: usize = 64 * 1024 * 1024;
+const MAX_CONNECTIONS: usize = 64;
+const HTTP_IO_TIMEOUT: Duration = Duration::from_secs(10);
+const HTTP_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
 const PREFILL_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(5);
 const TOOLS_PROMPT: &str = "## Tools\n\n\
 You have access to a set of tools to help answer the user question. You can invoke tools by writing a \"<｜DSML｜tool_calls>\" block like the following:\n\n\
@@ -71,6 +74,7 @@ struct State {
     models_path: PathBuf,
     cache_path: PathBuf,
     sequence: AtomicU64,
+    connections: Arc<AtomicUsize>,
     tool_memory: Mutex<HashMap<String, String>>,
     metrics: Arc<Metrics>,
 }
@@ -226,6 +230,7 @@ impl ServerHandle {
             models_path,
             cache_path,
             sequence: AtomicU64::new(0),
+            connections: Arc::new(AtomicUsize::new(0)),
             tool_memory: Mutex::new(HashMap::new()),
             metrics: Arc::clone(&metrics),
         });
@@ -257,10 +262,20 @@ fn serve(listener: TcpListener, state: Arc<State>, stop: Arc<AtomicBool>) {
     while !stop.load(Ordering::Relaxed) {
         match listener.accept() {
             Ok((stream, _)) => {
+                let Some(slot) = ConnectionSlot::acquire(&state.connections) else {
+                    continue;
+                };
                 let state = Arc::clone(&state);
-                let _ = thread::Builder::new()
+                if let Err(error) =
+                    thread::Builder::new()
                         .name("http-request".into())
-                    .spawn(move || handle(stream, &state));
+                        .spawn(move || {
+                            let _slot = slot;
+                            handle(stream, &state);
+                        })
+                {
+                    eprintln!("DS4Server: endpoint request thread failed: {error}");
+                }
             }
             Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
                 thread::sleep(Duration::from_millis(50));
@@ -273,10 +288,29 @@ fn serve(listener: TcpListener, state: Arc<State>, stop: Arc<AtomicBool>) {
     }
 }
 
+struct ConnectionSlot(Arc<AtomicUsize>);
+
+impl ConnectionSlot {
+    fn acquire(active: &Arc<AtomicUsize>) -> Option<Self> {
+        active
+            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |count| {
+                (count < MAX_CONNECTIONS).then_some(count + 1)
+            })
+            .ok()?;
+        Some(Self(Arc::clone(active)))
+    }
+}
+
+impl Drop for ConnectionSlot {
+    fn drop(&mut self) {
+        self.0.fetch_sub(1, Ordering::Relaxed);
+    }
+}
+
 fn handle(mut stream: TcpStream, state: &State) {
-    let _ = stream.set_read_timeout(Some(Duration::from_secs(30)));
+    let _ = stream.set_write_timeout(Some(HTTP_IO_TIMEOUT));
     let started = Instant::now();
-    let request = match read_request(&mut stream) {
+    let request = match read_request(&mut stream, started + HTTP_REQUEST_TIMEOUT) {
         Ok(request) => request,
         Err(error) => {
             state.metrics.http_started("", 0);
@@ -518,6 +552,41 @@ fn empty_arguments() -> String {
 mod tests {
     use super::*;
 
+    #[test]
+    fn connection_slots_enforce_the_limit_and_release_on_drop() {
+        let active = Arc::new(AtomicUsize::new(0));
+        let mut slots = (0..MAX_CONNECTIONS)
+            .map(|_| ConnectionSlot::acquire(&active).unwrap())
+            .collect::<Vec<_>>();
+
+        assert!(ConnectionSlot::acquire(&active).is_none());
+        slots.pop();
+        assert!(ConnectionSlot::acquire(&active).is_some());
+    }
+
+    #[test]
+    fn request_deadline_stops_slow_clients() {
+        let listener = match TcpListener::bind("127.0.0.1:0") {
+            Ok(listener) => listener,
+            Err(error) if error.kind() == std::io::ErrorKind::PermissionDenied => return,
+            Err(error) => panic!("could not start test server: {error}"),
+        };
+        let address = listener.local_addr().unwrap();
+        let client = thread::spawn(move || {
+            let mut stream = TcpStream::connect(address).unwrap();
+            stream.write_all(b"GET /v1/models HTTP/1.1\r\n").unwrap();
+            thread::sleep(Duration::from_millis(200));
+        });
+        let (mut stream, _) = listener.accept().unwrap();
+
+        let error = read_request(&mut stream, Instant::now() + Duration::from_millis(50))
+            .err()
+            .unwrap();
+
+        client.join().unwrap();
+        assert_eq!(error, "HTTP request timed out");
+    }
+
     #[test]
     fn tool_prompt_matches_ds4_server_text() {
         assert_eq!(TOOLS_PROMPT.len(), 1_183);
@@ -612,6 +681,7 @@ mod tests {
             models_path: PathBuf::new(),
             cache_path: PathBuf::new(),
             sequence: AtomicU64::new(0),
+            connections: Arc::new(AtomicUsize::new(0)),
             tool_memory: Mutex::new(HashMap::new()),
             metrics,
         };
@@ -647,6 +717,7 @@ mod tests {
             models_path: PathBuf::new(),
             cache_path: PathBuf::new(),
             sequence: AtomicU64::new(0),
+            connections: Arc::new(AtomicUsize::new(0)),
             tool_memory: Mutex::new(HashMap::new()),
             metrics,
         };
@@ -673,6 +744,7 @@ mod tests {
             models_path: PathBuf::new(),
             cache_path: PathBuf::new(),
             sequence: AtomicU64::new(0),
+            connections: Arc::new(AtomicUsize::new(0)),
             tool_memory: Mutex::new(HashMap::new()),
             metrics,
         };

src/server/http.rs

@@ -72,21 +72,27 @@ pub(super) fn send_response(
         .map_err(|error| error.to_string())
 }
 
-pub(super) fn read_request(stream: &mut TcpStream) -> Result<HttpRequest, String> {
+pub(super) fn read_request(
+    stream: &mut TcpStream,
+    deadline: Instant,
+) -> Result<HttpRequest, String> {
     let mut bytes = Vec::new();
     let header_end = loop {
-        if bytes.len() >= MAX_HEADER_BYTES {
-            return Err("HTTP headers are too large".into());
-        }
         let mut chunk = [0_u8; 4096];
-        let read = stream.read(&mut chunk).map_err(|error| error.to_string())?;
+        let read = read_with_deadline(stream, &mut chunk, deadline)?;
         if read == 0 {
             return Err("bad HTTP request".into());
         }
         bytes.extend_from_slice(&chunk[..read]);
         if let Some(end) = find_header_end(&bytes) {
+            if end > MAX_HEADER_BYTES {
+                return Err("HTTP headers are too large".into());
+            }
             break end;
         }
+        if bytes.len() >= MAX_HEADER_BYTES {
+            return Err("HTTP headers are too large".into());
+        }
     };
     let header = std::str::from_utf8(&bytes[..header_end])
         .map_err(|_| "HTTP headers are not UTF-8".to_owned())?;
@@ -115,7 +121,7 @@ pub(super) fn read_request(stream: &mut TcpStream) -> Result<HttpRequest, String
     }
     while bytes.len() < header_end + length {
         let mut chunk = [0_u8; 8192];
-        let read = stream.read(&mut chunk).map_err(|error| error.to_string())?;
+        let read = read_with_deadline(stream, &mut chunk, deadline)?;
         if read == 0 {
             return Err("incomplete HTTP body".into());
         }
@@ -128,6 +134,30 @@ pub(super) fn read_request(stream: &mut TcpStream) -> Result<HttpRequest, String
     })
 }
 
+fn read_with_deadline(
+    stream: &mut TcpStream,
+    buffer: &mut [u8],
+    deadline: Instant,
+) -> Result<usize, String> {
+    let remaining = deadline.saturating_duration_since(Instant::now());
+    if remaining.is_zero() {
+        return Err("HTTP request timed out".into());
+    }
+    stream
+        .set_read_timeout(Some(remaining.min(HTTP_IO_TIMEOUT)))
+        .map_err(|error| error.to_string())?;
+    stream.read(buffer).map_err(|error| {
+        if matches!(
+            error.kind(),
+            std::io::ErrorKind::TimedOut | std::io::ErrorKind::WouldBlock
+        ) {
+            "HTTP request timed out".into()
+        } else {
+            error.to_string()
+        }
+    })
+}
+
 fn find_header_end(bytes: &[u8]) -> Option<usize> {
     bytes
         .windows(4)