--- 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<RefCell<Connection>>;
type Request = Rc<RefCell<RequestState>>;
-fn ready(py: Python<'_>, value: Py<PyAny>, state: Option<&State>) -> PyResult<Py<PyAny>> {
- 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<PyAny>,
+ request: &Request,
+ state: Option<&State>,
+) -> PyResult<Py<PyAny>> {
+ completed::value(py, Some(value), cooperate(request), state)
}
fn message(py: Python<'_>, kind: &str, body: Option<(Py<PyBytes>, bool)>) -> PyResult<Py<PyAny>> {
let d = PyDict::new(py);
@@ -86,6 +100,7 @@
head_since: Option<Instant>,
peer: Py<PyAny>,
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<PyAny>,
lifespan: Py<PyDict>,
prototype: Py<PyDict>,
+ asgi: Py<PyDict>,
done: Py<PyAny>,
task_type: Py<PyAny>,
task_options: Py<PyDict>,
connections: RefCell<HashMap<u64, Link>>,
tasks: RefCell<HashMap<u64, Py<PyAny>>>,
next: Cell<u64>,
- sends: Cell<usize>,
+ eager_left: Cell<usize>,
+ polling: Cell<bool>,
+ flush_scheduled: Cell<bool>,
+ flushes: RefCell<Vec<Link>>,
stopping: Cell<bool>,
date: RefCell<Vec<u8>>,
date_at: Cell<u64>,
@@ -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<Self>, 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<Self>, 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<Py<PyAny>> {
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::<ClientDisconnected>(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<State>,
+}
+#[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::<Worker>()?;
+ m.add(
+ "ClientDisconnected",
+ m.py().get_type::<ClientDisconnected>(),
+ )?;
Ok(())
}