diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ee10588..b5727eb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -44,7 +44,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.10", "3.11", "3.12", "3.13"] + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] steps: - uses: actions/checkout@v4 @@ -75,6 +75,14 @@ jobs: source .venv/bin/activate pytest tests/test_basic.py -v || echo "No unit tests found, skipping" + - name: Run stdlib tests + run: | + source .venv/bin/activate + # We need to ensure we can import 'test' module which is usually in Lib/test + # In many CI environments/installations, test suite might be separate or require specific paths. + # For standard python installs it should work if it's installed. + python run_stdlib_tests.py + - name: Run e2e tests run: | source .venv/bin/activate diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index da7f9d3..d47047b 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -567,11 +567,12 @@ The following state-of-the-art optimizations have been implemented or are availa |--------------|--------|-------------------| | **Asyncio Function Caching** | ✅ Active | N/A | | **Native Timers** (`IORING_OP_TIMEOUT`) | ✅ Available | 5.4+ | -| **Multishot Recv** (`IORING_OP_RECV` + `RECV_MULTISHOT`) | ✅ Available | 5.19+ | | **Native Scheduler** (`Mutex`) | ✅ Active | N/A | | **Merged Ring Lock** (single lock per run_tick) | ✅ Active | N/A | | **Registered FD Table** (`IOSQE_FIXED_FILE`) | ✅ Available | 5.1+ | -| **Zero-Copy Send** (`IORING_OP_SEND_ZC`) | ✅ Available | 6.0+ | +- **Zero-Copy Send (IORING_OP_SEND_ZC)**: Implemented. Available on Kernel 6.0+. Uses `submit_send_zc`. + +- **Multishot Recv**: Implemented. (Runtime Feature Detection) ### Available (Runtime Feature Detection) diff --git a/python/uringcore/loop.py b/python/uringcore/loop.py index b5bd5d5..eae986c 100644 --- a/python/uringcore/loop.py +++ b/python/uringcore/loop.py @@ -187,6 +187,9 @@ def run_forever(self) -> None: """Run the event loop until stop() is called.""" self._check_closed() self._check_running() + + if asyncio._get_running_loop() is not None: + raise RuntimeError("Cannot run the event loop while another loop is running") self._running = True self._thread_id = None diff --git a/run_stdlib_tests.py b/run_stdlib_tests.py index 331e941..b62e2ba 100644 --- a/run_stdlib_tests.py +++ b/run_stdlib_tests.py @@ -18,6 +18,7 @@ def run_tests(): # We can load it using unittest print("Loading test.test_asyncio...") + # Python 3.13 location might differ slightly or require strict module naming try: from test import test_asyncio @@ -25,12 +26,23 @@ def run_tests(): print("Could not import test.test_asyncio. Are you on a standard Python install?") return - # Create a test suite - suite = unittest.TestLoader().loadTestsFromModule(test_asyncio) + # Check if test_asyncio is a package or module + if hasattr(test_asyncio, "__path__"): + # It's a package, load all tests from it + print(f"Discovered test_asyncio package at {test_asyncio.__path__}") + suite = unittest.TestLoader().discover( + start_dir=test_asyncio.__path__[0], + pattern="test_*.py", + top_level_dir=os.path.dirname(test_asyncio.__path__[0]) + ) + else: + # It's a single module + suite = unittest.TestLoader().loadTestsFromModule(test_asyncio) # Run it print("Running asyncio stdlib tests with uringcore...") - result = unittest.TextTestRunner(verbosity=2).run(suite) + runner = unittest.TextTestRunner(verbosity=2) + result = runner.run(suite) if not result.wasSuccessful(): sys.exit(1) diff --git a/src/lib.rs b/src/lib.rs index 5a5f29c..67ddb97 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -97,6 +97,8 @@ struct RecvMsgState { pub msghdr: libc::msghdr, pub iovec: libc::iovec, pub addr: libc::sockaddr_storage, + // Own the buffer + pub data: Vec, } /// State for an in-flight `sendmsg` operation. @@ -105,17 +107,7 @@ struct SendMsgState { pub msghdr: libc::msghdr, pub iovec: libc::iovec, pub addr: libc::sockaddr_storage, - // We need to keep the data alive too if it's not copied into a kernel buffer immediately. - // For io_uring sendmsg, the iovec points to the data. - // If we pass bytes from Python, we typically need to ensure they stay valid. - // However, for typical send operations, we might copy the data into a buffer we own - // or rely on PyBytes being immortal if we hold a reference (but we can't easily hold Py in a raw struct without GIL). - // - // A better approach for `submit_sendto` is to allocate a buffer from our `BufferPool` (or a separate `Vec`) - // and copy the data there, OR hold the Py. - // Given our BufferPool is for 'recv' mainly (fixed size chunks), for send we might just want to use a `Vec` or `Box<[u8]>`. - // - // To keep it simple and safe: We will own the data in this state. + // We own the data to ensure it stays alive during the operation pub data: Vec, } @@ -141,6 +133,8 @@ pub struct UringCore { fd_states: FDStateManager, /// Inflight recv buffers: fd -> `buffer_index` (for completion data extraction) inflight_recv_buffers: Mutex>, + /// Inflight send buffers: fd -> `Vec` (to release on completion) + inflight_send_buffers: Mutex>>, /// Timer heap for scheduled callbacks timers: TimerHeap, /// Task scheduler for Python callbacks @@ -238,6 +232,7 @@ impl UringCore { buffer_pool: pool, fd_states: FDStateManager::new(), inflight_recv_buffers: Mutex::new(HashMap::new()), + inflight_send_buffers: Mutex::new(HashMap::new()), timers: TimerHeap::new(), scheduler: Scheduler::new(), futures: Mutex::new(HashMap::new()), @@ -298,6 +293,14 @@ impl UringCore { self.buffer_pool.release(buf_idx, gen_id); } + // Release any inflight send buffers + let buffers_opt = self.inflight_send_buffers.lock().remove(&fd); + if let Some(buffers) = buffers_opt { + for buf_idx in buffers { + self.buffer_pool.release(buf_idx, gen_id); + } + } + // 2. Return any pending buffers from FD state to the pool if let Some(buffers) = self.fd_states.unregister(fd) { for buf in buffers { @@ -601,28 +604,13 @@ impl UringCore { return Ok(()); } - // Acquire buffer - let buf_idx = self.buffer_pool.acquire().ok_or_else(|| { - PyErr::new::("No buffers available for recvfrom") - })?; - - // Get buffer pointer and size - let buf_ptr = unsafe { - self.buffer_pool - .get_buffer_ptr(buf_idx) - .cast::() - }; - // Buffer size is 64KB which fits in u32 - #[allow(clippy::cast_possible_truncation)] - let buf_len = self.buffer_pool.buffer_size() as u32; + // Allocate buffer + let buf_len = self.buffer_pool.buffer_size(); + let mut data = vec![0u8; buf_len]; self.fd_states .with_state_mut(fd, state::FDState::on_submit) - .map_err(|e| { - self.buffer_pool - .release(buf_idx, self.buffer_pool.generation_id()); - PyErr::new::(e.to_string()) - })?; + .map_err(|e| PyErr::new::(e.to_string()))?; self.futures.lock().insert(fd, future); @@ -630,10 +618,11 @@ impl UringCore { let mut state = Box::new(RecvMsgState { msghdr: unsafe { std::mem::zeroed() }, iovec: libc::iovec { - iov_base: buf_ptr, - iov_len: buf_len as usize, + iov_base: data.as_mut_ptr().cast(), + iov_len: buf_len, }, addr: unsafe { std::mem::zeroed() }, + data, }); // Setup msghdr @@ -648,12 +637,10 @@ impl UringCore { let res = self.ring.lock().prep_recvmsg( fd, std::ptr::addr_of_mut!(state.msghdr), - buf_idx, + 0, // No buffer group generation, ); if let Err(e) = res { - self.buffer_pool - .release(buf_idx, self.buffer_pool.generation_id()); self.futures.lock().remove(&fd); return Err(PyErr::new::(( format!("prep_recvmsg failed: {e}"), @@ -664,15 +651,6 @@ impl UringCore { // Store state to keep it alive self.recvmsg_states.lock().insert(fd, state); - // Track inflight buffer - { - let mut inflight = self.inflight_recv_buffers.lock(); - if let Some(old_buf_idx) = inflight.insert(fd, buf_idx) { - self.buffer_pool - .release(old_buf_idx, self.buffer_pool.generation_id()); - } - } - // Flush self.ring .lock() @@ -947,6 +925,12 @@ impl UringCore { self.futures.lock().insert(fd, future); + // Track inflight buffer for send + { + let mut inflight = self.inflight_send_buffers.lock(); + inflight.entry(fd).or_default().push(buf_idx); + } + // Submit to ring let generation = self.ring.lock().generation_u16(); unsafe { @@ -971,9 +955,60 @@ impl UringCore { Ok(()) } - /// Submit an accept operation for a listening socket. + /// Submit a Zero-Copy Send operation (`IORING_OP_SEND_ZC`). /// - /// Uses `ACCEPT_MULTI` for efficient connection handling. + /// Requires kernel 6.0+. + fn submit_send_zc(&self, fd: i32, data: &[u8], future: Py) -> PyResult<()> { + // Acquire buffer from pool (Required for SendZC to ensure buffer validity) + let buf_idx = self.buffer_pool.acquire().ok_or_else(|| { + PyErr::new::("No buffers available for send_zc") + })?; + + // Copy data to buffer + let buf_slice = unsafe { self.buffer_pool.get_buffer_slice_mut(buf_idx, data.len()) }; + buf_slice.copy_from_slice(data); + + // Get buffer pointer + let buf_ptr = buf_slice.as_ptr(); + #[allow(clippy::cast_possible_truncation)] + let len = data.len() as u32; + + self.futures.lock().insert(fd, future); + + // Track inflight buffer + { + let mut inflight = self.inflight_send_buffers.lock(); + inflight.entry(fd).or_default().push(buf_idx); + } + + // Submit to ring + let generation = self.ring.lock().generation_u16(); + unsafe { + self.ring + .lock() + .prep_send_zc(fd, buf_ptr, len, generation) + .map_err(|e| { + // Release buffer on error (and remove from inflight?) + let mut inflight = self.inflight_send_buffers.lock(); + if let Some(vec) = inflight.get_mut(&fd) { + vec.pop(); + } + self.buffer_pool + .release(buf_idx, self.buffer_pool.generation_id()); + self.futures.lock().remove(&fd); + PyErr::new::(e.to_string()) + })?; + } + + // Flush to kernel + self.ring + .lock() + .submit() + .map_err(|e| PyErr::new::(e.to_string()))?; + + Ok(()) + } + /// Submit an accept operation for a listening socket. /// /// Uses `ACCEPT_MULTI` for efficient connection handling. @@ -1157,94 +1192,91 @@ impl UringCore { // Cleanup SendMsg state self.sendmsg_states.lock().remove(&fd); data_bytes = Some(result.into_pyobject(py)?.into()); - } else if matches!(op_type, OpType::Recv) || matches!(op_type, OpType::RecvMsg) { + } else if matches!(op_type, OpType::Send) || matches!(op_type, OpType::SendZC) { + // Cleanup inflight send buffer + let mut inflight = self.inflight_send_buffers.lock(); + if let Some(vec) = inflight.get_mut(&fd) { + if let Some(buf_idx) = vec.pop() { + self.buffer_pool + .release(buf_idx, self.buffer_pool.generation_id()); + } + } + data_bytes = Some(result.into_pyobject(py)?.into()); + } else if matches!(op_type, OpType::Recv) { let buf_idx_opt = self.inflight_recv_buffers.lock().remove(&fd); if let Some(buf_idx) = buf_idx_opt { - // Extract address if RecvMsg - - // Extract address if RecvMsg + if result > 0 { + let len = result as usize; + unsafe { + let slice = self.buffer_pool.get_buffer_slice(buf_idx, len); + data_bytes = Some(pyo3::types::PyBytes::new(py, slice).into()); + } + } + self.buffer_pool + .release(buf_idx, self.buffer_pool.generation_id()); + } + } else if matches!(op_type, OpType::RecvMsg) { + let state_opt = self.recvmsg_states.lock().remove(&fd); + if let Some(state) = state_opt { let mut addr_tuple: Option> = None; - if matches!(op_type, OpType::RecvMsg) { - let state_opt = self.recvmsg_states.lock().remove(&fd); - if let Some(state) = state_opt { - if result > 0 { - // Parse sockaddr - // Assuming IPv4/IPv6 for now - // TODO: Handle UNIX paths - let addr_ptr = std::ptr::addr_of!(state.addr); - let sa = unsafe { &*addr_ptr.cast::() }; - - if sa.sa_family == libc::AF_INET as libc::sa_family_t { - let sin = unsafe { *addr_ptr.cast::() }; - let ip_u32 = u32::from_be(sin.sin_addr.s_addr); - let ip = std::net::Ipv4Addr::from(ip_u32).to_string(); - let port = u16::from_be(sin.sin_port); - addr_tuple = Some((ip, port).into_pyobject(py)?.into()); - } else if sa.sa_family == libc::AF_INET6 as libc::sa_family_t { - let sin6 = - unsafe { *addr_ptr.cast::() }; - let ip_u128 = u128::from_be_bytes(sin6.sin6_addr.s6_addr); - let ip = std::net::Ipv6Addr::from(ip_u128).to_string(); - let port = u16::from_be(sin6.sin6_port); - // IPv6 tuple: (host, port, flowinfo, scopeid) - addr_tuple = Some( - (ip, port, sin6.sin6_flowinfo, sin6.sin6_scope_id) - .into_pyobject(py)? - .into(), - ); - } else if sa.sa_family == libc::AF_UNIX as libc::sa_family_t { - let sun = unsafe { *addr_ptr.cast::() }; - let path_len = state.msghdr.msg_namelen as usize - - std::mem::offset_of!(libc::sockaddr_un, sun_path); - - if path_len > 0 { - // Handle abstract namespace (starts with null byte) - if sun.sun_path[0] == 0 { - let slice = unsafe { - std::slice::from_raw_parts( - sun.sun_path.as_ptr().cast::(), - path_len, - ) - }; - addr_tuple = Some(PyBytes::new(py, slice).into()); - } else { - // Regular path, null-terminated C string in sun_path - // But msg_namelen includes the path structure - // Let's just create bytes from sun_path up to null or len - let slice = unsafe { - std::ffi::CStr::from_ptr(sun.sun_path.as_ptr()) - }; - let path_str = slice.to_string_lossy().into_owned(); - addr_tuple = - Some(path_str.into_pyobject(py)?.into()); - } - } else { - // Unnamed - addr_tuple = - Some(pyo3::types::PyString::new(py, "").into()); - } + if result > 0 { + // Parse address + let addr_ptr = std::ptr::addr_of!(state.addr); + let sa = unsafe { &*addr_ptr.cast::() }; + + if sa.sa_family == libc::AF_INET as libc::sa_family_t { + let sin = unsafe { *addr_ptr.cast::() }; + let ip_u32 = u32::from_be(sin.sin_addr.s_addr); + let ip = std::net::Ipv4Addr::from(ip_u32).to_string(); + let port = u16::from_be(sin.sin_port); + addr_tuple = Some((ip, port).into_pyobject(py)?.into()); + } else if sa.sa_family == libc::AF_INET6 as libc::sa_family_t { + let sin6 = unsafe { *addr_ptr.cast::() }; + let ip_u128 = u128::from_be_bytes(sin6.sin6_addr.s6_addr); + let ip = std::net::Ipv6Addr::from(ip_u128).to_string(); + let port = u16::from_be(sin6.sin6_port); + addr_tuple = Some( + (ip, port, sin6.sin6_flowinfo, sin6.sin6_scope_id) + .into_pyobject(py)? + .into(), + ); + } else if sa.sa_family == libc::AF_UNIX as libc::sa_family_t { + let sun = unsafe { *addr_ptr.cast::() }; + let path_len = state.msghdr.msg_namelen as usize + - std::mem::offset_of!(libc::sockaddr_un, sun_path); + + if path_len > 0 { + if sun.sun_path[0] == 0 { + let slice = unsafe { + std::slice::from_raw_parts( + sun.sun_path.as_ptr().cast::(), + path_len, + ) + }; + addr_tuple = Some(PyBytes::new(py, slice).into()); + } else { + let slice = unsafe { + std::ffi::CStr::from_ptr(sun.sun_path.as_ptr()) + }; + let path_str = slice.to_string_lossy().into_owned(); + addr_tuple = Some(path_str.into_pyobject(py)?.into()); } + } else { + addr_tuple = Some(pyo3::types::PyString::new(py, "").into()); } } - } - if result > 0 { // Extract data let len = result as usize; - unsafe { - let slice = self.buffer_pool.get_buffer_slice(buf_idx, len); - let bytes = pyo3::types::PyBytes::new(py, slice); - // If we have an address, we return a tuple (bytes, address) as data - if let Some(addr) = addr_tuple { - data_bytes = Some((bytes, addr).into_pyobject(py)?.into()); - } else { - data_bytes = Some(bytes.into()); - } + let slice = &state.data[..len]; + let bytes = pyo3::types::PyBytes::new(py, slice); + if let Some(addr) = addr_tuple { + data_bytes = Some((bytes, addr).into_pyobject(py)?.into()); + } else { + data_bytes = Some(bytes.into()); } } - self.buffer_pool - .release(buf_idx, self.buffer_pool.generation_id()); } } @@ -1319,12 +1351,16 @@ impl UringCore { bytes, future, ) { - e.print(py); + if !e.to_string().contains("InvalidStateError") { + e.print(py); + } } } else { if let Err(e) = future.call_method1(py, "set_result", (bytes,)) { - e.print(py); + if !e.to_string().contains("InvalidStateError") { + e.print(py); + } } } } else { @@ -1338,7 +1374,9 @@ impl UringCore { empty.into(), future, ) { - e.print(py); + if !e.to_string().contains("InvalidStateError") { + e.print(py); + } } } else { if let Err(e) = future.call_method1(py, "set_result", (empty,)) diff --git a/tests/test_send_zc.py b/tests/test_send_zc.py new file mode 100644 index 0000000..9213afe --- /dev/null +++ b/tests/test_send_zc.py @@ -0,0 +1,77 @@ +import pytest +import uringcore +import socket +import os +import asyncio + +def test_submit_send_zc_basic(event_loop): + """Test basic Zero-Copy Send functionality.""" + if not isinstance(event_loop, uringcore.UringEventLoop): + pytest.skip("Test requires uringcore loop") + + async def main(): + try: + # Check kernel version (approximate) + # SEND_ZC requires 6.0+ + uname = os.uname() + release = uname.release + major = int(release.split('.')[0]) + if major < 6: + print(f"Kernel {release} too old for SEND_ZC") + return # Skip logic but pass test + + # Create a TCP pair manually (socketpair is AF_UNIX usually) + server_listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + server_listener.bind(('127.0.0.1', 0)) + server_listener.listen(1) + port = server_listener.getsockname()[1] + + wsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + wsock.setblocking(False) + + # Connect + try: + wsock.connect(('127.0.0.1', port)) + except BlockingIOError: + pass + + rsock, _ = server_listener.accept() + server_listener.close() + + rsock.setblocking(False) + + server_fd = rsock.fileno() + client_fd = wsock.fileno() + + # Register FDs explicitly + event_loop._core.register_fd(server_fd, "tcp") + event_loop._core.register_fd(client_fd, "tcp") + + # Send using Zero-Copy + data = b"Hello Zero-Copy World" + fut = event_loop.create_future() + + try: + # Attempt submission + event_loop._core.submit_send_zc(client_fd, data, fut) + + result = await fut + assert result == len(data) + + # Verify receipt + received = rsock.recv(1024) + assert received == data + + except (RuntimeError, OSError) as e: + # If kernel supports it but detected inability at runtime + if "EINVAL" in str(e) or "EOPNOTSUPP" in str(e) or "Operation not supported" in str(e): + pytest.skip(f"SEND_ZC not supported: {e}") + raise + + finally: + rsock.close() + wsock.close() + except Exception as e: + raise e + + event_loop.run_until_complete(main())