Implement socket lifecycle

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",

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"

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()) };

src/ops/mod.rs

@@ -1 +1,2 @@
 pub mod no_op;
+pub mod socket;

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<Self::Output> {
-        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();
 }

src/ops/socket.rs

@@ -0,0 +1,36 @@
+use std::{io, mem};
+
+use io_uring::{
+    opcode::{Close, Socket as SocketOpcode},
+    types::Fd,
+};
+
+use crate::ring_op::{ring_op_result, submit_and_forget};
+
+pub async fn socket(domain: i32, socket_type: i32, protocol: i32) -> io::Result<Socket> {
+    let fd = ring_op_result(|| SocketOpcode::new(domain, socket_type, protocol).build()).await?;
+    Ok(Socket { fd: Some(fd) })
+}
+
+pub struct Socket {
+    fd: Option<i32>,
+}
+
+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());
+        }
+    }
+}

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<UncheckedRefCell<Inner>>);
 
+#[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<Complete> {
         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<E: FnMut() -> Entry> {
+    #[pin]
+    ring_op: RingOp,
+    with_entry: E,
+}
+
+impl<E: FnMut() -> Entry> Future for RingOpFut<E> {
+    type Output = Complete;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        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<i32> {
+    let Complete { result, .. } = ring_op(with_entry).await;
+    if result < 0 {
+        Err(io::Error::from_raw_os_error(-result))
+    } else {
+        Ok(result)
+    }
+}

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();
+}

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<Rc<RefCell<Slot>>>,
+    slots: Slab<Rc<UncheckedRefCell<Slot>>>,
     queue: BinaryHeap<Reverse<Entry>>,
 }
 
@@ -39,8 +39,8 @@ pub(crate) enum TimerPoll {
 }
 
 impl TimerHeap {
-    fn insert(&mut self, deadline: Instant, waker: LocalWaker) -> Rc<RefCell<Slot>> {
-        let slot = Rc::new(RefCell::new(Slot::Armed(waker)));
+    fn insert(&mut self, deadline: Instant, waker: LocalWaker) -> Rc<UncheckedRefCell<Slot>> {
+        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<RefCell<Slot>>);
+pub struct Queued(Rc<UncheckedRefCell<Slot>>);
 
 impl Future for Timer {
     type Output = ();

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)