From f0ac0405a5a970b02201d3130b8b4eb93f95d015 Mon Sep 17 00:00:00 2001 From: Kenneth Allen Date: Sun, 25 Oct 2026 00:51:20 +1100 Subject: [PATCH] Implement socket lifecycle --- Cargo.lock | 21 +++++++++++ Cargo.toml | 1 + src/lib.rs | 9 +++-- src/ops/mod.rs | 1 + src/ops/no_op.rs | 33 ++--------------- src/ops/socket.rs | 37 +++++++++++++++++++ src/ring_op.rs | 94 +++++++++++++++++++++++++++++++++++++---------- src/tests.rs | 19 +++++++++- src/timer.rs | 10 ++--- todo.txt | 5 ++- 10 files changed, 169 insertions(+), 61 deletions(-) create mode 100644 src/ops/socket.rs diff --git a/Cargo.lock b/Cargo.lock index e9f40c3..7deff3a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -217,6 +217,26 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[package]] +name = "pin-project" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "portable-atomic" version = "1.15.0" @@ -382,6 +402,7 @@ dependencies = [ "io-uring", "libc", "log", + "pin-project", "send_wrapper", "slab", "slotmap", diff --git a/Cargo.toml b/Cargo.toml index 817216f..0fc7040 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,6 +7,7 @@ edition = "2024" io-uring = "0.7.15" libc = "0.2.189" log = "0.4.34" +pin-project = "1.1.13" send_wrapper = "0.6.0" slab = "0.4.12" slotmap = "1.1.1" diff --git a/src/lib.rs b/src/lib.rs index 7748e7a..bd8933c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -11,7 +11,7 @@ use std::{ use io_uring::{ IoUring, - opcode::Nop, + opcode::{Close, Nop, Socket}, register::Restriction, types::{SubmitArgs, Timespec}, }; @@ -72,7 +72,11 @@ impl Runtime { ); let submitter = ring.submitter(); submitter - .register_restrictions(&mut [Restriction::sqe_op(Nop::CODE)]) + .register_restrictions(&mut [ + Restriction::sqe_op(Nop::CODE), + Restriction::sqe_op(Socket::CODE), + Restriction::sqe_op(Close::CODE), + ]) .expect("failed to register restrictions on ring"); submitter .register_enable_rings() @@ -214,7 +218,6 @@ impl Runtime { break; }; trace!("Processing CQE {:?}", &cqe); - debug_assert!(cqe.result() >= 0); drop(lock); { let handle = unsafe { RingOpHandle::from_user_data(cqe.user_data()) }; diff --git a/src/ops/mod.rs b/src/ops/mod.rs index 7fe2eaf..b4f9569 100644 --- a/src/ops/mod.rs +++ b/src/ops/mod.rs @@ -1 +1,2 @@ pub mod no_op; +pub mod socket; diff --git a/src/ops/no_op.rs b/src/ops/no_op.rs index e47c2e2..a25d5c0 100644 --- a/src/ops/no_op.rs +++ b/src/ops/no_op.rs @@ -1,34 +1,7 @@ -use std::{ - pin::Pin, - task::{Context, Poll}, -}; - use io_uring::opcode::Nop; -use crate::ring_op::RingOp; +use crate::ring_op::ring_op_result; -pub fn no_op() -> NoOp { - Default::default() -} - -#[derive(Default)] -pub struct NoOp { - ring_op: RingOp, -} - -impl Future for NoOp { - type Output = (); - - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - match self - .ring_op - .submit_and_poll(cx, |user_data| Nop::new().build().user_data(user_data)) - { - Poll::Ready(complete) => { - debug_assert_eq!(complete.result, 0); - Poll::Ready(()) - } - Poll::Pending => Poll::Pending, - } - } +pub async fn no_op() { + let _ = ring_op_result(|| Nop::new().build()).await.unwrap(); } diff --git a/src/ops/socket.rs b/src/ops/socket.rs new file mode 100644 index 0000000..da81752 --- /dev/null +++ b/src/ops/socket.rs @@ -0,0 +1,37 @@ +use std::{io, mem}; + +use io_uring::{ + opcode::{Close, Socket as SocketOpcode}, + types::Fd, +}; +use log::trace; + +use crate::ring_op::{ring_op_result, submit_and_forget}; + +pub async fn socket(domain: i32, socket_type: i32, protocol: i32) -> io::Result { + let fd = ring_op_result(|| SocketOpcode::new(domain, socket_type, protocol).build()).await?; + Ok(Socket { fd: Some(fd) }) +} + +pub struct Socket { + fd: Option, +} + +impl Socket { + pub async fn close(mut self) -> io::Result<()> { + let _ = ring_op_result(|| { + let fd = mem::take(&mut self.fd).unwrap(); + Close::new(Fd(fd)).build() + }) + .await?; + Ok(()) + } +} + +impl Drop for Socket { + fn drop(&mut self) { + if let Some(fd) = self.fd { + submit_and_forget(Close::new(Fd(fd)).build()); + } + } +} diff --git a/src/ring_op.rs b/src/ring_op.rs index 6c224ac..f0a2f73 100644 --- a/src/ring_op.rs +++ b/src/ring_op.rs @@ -1,6 +1,7 @@ use std::{ - mem, + io, mem, ops::DerefMut as _, + pin::Pin, ptr, rc::Rc, task::{Context, LocalWaker, Poll}, @@ -8,15 +9,20 @@ use std::{ use io_uring::{opcode::AsyncCancel, squeue::Entry}; use log::trace; +use pin_project::pin_project; use unchecked_refcell::UncheckedRefCell; use crate::{RT, Runtime}; +#[derive(Default)] pub(crate) struct RingOp(RingOpHandle); +#[derive(Default, Clone)] pub(crate) struct RingOpHandle(Rc>); +#[derive(Default)] enum Inner { + #[default] New, Live(LocalWaker), Complete(Complete), @@ -24,15 +30,9 @@ enum Inner { } #[derive(Clone, Copy)] -pub struct Complete { +pub(crate) struct Complete { pub result: i32, - pub flags: u32, -} - -impl Default for RingOp { - fn default() -> Self { - Self(RingOpHandle(Rc::new(UncheckedRefCell::new(Inner::New)))) - } + pub _flags: u32, } impl Drop for RingOp { @@ -60,10 +60,10 @@ impl Drop for RingOp { } impl RingOp { - pub fn submit_and_poll( - &mut self, + pub fn poll_submit( + self: Pin<&mut Self>, cx: &mut Context<'_>, - with_entry: impl FnOnce(u64) -> Entry, + with_entry: impl FnOnce() -> Entry, ) -> Poll { let mut lock = self.0.0.borrow_mut(); match &mut *lock { @@ -73,12 +73,7 @@ impl RingOp { { let rt: &Runtime = cx.ext().downcast_ref().unwrap(); // TODO Get this way or thread local? let mut lock = rt.0.borrow_mut(); - let entry = with_entry( - Rc::into_raw(self.0.0.clone()) - .expose_provenance() - .try_into() - .unwrap(), - ); + let entry = with_entry().user_data(self.0.clone().into_user_data()); trace!("Pushing SQE {:?}", &entry); let mut sq = lock.ring.submission(); unsafe { sq.push(&entry) }.expect("failed to submit no-op"); @@ -96,18 +91,25 @@ impl RingOp { } impl RingOpHandle { + pub fn into_user_data(self) -> u64 { + Rc::into_raw(self.0).expose_provenance().try_into().unwrap() + } + pub unsafe fn from_user_data(user_data: u64) -> Self { let ptr = ptr::with_exposed_provenance(user_data.try_into().unwrap()); Self(unsafe { Rc::from_raw(ptr) }) } pub fn complete(self, result: i32, flags: u32) { - // TODO don't drop for multishot accepts + // TODO don't drop self for multishot accepts let mut lock = self.0.borrow_mut(); match mem::replace( lock.deref_mut(), - Inner::Complete(Complete { result, flags }), + Inner::Complete(Complete { + result, + _flags: flags, + }), ) { Inner::New | Inner::Complete { .. } => unreachable!(), Inner::Cancelled => {} @@ -115,3 +117,55 @@ impl RingOpHandle { } } } + +pub fn submit_and_forget(entry: Entry) { + RT.with(|rt| { + let outer = rt.borrow(); + let lock = outer + .as_ref() + .expect("socket closed outside runtime thread") + .0 + .borrow_mut(); + + // Start cancelled to ignore result + let handle = RingOpHandle(Rc::new(UncheckedRefCell::new(Inner::Cancelled))); + let entry = entry.user_data(handle.into_user_data()); + trace!("Submitting and forgetting {:?}", &entry); + let mut sq = unsafe { lock.ring.submission_shared() }; + unsafe { sq.push(&entry) }.expect("failed to submit cancel"); + // TODO check for full + }); +} + +#[pin_project] +struct RingOpFut Entry> { + #[pin] + ring_op: RingOp, + with_entry: E, +} + +impl Entry> Future for RingOpFut { + type Output = Complete; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.project(); + this.ring_op.poll_submit(cx, this.with_entry) + } +} + +pub async fn ring_op(with_entry: impl FnMut() -> Entry) -> Complete { + RingOpFut { + ring_op: Default::default(), + with_entry, + } + .await +} + +pub async fn ring_op_result(with_entry: impl FnMut() -> Entry) -> io::Result { + let Complete { result, .. } = ring_op(with_entry).await; + if result < 0 { + Err(io::Error::from_raw_os_error(-result)) + } else { + Ok(result) + } +} diff --git a/src/tests.rs b/src/tests.rs index 4b4c6ce..09df5f2 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -3,8 +3,10 @@ use std::{ time::Duration, }; +use libc::{AF_INET, SOCK_DGRAM}; + use crate::{ - ops::no_op::no_op, + ops::{no_op::no_op, socket::socket}, timer::{delay, delay_until}, }; @@ -52,3 +54,18 @@ fn concurrency() { } runtime.run(); } + +#[test] +fn socket_lifecycle() { + init(); + let mut runtime = Runtime::new(); + runtime.spawn(Box::pin(async move { + let socket = socket(AF_INET, SOCK_DGRAM, 0).await.unwrap(); + socket.close().await.unwrap(); + })); + runtime.spawn(Box::pin(async move { + let socket = socket(AF_INET, SOCK_DGRAM, 0).await.unwrap(); + drop(socket); + })); + runtime.run(); +} diff --git a/src/timer.rs b/src/timer.rs index dc303c3..75309e1 100644 --- a/src/timer.rs +++ b/src/timer.rs @@ -1,5 +1,4 @@ use std::{ - cell::RefCell, cmp::Reverse, collections::BinaryHeap, mem, @@ -10,6 +9,7 @@ use std::{ }; use slab::Slab; +use unchecked_refcell::UncheckedRefCell; use crate::Runtime; @@ -28,7 +28,7 @@ enum Slot { #[derive(Default)] pub(crate) struct TimerHeap { - slots: Slab>>, + slots: Slab>>, queue: BinaryHeap>, } @@ -39,8 +39,8 @@ pub(crate) enum TimerPoll { } impl TimerHeap { - fn insert(&mut self, deadline: Instant, waker: LocalWaker) -> Rc> { - let slot = Rc::new(RefCell::new(Slot::Armed(waker))); + fn insert(&mut self, deadline: Instant, waker: LocalWaker) -> Rc> { + let slot = Rc::new(UncheckedRefCell::new(Slot::Armed(waker))); let handle = Rc::clone(&slot); let slot_id = self.slots.insert(slot); self.queue.push(Reverse(Entry { deadline, slot_id })); @@ -75,7 +75,7 @@ pub enum Timer { Queued(Queued), } -pub struct Queued(Rc>); +pub struct Queued(Rc>); impl Future for Timer { type Output = (); diff --git a/todo.txt b/todo.txt index 37e97a3..26067c1 100644 --- a/todo.txt +++ b/todo.txt @@ -1,7 +1,8 @@ Minimal wakers -Nop reactor -Sleep reactor +Sleep reactor? File reactor Network reactor Task soundness +Task handles Unchecked SQ/CQ without debug_assertions +Register ring fd (split)