diff --git a/ffi/README.md b/ffi/README.md index 2f8ae6824..dd4cf2e51 100644 --- a/ffi/README.md +++ b/ffi/README.md @@ -352,7 +352,8 @@ function in-process, so a call is one crossing, which releases the GIL while Rus - **Futures** are `async def`s, which run nothing until awaited, or for a function that is not `async`, an awaitable the call returns. They resume on the running event loop and never inside a wake: a wake from any thread queues the future's token and writes a byte - to a socket the loop reads, without the GIL. Cancelling the task frees the future and + to a socket the loop reads, without the GIL, or during a call on the loop's own thread, + writes nothing and schedules the loop's dispatch once the call returns. Cancelling the task frees the future and loses nothing. The call polls it once before suspending, and a future is awaited once, as a coroutine is. - A **waker** is a `Waker(callback)`, which Rust may wake from any thread, calling diff --git a/ffi/runtime/python.rs b/ffi/runtime/python.rs index a5002fc7a..85ebdc5fa 100644 --- a/ffi/runtime/python.rs +++ b/ffi/runtime/python.rs @@ -5,9 +5,11 @@ //! The glue holds the GIL only while it touches Python: it releases it around each call //! into the C ABI, so Rust may call Python back from any thread, as a lent closure or a //! trait method, while the caller waits. A wake never takes the GIL: it queues a token -//! and writes a byte to a socket the event loop reads, and the loop does the rest. +//! and writes a byte to a socket the event loop reads, and the loop does the rest. A wake +//! on the loop's own thread, during a call, writes nothing: the call schedules the loop's +//! dispatch once it returns. -use std::cell::UnsafeCell; +use std::cell::{Cell, RefCell, UnsafeCell}; use std::ffi::{CStr, c_char, c_int, c_void}; use std::io::{Read, Write}; use std::ptr::{NonNull, null, null_mut}; @@ -138,16 +140,95 @@ pub unsafe fn args<'a>(args: *mut *mut PyObject, nargs: Py_ssize_t, n: usize, na } /// Call into the C ABI with the GIL released, so Rust may call Python back from any -/// thread while this one waits. +/// thread while this one waits, then resume what it woke on this thread's loop. pub fn call(f: impl FnOnce() -> T) -> T { + // A call inside another keeps its loop: none starts on a thread one runs on. + let outer = RUNNING.get(); + if outer == 0 && LOOPS.get() { + let running = running_loop(); + // Until a loop polls or takes here again, which it does once it runs. + LOOPS.set(running != 0); + RUNNING.set(running); + } // SAFETY: this thread holds the GIL, and touches no Python object until it retakes it. let state = unsafe { PyEval_SaveThread() }; let result = f(); // SAFETY: released just above. unsafe { PyEval_RestoreThread(state) }; + RUNNING.set(outer); + if !RESUMES.with_borrow(Vec::is_empty) { + resume(); + } result } +thread_local! { + /// Whether an event loop may run on this thread, as one polled or took here since a + /// call last found none, so its calls look for the loop running. + static LOOPS: Cell = const { Cell::new(false) }; + /// The address of the event loop running on this thread, during a call, else `0`. + static RUNNING: Cell = const { Cell::new(0) }; + /// The queues woken on their loop's thread during the calls it is in. + static RESUMES: RefCell>> = const { RefCell::new(Vec::new()) }; +} + +static GET_RUNNING_LOOP: std::sync::OnceLock = std::sync::OnceLock::new(); + +/// The address of the event loop running on this thread, or `0`: none is while an +/// exception is set, which a call must not run Python over. +fn running_loop() -> usize { + // SAFETY: the GIL is held, and `asyncio._get_running_loop` neither raises nor runs + // Python code. + unsafe { + if !PyErr_Occurred().is_null() { + return 0; + } + let get = GET_RUNNING_LOOP.get().expect("loaded").ptr(); + let Ok(running) = Obj::new(PyObject_CallNoArgs(get)) else { + PyErr_Clear(); + return 0; + }; + match is_none(running.ptr()) { + true => 0, + false => running.ptr() as usize, + } + } +} + +/// Schedule the dispatch of each queue a call woke on its loop's thread, as the loop's +/// next callback, or else through its socket as a wake from another thread is. +#[cold] +fn resume() { + let woken = RESUMES.with_borrow_mut(std::mem::take); + let resume = ty(T_RESUME); + // SAFETY: the GIL is held. An exception the caller has set is kept aside while Python + // runs, and put back after. + unsafe { + let (mut kind, mut value, mut trace) = (null_mut(), null_mut(), null_mut()); + PyErr_Fetch(&mut kind, &mut value, &mut trace); + let running = match Obj::new(PyObject_CallNoArgs(resume)) { + Ok(queue) => is_type(queue.ptr(), &QUEUE).then(|| queue_of(queue.ptr()).ok()).flatten(), + Err(Raised) => { + // Nothing is lost: each queue's socket signals it instead, and an interrupt + // its Python ran into is raised again at the thread's next chance. + if PyErr_ExceptionMatches(PyExc_KeyboardInterrupt) != 0 { + PyErr_Clear(); + PyErr_SetInterrupt(); + } else { + PyErr_WriteUnraisable(resume); + } + None + } + }; + for queue in woken { + if !running.as_ref().is_some_and(|running| Arc::ptr_eq(running, &queue)) { + queue.signal(); + } + } + PyErr_Restore(kind, value, trace); + } +} + /// Run `f` calling into Python, from whatever thread Rust calls on, holding the GIL; an /// exception it raises aborts with what it was, since Rust has nowhere to unwind it to. pub fn upcall(what: &str, f: impl FnOnce() -> R) -> T { @@ -256,17 +337,18 @@ impl Drop for Scope { } } -/// The classes the generated Python registers, by the index the glue names each with: -/// the runtime's, then the API's. +/// The objects the generated Python registers, by the index the glue names each with: +/// the runtime's, then the API's classes. static TYPES: Gil> = Gil::new(Vec::new()); /// The plain enums' members, by class, with their discriminants. static MEMBERS: Gil>>> = Gil::new(Vec::new()); -/// The runtime's classes, at the start of the registry: the awaitable a pending future -/// is, and `datetime.timedelta`. +/// The runtime's entries, at the start of the registry: the awaitable a pending future +/// is, `datetime.timedelta`, and `_resume`, which schedules the running loop's dispatch. pub const T_FUTURE: usize = 0; const T_TIMEDELTA: usize = 1; +const T_RESUME: usize = 2; /// The class registered at `index`. pub fn ty(index: usize) -> *mut PyObject { @@ -1153,10 +1235,15 @@ unsafe extern "C" fn bytes_dealloc(obj: *mut PyObject) { } /// What a wake does: queue a token for the event loop, and write a byte it reads, once -/// until it takes them. +/// until it takes them; or, during a call on the thread running the loop, note the queue +/// for the call to resume once it returns. pub struct Queue { tokens: Mutex>, signaled: AtomicBool, + /// Whether a call holds the queue to resume, until it is taken. + woken: AtomicBool, + /// The address of its event loop, which [`RUNNING`] matches while it runs a call. + event_loop: usize, tx: Socket, rx: Socket, } @@ -1189,20 +1276,36 @@ fn socket_pair() -> std::io::Result<(Socket, Socket)> { } impl Queue { - fn push(&self, token: i64) { + fn push(self: &Arc, token: i64) { self.tokens.lock().unwrap_or_else(|e| e.into_inner()).push(token); + // During a call on the thread running the loop, which runs nothing until it returns. + if RUNNING.get() == self.event_loop { + if !self.woken.swap(true, Ordering::SeqCst) { + RESUMES.with_borrow_mut(|resumes| resumes.push(self.clone())); + } + return; + } + self.signal(); + } + + fn signal(&self) { if !self.signaled.swap(true, Ordering::SeqCst) { // A full socket already signals. let _ = (&self.tx).write(&[1]); } } - /// The tokens queued, each once, in the order first queued: drained, then the signal - /// reset, then the queue taken, so a token queued after is signaled again. - fn take(&self) -> Vec { - let mut sink = [0u8; 64]; - while matches!((&self.rx).read(&mut sink), Ok(n) if n > 0) {} - self.signaled.store(false, Ordering::SeqCst); + /// The tokens queued, each once, in the order first queued: the socket drained if + /// `read`, as the loop does once it signals, then each signal reset, then the queue + /// taken, so a token queued after is signaled again. + fn take(&self, read: bool) -> Vec { + LOOPS.set(true); + if read { + let mut sink = [0u8; 64]; + while matches!((&self.rx).read(&mut sink), Ok(n) if n > 0) {} + self.signaled.store(false, Ordering::SeqCst); + } + self.woken.store(false, Ordering::SeqCst); let tokens = std::mem::take(&mut *self.tokens.lock().unwrap_or_else(|e| e.into_inner())); let mut seen = std::collections::HashSet::new(); tokens.into_iter().filter(|t| seen.insert(*t)).collect() @@ -1282,38 +1385,36 @@ unsafe extern "C" fn node_drop(data: *mut c_void) { drop(unsafe { Arc::from_raw(data.cast::()) }); } -/// `Queue`: the queue of one event loop, with its socket's file descriptor to read. +/// `Queue(loop)`: the queue of one event loop, with its socket's file descriptor to read. #[repr(C)] struct QueueObj { ob_base: PyObject, queue: Option>, } -unsafe extern "C" fn queue_new(class: *mut PyTypeObject, _: *mut PyObject, _: *mut PyObject) -> *mut PyObject { - let (tx, rx) = match socket_pair() { - Ok(pair) => pair, - Err(err) => { +unsafe extern "C" fn queue_new(class: *mut PyTypeObject, args: *mut PyObject, _: *mut PyObject) -> *mut PyObject { + wrap(|| { + let event_loop = tuple_item(args, 0, 1)? as usize; + let (tx, rx) = match socket_pair() { + Ok(pair) => pair, // SAFETY: a static exception type. - let _ = raise::<()>(unsafe { PyExc_OSError }, &err.to_string()); - return null_mut(); - } - }; - // SAFETY: the GIL is held, and `class` is `Queue`. - unsafe { - let obj = PyType_GenericAlloc(class, 0); - if !obj.is_null() { - std::ptr::write( - &raw mut (*obj.cast::()).queue, - Some(Arc::new(Queue { - tokens: Mutex::new(Vec::new()), - signaled: AtomicBool::new(false), - tx, - rx, - })), - ); + Err(err) => return raise(unsafe { PyExc_OSError }, &err.to_string()), + }; + // SAFETY: the GIL is held, and `class` is `Queue`. + unsafe { + let obj = Obj::new(PyType_GenericAlloc(class, 0))?; + let queue = Queue { + tokens: Mutex::new(Vec::new()), + signaled: AtomicBool::new(false), + woken: AtomicBool::new(false), + event_loop, + tx, + rx, + }; + std::ptr::write(&raw mut (*obj.ptr().cast::()).queue, Some(Arc::new(queue))); + Ok(obj) } - obj - } + }) } unsafe extern "C" fn queue_dealloc(obj: *mut PyObject) { @@ -1343,10 +1444,12 @@ unsafe extern "C" fn queue_fileno(obj: *mut PyObject, _: *mut PyObject) -> *mut int_out(fd).map_or(null_mut(), Obj::into_ptr) } -unsafe extern "C" fn queue_take(obj: *mut PyObject, _: *mut PyObject) -> *mut PyObject { +/// `take(read)`: the tokens queued, the socket drained first if `read`. +unsafe extern "C" fn queue_take(obj: *mut PyObject, read: *mut PyObject) -> *mut PyObject { // SAFETY: Python calls it on a `Queue`. let queue = unsafe { (*obj.cast::()).queue.as_ref() }.expect("a live queue"); - let tokens = queue.take(); + // SAFETY: comparing addresses. + let tokens = queue.take(read == unsafe { Py_True() }); wrap(|| { let list = rt::List::owned(tokens); list_out(list, int_out) @@ -1522,6 +1625,7 @@ unsafe extern "C" fn future_poll(obj: *mut PyObject, args: *mut *mut PyObject, n node } }; + LOOPS.set(true); loop { let first = node.state.swap(POLLING, Ordering::AcqRel) == IDLE; // The first poll with the node gives the C ABI a reference, which it keeps @@ -1611,6 +1715,15 @@ const fn noargs(name: &'static CStr, f: PyCFunction) -> PyMethodDef { } } +const fn onearg(name: &'static CStr, f: PyCFunction) -> PyMethodDef { + PyMethodDef { + ml_name: name.as_ptr(), + ml_meth: PyMethodDefPointer { PyCFunction: f }, + ml_flags: METH_O, + ml_doc: null(), + } +} + pub const END: PyMethodDef = PyMethodDef::zeroed(); /// A table of methods Python reads for the interpreter's life. @@ -1634,7 +1747,7 @@ static FUTURE_METHODS: Methods<4> = Methods(UnsafeCell::new([ static QUEUE_METHODS: Methods<3> = Methods(UnsafeCell::new([ noargs(c"fileno", queue_fileno), - noargs(c"take", queue_take), + onearg(c"take", queue_take), END, ])); @@ -1745,6 +1858,9 @@ unsafe fn runtime(module: *mut PyObject) -> R<()> { return Err(Raised); } let _ = PENDING.set(Static(pending)); + let asyncio = Obj::new(PyImport_ImportModule(c"asyncio".as_ptr()))?; + let get = Obj::new(PyObject_GetAttrString(asyncio.ptr(), c"_get_running_loop".as_ptr()))?; + let _ = GET_RUNNING_LOOP.set(Static(get)); } Ok(()) } diff --git a/ffi/runtime/rt.py b/ffi/runtime/rt.py index 3b66d0dd8..ceb63ac62 100644 --- a/ffi/runtime/rt.py +++ b/ffi/runtime/rt.py @@ -29,19 +29,20 @@ def __await__(self): class _Wakeup: """One event loop's wakes: Rust queues a token and writes a byte to a socket the loop - reads, from any thread and without the GIL, and the loop resumes what each names.""" + reads, from any thread and without the GIL, and the loop resumes what each names. A wake + during a call on the loop's own thread writes nothing: `_resume` dispatches it.""" __slots__ = ("queue", "parked", "wakers", "pump", "__weakref__") def __init__(self, loop): - self.queue = _n.Queue() + self.queue = _n.Queue(loop) # The future each parked poll waits on, and each `Waker`'s callback, by token. self.parked = {} self.wakers = {} self.pump = None fd = self.queue.fileno() try: - loop.add_reader(fd, self.dispatch) + loop.add_reader(fd, self.dispatch, True) except NotImplementedError: # A proactor loop, Windows' default, reads only through its own calls. sock = _socket.fromfd(fd, _socket.AF_INET if _sys.platform == "win32" else _socket.AF_UNIX, _socket.SOCK_STREAM) @@ -52,11 +53,11 @@ async def _pump(self, loop, sock): with sock: while True: await loop.sock_recv(sock, 64) - self.dispatch() + self.dispatch(True) - def dispatch(self): + def dispatch(self, read): loop = None - for token in self.queue.take(): + for token in self.queue.take(read): if token < 0: self.wakers.pop(-token, None) continue @@ -73,6 +74,18 @@ def dispatch(self): _wakeups = _weakref.WeakKeyDictionary() +def _resume(): + """Called once a call on the running loop's thread returns having woken some of its + futures or wakers: schedule their dispatch as the loop's next callback, never inside + the call, and return the queue it takes from.""" + loop = _get_running_loop() + w = _wakeups.get(loop) + if w is None: + return None + loop.call_soon(w.dispatch, False) + return w.queue + + def _wakeup(loop): w = _wakeups.get(loop) if w is None: diff --git a/ffi/src/python.rs b/ffi/src/python.rs index 4b5ecdd08..41b4321cb 100644 --- a/ffi/src/python.rs +++ b/ffi/src/python.rs @@ -48,6 +48,7 @@ const TAKEN: &[&str] = &[ "_Wakeup", "_wakeups", "_wakeup", + "_resume", "_Const", "_PENDING", "_consts", @@ -69,8 +70,9 @@ const TAKEN: &[&str] = &[ "_AttributeError", ]; -/// The registry's first classes, the runtime's: the glue's `T_FUTURE` and `T_TIMEDELTA`. -const RUNTIME_TYPES: &[&str] = &["_Future", "_datetime.timedelta"]; +/// The registry's first entries, the runtime's: the glue's `T_FUTURE`, `T_TIMEDELTA`, and +/// `T_RESUME`. +const RUNTIME_TYPES: &[&str] = &["_Future", "_datetime.timedelta", "_resume"]; /// The Python for `abi`, whose header's contract is `contract`, or the names two items /// claim alike. The Python's own contract hashes the header's with everything it and its diff --git a/ffi/tests/api.py b/ffi/tests/api.py index b715ce0a1..6a4df8396 100644 --- a/ffi/tests/api.py +++ b/ffi/tests/api.py @@ -8,6 +8,7 @@ import datetime import gc import os +import select import sys import threading @@ -561,7 +562,7 @@ def log(self, level, line): def close(self): try: - later.poll(api._n.Queue()) + later.poll(api._n.Queue(loop)) check(False, "a future polled while it is freed raises") except ValueError: pass @@ -579,10 +580,94 @@ def close(self): sender.join() until(loop, lambda: wakes) check(wakes == [threading.get_ident()], "a waker calls back on its loop's thread") + loop.run_until_complete(asyncio.wait_for(on_loop(), 5)) loop.close() print("futures ok") +async def on_loop(): + """Wakes during a call on the running loop's own thread, which write no byte to its + socket and dispatch once the call returns, beside wakes from other threads, which do.""" + channel = api.Channel.new() + got = [] + + async def recv(): + got.append(text(await channel.recv())) + + async def parked(): + task = asyncio.ensure_future(recv()) + await asyncio.sleep(0) + return task + + def signaled(): + return bool(select.select([wakeup.queue.fileno()], [], [], 0)[0]) + + async def settle(): + for _ in range(4): + await asyncio.sleep(0) + + task = await parked() + wakeup = api._wakeups[asyncio.get_running_loop()] + channel.send(b"near") + check(not got, "a wake on the loop's thread never resumes inside the call") + check(not signaled(), "a wake on the loop's thread writes no byte") + await task + check(got == ["near"], "a wake on the loop's thread resumes the receive") + + task = await parked() + sender = threading.Thread(target=lambda: channel.send(b"far")) + sender.start() + sender.join() + check(signaled(), "a wake from another thread writes a byte") + await task + check(got == ["near", "far"], "a wake from another thread resumes the receive") + + # Before the first await, the call's poll with no waker leaves nothing to wake. + early = channel.recv() + channel.send(b"early") + check(text(await early) == "early", "a send before the first await is received") + + # Cancelled after its wake, or before it, in the same step, a receive never resumes and + # leaves the payload. + for send_first in (True, False): + task = await parked() + if send_first: + channel.send(b"kept") + task.cancel() + else: + task.cancel() + channel.send(b"kept") + await settle() + check(task.cancelled() and len(got) == 2, "a cancelled receive never resumes") + check(text(await channel.recv()) == "kept", "a cancelled receive loses nothing") + check(not signaled(), "wakes on the loop's thread wrote no byte") + + # Wakes that land before the loop gets to them call the waker once. + wakes = [] + waker = api.Waker(lambda: wakes.append(1)) + for _ in range(2): + channel.notify(waker) + channel.send(b"rung") + check(not wakes, "a waker is never called inside the call") + await settle() + check(wakes == [1], "wakes on the loop's thread call the waker once") + await channel.recv() + + # Many receives, each woken once from the loop's thread or another, while both wake. + channels = [api.Channel.new() for _ in range(64)] + for _ in range(20): + tasks = [asyncio.ensure_future(c.recv()) for c in channels] + await asyncio.sleep(0) + sender = threading.Thread(target=lambda: [c.send(b"far") for c in channels[1::2]]) + sender.start() + for c in channels[::2]: + c.send(b"near") + payloads = await asyncio.gather(*tasks) + sender.join() + expected = ["near", "far"] * (len(channels) // 2) + check([text(p) for p in payloads] == expected, "every receive resumes once with its payload") + + def owned(): """What frees an object or bytes: its last reference going, or the cycle collector. Rust's bytes cross lent, until they and every view of them are released.""" diff --git a/quest/m3/README.md b/quest/m3/README.md index c9ef7bdda..403c6db04 100644 --- a/quest/m3/README.md +++ b/quest/m3/README.md @@ -29,7 +29,6 @@ generic std gap is planned here even when moq never calls it. - [A vector's exact capacity as it shrinks](/quest/m3/ts-vec-exact-capacity.md) - a deque made from a `Vec` splits its slices where std's does - [A string's exact capacity](/quest/m3/ts-string-capacity.md) - a `String` holds std's capacity, so a program making a `Bytes` of one reads its capacity - [A B-tree's searches call Ord as std's do](/quest/m3/ts-btree-search-order.md) - a B-tree's lookups and inserts call a key's own `Ord` in std's order and number -- [Python's wake on the loop's own thread](/quest/m3/ffi-python-wake.md) - a wake on the event loop's thread resumes the future with no socket pair - [A panic hook](/quest/m3/ts-panic-hook.md) - `set_hook`, `take_hook`, and `PanicHookInfo` translate, every panic running the hook - [A string Cow's variant](/quest/m3/ts-cow-str-variant.md) - matching a `Cow`'s `Borrowed` or `Owned`, and `to_mut`'s change of it, translate as std's - [An Arc's raw pointer to a value JS copies](/quest/m3/ts-arc-raw-place.md) - `from_raw` of a pointer remade from a reference to an `Arc`'s number, and an `Arc`'s raw pointer cast to `*const ()`, translate diff --git a/quest/m3/ffi-python-wake.md b/quest/m3/ffi-python-wake.md deleted file mode 100644 index 58ad5b613..000000000 --- a/quest/m3/ffi-python-wake.md +++ /dev/null @@ -1,15 +0,0 @@ -# [S] Python's wake on the loop's own thread - -## Goal - -A future the exporter's [Python](/ffi/README.md) awaits resumes directly when Rust wakes it -on the event loop's own thread, skipping the socket pair, and `just ffi bench` measures the -woken read against [BoltFFI's](/bench/ffi/README.md#boltffi-comparison) 10 µs. - -## Plan - -Found by [#407](https://github.com/moq-dev/rust2/pull/407): a woken read costs about 21 µs, -about 5 of them the socket pair's write and read. A wake from another thread still goes -through it, and the wake-latch and cancel tests keep passing. - -Public API: none.