Client tracing preparations

build/wire.rs

@@ -4,6 +4,7 @@ use crate::open;
 use crate::wire::parser::Field;
 use crate::wire::parser::Lined;
 use crate::wire::parser::Message;
+use crate::wire::parser::ParseResult;
 use crate::wire::parser::Type;
 use crate::wire::parser::parse_messages;
 use crate::wire::parser::to_camel;
@@ -396,11 +397,13 @@ fn write_request_handler<W: Write>(
     Ok(())
 }
 
-fn write_file<W: Write>(
-    f: &mut W,
-    file: &DirEntry,
-    interface_names: &mut Vec<String>,
-) -> Result<()> {
+struct ParsedFile {
+    obj_name: String,
+    camel_obj_name: String,
+    messages: ParseResult,
+}
+
+fn parse_file(file: &DirEntry, interface_names: &mut Vec<String>) -> Result<ParsedFile> {
     let file_name = file.file_name();
     let file_name = std::str::from_utf8(file_name.as_bytes())?;
     println!("cargo:rerun-if-changed=wire/{}", file_name);
@@ -409,6 +412,19 @@ fn write_file<W: Write>(
     interface_names.push(camel_obj_name.clone());
     let contents = std::fs::read(file.path())?;
     let messages = parse_messages(&contents)?;
+    Ok(ParsedFile {
+        obj_name: obj_name.to_string(),
+        camel_obj_name,
+        messages,
+    })
+}
+
+fn write_file(f: &mut impl Write, file: &ParsedFile) -> Result<()> {
+    let ParsedFile {
+        obj_name,
+        camel_obj_name,
+        messages,
+    } = file;
     writeln!(f)?;
     writeln!(f, "id!({}Id);", camel_obj_name)?;
     writeln!(f)?;
@@ -422,19 +438,19 @@ fn write_file<W: Write>(
     {
         let f = &mut open(&format!("wire/{obj_name}.rs"))?;
         writeln!(f, "use super::*;")?;
-        for message in messages.requests.iter().chain(messages.events.iter()) {
-            write_message(f, &camel_obj_name, &message.val)?;
+        for message in messages.messages() {
+            write_message(f, camel_obj_name, &message.val)?;
         }
         write_request_handler(
             f,
-            &camel_obj_name,
+            camel_obj_name,
             &messages.requests,
             RequestHandlerDirection::Request,
             messages.dead,
         )?;
         write_request_handler(
             f,
-            &camel_obj_name,
+            camel_obj_name,
             &messages.events,
             RequestHandlerDirection::Event,
             messages.dead,
@@ -463,9 +479,14 @@ pub fn main() -> Result<()> {
     }
     files.sort_by_key(|f| f.file_name());
     let mut interface_names = vec![];
+    let mut parsed_files = vec![];
     for file in files {
-        write_file(&mut f, &file, &mut interface_names)
+        let parsed = parse_file(&file, &mut interface_names)
             .with_context(|| format!("While processing {}", file.path().display()))?;
+        parsed_files.push(parsed);
+    }
+    for file in parsed_files {
+        write_file(&mut f, &file)?;
     }
     writeln!(f)?;
     writeln!(f, "#[doc(hidden)]")?;

build/wire/parser.rs

@@ -238,6 +238,8 @@ pub struct Message {
     pub attribs: MessageAttribs,
     pub has_reference_type: bool,
     pub is_fixed_size: bool,
+    #[expect(dead_code)]
+    pub is_request: bool,
 }
 
 #[derive(Debug, Default)]
@@ -264,23 +266,33 @@ pub struct ParseResult {
     pub events: Vec<Lined<Message>>,
 }
 
+impl ParseResult {
+    pub fn messages(&self) -> impl Iterator<Item = &Lined<Message>> {
+        self.requests.iter().chain(self.events.iter())
+    }
+}
+
 impl<'a> Parser<'a> {
     fn parse(&mut self) -> Result<ParseResult> {
         let mut dead = false;
         let mut requests = vec![];
         let mut events = vec![];
         while !self.eof() {
             let (line, ty) = self.expect_ident()?;
-            let res = match ty.as_bytes() {
+            let is_request = match ty.as_bytes() {
                 b"dead" => {
                     dead = true;
                     continue;
                 }
-                b"request" => &mut requests,
-                b"event" => &mut events,
+                b"request" => true,
+                b"event" => false,
                 _ => bail!("In line {}: Unexpected entry {:?}", line, ty),
             };
-            res.push(self.parse_message(res.len() as _)?);
+            let res = match is_request {
+                true => &mut requests,
+                false => &mut events,
+            };
+            res.push(self.parse_message(res.len() as _, is_request)?);
         }
         Ok(ParseResult {
             dead,
@@ -332,7 +344,7 @@ impl<'a> Parser<'a> {
         Ok(())
     }
 
-    fn parse_message(&mut self, id: u32) -> Result<Lined<Message>> {
+    fn parse_message(&mut self, id: u32, is_request: bool) -> Result<Lined<Message>> {
         let (line, name) = self.expect_ident()?;
         let res: Result<_> = (|| {
             self.not_eof()?;
@@ -380,6 +392,7 @@ impl<'a> Parser<'a> {
                     attribs,
                     has_reference_type,
                     is_fixed_size,
+                    is_request,
                 },
             })
         })();

src/io_uring.rs

@@ -3,6 +3,8 @@ use crate::io_uring::debounce::Debouncer;
 use crate::io_uring::ops::accept::AcceptTask;
 use crate::io_uring::ops::async_cancel::AsyncCancelTask;
 use crate::io_uring::ops::connect::ConnectTask;
+use crate::io_uring::ops::futex::FutexWaitTask;
+use crate::io_uring::ops::futex::FutexWakeTask;
 use crate::io_uring::ops::poll::PollTask;
 use crate::io_uring::ops::poll_external::PollExternalTask;
 use crate::io_uring::ops::read_write::ReadWriteTask;
@@ -42,6 +44,7 @@ use crate::utils::ptr_ext::PtrExt;
 use crate::utils::stack::Stack;
 use crate::utils::syncqueue::SyncQueue;
 pub use ops::TaskResultExt;
+pub use ops::futex::FutexObj;
 pub use ops::poll_external::PendingPoll;
 pub use ops::poll_external::PollCallback;
 pub use ops::timeout_external::PendingTimeout;
@@ -251,6 +254,8 @@ impl IoUring {
             cached_cmsg_bufs: Default::default(),
             cached_connects: Default::default(),
             cached_accepts: Default::default(),
+            cached_futex_wakes: Default::default(),
+            cached_futex_waits: Default::default(),
             fd_ids_scratch: Default::default(),
             iteration: Default::default(),
             yields: Default::default(),
@@ -325,6 +330,8 @@ struct IoUringData {
     cached_cmsg_bufs: Stack<Buf>,
     cached_connects: Stack<Box<ConnectTask>>,
     cached_accepts: Stack<Box<AcceptTask>>,
+    cached_futex_wakes: Stack<Box<FutexWakeTask>>,
+    cached_futex_waits: Stack<Box<FutexWaitTask>>,
 
     fd_ids_scratch: RefCell<Vec<c::c_int>>,
 

src/io_uring/ops.rs

@@ -4,6 +4,7 @@ use crate::utils::oserror::OsError;
 pub mod accept;
 pub mod async_cancel;
 pub mod connect;
+pub mod futex;
 pub mod poll;
 pub mod poll_external;
 pub mod read_write;

src/io_uring/ops/futex.rs

@@ -0,0 +1,155 @@
+use crate::io_uring::IoUring;
+use crate::io_uring::IoUringData;
+use crate::io_uring::IoUringError;
+use crate::io_uring::IoUringTaskId;
+use crate::io_uring::Task;
+use crate::io_uring::TaskResultExt;
+use crate::io_uring::pending_result::PendingResult;
+use crate::io_uring::sys::IORING_OP_FUTEX_WAIT;
+use crate::io_uring::sys::IORING_OP_FUTEX_WAKE;
+use crate::io_uring::sys::io_uring_sqe;
+use std::rc::Rc;
+use std::sync::atomic::AtomicU32;
+
+pub trait FutexObj: 'static {
+    fn get(&self) -> &AtomicU32;
+}
+
+#[expect(dead_code)]
+const FUTEX2_SIZE_U8: i32 = 0x00;
+#[expect(dead_code)]
+const FUTEX2_SIZE_U16: i32 = 0x01;
+const FUTEX2_SIZE_U32: i32 = 0x02;
+#[expect(dead_code)]
+const FUTEX2_SIZE_U64: i32 = 0x03;
+#[expect(dead_code)]
+const FUTEX2_NUMA: i32 = 0x04;
+#[expect(dead_code)]
+const FUTEX2_MPOL: i32 = 0x08;
+const FUTEX2_PRIVATE: i32 = 128;
+
+impl IoUring {
+    #[expect(dead_code)]
+    pub fn futex_wake(
+        &self,
+        futex: &Rc<impl FutexObj>,
+        n: i32,
+        private: bool,
+    ) -> Result<(), IoUringError> {
+        self.ring.check_destroyed()?;
+        {
+            let mut pw = self.ring.cached_futex_wakes.pop().unwrap_or_default();
+            pw.id = self.ring.id_raw();
+            pw.addr = futex.get().as_ptr();
+            pw.n = n;
+            pw.private = private;
+            pw.futex = Some(futex.clone());
+            self.ring.schedule(pw);
+        }
+        Ok(())
+    }
+
+    #[expect(dead_code)]
+    pub async fn futex_wait(
+        &self,
+        futex: &Rc<impl FutexObj>,
+        val: u32,
+        private: bool,
+    ) -> Result<(), IoUringError> {
+        self.ring.check_destroyed()?;
+        let id = self.ring.id();
+        let pr = self.ring.pending_results.acquire();
+        {
+            let mut pw = self.ring.cached_futex_waits.pop().unwrap_or_default();
+            pw.id = id.id;
+            pw.addr = futex.get().as_ptr();
+            pw.val = val;
+            pw.private = private;
+            pw.data = Some(FutexWaitData {
+                pr: pr.clone(),
+                _futex: futex.clone(),
+            });
+            self.ring.schedule(pw);
+        }
+        Ok(pr.await.map(drop)).merge()
+    }
+}
+
+#[derive(Default)]
+pub struct FutexWakeTask {
+    id: IoUringTaskId,
+    addr: *mut u32,
+    n: i32,
+    private: bool,
+    futex: Option<Rc<dyn FutexObj>>,
+}
+
+struct FutexWaitData {
+    pr: PendingResult,
+    _futex: Rc<dyn FutexObj>,
+}
+
+#[derive(Default)]
+pub struct FutexWaitTask {
+    id: IoUringTaskId,
+    addr: *mut u32,
+    val: u32,
+    private: bool,
+    data: Option<FutexWaitData>,
+}
+
+unsafe impl Task for FutexWakeTask {
+    fn id(&self) -> IoUringTaskId {
+        self.id
+    }
+
+    fn complete(mut self: Box<Self>, ring: &IoUringData, _res: i32) {
+        self.futex.take();
+        ring.cached_futex_wakes.push(self);
+    }
+
+    fn encode(&self, sqe: &mut io_uring_sqe) {
+        encode_futex(
+            sqe,
+            IORING_OP_FUTEX_WAKE,
+            self.addr,
+            self.n as _,
+            self.private,
+        );
+    }
+}
+
+unsafe impl Task for FutexWaitTask {
+    fn id(&self) -> IoUringTaskId {
+        self.id
+    }
+
+    fn complete(mut self: Box<Self>, ring: &IoUringData, res: i32) {
+        if let Some(data) = self.data.take() {
+            data.pr.complete(res);
+        }
+        ring.cached_futex_waits.push(self);
+    }
+
+    fn encode(&self, sqe: &mut io_uring_sqe) {
+        encode_futex(
+            sqe,
+            IORING_OP_FUTEX_WAIT,
+            self.addr,
+            self.val as _,
+            self.private,
+        );
+    }
+}
+
+fn encode_futex(sqe: &mut io_uring_sqe, op: u8, addr: *mut u32, val: u64, private: bool) {
+    let mut flags = FUTEX2_SIZE_U32;
+    if private {
+        flags |= FUTEX2_PRIVATE;
+    }
+    sqe.opcode = op;
+    sqe.fd = flags;
+    sqe.u2.addr = addr as _;
+    sqe.u1.addr2 = val as _;
+    sqe.u6.s1.addr3 = !0u32 as u64;
+}