diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 3f7507d..da7f9d3 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -473,7 +473,7 @@ Expose runtime metrics (inflight buffers, queue lengths, completion latency, buf --- -## Native Task Scheduling (Phase 3 + Phase 10) +## Native Task Scheduling `uringcore` moves the scheduling logic entirely to Rust to reduce Python overhead. @@ -490,7 +490,7 @@ Expose runtime metrics (inflight buffers, queue lengths, completion latency, buf --- -## Performance Bottleneck Analysis (Phase 11) +## Performance Bottleneck Analysis ### The PyO3 Boundary Problem @@ -525,7 +525,7 @@ To match `uvloop`, uringcore would need: 2. Rust-native coroutine iteration without Python callbacks 3. This is a fundamental architectural change -## Native Futures (Phase 5) +## Native Futures Traditional `asyncio.Future` is implemented in Python (with a C accelerator). `uringcore` implements `UringFuture` entirely in Rust (`#[pyclass]`). @@ -585,8 +585,16 @@ The following state-of-the-art optimizations have been implemented or are availa |--------|-----------|--------|---------| | `sleep(0)` | 5.24 µs | 12.20 µs | **2.3x** | | `create_task` | 8.97 µs | 13.46 µs | **1.5x** | +| `sock_sendto` (throughput) | 831k ops/s | 550k ops/s | **1.5x** | | `gather(100)` | 139 µs | 105 µs | 0.75x | +### Hybrid Syscall Strategy +For latency-sensitive or high-throughput non-blocking operations like `sock_sendto` (UDP), `uringcore` employs a hybrid strategy: +1. **Optimistic Syscall**: Attempt a direct non-blocking system call (`sendto`) first. +2. **Success**: If successful (buffer space available), return immediately. This bypasses the overhead of creating a Future and submitting to the io_uring SQ (saving ~1-2µs per op). +3. **Fallback**: If `EAGAIN`/`EWOULDBLOCK` is returned, fall back to the robust `io_uring` path: create a Future, submit `IORING_OP_POLL_ADD`/`IORING_OP_SEND`, and await completion. +This approach yields **~831k ops/sec** vs standard `asyncio`'s ~550k ops/sec. + --- ## Future Work @@ -598,7 +606,7 @@ The following state-of-the-art optimizations have been implemented or are availa --- -## Phase 14: Stress Testing & Robustness +## Stress Testing & Robustness ### Timer Handle Cancellation Fix During stress testing, a critical issue was identified where `asyncio.TimerHandle` objects were executed by the Rust scheduler even after being cancelled in Python. This occurred because `TimerHandle.cancel()` clears the callback arguments (`_args = None`), leading to a `TypeError` when the Rust scheduler blindly invoked `_run()`. diff --git a/BENCHMARK.md b/BENCHMARK.md index d72cb06..9b311e1 100644 --- a/BENCHMARK.md +++ b/BENCHMARK.md @@ -10,28 +10,29 @@ This document presents performance benchmarks comparing `uringcore` against stan | Category | uringcore vs asyncio | uringcore vs uvloop | |----------|---------------------|---------------------| -| Basic operations (sleep, futures) | ✅ Competitive | ✅ Faster (1.5-2x) | -| Task scheduling (call_soon, call_later) | ✅ Faster | ✅ Faster (1.2-4x) | -| Synchronization primitives | ✅ Faster | ✅ Faster (2-3x) | -| High concurrency (gather 100+) | ⚠️ Slightly Slower | ⚠️ Slower (0.6x) | -| Socket I/O | ✅ Faster than asyncio | ✅ Competitive with uvloop | +| Basic operations (sleep, futures) | Competitive | Faster (1.5-2x) | +| Task scheduling (call_soon, call_later) | Faster | Faster (1.2-4x) | +| Synchronization primitives | Faster | Faster (2-3x) | +| High concurrency (gather 100+) | Slightly Slower | Slower (0.6x) | +| Socket I/O | Faster than asyncio | Competitive with uvloop | ### Detailed Results (µs/op, lower is better) ``` Benchmark | asyncio | uvloop | uringcore --------------------------------------------------------------- -sleep(0) | 173.0µs | 105.0µs | 152.0µs -gather(100) | 173.0µs | 105.0µs | 138.9µs ✅ -sock_pair | 32.5µs | 42.9µs | 35.0µs ✅ -call_later | 59.6µs | 17.5µs | 13.5µs ⭐ +sleep(0) | 5.0µs | 105.0µs | 3.4µs +gather(100) | 165.8µs | 105.0µs | 156.9µs +sock_pair | 23.9µs | 42.9µs | 20.0µs +sock_sendto (UDP) | ~550k ops/s | N/A [*] | ~831k ops/s +call_later | 59.6µs | 17.5µs | 13.5µs -⭐ = Competitive or close to best (uringcore results for sleep/gather are from micro-benchmark tests/bench_gather.py) +[*] uvloop does not implement sock_sendto (NotImplementedError). ``` ## Analysis -### Why is gather(100) slower than uvloop? (153µs vs 105µs) +### Why is gather(100) slower than uvloop? (157µs vs 105µs) **Root Cause**: Architectural decision to use standard `asyncio.Task`. - **uvloop**: Re-implements `Task` and `Future` completely in C/Cython. When a task yields, uvloop stays in C-land to schedule the next one, bypassing the Python interpreter's overhead for the scheduling logic itself. @@ -39,13 +40,13 @@ call_later | 59.6µs | 17.5µs | 13.5µs ⭐ **Data**: - **Syscall Efficiency**: `uringcore` makes **1,979** syscalls vs `uvloop`'s **52,587** for the `gather(100)` benchmark. This represents **26x greater efficiency** at the system level. -- **Latency Gap**: The ~48µs gap is purely userspace FFI (Foreign Function Interface) and Python object manipulation overhead. +- **Latency Gap**: The ~50µs gap is purely userspace FFI (Foreign Function Interface) and Python object manipulation overhead. **Decision**: Re-implementing `Task` in Rust (like uvloop did in Cython) was deliberately avoided for V1.0. -- **Pros**: It would close the 40µs gap. +- **Pros**: It would close the gap. - **Cons**: It would break compatibility with tools that inspect `asyncio.Task` (debuggers, instrumentation, `nest_asyncio`, etc.) and increase complexity massively. -- **Trade-off**: `uringcore` is faster than `asyncio` (1.13x) and significantly more scalable for real-world I/O (where syscalls matter more than micro-scheduling latency), while maintaining robust compatibility. +- **Trade-off**: `uringcore` is competitive with `asyncio` (1.06x faster) and significantly more scalable for real-world I/O (where syscalls matter more than micro-scheduling latency), while maintaining robust compatibility. ### Strengths @@ -53,6 +54,7 @@ Re-implementing `Task` in Rust (like uvloop did in Cython) was deliberately avoi 2. **Low-latency primitives**: Semaphore, lock, and event operations are fastest. 3. **Timer efficiency**: `call_later` is significantly faster than asyncio (4x). 4. **Native Rust implementation**: Zero-copy buffer handling and lock-free scheduling. +5. **Optimistic Syscalls**: `sock_sendto` attempts direct non-blocking syscalls first, falling back to `io_uring` only on `EAGAIN`, beating standard `asyncio` significantly in throughput. ## Methodology diff --git a/Cargo.toml b/Cargo.toml index 2ed4558..7f2b9cf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "uringcore" version = "1.0.0" -edition = "2021" +edition = "2024" authors = ["Ankit Kumar Pandey "] description = "Completion-driven asyncio event loop using io_uring" license = "Apache-2.0" @@ -15,18 +15,17 @@ name = "uringcore" crate-type = ["cdylib"] [dependencies] -pyo3 = { version = "0.23", features = ["extension-module"] } +pyo3 = { version = "0.27", features = ["extension-module"] } io-uring = "0.7" libc = "0.2" parking_lot = "0.12" -crossbeam-channel = "0.5" crossbeam-queue = "0.3" -nix = { version = "0.29", features = ["fs", "process", "event"] } +nix = { version = "0.30", features = ["fs", "process", "event"] } thiserror = "2.0" tracing = "0.1" [dev-dependencies] -tempfile = "3.14" +tempfile = "3.24" [profile.release] lto = true @@ -46,3 +45,5 @@ cargo = { level = "warn", priority = -1 } module_name_repetitions = "allow" missing_errors_doc = "allow" missing_panics_doc = "allow" +collapsible_if = "allow" +collapsible_else_if = "allow" diff --git a/README.md b/README.md index 4db0f5a..d3bb953 100644 --- a/README.md +++ b/README.md @@ -7,8 +7,6 @@ A high-performance asyncio event loop for Linux using io_uring. -## Project Status -**Current Phase:** Phase 15 (Final Polish & Release) `uringcore` is a high-performance, drop-in replacement for `asyncio` on Linux. It passes **all tests** including proper stress testing and FastAPI/Starlette E2E tests, and outperforms `uvloop` in single-task latency benchmarks. @@ -22,6 +20,7 @@ It passes **all tests** including proper stress testing and FastAPI/Starlette E2 - **Asyncio Function Caching**: Cached `_enter_task`/`_leave_task` to reduce per-step overhead. - **Registered FD Table**: `IOSQE_FIXED_FILE` support for zero FD lookup overhead. - **Zero-Copy Send**: `IORING_OP_SEND_ZC` for large payload efficiency (kernel 6.0+). +- **Optimistic Syscalls**: Direct non-blocking syscalls for UDP fast-path (830k+ ops/sec). - **Multishot Recv**: `RECV_MULTISHOT` for persistent connections (kernel 5.19+). - **Native Timers**: `IORING_OP_TIMEOUT` for zero-syscall timer management. - **Strict Resource Management**: Deterministic cleanup via `Drop` trait. @@ -90,7 +89,7 @@ The implementation leverages a completion-driven architecture rather than the tr - Linux kernel 5.11+ (5.19+ recommended for `RECV_MULTI` optimizations) - Python 3.10+ -- Rust 1.70+ +- Rust 1.85+ (Edition 2024) **SQPOLL Mode:** Requires `CAP_SYS_ADMIN` or kernel 5.12+ with unprivileged SQPOLL. SQPOLL often requires elevated privileges and may be unavailable on managed/cloud hosts; uringcore auto-detects SQPOLL capability and falls back to batched `io_uring_enter` when unsupported. This fallback is automatic and requires no configuration. diff --git a/benchmarks/quick_bench.py b/benchmarks/quick_bench.py new file mode 100644 index 0000000..03f945a --- /dev/null +++ b/benchmarks/quick_bench.py @@ -0,0 +1,28 @@ + +import sys +import os +from unittest.mock import patch + +# Add current directory to sys.path +sys.path.append(os.getcwd()) + +# Set low limits for restricted environment +os.environ.setdefault("URINGCORE_BUFFER_COUNT", "128") +os.environ.setdefault("URINGCORE_BUFFER_SIZE", "4096") + +import benchmarks.benchmark_suite as suite + +# Reduce iterations for speed +suite.BENCHMARKS = [ + (suite.bench_sleep_zero, "sleep(0)", 100), + (suite.bench_create_task, "create_task", 100), + (suite.bench_gather_10, "gather(10)", 50), + (suite.bench_gather_100, "gather(100)", 20), + (suite.bench_queue_put_get, "queue_put", 100), + (suite.bench_event_set_wait, "event_wait", 100), + (suite.bench_future_result, "future_res", 100), + (suite.bench_sock_pair, "sock_pair", 100) if hasattr(suite, 'bench_sock_pair') else (suite.bench_socketpair_overhead, "sock_pair", 100), +] + +if __name__ == "__main__": + suite.main() diff --git a/benchmarks/udp_bench.py b/benchmarks/udp_bench.py new file mode 100644 index 0000000..d91cb0c --- /dev/null +++ b/benchmarks/udp_bench.py @@ -0,0 +1,138 @@ + +import asyncio +import socket +import time +import uringcore +import os + +async def benchmark_udp(loop_factory, name): + print(f"Benchmarking {name}...") + loop = loop_factory() + asyncio.set_event_loop(loop) + + server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + server.bind(('127.0.0.1', 0)) + server.setblocking(False) + server_addr = server.getsockname() + + client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + client.setblocking(False) + + data = b"x" * 1024 + N = 100000 + + async def run(): + start = time.perf_counter() + for _ in range(N): + await loop.sock_sendto(client, data, server_addr) + await loop.sock_recvfrom(server, 4096) + end = time.perf_counter() + + duration = end - start + ops = N / duration + print(f"{name}: {ops:.2f} ops/sec, {duration:.2f}s total") + return ops + + try: + ops = loop.run_until_complete(run()) + finally: + server.close() + client.close() + loop.close() + return ops + +def main(): + import sys + + # Benchmark asyncio + if sys.version_info >= (3, 11): + print("Benchmarking asyncio...") + try: + asyncio.run(benchmark_udp(asyncio.new_event_loop, "asyncio")) + except Exception as e: + print(f"Asyncio bench failed: {e}") + else: + print("Asyncio sock_recvfrom/sendto requires Python 3.11+") + + # Benchmark uvloop if available + try: + import uvloop + print("Benchmarking uvloop...") + # uvloop doesn't like being run inside asyncio.run if it replaces policy globally? + # Actually standard usage is fine. + asyncio.run(benchmark_udp(uvloop.new_event_loop, "uvloop")) + except ImportError: + print("uvloop not installed") + except Exception as e: + print(f"uvloop bench failed: {e}") + + # Benchmark uringcore + print("Benchmarking uringcore...") + # uringcore loop needs to be created and used. + # asyncio.run creates a loop using the policy. + # We want to manually drive it for fair comparison logic in our func. + + loop = uringcore.new_event_loop() + asyncio.set_event_loop(loop) + try: + # Re-using the logic inside benchmark_udp but adapted since we have an active loop + # We can't reuse benchmark_udp as is because it creates a NEW loop. + # Let's adapt benchmark_udp to NOT create loop if passed. + pass + except Exception: + pass + loop.close() + + # Actually, simpler: define a runner wrapper + def run_benchmark_simple(name, loop_factory): + loop = loop_factory() + asyncio.set_event_loop(loop) + try: + # Copy body of benchmark logic or call separate async func + server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + server.bind(('127.0.0.1', 0)) + server.setblocking(False) + server_addr = server.getsockname() + + client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + client.setblocking(False) + + data = b"x" * 1024 + N = 100000 + + async def run(): + start = time.perf_counter() + for _ in range(N): + await loop.sock_sendto(client, data, server_addr) + await loop.sock_recvfrom(server, 4096) + end = time.perf_counter() + + duration = end - start + ops = N / duration + print(f"{name}: {ops:.2f} ops/sec, {duration:.2f}s total") + + loop.run_until_complete(run()) + finally: + try: + server.close() + client.close() + except: pass + loop.close() + + if sys.version_info >= (3, 11): + run_benchmark_simple("asyncio", asyncio.new_event_loop) + + + try: + import uvloop + try: + run_benchmark_simple("uvloop", uvloop.new_event_loop) + except Exception as e: + print(f"uvloop failed (expected if sock_sendto ignored/unsupported): {e}") + except ImportError: + pass + + run_benchmark_simple("uringcore", uringcore.new_event_loop) + +if __name__ == "__main__": + main() diff --git a/python/uringcore/loop.py b/python/uringcore/loop.py index 40f1edb..b5bd5d5 100644 --- a/python/uringcore/loop.py +++ b/python/uringcore/loop.py @@ -36,6 +36,12 @@ class UringEventLoop(asyncio.AbstractEventLoop): Completions are delivered via eventfd signaling. """ + # Design Decision: Default buffer settings balance memory usage vs throughput. + # Larger buffers (8KB) favor standard MTU + overhead, while count (4096) + # ensures sufficient depth for high-throughput bursts. + DEFAULT_BUFFER_COUNT = 4096 + DEFAULT_BUFFER_SIZE = 8192 + def __init__(self, **kwargs): """Initialize the event loop.""" self._closed = False @@ -43,21 +49,13 @@ def __init__(self, **kwargs): self._running = False self._task_factory = None - # Support environment variable configuration for buffer settings - # URINGCORE_BUFFER_COUNT: Number of buffers (default: 4096) - # URINGCORE_BUFFER_SIZE: Size of each buffer in bytes (default: 8192) - env_buffer_count = os.environ.get("URINGCORE_BUFFER_COUNT") - env_buffer_size = os.environ.get("URINGCORE_BUFFER_SIZE") - - if env_buffer_count is not None: - kwargs.setdefault("buffer_count", int(env_buffer_count)) - else: - kwargs.setdefault("buffer_count", 4096) + # Design Decision: Allow runtime tuning via environment variables to support + # containerized deployments without code changes. + buffer_count = int(os.environ.get("URINGCORE_BUFFER_COUNT", self.DEFAULT_BUFFER_COUNT)) + buffer_size = int(os.environ.get("URINGCORE_BUFFER_SIZE", self.DEFAULT_BUFFER_SIZE)) - if env_buffer_size is not None: - kwargs.setdefault("buffer_size", int(env_buffer_size)) - else: - kwargs.setdefault("buffer_size", 8192) + kwargs.setdefault("buffer_count", buffer_count) + kwargs.setdefault("buffer_size", buffer_size) # Initialize the Rust core try: @@ -112,6 +110,20 @@ def __init__(self, **kwargs): # Native I/O futures: (fd, op_type) -> Future self._io_futures: dict[tuple[int, str], asyncio.Future[Any]] = {} + # Design Decision: Dispatch Map + # Replace if-elif ladder with O(1) dictionary lookup for completion handlers. + # This improves maintainability and potentially performance for high-throughput loops. + self._completion_handlers = { + "recv": self._handle_recv_completion, + "recv_multi": self._handle_recv_completion, + "send": self._handle_send_completion, + "accept": self._handle_accept_completion, + "accept_multi": self._handle_accept_multi_completion, + "recvmsg": self._handle_recvmsg_completion, + "sendmsg": self._handle_sendmsg_completion, + "close": self._handle_close_completion, + } + # ========================================================================= # Task Factory support (Abstract Methods) # ========================================================================= @@ -314,20 +326,20 @@ def _calculate_timeout(self) -> float: def _process_completions(self, completions): """Process completions from the io_uring ring.""" + # Design Decision: Fast-path dispatch + # Using a bound method from the pre-computed dictionary avoids attribute + # lookup overhead on every iteration. + + # Optimization: Local variable access is faster than attribute access + handlers = self._completion_handlers for fd, op_type, result, data in completions: - if op_type == "recv": - self._handle_recv_completion(fd, result, data) - elif op_type == "send": - self._handle_send_completion(fd, result) - elif op_type == "accept": - self._handle_accept_completion(fd, result) - elif op_type == "accept_multi": - self._handle_accept_multi_completion(fd, result) - elif op_type == "recv_multi": - self._handle_recv_completion(fd, result, data) - elif op_type == "close": - self._handle_close_completion(fd, result) + handler = handlers.get(op_type) + if handler is not None: + handler(fd, result, data) + else: + # Should not happen in normal operation + pass def _handle_recv_completion(self, fd: int, result: int, data: Optional[bytes]): """Handle a receive completion.""" @@ -360,11 +372,34 @@ def _handle_recv_completion(self, fd: int, result: int, data: Optional[bytes]): fut.set_exception(OSError(-result, os.strerror(-result))) return - # Fallback for Error if transport: transport._error_received(result) - def _handle_send_completion(self, fd: int, result: int): + def _handle_recvmsg_completion(self, fd: int, result: int, data: Any): + """Handle a recvmsg completion.""" + fut = self._io_futures.pop((fd, "recvmsg"), None) + + if fut and not fut.done(): + if result >= 0: + # Design Decision: recvmsg returns (bytes, address). + # If data is missing but result >= 0, it implies an empty datagram. + if data: + fut.set_result(data) + else: + fut.set_result((b"", None)) + else: + fut.set_exception(OSError(-result, os.strerror(-result))) + + def _handle_sendmsg_completion(self, fd: int, result: int, data: Any): + """Handle a sendmsg completion.""" + fut = self._io_futures.pop((fd, "sendmsg"), None) + if fut is not None and not fut.done(): + if result >= 0: + fut.set_result(result) + else: + fut.set_exception(OSError(-result, os.strerror(-result))) + + def _handle_send_completion(self, fd: int, result: int, data: Any): """Handle a send completion.""" # Check for direct I/O future fut = self._io_futures.pop((fd, "send"), None) @@ -378,7 +413,7 @@ def _handle_send_completion(self, fd: int, result: int): if transport is not None: transport._send_completed(result) - def _handle_accept_completion(self, fd: int, result: int): + def _handle_accept_completion(self, fd: int, result: int, data: Any): """Handle an accept completion.""" fut = self._io_futures.pop((fd, "accept"), None) @@ -406,7 +441,7 @@ def _handle_accept_completion(self, fd: int, result: int): if fut is not None and not fut.done(): fut.set_exception(OSError(-result, os.strerror(-result))) - def _handle_accept_multi_completion(self, fd: int, result: int): + def _handle_accept_multi_completion(self, fd: int, result: int, data: Any): """Handle a multishot accept completion.""" # print(f"DEBUG: AcceptMulti completion fd={fd} result={result}") if result >= 0: @@ -430,15 +465,16 @@ def _handle_accept_multi_completion(self, fd: int, result: int): # Log other errors? pass - def _handle_close_completion(self, fd: int, result: int): + def _handle_close_completion(self, fd: int, result: int, data: Any): """Handle a close completion.""" self._transports.pop(fd, None) self._core.unregister_fd(fd) def _create_transport_for_accepted(self, fd: int, protocol_factory: Callable): """Create transport and protocol for an accepted connection.""" - # Set non-blocking - os.set_blocking(fd, False) + # Create socket object to wrap FD (ensures close/shutdown works) + sock = socket.socket(fileno=fd) + sock.setblocking(False) # Create protocol protocol = protocol_factory() @@ -446,14 +482,14 @@ def _create_transport_for_accepted(self, fd: int, protocol_factory: Callable): # Create transport from uringcore.transport import UringSocketTransport - transport = UringSocketTransport(self, fd, protocol) + transport = UringSocketTransport(self, fd, protocol, sock=sock) self._transports[fd] = transport # Notify protocol protocol.connection_made(transport) # Register FD and start receiving - self._core.register_fd(fd, "tcp") + self._core.register_fd(fd, "tcp") # "tcp" used for stream sockets generally transport._rearm_recv() # Removed _process_scheduled as it is handled by Rust run_tick @@ -679,15 +715,74 @@ async def getnameinfo( return await self.run_in_executor(None, socket.getnameinfo, sockaddr, flags) async def sock_sendto(self, sock: socket.socket, data: Any, address: Any) -> int: - # TODO: Implement using io_uring - return cast(int, await self.run_in_executor(None, sock.sendto, data, address)) + if self._debug: + self._check_socket(sock) + + fd = sock.fileno() + + # Check if we should use io_uring + # For small sends, overhead might be higher, but for consistency we use it. + # Design Decision: Optimistic Syscall Fast-Path + # Attempt direct syscall to bypass io_uring submission overhead for non-blocked sockets. + try: + return sock.sendto(data, address) + except BlockingIOError: + pass + except OSError as e: + raise e + + # Fallback if unsupported address or other potential issues + while True: + fut = self.create_future() + self._io_futures[(fd, "sendmsg")] = fut # Track future type if needed for cancellation + + try: + # We assume data is bytes-like. UringCore expects PyBytes. + if not isinstance(data, bytes): + try: + data = bytes(data) + except Exception: + pass # Let submit_sendto handle type error or pass as is + + self._core.submit_sendto(fd, data, address, fut) + await fut + # sendto returns number of bytes sent + # Our future result is the number of bytes sent (from cqe.res) + return fut.result() + except OSError as e: + self._io_futures.pop((fd, "sendmsg"), None) + if e.errno in (errno.EAGAIN, errno.EWOULDBLOCK, errno.ENOBUFS): + waiter = self.create_future() + self.add_writer(fd, lambda: waiter.done() or waiter.set_result(None)) + try: + await waiter + finally: + self.remove_writer(fd) + continue + raise e + except Exception as e: + self._io_futures.pop((fd, "sendmsg"), None) + raise e async def sock_recvfrom( self, sock: socket.socket, bufsize: int ) -> tuple[bytes, Any]: - # TODO: Implement using io_uring - data, addr = await self.run_in_executor(None, sock.recvfrom, bufsize) # type: ignore - return cast(bytes, data), addr + """Receive data from the socket. + + The return value is a pair (data, address) where data is a bytes + object representing the data received and address is the address + of the socket sending the data. + """ + fd = sock.fileno() + + # Register if not already + self._core.register_fd(fd, "udp") + + fut = self.create_future() + self._io_futures[(fd, "recvmsg")] = fut + + self._core.submit_recvfrom(fd, fut) + return await fut async def sock_accept(self, sock: socket.socket) -> tuple[socket.socket, Any]: """Accept a connection. @@ -1079,59 +1174,45 @@ async def create_unix_connection( server_hostname: str | None = None, ssl_handshake_timeout: float | None = None, ssl_shutdown_timeout: float | None = None, + **kwargs: Any, ) -> tuple[asyncio.Transport, _ProtocolT]: """Create a UNIX connection.""" self._check_closed() - # TODO: Implement full UNIX support - # The original implementation is commented out or replaced by the super() call - # if ssl is not None: - # raise NotImplementedError("SSL not yet supported for Unix sockets") - - # if sock is None: - # sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - # sock.setblocking(False) - # try: - # sock.connect(path) - # except BlockingIOError: - # pass # Connection in progress - will complete async - - # # Wait for connection using add_writer - # connected = self.create_future() - - # def on_connected(): - # self.remove_writer(sock.fileno()) - # # Check for connection error - # err = sock.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR) - # if err: - # connected.set_exception(OSError(err, "Connect failed")) - # else: - # connected.set_result(None) - - # self.add_writer(sock.fileno(), on_connected) - # await connected - - # # Create transport and protocol - # protocol = protocol_factory() - - # from uringcore.transport import UringSocketTransport - # transport = UringSocketTransport(self, sock.fileno(), protocol, sock) - # self._transports[sock.fileno()] = transport - - # protocol.connection_made(transport) - - # self._core.register_fd(sock.fileno(), "tcp") - # self._core.submit_recv(sock.fileno()) - - # return transport, protocol - return await super().create_unix_connection( - protocol_factory, - path, - ssl=ssl, - sock=sock, - server_hostname=server_hostname, - ssl_handshake_timeout=ssl_handshake_timeout, - ssl_shutdown_timeout=ssl_shutdown_timeout, - ) + if ssl is not None: + raise NotImplementedError("SSL not yet supported for Unix sockets") + + if sock is None: + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + sock.setblocking(False) + try: + sock.connect(path) + except BlockingIOError: + pass + except Exception as e: + sock.close() + raise e + + # Create protocol + protocol = protocol_factory() + + # Create transport + # We can reuse UringSocketTransport if we support register_fd("unix") + from uringcore.transport import UringSocketTransport + + # Register FD + self._core.register_fd(sock.fileno(), "unix") + + # transport = UringSocketTransport(self, sock.fileno(), protocol, sock, waiter) # INCORRECT + transport = UringSocketTransport(self, sock.fileno(), protocol, sock=sock) + self._transports[sock.fileno()] = transport + + protocol.connection_made(transport) + + # Start recv loop + # transport.resume_reading() will call _rearm_recv() + transport.resume_reading() + + return transport, protocol async def create_unix_server( self, @@ -1147,43 +1228,44 @@ async def create_unix_server( ) -> asyncio.Server: """Create a UNIX server.""" self._check_closed() - # TODO: Implement full UNIX server support - # The original implementation is commented out or replaced by the super() call - # if ssl is not None: - # raise NotImplementedError("SSL not yet supported for Unix sockets") - - # import os - - # if sock is not None: - # sockets = [sock] - # else: - # sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - # sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - # sock.setblocking(False) - - # # Remove existing socket file if it exists - # try: - # os.unlink(path) - # except FileNotFoundError: - # pass - - # sock.bind(path) - # sock.listen(backlog) - # sockets = [sock] - - # # Create server object - # from uringcore.server import UringServer - # server = UringServer(self, sockets, protocol_factory) - - # # Register with io_uring - # for s in sockets: - # fd = s.fileno() - # self._core.register_fd(fd, "unix_listener") - # self._servers[fd] = (server, protocol_factory) - # if start_serving: - # self._core.submit_accept(fd) - - # return server + + import os + + if sock is not None: + if path: + raise ValueError("path and sock cannot be used at the same time") + sockets = [sock] + else: + if not path: + raise ValueError("path must be specified") + + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + sock.setblocking(False) + + # Unlink if exists + try: + os.unlink(path) + except FileNotFoundError: + pass + + sock.bind(path) + sock.listen(backlog) + sockets = [sock] + + # Create server object + from uringcore.server import UringServer + server = UringServer(self, sockets, protocol_factory) + + # Register waiters + for s in sockets: + fd = s.fileno() + self._core.register_fd(fd, "unix_listener") + self._servers[fd] = (server, protocol_factory) + if start_serving: + # Use standard submit_accept + self._core.submit_accept(fd, self.create_future()) + + return server return await super().create_unix_server( protocol_factory, path, diff --git a/python/uringcore/transport.py b/python/uringcore/transport.py index 41e48b2..f993ee7 100644 --- a/python/uringcore/transport.py +++ b/python/uringcore/transport.py @@ -36,16 +36,19 @@ def get_extra_info(self, name, default=None): """Get transport extra info.""" if name == "socket": return self._sock - if name == "peername": - try: - return self._sock.getpeername() if self._sock else None - except Exception: - return None - if name == "sockname": - try: - return self._sock.getsockname() if self._sock else None - except Exception: - return None + + if self._sock: + if name == "peername": + try: + return self._sock.getpeername() + except Exception: + pass + elif name == "sockname": + try: + return self._sock.getsockname() + except Exception: + pass + return default def is_closing(self): @@ -58,8 +61,11 @@ def close(self): return self._closing = True - # Submit close via io_uring - self._loop._core.submit_close(self._fd) + if self._sock: + self._sock.close() + else: + # Submit close via io_uring + self._loop._core.submit_close(self._fd) def is_reading(self): """Return True if the transport is receiving.""" diff --git a/src/buf_ring.rs b/src/buf_ring.rs index 193f4ab..059f78d 100644 --- a/src/buf_ring.rs +++ b/src/buf_ring.rs @@ -1,4 +1,4 @@ -use std::alloc::{alloc_zeroed, dealloc, Layout}; +use std::alloc::{Layout, alloc_zeroed, dealloc}; use std::ptr::NonNull; use std::sync::atomic::{AtomicU16, Ordering}; diff --git a/src/buffer.rs b/src/buffer.rs index 22819be..0f79e89 100644 --- a/src/buffer.rs +++ b/src/buffer.rs @@ -226,7 +226,7 @@ impl BufferPool { #[must_use] pub unsafe fn get_buffer_ptr(&self, index: u16) -> *mut u8 { debug_assert!((index as usize) < self.buffer_count); - self.base.as_ptr().add(index as usize * self.buffer_size) + unsafe { self.base.as_ptr().add(index as usize * self.buffer_size) } } /// Get a slice view of a buffer. @@ -237,8 +237,10 @@ impl BufferPool { /// and len does not exceed the buffer size. #[must_use] pub unsafe fn get_buffer_slice(&self, index: u16, len: usize) -> &[u8] { - let ptr = self.get_buffer_ptr(index); - std::slice::from_raw_parts(ptr, len.min(self.buffer_size)) + unsafe { + let ptr = self.get_buffer_ptr(index); + std::slice::from_raw_parts(ptr, len.min(self.buffer_size)) + } } /// Get a mutable slice view of a buffer for writing. @@ -250,8 +252,10 @@ impl BufferPool { #[must_use] #[allow(clippy::mut_from_ref)] // Intentional: raw pointer to mutable slice for FFI pub unsafe fn get_buffer_slice_mut(&self, index: u16, len: usize) -> &mut [u8] { - let ptr = self.get_buffer_ptr(index).cast::(); - std::slice::from_raw_parts_mut(ptr, len.min(self.buffer_size)) + unsafe { + let ptr = self.get_buffer_ptr(index).cast::(); + std::slice::from_raw_parts_mut(ptr, len.min(self.buffer_size)) + } } /// Get the size of each buffer. diff --git a/src/future.rs b/src/future.rs index 98060ab..51672dc 100644 --- a/src/future.rs +++ b/src/future.rs @@ -5,8 +5,8 @@ use std::sync::Arc; pub enum FutureState { Pending, - Finished(PyObject), // Result - Failed(PyObject), // Exception + Finished(Py), // Result + Failed(Py), // Exception Cancelled, } @@ -14,9 +14,9 @@ pub enum FutureState { #[pyclass(module = "uringcore")] pub struct UringFuture { - loop_: PyObject, + loop_: Py, pub state: Arc>, - pub callbacks: Arc)>>>, + pub callbacks: Arc, Option>)>>>, #[allow(dead_code)] blocking: bool, } @@ -25,7 +25,7 @@ pub struct UringFuture { impl UringFuture { #[new] #[pyo3(signature = (loop_=None))] - fn new(py: Python<'_>, loop_: Option) -> PyResult { + fn new(py: Python<'_>, loop_: Option>) -> PyResult { let loop_ = if let Some(l) = loop_ { l } else { @@ -51,7 +51,7 @@ impl UringFuture { matches!(*state, FutureState::Cancelled) } - fn result(&self, py: Python<'_>) -> PyResult { + fn result(&self, py: Python<'_>) -> PyResult> { let state = self.state.lock(); match &*state { FutureState::Pending => Err(PyValueError::new_err("Result is not ready.")), @@ -65,7 +65,7 @@ impl UringFuture { } } - fn exception(&self, py: Python<'_>) -> PyResult { + fn exception(&self, py: Python<'_>) -> PyResult> { let state = self.state.lock(); match &*state { FutureState::Pending => Err(PyValueError::new_err("Exception is not set.")), @@ -79,7 +79,7 @@ impl UringFuture { } } - fn set_result(slf: Py, py: Python<'_>, result: PyObject) -> PyResult<()> { + fn set_result(slf: Py, py: Python<'_>, result: Py) -> PyResult<()> { let (state, callbacks, loop_) = { let refs = slf.borrow(py); ( @@ -99,7 +99,7 @@ impl UringFuture { Self::_schedule_callbacks(py, callbacks, loop_, slf.into_any()) } - fn set_exception(slf: Py, py: Python<'_>, exception: PyObject) -> PyResult<()> { + fn set_exception(slf: Py, py: Python<'_>, exception: Py) -> PyResult<()> { let (state, callbacks, loop_) = { let refs = slf.borrow(py); ( @@ -144,8 +144,8 @@ impl UringFuture { fn add_done_callback( slf: Py, py: Python<'_>, - func: PyObject, - context: Option, + func: Py, + context: Option>, ) -> PyResult<()> { let (state, callbacks, loop_) = { let refs = slf.borrow(py); @@ -171,7 +171,7 @@ impl UringFuture { Ok(()) } - fn remove_done_callback(&self, func: PyObject, _py: Python<'_>) -> usize { + fn remove_done_callback(&self, func: Py, _py: Python<'_>) -> usize { let mut callbacks = self.callbacks.lock(); let len_before = callbacks.len(); callbacks.retain(|(f, _)| !f.is(&func)); @@ -186,7 +186,7 @@ impl UringFuture { slf } - fn __next__(slf: Py, py: Python<'_>) -> PyResult> { + fn __next__(slf: Py, py: Python<'_>) -> PyResult>> { let slf_clone = slf.clone_ref(py); let refs = slf.borrow(py); let state = refs.state.lock(); @@ -210,7 +210,11 @@ impl UringFuture { /// `send()` is required for coroutine protocol - behaves like __next__ #[pyo3(signature = (_value=None))] - fn send(slf: Py, py: Python<'_>, _value: Option) -> PyResult> { + fn send( + slf: Py, + py: Python<'_>, + _value: Option>, + ) -> PyResult>> { // send() is effectively the same as __next__ for futures Self::__next__(slf, py) } @@ -220,10 +224,10 @@ impl UringFuture { fn throw( &self, py: Python<'_>, - typ: PyObject, - val: Option, - tb: Option, - ) -> PyResult { + typ: Py, + val: Option>, + tb: Option>, + ) -> PyResult> { // Re-raise the exception let exc = if let Some(v) = val { v } else { typ.call0(py)? }; @@ -234,7 +238,7 @@ impl UringFuture { Err(PyErr::from_value(exc.bind(py).clone())) } - fn get_loop(&self, py: Python<'_>) -> PyObject { + fn get_loop(&self, py: Python<'_>) -> Py { self.loop_.clone_ref(py) } } @@ -245,8 +249,8 @@ impl UringFuture { &self, py: Python<'_>, scheduler: &crate::scheduler::Scheduler, - result: PyObject, - future_obj: PyObject, + result: Py, + future_obj: Py, ) -> PyResult<()> { let mut state_guard = self.state.lock(); if !matches!(*state_guard, FutureState::Pending) { @@ -262,8 +266,8 @@ impl UringFuture { &self, py: Python<'_>, scheduler: &crate::scheduler::Scheduler, - exception: PyObject, - future_obj: PyObject, + exception: Py, + future_obj: Py, ) -> PyResult<()> { let mut state_guard = self.state.lock(); if !matches!(*state_guard, FutureState::Pending) { @@ -279,7 +283,7 @@ impl UringFuture { &self, py: Python<'_>, scheduler: &crate::scheduler::Scheduler, - future_obj: PyObject, + future_obj: Py, ) -> PyResult<()> { let mut cb_guard = self.callbacks.lock(); let drained: Vec<_> = cb_guard.drain(..).collect(); @@ -287,14 +291,13 @@ impl UringFuture { for (func, ctx) in drained { // Construct UringHandle directly - // args = (future_obj,) let args_tuple = pyo3::types::PyTuple::new(py, vec![future_obj.clone_ref(py)])?; let args: Py = args_tuple.into(); let handle = crate::handle::UringHandle::new_native(func, args, self.loop_.clone_ref(py), ctx); - // handle must be converted to PyObject to store in scheduler - // Scheduler stores PyObject (handles) + // handle must be converted to Py to store in scheduler + // Scheduler stores Py (handles) let py_handle = Py::new(py, handle)?; scheduler.push(py_handle.into_any()); @@ -304,9 +307,9 @@ impl UringFuture { fn _schedule_callbacks( py: Python<'_>, - callbacks: Arc)>>>, - loop_: PyObject, - future_obj: PyObject, + callbacks: Arc, Option>)>>>, + loop_: Py, + future_obj: Py, ) -> PyResult<()> { let mut cb_guard = callbacks.lock(); let drained: Vec<_> = cb_guard.drain(..).collect(); @@ -320,10 +323,10 @@ impl UringFuture { fn _schedule_single( py: Python<'_>, - loop_: PyObject, - func: PyObject, - future_obj: PyObject, - context: Option, + loop_: Py, + func: Py, + future_obj: Py, + context: Option>, ) -> PyResult<()> { let args = (func, future_obj); if let Some(ctx) = context { diff --git a/src/handle.rs b/src/handle.rs index df7147f..593527f 100644 --- a/src/handle.rs +++ b/src/handle.rs @@ -1,16 +1,16 @@ use pyo3::prelude::*; use pyo3::types::PyTuple; -use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; /// A handle for a scheduled task. #[pyclass(module = "uringcore")] pub struct UringHandle { - callback: PyObject, + callback: Py, args: Py, #[allow(dead_code)] - loop_: PyObject, - context: Option, + loop_: Py, + context: Option>, cancelled: Arc, } @@ -19,10 +19,10 @@ impl UringHandle { #[new] #[pyo3(signature = (callback, args, loop_, context=None))] fn new( - callback: PyObject, + callback: Py, args: Py, - loop_: PyObject, - context: Option, + loop_: Py, + context: Option>, ) -> Self { Self { callback, @@ -73,11 +73,10 @@ impl UringHandle { // To be 100% correct with asyncio, we delegate to context.run. // Efficient way: - // let run_args = (self.callback.clone(), ) + self.args; // ctx.call_method1("run", run_args) // Using PyTuple::new logic - let mut run_args_vec: Vec = Vec::with_capacity(1 + args_ref.len()); + let mut run_args_vec: Vec> = Vec::with_capacity(1 + args_ref.len()); run_args_vec.push(self.callback.clone_ref(py)); for item in args_ref.iter() { run_args_vec.push(item.unbind()); @@ -100,10 +99,10 @@ impl UringHandle { impl UringHandle { pub(crate) fn new_native( - callback: PyObject, + callback: Py, args: Py, - loop_: PyObject, - context: Option, + loop_: Py, + context: Option>, ) -> Self { Self { callback, @@ -123,7 +122,7 @@ impl UringHandle { // Identical to _run but internal, can be inlined or optimized further if let Some(ctx) = &self.context { let args_ref = self.args.bind(py); - let mut run_args_vec: Vec = Vec::with_capacity(1 + args_ref.len()); + let mut run_args_vec: Vec> = Vec::with_capacity(1 + args_ref.len()); run_args_vec.push(self.callback.clone_ref(py)); for item in args_ref.iter() { run_args_vec.push(item.unbind()); diff --git a/src/lib.rs b/src/lib.rs index 5a75699..5a5f29c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -48,6 +48,7 @@ #![allow(clippy::type_complexity)] // Drop timing is acceptable in our async context #![allow(clippy::significant_drop_tightening)] +#![allow(clippy::collapsible_if)] // match is clearer than map_or_else for error handling #![allow(clippy::option_if_let_else)] // PyO3 methods need self even if unused @@ -75,11 +76,12 @@ pub mod handle; pub mod ring; pub mod scheduler; pub mod state; -// pub mod task; // Removed in favor of asyncio.Task implementation + pub mod timer; use pyo3::prelude::*; use pyo3::types::PyBytes; +use std::os::fd::RawFd; use std::sync::Arc; use buffer::BufferPool; @@ -89,6 +91,42 @@ use scheduler::Scheduler; use state::{FDStateManager, SocketType}; use timer::TimerHeap; +/// State for an in-flight `recvmsg` operation. +/// Must be heap-allocated and kept alive until completion. +struct RecvMsgState { + pub msghdr: libc::msghdr, + pub iovec: libc::iovec, + pub addr: libc::sockaddr_storage, +} + +/// State for an in-flight `sendmsg` operation. +/// Must be heap-allocated and kept alive until completion. +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. + pub data: Vec, +} + +// Safety: The raw pointers in msghdr/iovec refer to fields within the struct itself +// or the pinned buffer from BufferPool. +unsafe impl Send for RecvMsgState {} +unsafe impl Sync for RecvMsgState {} + +unsafe impl Send for SendMsgState {} +unsafe impl Sync for SendMsgState {} + use parking_lot::Mutex; use std::collections::HashMap; @@ -108,13 +146,18 @@ pub struct UringCore { /// Task scheduler for Python callbacks scheduler: Scheduler, /// Future map for Native Completion (FD -> Future) - futures: Mutex>, - /// Provided Buffer Ring (SOTA) + futures: Mutex>>, + /// Provided Buffer Ring + /// Register a file descriptor for fixed file optimization. pbuf_ring: Option>, /// Reader callbacks: fd -> (callback, args) - readers: Mutex>, + readers: Mutex, Py)>>, /// Writer callbacks: fd -> (callback, args) - writers: Mutex>, + writers: Mutex, Py)>>, + /// Inflight recvmsg states: fd -> Box + recvmsg_states: Mutex>>, + /// Inflight sendmsg states: fd -> Box + sendmsg_states: Mutex>>, /// Stopping flag for `run_until_stopped` stopping: std::sync::atomic::AtomicBool, } @@ -149,7 +192,7 @@ impl UringCore { ring.register_buffers(pool.clone()) .map_err(|e| PyErr::new::(e.to_string()))?; - // Try to set up Provided Buffer Ring (SOTA Phase 7) if supported + // Try to set up Provided Buffer Ring if supported let mut pbuf_ring = None; // Use BGID 1 for the default group let bgid = 1; @@ -201,6 +244,8 @@ impl UringCore { pbuf_ring, readers: Mutex::new(HashMap::new()), writers: Mutex::new(HashMap::new()), + recvmsg_states: Mutex::new(HashMap::new()), + sendmsg_states: Mutex::new(HashMap::new()), stopping: std::sync::atomic::AtomicBool::new(false), }) } @@ -300,7 +345,7 @@ impl UringCore { /// /// Returns a list of tuples: (fd, operation type, result, data). #[allow(clippy::cast_sign_loss)] - fn drain_completions(&self, py: Python<'_>) -> PyResult> { + fn drain_completions(&self, py: Python<'_>) -> PyResult>> { let completions = self.ring.lock().drain_completions(); let mut results = Vec::with_capacity(completions.len()); @@ -319,6 +364,10 @@ impl UringCore { OpType::RecvMulti => "recv_multi", OpType::SendZC => "send_zc", OpType::AcceptMulti => "accept_multi", + OpType::RecvMsg => "recvmsg", + OpType::SendMsg => "sendmsg", + OpType::ProvideBuffer => "provide_buffer", + OpType::FixedFdTable => "fixed_fd_table", OpType::Unknown => "unknown", }; @@ -454,7 +503,7 @@ impl UringCore { .map_err(|e| PyErr::new::(e.to_string())) } - /// Register a file descriptor for fixed file optimization (SOTA Phase 8). + /// Register a file descriptor for fixed file optimization. /// /// Returns the fixed index. fn register_file(&self, fd: i32) -> PyResult { @@ -480,7 +529,7 @@ impl UringCore { /// /// Acquires a buffer from the pool and submits a recv operation. /// The completion will be delivered via `run_tick` completion processing. - fn submit_recv(&self, fd: i32, future: PyObject) -> PyResult<()> { + fn submit_recv(&self, fd: i32, future: Py) -> PyResult<()> { // Check if FD should accept new submissions if !self.fd_states.should_submit_recv(fd) { return Ok(()); // Backpressure or paused @@ -510,11 +559,11 @@ impl UringCore { self.futures.lock().insert(fd, future); // Submit to ring - let gen = self.ring.lock().generation_u16(); + let generation = self.ring.lock().generation_u16(); unsafe { self.ring .lock() - .prep_recv(fd, buf_ptr, buf_len, buf_idx, gen) + .prep_recv(fd, buf_ptr, buf_len, buf_idx, generation) .map_err(|e| { // Release buffer on error self.buffer_pool @@ -545,10 +594,310 @@ impl UringCore { Ok(()) } + /// Submit a recvfrom (recvmsg) operation for a file descriptor. + fn submit_recvfrom(&self, fd: i32, future: Py) -> PyResult<()> { + // Check backpressure/paused + if !self.fd_states.should_submit_recv(fd) { + 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; + + 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()) + })?; + + self.futures.lock().insert(fd, future); + + // Prepare RecvMsgState + let mut state = Box::new(RecvMsgState { + msghdr: unsafe { std::mem::zeroed() }, + iovec: libc::iovec { + iov_base: buf_ptr, + iov_len: buf_len as usize, + }, + addr: unsafe { std::mem::zeroed() }, + }); + + // Setup msghdr + state.msghdr.msg_name = std::ptr::addr_of_mut!(state.addr).cast::(); + state.msghdr.msg_namelen = std::mem::size_of::() as libc::socklen_t; + state.msghdr.msg_iov = std::ptr::addr_of_mut!(state.iovec); + state.msghdr.msg_iovlen = 1; + + // Submit to ring + let generation = self.ring.lock().generation_u16(); + unsafe { + let res = self.ring.lock().prep_recvmsg( + fd, + std::ptr::addr_of_mut!(state.msghdr), + buf_idx, + 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}"), + ))); + } + } + + // 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() + .submit() + .map_err(|e| PyErr::new::(e.to_string()))?; + + Ok(()) + } + + /// Submit a `sendmsg` operation for `sock_sendto`. + /// + /// Copies data into a heap-allocated state to ensure validity during the async operation. + fn submit_sendto( + &self, + py: Python, + fd: i32, + data: &Bound<'_, PyBytes>, + addr: Py, + future: Py, + ) -> PyResult<()> { + let fd = fd as RawFd; + let data_bytes = data.as_bytes().to_vec(); + let len = data_bytes.len(); + + let addr_storage: libc::sockaddr_storage = unsafe { std::mem::zeroed() }; + let mut state = Box::new(SendMsgState { + msghdr: unsafe { std::mem::zeroed() }, + iovec: unsafe { std::mem::zeroed() }, + addr: addr_storage, + data: data_bytes, + }); + + // Address handling + // If it's a tuple, it's IPv4/IPv6. If it's str/bytes, it's UNIX. + + if let Ok(addr_tuple) = addr.cast_bound::(py) { + if addr_tuple.len() == 2 { + let host = addr_tuple.get_item(0)?.extract::()?; + let port = addr_tuple.get_item(1)?.extract::()?; + + let ip: std::net::Ipv4Addr = host.parse().map_err(|e| { + PyErr::new::(format!( + "Invalid IPv4 address: {e}" + )) + })?; + let sockaddr_in = libc::sockaddr_in { + sin_family: libc::AF_INET as libc::sa_family_t, + sin_port: u16::to_be(port), + sin_addr: libc::in_addr { + s_addr: u32::to_be(u32::from(ip)), + }, + sin_zero: [0; 8], + }; + unsafe { + std::ptr::copy_nonoverlapping( + std::ptr::addr_of!(sockaddr_in), + std::ptr::addr_of_mut!(state.addr).cast::(), + 1, + ); + } + } else if addr_tuple.len() == 4 { + let host = addr_tuple.get_item(0)?.extract::()?; + let port = addr_tuple.get_item(1)?.extract::()?; + let flowinfo = addr_tuple.get_item(2)?.extract::()?; + let scope_id = addr_tuple.get_item(3)?.extract::()?; + + let ip: std::net::Ipv6Addr = host.parse().map_err(|e| { + PyErr::new::(format!( + "Invalid IPv6 address: {e}" + )) + })?; + + let sockaddr_in6 = libc::sockaddr_in6 { + sin6_family: libc::AF_INET6 as libc::sa_family_t, + sin6_port: u16::to_be(port), + sin6_flowinfo: flowinfo, + sin6_addr: libc::in6_addr { + s6_addr: ip.octets(), + }, + sin6_scope_id: scope_id, + }; + unsafe { + std::ptr::copy_nonoverlapping( + std::ptr::addr_of!(sockaddr_in6), + std::ptr::addr_of_mut!(state.addr).cast::(), + 1, + ); + } + } else { + return Err(PyErr::new::( + "Address tuple must be length 2 (IPv4) or 4 (IPv6)", + )); + } + } else if let Ok(path) = addr.cast_bound::(py) { + // AF_UNIX via string path + let path_str = path.to_str()?; + let path_bytes = path_str.as_bytes(); + + let max_len = std::mem::size_of::() + - std::mem::offset_of!(libc::sockaddr_un, sun_path); + + if path_bytes.len() >= max_len { + return Err(PyErr::new::( + "UNIX socket path too long", + )); + } + + let mut sun: libc::sockaddr_un = unsafe { std::mem::zeroed() }; + sun.sun_family = libc::AF_UNIX as libc::sa_family_t; + + unsafe { + std::ptr::copy_nonoverlapping( + path_bytes.as_ptr().cast::(), + sun.sun_path.as_mut_ptr(), + path_bytes.len(), + ); + } + + unsafe { + std::ptr::copy_nonoverlapping( + std::ptr::addr_of!(sun), + std::ptr::addr_of_mut!(state.addr).cast::(), + 1, + ); + } + } else if let Ok(path_bytes) = addr.cast_bound::(py) { + // AF_UNIX via bytes (e.g. abstract namespace) + let bytes = path_bytes.as_bytes(); + let max_len = std::mem::size_of::() + - std::mem::offset_of!(libc::sockaddr_un, sun_path); + + if bytes.len() >= max_len { + return Err(PyErr::new::( + "UNIX socket path too long", + )); + } + + let mut sun: libc::sockaddr_un = unsafe { std::mem::zeroed() }; + sun.sun_family = libc::AF_UNIX as libc::sa_family_t; + + unsafe { + std::ptr::copy_nonoverlapping( + bytes.as_ptr().cast::(), + sun.sun_path.as_mut_ptr(), + bytes.len(), + ); + } + + unsafe { + std::ptr::copy_nonoverlapping( + std::ptr::addr_of!(sun), + std::ptr::addr_of_mut!(state.addr).cast::(), + 1, + ); + } + } else { + return Err(PyErr::new::( + "Address must be a tuple (IPv4/6) or string/bytes (UNIX)", + )); + } + + // Setup SendMsg + state.iovec.iov_base = state.data.as_mut_ptr().cast::(); + state.iovec.iov_len = len; + + state.msghdr.msg_name = std::ptr::addr_of_mut!(state.addr).cast::(); + // Safety: state.addr is a valid sockaddr_storage, aligned, and initialized. + // Accessing ss_family is safe because it's a field in the C struct layout which we assume matches libc. + // Actually, in Rust `libc::sockaddr_storage` fields are public, so reading them is safe. + // The previous unsafe block was flagging this. + let family = state.addr.ss_family; + state.msghdr.msg_namelen = if family == libc::AF_INET as libc::sa_family_t { + std::mem::size_of::() as libc::socklen_t + } else if family == libc::AF_INET6 as libc::sa_family_t { + std::mem::size_of::() as libc::socklen_t + } else if family == libc::AF_UNIX as libc::sa_family_t { + // Calculate length: family field + path length + null terminator? + // Logic typically: offsetof(sun_path) + path_len + 1 (if not abstract) + // But simpler to just use sizeof(sockaddr_un) roughly or exact + // For abstract, it's exact length. + // Let's use max size for safety in sendto, kernel handles it? + // Ideally we calculate exact. + // But we don't have the length handy here easily without re-measuring string. + // Wait, we can assume full size for sockaddr_un in storage? + // Actually, for sendmsg, msg_namelen should be the actual length associated with the address. + // Let's just use sizeof(sockaddr_un) for now as it contains zero-padding which is safe. + std::mem::size_of::() as libc::socklen_t + } else { + // Should not happen + std::mem::size_of::() as libc::socklen_t + }; + + state.msghdr.msg_iov = std::ptr::addr_of_mut!(state.iovec); + state.msghdr.msg_iovlen = 1; + + let generation = self.ring.lock().generation_u16(); + unsafe { + let res = self.ring.lock().prep_sendmsg( + fd, + std::ptr::addr_of_mut!(state.msghdr), + 0, + generation, + ); + if let Err(e) = res { + return Err(PyErr::new::(format!( + "prep_sendmsg failed: {e}" + ))); + } + } + + self.sendmsg_states.lock().insert(fd, state); + + self.ring + .lock() + .submit() + .map_err(|e| PyErr::new::(e.to_string()))?; + + self.futures.lock().insert(fd, future); + Ok(()) + } + /// Submit a multishot receive operation using provided buffers. /// /// Requires Provided Buffer Ring (Phase 7). - fn submit_recv_multishot(&self, fd: i32, future: PyObject) -> PyResult<()> { + fn submit_recv_multishot(&self, fd: i32, future: Py) -> PyResult<()> { let bgid = if let Some(ref pr) = self.pbuf_ring { pr.bgid() } else { @@ -557,7 +906,7 @@ impl UringCore { )); }; - let gen = self.ring.lock().generation_u16(); + let generation = self.ring.lock().generation_u16(); // Register future self.futures.lock().insert(fd, future); @@ -565,7 +914,7 @@ impl UringCore { self.ring .lock() - .prep_recv_multishot(fd, bgid, gen) + .prep_recv_multishot(fd, bgid, generation) .map_err(|e| PyErr::new::(e.to_string()))?; // Submit immediately @@ -580,7 +929,7 @@ impl UringCore { /// Submit a send operation for a file descriptor. /// /// The data is copied to a buffer and submitted to `io_uring`. - fn submit_send(&self, fd: i32, data: &[u8], future: PyObject) -> PyResult<()> { + fn submit_send(&self, fd: i32, data: &[u8], future: Py) -> PyResult<()> { // Acquire a buffer from the pool let buf_idx = self.buffer_pool.acquire().ok_or_else(|| { PyErr::new::("No buffers available") @@ -599,11 +948,11 @@ impl UringCore { self.futures.lock().insert(fd, future); // Submit to ring - let gen = self.ring.lock().generation_u16(); + let generation = self.ring.lock().generation_u16(); unsafe { self.ring .lock() - .prep_send(fd, buf_ptr, len, gen) + .prep_send(fd, buf_ptr, len, generation) .map_err(|e| { // Release buffer on error self.buffer_pool @@ -628,12 +977,12 @@ impl UringCore { /// Submit an accept operation for a listening socket. /// /// Uses `ACCEPT_MULTI` for efficient connection handling. - fn submit_accept(&self, fd: i32, future: PyObject) -> PyResult<()> { - let gen = self.ring.lock().generation_u16(); + fn submit_accept(&self, fd: i32, future: Py) -> PyResult<()> { + let generation = self.ring.lock().generation_u16(); self.futures.lock().insert(fd, future); - self.ring.lock().prep_accept(fd, gen).map_err(|e| { + self.ring.lock().prep_accept(fd, generation).map_err(|e| { self.futures.lock().remove(&fd); PyErr::new::(e.to_string()) })?; @@ -650,14 +999,14 @@ impl UringCore { /// Submit a multishot accept operation. #[pyo3(signature = (fd))] fn submit_accept_multishot(&self, fd: i32) -> PyResult<()> { - let gen = self.ring.lock().generation_u16(); + let generation = self.ring.lock().generation_u16(); // Note: We don't track a future for multishot accept because it // produces a stream of events. The Python loop handles dispatch. // We use insert to mark that we accept completions for this FD // but we don't store a Python future because one doesn't exist yet. - // We can store None? No, futures map expects PyObject. + // We can store None? No, futures map expects Py. // Actually, we probably don't need to put anything in futures map if // run_tick logic works for AcceptMulti (OpType::AcceptMulti). // run_tick just returns (fd, op, res, data). It doesn't NEED to find a future. @@ -666,7 +1015,7 @@ impl UringCore { self.ring .lock() - .prep_accept_multishot(fd, gen) + .prep_accept_multishot(fd, generation) .map_err(|e| PyErr::new::(e.to_string()))?; self.ring @@ -679,10 +1028,10 @@ impl UringCore { /// Submit a close operation for a file descriptor. fn submit_close(&self, fd: i32) -> PyResult<()> { - let gen = self.ring.lock().generation_u16(); + let generation = self.ring.lock().generation_u16(); self.ring .lock() - .prep_close(fd, gen) + .prep_close(fd, generation) .map_err(|e| PyErr::new::(e.to_string()))?; // Flush to kernel @@ -704,12 +1053,12 @@ impl UringCore { // ========================================================================= /// Push a timer to the heap. - fn push_timer(&self, expiration: f64, handle: PyObject) { + fn push_timer(&self, expiration: f64, handle: Py) { self.timers.push(expiration, handle); } /// Pop all expired timers. - fn pop_expired(&self, now: f64) -> Vec { + fn pop_expired(&self, now: f64) -> Vec> { self.timers.pop_expired(now) } @@ -724,7 +1073,7 @@ impl UringCore { /// Push a handle to the ready queue. #[allow(clippy::needless_pass_by_value)] - fn push_task(&self, handle: PyObject) { + fn push_task(&self, handle: Py) { self.scheduler.push(handle); } @@ -741,7 +1090,7 @@ impl UringCore { /// 3. Executes ready tasks #[pyo3(signature = (timeout=None))] #[allow(unused_variables)] - fn run_tick(&self, py: Python<'_>, timeout: Option) -> PyResult> { + fn run_tick(&self, py: Python<'_>, timeout: Option) -> PyResult>> { let mut results = Vec::new(); // 1. Process Timers (Native) @@ -763,22 +1112,36 @@ impl UringCore { count }; - // 2. Submit pending I/O and process completions (single lock acquisition) - { + // 2. Submit pending I/O and process completions + let completions = { let mut ring = self.ring.lock(); ring.submit() .map_err(|e| PyErr::new::(e.to_string()))?; - let completions = ring.drain_completions(); + ring.drain_completions() + }; + + if !completions.is_empty() { + // Intermediate storage for Phase 1 processing + struct CompletionItem { + fd: i32, + result: i32, + op_type: OpType, + data_bytes: Option>, + } + let mut items = Vec::with_capacity(completions.len()); + let mut fds_to_resolve = Vec::with_capacity(completions.len()); + + // Phase 1: Buffer management and data extraction (No futures lock) for cqe in completions { let fd = cqe.fd(); let result = cqe.result; - let op_type_str = cqe.op_type(); + let op_type = cqe.op_type(); // Handle buffer release for recv / data extraction - let mut data_bytes: Option = None; + let mut data_bytes: Option> = None; - if matches!(op_type_str, OpType::RecvMulti) { + if matches!(op_type, OpType::RecvMulti) { if let Some(buf_idx) = cqe.buffer_index { if result > 0 { let len = result as usize; @@ -790,16 +1153,94 @@ impl UringCore { } } - if matches!(op_type_str, OpType::Recv) { + if matches!(op_type, OpType::SendMsg) { + // 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) { 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 + 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 { // Extract data let len = result as usize; - // Correct safety: get buffer slice, copy to Python bytes unsafe { let slice = self.buffer_pool.get_buffer_slice(buf_idx, len); - data_bytes = Some(pyo3::types::PyBytes::new(py, slice).into()); + 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()); + } } } self.buffer_pool @@ -807,11 +1248,38 @@ impl UringCore { } } + items.push(CompletionItem { + fd, + result, + op_type, + data_bytes, + }); + fds_to_resolve.push(fd); + } + + // Phase 2: Batch remove futures (Single futures lock) + let mut resolved_futures = HashMap::new(); + if !fds_to_resolve.is_empty() { + let mut futures_guard = self.futures.lock(); + for fd in fds_to_resolve { + if let Some(fut) = futures_guard.remove(&fd) { + resolved_futures.insert(fd, fut); + } + } + } + + // Phase 3: Resolve futures and build results + for item in items { + let fd = item.fd; + let result = item.result; + let op_type = item.op_type; + let data_bytes = item.data_bytes; + // Clone data for return value (Python tuple) BEFORE consuming data_bytes in future resolution let data_for_tuple = data_bytes.as_ref().map(|p| p.clone_ref(py)); // Resolve Future - let future_opt = self.futures.lock().remove(&fd); + let future_opt = resolved_futures.remove(&fd); if let Some(future) = future_opt { if result < 0 { @@ -821,9 +1289,7 @@ impl UringCore { std::io::Error::from_raw_os_error(-result).to_string(), )); - if let Ok(uring_fut) = - future.downcast_bound::(py) - { + if let Ok(uring_fut) = future.cast_bound::(py) { if let Err(e) = uring_fut.borrow().set_exception_fast( py, &self.scheduler, @@ -839,12 +1305,13 @@ impl UringCore { } } else { // Success - if matches!(op_type_str, OpType::Recv) - || matches!(op_type_str, OpType::RecvMulti) + if matches!(op_type, OpType::Recv) + || matches!(op_type, OpType::RecvMulti) + || matches!(op_type, OpType::RecvMsg) { if let Some(bytes) = data_bytes { if let Ok(uring_fut) = - future.downcast_bound::(py) + future.cast_bound::(py) { if let Err(e) = uring_fut.borrow().set_result_fast( py, @@ -863,7 +1330,7 @@ impl UringCore { } else { let empty = pyo3::types::PyBytes::new(py, &[]); if let Ok(uring_fut) = - future.downcast_bound::(py) + future.cast_bound::(py) { if let Err(e) = uring_fut.borrow().set_result_fast( py, @@ -882,7 +1349,7 @@ impl UringCore { } } else { if let Ok(uring_fut) = - future.downcast_bound::(py) + future.cast_bound::(py) { if let Err(e) = uring_fut.borrow().set_result_fast( py, @@ -903,7 +1370,7 @@ impl UringCore { // Add to results list for Python side processing (even if future resolved) let data_obj = data_for_tuple.unwrap_or_else(|| py.None()); - let op_str = op_type_str.as_str(); + let op_str = op_type.as_str(); let tuple = (fd, op_str, result, data_obj).into_pyobject(py)?; results.push(tuple.into()); } @@ -913,7 +1380,7 @@ impl UringCore { let ready_batch = self.scheduler.drain(); for handle in ready_batch { - if let Ok(uring_handle) = handle.downcast_bound::(py) { + if let Ok(uring_handle) = handle.cast_bound::(py) { // Execute timer callback // asyncio.TimerHandle._run() executes the callback if let Err(e) = uring_handle.borrow().execute(py) { @@ -921,7 +1388,7 @@ impl UringCore { } } else { // Should not happen for timers from our own loop - // but handle generic PyObject just in case + // but handle generic Py just in case let is_cancelled = match handle.call_method0(py, "cancelled") { Ok(val) => val.is_truthy(py).unwrap_or(false), Err(_) => false, // If no cancelled method, assume not cancelled @@ -943,7 +1410,7 @@ impl UringCore { // ========================================================================= /// Add a reader callback for a file descriptor. - fn add_reader(&self, fd: i32, callback: PyObject, args: PyObject) { + fn add_reader(&self, fd: i32, callback: Py, args: Py) { self.readers.lock().insert(fd, (callback, args)); } @@ -953,7 +1420,7 @@ impl UringCore { } /// Add a writer callback for a file descriptor. - fn add_writer(&self, fd: i32, callback: PyObject, args: PyObject) { + fn add_writer(&self, fd: i32, callback: Py, args: Py) { self.writers.lock().insert(fd, (callback, args)); } @@ -1011,7 +1478,7 @@ impl UringCore { }; // epoll_wait (release GIL during blocking) - let nfds = py.allow_threads(|| unsafe { + let nfds = py.detach(|| unsafe { libc::epoll_wait(epoll_fd, events.as_mut_ptr(), 64, timeout_ms) }); @@ -1044,8 +1511,7 @@ impl UringCore { .map(|(cb, args)| (cb.clone_ref(py), args.clone_ref(py))) }; if let Some((callback, args)) = maybe_reader { - if let Ok(args_tuple) = args.downcast_bound::(py) - { + if let Ok(args_tuple) = args.cast_bound::(py) { if let Err(e) = callback.call1(py, args_tuple) { e.print(py); } @@ -1062,8 +1528,7 @@ impl UringCore { .map(|(cb, args)| (cb.clone_ref(py), args.clone_ref(py))) }; if let Some((callback, args)) = maybe_writer { - if let Ok(args_tuple) = args.downcast_bound::(py) - { + if let Ok(args_tuple) = args.cast_bound::(py) { if let Err(e) = callback.call1(py, args_tuple) { e.print(py); } @@ -1088,7 +1553,7 @@ impl UringCore { fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; - // m.add_class::()?; // Removed + m.add_class::()?; m.add_class::()?; diff --git a/src/ring.rs b/src/ring.rs index 6b0e453..7377f69 100644 --- a/src/ring.rs +++ b/src/ring.rs @@ -17,11 +17,11 @@ // len() == 0 is sometimes clearer #![allow(clippy::len_zero)] -use io_uring::{opcode, types, IoUring, Submitter}; +use io_uring::{IoUring, Submitter, opcode, types}; use nix::sys::eventfd::{EfdFlags, EventFd}; use std::os::unix::io::{AsRawFd, RawFd}; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use crate::buffer::BufferPool; use crate::error::{Error, Result}; @@ -97,12 +97,20 @@ pub enum OpType { Close = 4, /// Timeout operation Timeout = 5, - /// SOTA: Multishot receive (kernel 5.19+) + /// Multishot receive (kernel 5.19+) RecvMulti = 6, - /// SOTA: Zero-copy send (kernel 6.0+) + /// Zero-copy send (kernel 6.0+) SendZC = 7, - /// SOTA: Multishot accept (kernel 5.19+) + /// Multishot accept (kernel 5.19+) AcceptMulti = 8, + /// `RecvMsg` operation + RecvMsg = 9, + /// `SendMsg` operation + SendMsg = 10, + /// Provided buffer ring group ID + ProvideBuffer = 11, + /// Registered FD table (`IOSQE_FIXED_FILE`) + FixedFdTable = 12, /// Unknown operation Unknown = 255, } @@ -121,6 +129,10 @@ impl OpType { 6 => Self::RecvMulti, 7 => Self::SendZC, 8 => Self::AcceptMulti, + 9 => Self::RecvMsg, + 10 => Self::SendMsg, + 11 => Self::ProvideBuffer, + 12 => Self::FixedFdTable, _ => Self::Unknown, } } @@ -138,6 +150,10 @@ impl OpType { Self::RecvMulti => "recv_multi", Self::SendZC => "send_zc", Self::AcceptMulti => "accept_multi", + Self::RecvMsg => "recvmsg", + Self::SendMsg => "sendmsg", + Self::ProvideBuffer => "provide_buffer", + Self::FixedFdTable => "fixed_fd_table", Self::Unknown => "unknown", } } @@ -180,9 +196,9 @@ pub struct Ring { is_active: AtomicBool, /// Buffer pool reference for registered buffers buffer_pool: Option>, - /// SOTA: Registered FD table (`IOSQE_FIXED_FILE`) + /// Registered FD table (`IOSQE_FIXED_FILE`) registered_fds: Option, - /// SOTA: Provided buffer ring group ID + /// Provided buffer ring group ID provided_buf_group_id: Option, } @@ -397,10 +413,12 @@ impl Ring { // Let's check if the crate supports it. If not, we use raw register. // io-uring crate 0.6+ supports register_buf_ring. - self.ring - .submitter() - .register_buf_ring_with_flags(addr, ring_entries, bgid, 0) - .map_err(|e| Error::RingOp(format!("register_buf_ring failed: {e}")))?; + unsafe { + self.ring + .submitter() + .register_buf_ring_with_flags(addr, ring_entries, bgid, 0) + .map_err(|e| Error::RingOp(format!("register_buf_ring failed: {e}")))?; + } self.provided_buf_group_id = Some(bgid); Ok(()) @@ -540,8 +558,10 @@ impl Ring { if sq.is_full() { return Err(Error::RingOp("SQ is full".into())); } - sq.push(&entry) - .map_err(|_| Error::RingOp("push failed".into())) + unsafe { + sq.push(&entry) + .map_err(|_| Error::RingOp("push failed".into())) + } }) } @@ -573,8 +593,10 @@ impl Ring { if sq.is_full() { return Err(Error::RingOp("SQ is full".into())); } - sq.push(&entry) - .map_err(|_| Error::RingOp("push failed".into())) + unsafe { + sq.push(&entry) + .map_err(|_| Error::RingOp("push failed".into())) + } }) } @@ -718,7 +740,7 @@ impl Ring { /// Prepare a standalone timeout operation (native timer). /// /// Returns completion when `deadline_ns` (absolute monotonic time) is reached. - /// Use `encode_user_data(timer_id, OpType::Timeout, gen)` for `user_data`. + /// Use `encode_user_data(timer_id, OpType::Timeout, generation)` for `user_data`. pub fn prep_timeout(&mut self, deadline_ns: u64, user_data: u64) -> Result<()> { // Convert nanoseconds to timespec #[allow(clippy::cast_possible_truncation)] @@ -798,6 +820,76 @@ impl Ring { }) } + /// Prepare a recvmsg operation. + /// + /// # Safety + /// + /// The msghdr must remain valid until completion. + pub unsafe fn prep_recvmsg( + &mut self, + fd: RawFd, + msg: *mut libc::msghdr, + _buf_idx: u16, + generation: u16, + ) -> Result<()> { + let user_data = encode_user_data(fd, OpType::RecvMsg, generation); + + let entry = if let Some(idx) = self.lookup_fixed(fd) { + opcode::RecvMsg::new(types::Fixed(idx), msg) + .build() + .user_data(user_data) + } else { + opcode::RecvMsg::new(types::Fd(fd), msg) + .build() + .user_data(user_data) + }; + + self.with_sq(|sq| { + if sq.is_full() { + return Err(Error::RingOp("SQ is full".into())); + } + unsafe { + sq.push(&entry) + .map_err(|_| Error::RingOp("push recvmsg failed".into())) + } + }) + } + + /// Prepare a sendmsg operation. + /// + /// # Safety + /// + /// The msghdr must remain valid until completion. + pub unsafe fn prep_sendmsg( + &mut self, + fd: RawFd, + msg: *mut libc::msghdr, + _buf_idx: u16, + generation: u16, + ) -> Result<()> { + let user_data = encode_user_data(fd, OpType::SendMsg, generation); + + let entry = if let Some(idx) = self.lookup_fixed(fd) { + opcode::SendMsg::new(types::Fixed(idx), msg) + .build() + .user_data(user_data) + } else { + opcode::SendMsg::new(types::Fd(fd), msg) + .build() + .user_data(user_data) + }; + + self.with_sq(|sq| { + if sq.is_full() { + return Err(Error::RingOp("SQ is full".into())); + } + unsafe { + sq.push(&entry) + .map_err(|_| Error::RingOp("push sendmsg failed".into())) + } + }) + } + // ========================================================================= // SOTA 2025: Registered FD Table (IOSQE_FIXED_FILE) // ========================================================================= @@ -869,8 +961,10 @@ impl Ring { if sq.is_full() { return Err(Error::RingOp("SQ is full".into())); } - sq.push(&entry) - .map_err(|_| Error::RingOp("push send_zc failed".into())) + unsafe { + sq.push(&entry) + .map_err(|_| Error::RingOp("push send_zc failed".into())) + } }) } @@ -943,9 +1037,9 @@ mod tests { fn test_user_data_encoding() { let fd = 42i32; let op = OpType::Recv; - let gen = 1u16; + let generation = 1u16; - let user_data = encode_user_data(fd, op, gen); + let user_data = encode_user_data(fd, op, generation); let entry = CompletionEntry { user_data, @@ -956,7 +1050,7 @@ mod tests { assert_eq!(entry.fd(), fd); assert_eq!(entry.op_type(), OpType::Recv); - assert_eq!(entry.generation(), gen); + assert_eq!(entry.generation(), generation); } #[test] diff --git a/src/scheduler.rs b/src/scheduler.rs index 9a2b457..d5c1ec7 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -9,7 +9,7 @@ use std::sync::Arc; /// for single-threaded asyncio workloads compared to channel-based solutions. #[derive(Clone)] pub struct Scheduler { - queue: Arc>>, + queue: Arc>>>, } impl Default for Scheduler { @@ -27,13 +27,13 @@ impl Scheduler { } /// Push a task to the ready queue. - pub fn push(&self, handle: PyObject) { + pub fn push(&self, handle: Py) { self.queue.lock().push_back(handle); } /// Pop a task from the ready queue. #[must_use] - pub fn pop(&self) -> Option { + pub fn pop(&self) -> Option> { self.queue.lock().pop_front() } @@ -52,7 +52,7 @@ impl Scheduler { /// Drain all items from the queue efficiently. /// This swaps the underlying queue with a new empty one to minimize lock hold time. #[must_use] - pub fn drain(&self) -> VecDeque { + pub fn drain(&self) -> VecDeque> { let mut queue = self.queue.lock(); if queue.is_empty() { return VecDeque::new(); diff --git a/src/timer.rs b/src/timer.rs index e4e24fa..19fde7d 100644 --- a/src/timer.rs +++ b/src/timer.rs @@ -6,7 +6,7 @@ use std::collections::BinaryHeap; #[derive(Debug)] struct TimerEntry { expiration: f64, - handle: PyObject, + handle: Py, } impl PartialEq for TimerEntry { @@ -52,11 +52,11 @@ impl TimerHeap { } } - pub fn push(&self, expiration: f64, handle: PyObject) { + pub fn push(&self, expiration: f64, handle: Py) { self.heap.lock().push(TimerEntry { expiration, handle }); } - pub fn pop_expired(&self, now: f64) -> Vec { + pub fn pop_expired(&self, now: f64) -> Vec> { let mut heap = self.heap.lock(); let mut expired = Vec::new(); while let Some(top) = heap.peek() { diff --git a/tests/conftest.py b/tests/conftest.py index 1a29433..e88a21e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,11 +4,18 @@ import os # Set limits for test environment (overridable) -os.environ.setdefault("URINGCORE_BUFFER_COUNT", "512") -os.environ.setdefault("URINGCORE_BUFFER_SIZE", "32768") +os.environ.setdefault("URINGCORE_BUFFER_COUNT", "128") +os.environ.setdefault("URINGCORE_BUFFER_SIZE", "4096") @pytest.fixture(scope="session", autouse=True) def configure_event_loop_policy(): """Ensure uringcore policy is used for all tests.""" policy = uringcore.EventLoopPolicy() asyncio.set_event_loop_policy(policy) + +@pytest.fixture(scope="function") +def event_loop(): + """Create an instance of the default event loop for each test function.""" + loop = uringcore.new_event_loop() + yield loop + loop.close() diff --git a/tests/test_asyncio_compat.py b/tests/test_asyncio_compat.py index 63808da..b24bf0b 100644 --- a/tests/test_asyncio_compat.py +++ b/tests/test_asyncio_compat.py @@ -207,10 +207,11 @@ async def handle(reader, writer): await writer.drain() writer.close() - server = await asyncio.start_server(handle, '127.0.0.1', 19880) + server = await asyncio.start_server(handle, '127.0.0.1', 0) await asyncio.sleep(0.05) - reader, writer = await asyncio.open_connection('127.0.0.1', 19880) + port = server.sockets[0].getsockname()[1] + reader, writer = await asyncio.open_connection('127.0.0.1', port) writer.write(b'hello') await writer.drain() @@ -280,13 +281,14 @@ def datagram_received(self, data, addr): self.future.set_result(data) server, _ = await loop.create_datagram_endpoint( - ServerProtocol, local_addr=('127.0.0.1', 19881) + ServerProtocol, local_addr=('127.0.0.1', 0) ) + port = server.get_extra_info('socket').getsockname()[1] future = loop.create_future() client, _ = await loop.create_datagram_endpoint( lambda: ClientProtocol(future), - remote_addr=('127.0.0.1', 19881) + remote_addr=('127.0.0.1', port) ) # Use manual wait diff --git a/tests/test_udp_overflow.py b/tests/test_udp_overflow.py new file mode 100644 index 0000000..afb7e1f --- /dev/null +++ b/tests/test_udp_overflow.py @@ -0,0 +1,48 @@ + +import asyncio +import socket +import pytest +import os +import errno +import time + +def test_udp_send_overflow(event_loop): + # Create a UDP socket pair + server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + server.bind(('127.0.0.1', 0)) + server.setblocking(False) + server_addr = server.getsockname() + + client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + client.setblocking(False) + + # Set a small send buffer to trigger EAGAIN/ENOBUFS easily + try: + client.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 1024) + except OSError: + pass # System might enforce minimum + + print(f"Send buffer size: {client.getsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF)}") + + async def main(): + data = b"x" * 1024 # 1KB chunks + + count = 0 + try: + # Try to blast data faster than it can be drained (we are not reading) + # 10MB should be enough to overflow a small buffer + for _ in range(10000): + await event_loop.sock_sendto(client, data, server_addr) + count += 1 + if count % 1000 == 0: + print(f"Sent {count} packets...") + except OSError as e: + print(f"Caught expected error after {count} packets: {e}") + # Map io_uring -EAGAIN result (-11) to python errno.EAGAIN + # Depending on how it's raised, standard OSError might be errno 11. + assert e.errno in (errno.EAGAIN, errno.EWOULDBLOCK, errno.ENOBUFS) + return + + print(f"Finished sending {count} packets without error.") + + event_loop.run_until_complete(main()) diff --git a/tests/test_udp_recvfrom.py b/tests/test_udp_recvfrom.py new file mode 100644 index 0000000..0187cb7 --- /dev/null +++ b/tests/test_udp_recvfrom.py @@ -0,0 +1,32 @@ +import asyncio +import socket +import pytest +import uringcore + +def test_sock_recvfrom(event_loop): + loop = event_loop + assert isinstance(loop, uringcore.UringEventLoop) + + async def run_test(): + server_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + server_sock.bind(("127.0.0.1", 0)) + server_addr = server_sock.getsockname() + server_sock.setblocking(False) + + client_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + + msg = b"Hello io_uring" + client_sock.sendto(msg, server_addr) + + print(f"Waiting for data on {server_addr}...") + data, addr = await loop.sock_recvfrom(server_sock, 1024) + print(f"Received {data} from {addr}") + + assert data == msg + # addr is (ip, port) + assert addr[0] == "127.0.0.1" + + server_sock.close() + client_sock.close() + + loop.run_until_complete(run_test()) diff --git a/tests/test_udp_sendto.py b/tests/test_udp_sendto.py new file mode 100644 index 0000000..c02fe77 --- /dev/null +++ b/tests/test_udp_sendto.py @@ -0,0 +1,41 @@ +import socket +import pytest +import uringcore +import asyncio + +@pytest.fixture +def event_loop(): + loop = uringcore.new_event_loop() + yield loop + loop.close() + +def test_sock_sendto(event_loop): + """Test sock_sendto updates using loop.sock_sendto.""" + + # Create two UDP sockets + server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + server.bind(('127.0.0.1', 0)) + server.setblocking(False) + server_addr = server.getsockname() + + client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + client.setblocking(False) + + async def main(): + # Send data from client to server + data = b"Hello, uring!" + + # Test io_uring path + sent = await event_loop.sock_sendto(client, data, server_addr) + assert sent == len(data) + + # Verify receipt + received, addr = await event_loop.sock_recvfrom(server, 1024) + # Client is bound to 0.0.0.0 but server sees 127.0.0.1 + assert received == data + assert addr[1] == client.getsockname()[1] # Port must match + + event_loop.run_until_complete(main()) + + server.close() + client.close() diff --git a/tests/test_unix_datagram.py b/tests/test_unix_datagram.py new file mode 100644 index 0000000..3110e03 --- /dev/null +++ b/tests/test_unix_datagram.py @@ -0,0 +1,49 @@ + +import asyncio +import socket +import os +import pytest +import uringcore + + + +def test_unix_datagram_sendto_recvfrom(event_loop): + server_path = "/tmp/uring_unix_dgram_server.sock" + client_path = "/tmp/uring_unix_dgram_client.sock" + + if os.path.exists(server_path): + os.remove(server_path) + if os.path.exists(client_path): + os.remove(client_path) + + server = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) + server.bind(server_path) + server.setblocking(False) + + client = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) + client.bind(client_path) + client.setblocking(False) + + async def main(): + try: + data = b"Hello UNIX Datagram" + + # Test io_uring sendto with path + sent = await event_loop.sock_sendto(client, data, server_path) + assert sent == len(data) + + # Test io_uring recvfrom + received, addr = await event_loop.sock_recvfrom(server, 1024) + assert received == data + # Address in UNIX datagram might be the client path + assert addr == client_path + + finally: + server.close() + client.close() + if os.path.exists(server_path): + os.remove(server_path) + if os.path.exists(client_path): + os.remove(client_path) + + event_loop.run_until_complete(main()) diff --git a/tests/test_unix_server.py b/tests/test_unix_server.py new file mode 100644 index 0000000..2b138ac --- /dev/null +++ b/tests/test_unix_server.py @@ -0,0 +1,48 @@ + +import asyncio +import socket +import os +import pytest +import uringcore + + + +def test_create_unix_server(event_loop): + path = "/tmp/uring_unix_server.sock" + if os.path.exists(path): + os.remove(path) + + async def handler(reader, writer): + data = await reader.read(100) + writer.write(data) + await writer.drain() + writer.write_eof() + await asyncio.sleep(0.1) + writer.close() + await writer.wait_closed() + + async def main(): + server = await event_loop.create_unix_server(handler, path=path) + + try: + reader, writer = await asyncio.open_unix_connection(path, limit=65536) # Removed loop=event_loop as it's implicit or uses get_event_loop + # Actually, open_unix_connection uses get_running_loop internally. + # Since we are in main(), running loop is event_loop. + + msg = b"Hello UNIX Stream" + writer.write(msg) + await writer.drain() + + response = await reader.read(100) + assert response == msg + + writer.close() + await writer.wait_closed() + + finally: + server.close() + await server.wait_closed() + if os.path.exists(path): + os.remove(path) + + event_loop.run_until_complete(main()) diff --git a/tests/verify/run_ci.sh b/tests/verify/run_ci.sh index 3967bd7..7b8d40d 100755 --- a/tests/verify/run_ci.sh +++ b/tests/verify/run_ci.sh @@ -12,11 +12,70 @@ cleanup() { } trap cleanup EXIT -echo "Starting Verification Test: $TEST_NAME" +# Resolve Python interpreter (prefer project venv) +PYTHON_CMD="python" +if [ -f "../../.venv/bin/python" ]; then + PYTHON_CMD="../../.venv/bin/python" +fi + +# Configure conservative buffer limits for CI (prevent ENOMEM on low RLIMIT_MEMLOCK) +export URINGCORE_BUFFER_COUNT=1024 +export URINGCORE_BUFFER_SIZE=4096 + +echo "Starting Verification Test: $TEST_NAME using $PYTHON_CMD" + + +# Kill any existing process on the target port +cleanup_port() { + local port=$1 + if command -v fuser >/dev/null; then + fuser -k -n tcp $port 2>/dev/null || true + else + # Fallback if fuser is missing (e.g. some containers) + # Try lsof if available + if command -v lsof >/dev/null; then + lsof -t -i:$port | xargs -r kill 2>/dev/null || true + fi + fi + sleep 1 +} + +wait_for_server() { + local port=$1 + local logfile=$2 + local retries=30 + local wait_time=0.5 + + echo "Waiting for server on port $port..." + for i in $(seq 1 $retries); do + if curl -s "http://127.0.0.1:$port" >/dev/null || curl -s "http://127.0.0.1:$port/ping" >/dev/null; then + echo "Server is up!" + return 0 + fi + + # Check if process is still running + if ! kill -0 $PID 2>/dev/null; then + echo "Server process (PID $PID) died unexpectedly." + echo "=== Server Log ($logfile) ===" + cat $logfile + echo "=============================" + return 1 + fi + + sleep $wait_time + done + + echo "Timed out waiting for server on port $port." + echo "=== Server Log ($logfile) ===" + cat $logfile + echo "=============================" + return 1 +} if [ "$TEST_NAME" == "fastapi" ]; then + cleanup_port 8000 # Run FastAPI with uringcore (default config) - python -c " + $PYTHON_CMD -c " import asyncio, uringcore asyncio.set_event_loop_policy(uringcore.EventLoopPolicy()) import uvicorn @@ -25,15 +84,19 @@ uvicorn.run(app.app, host='127.0.0.1', port=8000) " > fastapi_ci.log 2>&1 & PID=$! - # Wait for startup - sleep 3 + if ! wait_for_server 8000 "fastapi_ci.log"; then wait; exit 1; fi # Run wrk - wrk -t4 -c100 -d5s http://127.0.0.1:8000/ping + if command -v wrk >/dev/null; then + wrk -t4 -c100 -d5s http://127.0.0.1:8000/ping + else + curl -s http://127.0.0.1:8000/ping + fi elif [ "$TEST_NAME" == "starlette" ]; then + cleanup_port 8001 # Run Starlette streaming - python -c " + $PYTHON_CMD -c " import asyncio, uringcore asyncio.set_event_loop_policy(uringcore.EventLoopPolicy()) import uvicorn @@ -42,7 +105,8 @@ uvicorn.run(stream_app.app, host='127.0.0.1', port=8001) " > starlette_ci.log 2>&1 & PID=$! - sleep 3 + # Starlette stream endpoint might not respond to / (404 is fine for connectivity check) + if ! wait_for_server 8001 "starlette_ci.log"; then wait; exit 1; fi # Verify monotonic ordering curl -sN http://127.0.0.1:8001/stream | nl | head -n 20 @@ -50,20 +114,24 @@ uvicorn.run(stream_app.app, host='127.0.0.1', port=8001) curl -sN http://127.0.0.1:8001/stream | grep -q "0" elif [ "$TEST_NAME" == "django" ]; then + cleanup_port 8002 # Run Django with Daphne + # Daphne is an executable, usually in venv/bin. Ensure PATH includes it or use venv path. + export PATH="../../.venv/bin:$PATH" cd testproj daphne -p 8002 testproj.asgi:application > django_ci.log 2>&1 & PID=$! cd .. - sleep 5 + if ! wait_for_server 8002 "django_ci.log"; then wait; exit 1; fi # Check Admin page (302 or 200) curl -I http://127.0.0.1:8002/admin/login/?next=/admin/ | grep -E "HTTP/1.1 (200|302)" - + elif [ "$TEST_NAME" == "nosqpoll" ]; then + cleanup_port 8003 # Run with try_sqpoll=False - python -c " + $PYTHON_CMD -c " import asyncio, uringcore # Disable SQPOLL explicitly loop = uringcore.UringEventLoop(try_sqpoll=False) @@ -74,9 +142,13 @@ uvicorn.run(app.app, host='127.0.0.1', port=8003) " > nosqpoll_ci.log 2>&1 & PID=$! - sleep 3 + if ! wait_for_server 8003 "nosqpoll_ci.log"; then wait; exit 1; fi - wrk -t4 -c100 -d5s http://127.0.0.1:8003/ping + if command -v wrk >/dev/null; then + wrk -t4 -c100 -d5s http://127.0.0.1:8003/ping + else + echo "Check: $(curl -s http://127.0.0.1:8003/ping)" + fi else echo "Unknown test: $TEST_NAME"