--- src/lib.rs.before +++ src/lib.rs @@ -13,19 +13,33 @@ use std::rc::{Rc, Weak}; use std::time::{Duration, Instant}; +pyo3::create_exception!(skirr._native, ClientDisconnected, PyOSError); + +const EAGER_LIMIT: usize = net::EVENT_BATCH; const READ_HWM: usize = 64 * 1024; const WRITE_HWM: usize = 512 * 1024; const INPUT_LIMIT: usize = 1024 * 1024; type Link = Rc>; type Request = Rc>; -fn ready(py: Python<'_>, value: Py, state: Option<&State>) -> PyResult> { - let yield_first = state.is_some_and(|s| { - let n = s.sends.get() + 1; - s.sends.set(if n >= 128 { 0 } else { n }); - n >= 128 - }); - completed::value(py, Some(value), yield_first, state) +// Fairness belongs to the producing request. Unrelated short requests must +// not spend each other's quota and delay cleanup after an already-written body. +fn cooperate(request: &Request) -> bool { + let mut r = request.borrow_mut(); + r.operations += 1; + if r.operations < 128 { + return false; + } + r.operations = 0; + true +} +fn ready( + py: Python<'_>, + value: Py, + request: &Request, + state: Option<&State>, +) -> PyResult> { + completed::value(py, Some(value), cooperate(request), state) } fn message(py: Python<'_>, kind: &str, body: Option<(Py, bool)>) -> PyResult> { let d = PyDict::new(py); @@ -86,6 +100,7 @@ head_since: Option, peer: Py, resume_scheduled: bool, + flush_queued: bool, } /// Coalesce small messages, but send a large immutable Python body without an /// intermediate copy. A nonblocking partial write copies only the unsent suffix; @@ -156,6 +171,7 @@ head: bool, minor: u8, task_done: bool, + operations: usize, trailers: bool, body_finished: bool, trailers_started: bool, @@ -167,13 +183,17 @@ loop_: Py, lifespan: Py, prototype: Py, + asgi: Py, done: Py, task_type: Py, task_options: Py, connections: RefCell>, tasks: RefCell>>, next: Cell, - sends: Cell, + eager_left: Cell, + polling: Cell, + flush_scheduled: Cell, + flushes: RefCell>, stopping: Cell, date: RefCell>, date_at: Cell, @@ -226,7 +246,7 @@ // restoring EPOLLIN costs two syscalls per response. If pipelined bytes // arrive before the app finishes, buffer once and pause until dispatch. (if r.body_done { - c.pos == c.input.len() + c.pos == c.input.len() || (r.complete && c.input.len() - c.pos < READ_HWM) } else { r.queued < READ_HWM }) && c.output.len() - c.written < WRITE_HWM @@ -274,7 +294,7 @@ { f.bind(py).call_method1( "set_exception", - (PyOSError::new_err("Client disconnected").value(py),), + (ClientDisconnected::new_err("Client disconnected").value(py),), )?; } let task = self.tasks.borrow().get(&id).map(|t| t.clone_ref(py)); @@ -324,7 +344,7 @@ c.output.clear(); c.written = 0; } - let waiter = if c.output.len() - c.written < READ_HWM { + let waiter = if !failed && c.output.len() - c.written < READ_HWM { c.active.as_ref().and_then(|r| r.borrow_mut().drain.take()) } else { None @@ -338,6 +358,67 @@ self.close(py, c) } else { self.interest(c) + } + } + // Group small completed responses so their writes wake reader CPUs together. + // Streaming chunks and large buffers keep the immediate flush path. + fn flush_or_queue(self: &Rc, py: Python<'_>, c: &Link, complete: bool) -> PyResult<()> { + let size = { + let c = c.borrow(); + c.output.len() - c.written + }; + if !complete || size == 0 || size >= READ_HWM || self.stopping.get() { + return self.flush(py, c); + } + // Check a pending reset before accepting a completed response into the + // transport buffer. A later disconnect remains an unavoidable I/O race. + let error = c.borrow().socket.take_error(); + if !matches!(error, Ok(None)) { + self.close(py, c)?; + return Err(ClientDisconnected::new_err("Client disconnected")); + } + if !c.borrow().flush_queued { + c.borrow_mut().flush_queued = true; + self.flushes.borrow_mut().push(c.clone()); + } + if !self.polling.get() && !self.flush_scheduled.replace(true) { + let queued = (|| -> PyResult<()> { + let callback = Py::new( + py, + Flush { + state: Rc::downgrade(self), + }, + )?; + self.loop_ + .bind(py) + .call_method1(pyo3::intern!(py, "call_soon"), (callback,))?; + Ok(()) + })(); + if let Err(error) = queued { + self.flush_scheduled.set(false); + self.flush_batch(py)?; + return Err(error); + } + } + Ok(()) + } + fn flush_batch(&self, py: Python<'_>) -> PyResult<()> { + let mut queued = self.flushes.replace(Vec::new()); + let mut error = None; + for c in queued.drain(..) { + c.borrow_mut().flush_queued = false; + if let Err(e) = self.flush(py, &c) + && error.is_none() + { + error = Some(e); + } + } + if self.flushes.borrow().is_empty() { + *self.flushes.borrow_mut() = queued; + } + match error { + Some(e) => Err(e), + None => Ok(()), } } fn error(&self, py: Python<'_>, c: &Link, status: u16) -> PyResult<()> { @@ -462,27 +543,20 @@ } fn dispatch(self: &Rc, py: Python<'_>, c: &Link, h: http::Head) -> PyResult<()> { let scope = self.prototype.bind(py).copy()?; - let spec = PyDict::new(py); - spec.set_item(pyo3::intern!(py, "version"), pyo3::intern!(py, "3.0"))?; - spec.set_item(pyo3::intern!(py, "spec_version"), pyo3::intern!(py, "2.5"))?; - scope.set_item(pyo3::intern!(py, "asgi"), spec)?; - scope.set_item( - pyo3::intern!(py, "http_version"), - if h.minor == 1 { - pyo3::intern!(py, "1.1") - } else { - pyo3::intern!(py, "1.0") - }, - )?; - scope.set_item( - pyo3::intern!(py, "method"), - match h.method.as_str() { - "GET" => pyo3::intern!(py, "GET").as_any().clone(), - "POST" => pyo3::intern!(py, "POST").as_any().clone(), - "HEAD" => pyo3::intern!(py, "HEAD").as_any().clone(), - _ => pyo3::types::PyString::new(py, &h.method).into_any(), - }, - )?; + scope.set_item(pyo3::intern!(py, "asgi"), self.asgi.bind(py).copy()?)?; + if h.minor != 1 { + scope.set_item(pyo3::intern!(py, "http_version"), pyo3::intern!(py, "1.0"))?; + } + if h.method != "GET" { + scope.set_item( + pyo3::intern!(py, "method"), + match h.method.as_ref() { + "POST" => pyo3::intern!(py, "POST").as_any().clone(), + "HEAD" => pyo3::intern!(py, "HEAD").as_any().clone(), + _ => pyo3::types::PyString::new(py, &h.method).into_any(), + }, + )?; + } let target = h.target.as_slice(); let q = target .iter() @@ -525,22 +599,24 @@ }, ), )?; - let headers = PyList::empty(py); - for (n, v) in h.headers() { - let cached = self.header_names.borrow().get(n).map(|v| v.clone_ref(py)); - let name = if let Some(name) = cached { - name - } else { - let name = PyBytes::new(py, n).unbind(); - if n.len() <= 128 && self.header_names.borrow().len() < 512 { - self.header_names - .borrow_mut() - .insert(n.to_vec(), name.clone_ref(py)); - } - name - }; - headers.append((name, PyBytes::new(py, v)))?; - } + let headers = PyList::new( + py, + h.headers().map(|(n, v)| { + let cached = self.header_names.borrow().get(n).map(|v| v.clone_ref(py)); + let name = if let Some(name) = cached { + name + } else { + let name = PyBytes::new(py, n).unbind(); + if n.len() <= 128 && self.header_names.borrow().len() < 512 { + self.header_names + .borrow_mut() + .insert(n.to_vec(), name.clone_ref(py)); + } + name + }; + (name, PyBytes::new(py, v)) + }), + )?; scope.set_item(pyo3::intern!(py, "headers"), headers)?; scope.set_item(pyo3::intern!(py, "state"), self.lifespan.bind(py).copy()?)?; scope.set_item(pyo3::intern!(py, "client"), c.borrow().peer.bind(py))?; @@ -570,6 +646,7 @@ head: h.method == "HEAD", minor: h.minor, task_done: false, + operations: 0, trailers: false, body_finished: false, trailers_started: false, @@ -598,7 +675,9 @@ .bind(py) .call_method0(pyo3::intern!(py, "get_task_factory"))? .is_none() + && self.eager_left.get() > 0 { + self.eager_left.set(self.eager_left.get() - 1); self.task_type .bind(py) .call((coro,), Some(self.task_options.bind(py)))? @@ -730,15 +809,15 @@ let s = self .state .upgrade() - .ok_or_else(|| PyOSError::new_err("Server closed"))?; + .ok_or_else(|| ClientDisconnected::new_err("Server closed"))?; let c = self .request .borrow() .connection .upgrade() - .ok_or_else(|| PyOSError::new_err("Client disconnected"))?; + .ok_or_else(|| ClientDisconnected::new_err("Client disconnected"))?; if c.borrow().dead { - return Err(PyOSError::new_err("Client disconnected")); + return Err(ClientDisconnected::new_err("Client disconnected")); } if self.request.borrow().task_done || self.request.borrow().complete { return Err(PyRuntimeError::new_err("Response already completed")); @@ -805,7 +884,9 @@ continue; } has_date |= name.eq_ignore_ascii_case(b"date"); - headers.extend(name.iter().map(u8::to_ascii_lowercase)); + let start = headers.len(); + headers.extend_from_slice(name); + headers[start..].make_ascii_lowercase(); headers.extend_from_slice(b": "); headers.extend_from_slice(value); headers.extend_from_slice(b"\r\n"); @@ -926,12 +1007,13 @@ drop(date); if failed { s.close(py, &c)?; - return Err(PyOSError::new_err("Client disconnected")); + return Err(ClientDisconnected::new_err("Client disconnected")); } response_complete(py, &self.request)?; - s.flush(py, &c)?; + let complete = self.request.borrow().complete; + s.flush_or_queue(py, &c, complete)?; if c.borrow().dead { - return Err(PyOSError::new_err("Client disconnected")); + return Err(ClientDisconnected::new_err("Client disconnected")); } if c.borrow().output.len() - c.borrow().written > WRITE_HWM { if self.request.borrow().drain.is_none() { @@ -996,14 +1078,12 @@ } c.borrow_mut().output.extend_from_slice(&output); response_complete(py, &self.request)?; - s.flush(py, &c)?; + let complete = self.request.borrow().complete; + s.flush_or_queue(py, &c, complete)?; } else { return Err(PyValueError::new_err("Unsupported ASGI message")); } - let n = s.sends.get() + 1; - s.sends.set(n); - if n >= 128 { - s.sends.set(0); + if cooperate(&self.request) { return completed::value(py, None, true, Some(&s)); } Ok(s.done.clone_ref(py)) @@ -1019,13 +1099,28 @@ impl Receiver { fn __call__(&self, py: Python<'_>) -> PyResult> { let Some(s) = self.state.upgrade() else { - return ready(py, message(py, "http.disconnect", None)?, None); + return ready( + py, + message(py, "http.disconnect", None)?, + &self.request, + None, + ); }; let Some(c) = self.request.borrow().connection.upgrade() else { - return ready(py, message(py, "http.disconnect", None)?, Some(&s)); + return ready( + py, + message(py, "http.disconnect", None)?, + &self.request, + Some(&s), + ); }; if c.borrow().dead || self.request.borrow().complete || self.request.borrow().task_done { - return ready(py, message(py, "http.disconnect", None)?, Some(&s)); + return ready( + py, + message(py, "http.disconnect", None)?, + &self.request, + Some(&s), + ); } s.pump(py, &c)?; if !self.request.borrow().body.is_empty() @@ -1033,7 +1128,7 @@ { let msg = take_message(py, &self.request)?; s.pump(py, &c)?; - return ready(py, msg, Some(&s)); + return ready(py, msg, &self.request, Some(&s)); } let old = self.request.borrow_mut().receive.take(); if let Some(f) = old @@ -1072,6 +1167,9 @@ #[pymethods] impl Done { fn __call__(&self, py: Python<'_>, task: &Bound<'_, PyAny>) -> PyResult<()> { + if let Some(s) = self.state.upgrade() { + s.eager_left.set(EAGER_LIMIT); + } self.finish(py, task, true) } } @@ -1084,23 +1182,22 @@ if drive { s.tasks.borrow_mut().remove(&id); } - let cancelled = task - .call_method0(pyo3::intern!(py, "cancelled"))? - .is_truthy()?; - let failed = if cancelled { - false - } else { - let error = task.call_method0(pyo3::intern!(py, "exception"))?; - if !error.is_none() { - let ctx = PyDict::new(py); - ctx.set_item(pyo3::intern!(py, "message"), "ASGI application failed")?; - ctx.set_item(pyo3::intern!(py, "exception"), error)?; - s.loop_ - .bind(py) - .call_method1(pyo3::intern!(py, "call_exception_handler"), (ctx,))?; - true - } else { - false + let (cancelled, failed) = match task.call_method0(pyo3::intern!(py, "result")) { + Ok(_) => (false, false), + Err(error) => { + if error.is_instance_of::(py) { + (false, true) + } else if error.is_instance(py, &py.import("asyncio")?.getattr("CancelledError")?) { + (true, false) + } else { + let ctx = PyDict::new(py); + ctx.set_item(pyo3::intern!(py, "message"), "ASGI application failed")?; + ctx.set_item(pyo3::intern!(py, "exception"), error.value(py))?; + s.loop_ + .bind(py) + .call_method1(pyo3::intern!(py, "call_exception_handler"), (ctx,))?; + (false, true) + } } }; self.request.borrow_mut().task_done = true; @@ -1134,9 +1231,24 @@ c.closing |= !keep; c.activity = Instant::now(); } - s.flush(py, &c)?; + s.flush_or_queue(py, &c, true)?; if drive && !c.borrow().dead { s.pump(py, &c)?; + } + Ok(()) + } +} + +#[pyclass(unsendable, frozen)] +struct Flush { + state: Weak, +} +#[pymethods] +impl Flush { + fn __call__(&self, py: Python<'_>) -> PyResult<()> { + if let Some(s) = self.state.upgrade() { + s.flush_scheduled.set(false); + s.flush_batch(py)?; } Ok(()) } @@ -1151,6 +1263,7 @@ impl Resume { fn __call__(&self, py: Python<'_>) -> PyResult<()> { if let (Some(s), Some(c)) = (self.state.upgrade(), self.connection.upgrade()) { + s.eager_left.set(EAGER_LIMIT); c.borrow_mut().resume_scheduled = false; s.pump(py, &c)?; } @@ -1189,6 +1302,11 @@ ] { prototype.set_item(key, py.None())?; } + prototype.set_item(pyo3::intern!(py, "method"), pyo3::intern!(py, "GET"))?; + prototype.set_item(pyo3::intern!(py, "http_version"), pyo3::intern!(py, "1.1"))?; + let asgi = PyDict::new(py); + asgi.set_item(pyo3::intern!(py, "version"), pyo3::intern!(py, "3.0"))?; + asgi.set_item(pyo3::intern!(py, "spec_version"), pyo3::intern!(py, "2.5"))?; prototype.set_item(pyo3::intern!(py, "type"), "http")?; prototype.set_item(pyo3::intern!(py, "scheme"), "http")?; prototype.set_item(pyo3::intern!(py, "root_path"), "")?; @@ -1212,13 +1330,17 @@ loop_, lifespan, prototype: prototype.unbind(), + asgi: asgi.unbind(), done, task_type, task_options: task_options.unbind(), connections: RefCell::new(HashMap::new()), tasks: RefCell::new(HashMap::new()), next: Cell::new(0), - sends: Cell::new(0), + eager_left: Cell::new(EAGER_LIMIT), + polling: Cell::new(false), + flush_scheduled: Cell::new(false), + flushes: RefCell::new(Vec::new()), stopping: Cell::new(false), date: RefCell::new(Vec::new()), date_at: Cell::new(0), @@ -1236,71 +1358,79 @@ } fn poll(&self, py: Python<'_>) -> PyResult<()> { let s = &self.state; - s.sends.set(0); - for (token, mask) in s.poller.events()? { - if token == 0 { - if s.stopping.get() { - continue; - } - for _ in 0..64 { - match s.listener.accept() { - Ok((socket, peer)) => { - if s.connections.borrow().len() >= 10000 { - continue; + s.eager_left.set(EAGER_LIMIT); + s.polling.set(true); + let result = (|| -> PyResult<()> { + for (token, mask) in s.poller.events()? { + if token == 0 { + if s.stopping.get() { + continue; + } + for _ in 0..64 { + match s.listener.accept() { + Ok((socket, peer)) => { + if s.connections.borrow().len() >= 10000 { + continue; + } + socket.set_nonblocking(true)?; + socket.set_nodelay(true)?; + let token = s.id(); + let mask = libc::EPOLLIN as u32 | libc::EPOLLRDHUP as u32; + s.poller.watch(socket.as_raw_fd(), token, mask, true)?; + s.connections.borrow_mut().insert( + token, + Rc::new(RefCell::new(Connection { + socket, + token, + input: Vec::new(), + pos: 0, + active: None, + output: Vec::new(), + headers: Vec::new(), + written: 0, + closing: false, + dead: false, + mask, + activity: Instant::now(), + head_since: None, + resume_scheduled: false, + flush_queued: false, + peer: (peer.ip().to_string(), peer.port()) + .into_pyobject(py)? + .into_any() + .unbind(), + })), + ); } - socket.set_nonblocking(true)?; - socket.set_nodelay(true)?; - let token = s.id(); - let mask = libc::EPOLLIN as u32 | libc::EPOLLRDHUP as u32; - s.poller.watch(socket.as_raw_fd(), token, mask, true)?; - s.connections.borrow_mut().insert( - token, - Rc::new(RefCell::new(Connection { - socket, - token, - input: Vec::new(), - pos: 0, - active: None, - output: Vec::new(), - headers: Vec::new(), - written: 0, - closing: false, - dead: false, - mask, - activity: Instant::now(), - head_since: None, - resume_scheduled: false, - peer: (peer.ip().to_string(), peer.port()) - .into_pyobject(py)? - .into_any() - .unbind(), - })), - ); + Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => break, + Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue, + Err(e) => return Err(e.into()), } - Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => break, - Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue, - Err(e) => return Err(e.into()), - } - } - } else { - let conn = s.connections.borrow().get(&token).cloned(); - if let Some(c) = conn { - if mask & (libc::EPOLLERR | libc::EPOLLHUP) as u32 != 0 { - s.close(py, &c)?; - continue; - } - if mask & (libc::EPOLLIN | libc::EPOLLRDHUP) as u32 != 0 { - s.read(py, &c)?; - } - if mask & libc::EPOLLOUT as u32 != 0 { - s.flush(py, &c)?; - s.pump(py, &c)?; - } - } - } - } - Ok(()) - } + } + } else { + let conn = s.connections.borrow().get(&token).cloned(); + if let Some(c) = conn { + if mask & (libc::EPOLLERR | libc::EPOLLHUP) as u32 != 0 { + s.close(py, &c)?; + continue; + } + if mask & (libc::EPOLLIN | libc::EPOLLRDHUP) as u32 != 0 { + s.read(py, &c)?; + } + if mask & libc::EPOLLOUT as u32 != 0 { + s.flush(py, &c)?; + s.pump(py, &c)?; + } + } + } + } + Ok(()) + })(); + s.polling.set(false); + let flushed = s.flush_batch(py); + result.and(flushed) + } + fn tick(&self, py: Python<'_>) -> PyResult<()> { self.state.date(py)?; let conns: Vec<_> = self.state.connections.borrow().values().cloned().collect(); @@ -1381,5 +1511,9 @@ #[pymodule] fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; + m.add( + "ClientDisconnected", + m.py().get_type::(), + )?; Ok(()) }