From ad85d89641f67c82e1a35f3c761ffe4c4577c881 Mon Sep 17 00:00:00 2001
From: Julian Orth <ju.orth@gmail.com>
Date: Fri, 13 May 2022 18:24:12 +0200
Subject: [PATCH] io-uring: add timeout argument to write

---
 src/io_uring/ops/write.rs | 48 ++++++++++++++++++++++-----------------
 src/xwayland/xwm.rs       | 33 ++++++++++++---------------
 2 files changed, 42 insertions(+), 39 deletions(-)

diff --git a/src/io_uring/ops/write.rs b/src/io_uring/ops/write.rs
index d85e8d7c..13869fc4 100644
--- a/src/io_uring/ops/write.rs
+++ b/src/io_uring/ops/write.rs
@@ -3,37 +3,39 @@ use {
         io_uring::{
             ops::TaskResult,
             pending_result::PendingResult,
-            sys::{io_uring_sqe, IORING_OP_WRITE},
+            sys::{io_uring_sqe, IORING_OP_WRITE, IOSQE_IO_LINK},
             IoUring, IoUringData, Task,
         },
+        time::Time,
         utils::buf::Buf,
     },
-    std::{
-        cell::{Cell, RefCell},
-        rc::Rc,
-    },
+    std::rc::Rc,
     uapi::OwnedFd,
 };
 
 impl IoUring {
-    pub async fn write(&self, fd: &Rc<OwnedFd>, buf: Buf) -> TaskResult<usize> {
+    pub async fn write(
+        &self,
+        fd: &Rc<OwnedFd>,
+        buf: Buf,
+        timeout: Option<Time>,
+    ) -> TaskResult<usize> {
         self.ring.check_destroyed()?;
         let id = self.ring.id();
         let pr = self.ring.pending_results.acquire();
         {
-            let pw = self.ring.cached_writes.pop().unwrap_or_else(|| {
-                Box::new(WriteTask {
-                    id: Cell::new(0),
-                    data: Default::default(),
-                })
-            });
-            pw.id.set(id.id);
-            *pw.data.borrow_mut() = Some(WriteTaskData {
+            let mut pw = self.ring.cached_writes.pop().unwrap_or_default();
+            pw.id = id.id;
+            pw.has_timeout = timeout.is_some();
+            pw.data = Some(WriteTaskData {
                 fd: fd.clone(),
                 buf,
                 res: pr.clone(),
             });
             self.ring.schedule(pw);
+            if let Some(time) = timeout {
+                self.schedule_timeout(time);
+            }
         }
         Ok(pr.await.map(|v| v as usize))
     }
@@ -45,31 +47,35 @@ struct WriteTaskData {
     res: PendingResult,
 }
 
+#[derive(Default)]
 pub struct WriteTask {
-    id: Cell<u64>,
-    data: RefCell<Option<WriteTaskData>>,
+    id: u64,
+    has_timeout: bool,
+    data: Option<WriteTaskData>,
 }
 
 unsafe impl Task for WriteTask {
     fn id(&self) -> u64 {
-        self.id.get()
+        self.id
     }
 
-    fn complete(self: Box<Self>, ring: &IoUringData, res: i32) {
-        if let Some(data) = self.data.borrow_mut().take() {
+    fn complete(mut self: Box<Self>, ring: &IoUringData, res: i32) {
+        if let Some(data) = self.data.take() {
             data.res.complete(res);
         }
         ring.clone().cached_writes.push(self);
     }
 
     fn encode(&self, sqe: &mut io_uring_sqe) {
-        let data = self.data.borrow_mut();
-        let data = data.as_ref().unwrap();
+        let data = self.data.as_ref().unwrap();
         sqe.opcode = IORING_OP_WRITE;
         sqe.fd = data.fd.raw();
         sqe.u1.off = !0;
         sqe.u2.addr = data.buf.as_ptr() as _;
         sqe.u3.rw_flags = 0;
         sqe.len = data.buf.len() as _;
+        if self.has_timeout {
+            sqe.flags = IOSQE_IO_LINK;
+        }
     }
 }
diff --git a/src/xwayland/xwm.rs b/src/xwayland/xwm.rs
index b9392296..77157afe 100644
--- a/src/xwayland/xwm.rs
+++ b/src/xwayland/xwm.rs
@@ -18,9 +18,10 @@ use {
                 WlSurface,
             },
         },
-        io_uring::{IoUring, TaskResultExt},
+        io_uring::{IoUring, IoUringError, TaskResultExt},
         rect::Rect,
         state::State,
+        time::Time,
         tree::ToplevelNode,
         utils::{
             bitflags::BitflagsExt, buf::Buf, clonecell::CloneCell, copyhashmap::CopyHashMap,
@@ -59,10 +60,7 @@ use {
     },
     ahash::{AHashMap, AHashSet},
     bstr::ByteSlice,
-    futures_util::{
-        future::{self, Either},
-        pin_mut, select, FutureExt,
-    },
+    futures_util::{select, FutureExt},
     smallvec::SmallVec,
     std::{
         borrow::Cow,
@@ -2380,24 +2378,23 @@ struct XToWaylandTransfer {
 
 impl XToWaylandTransfer {
     async fn run(mut self) {
-        let timeout = self.state.wheel.timeout(5000);
-        pin_mut!(timeout);
+        let timeout = Time::in_ms(5000).unwrap();
         let mut pos = 0;
         while pos < self.data.len() {
-            let f1 = self.state.ring.write(&self.fd, self.data.slice(pos..));
-            pin_mut!(f1);
-            match future::select(f1, &mut timeout).await {
-                Either::Left((res, _)) => match res.merge() {
-                    Ok(n) => pos += n,
-                    Err(e) => {
-                        log::error!("Could not write to wayland client: {}", ErrorFmt(e));
-                        break;
-                    }
-                },
-                Either::Right(_) => {
+            let res = self
+                .state
+                .ring
+                .write(&self.fd, self.data.slice(pos..), Some(timeout));
+            match res.await.merge() {
+                Ok(n) => pos += n,
+                Err(IoUringError::OsError(OsError(c::ECANCELED))) => {
                     log::error!("Transfer timed out");
                     break;
                 }
+                Err(e) => {
+                    log::error!("Could not write to wayland client: {}", ErrorFmt(e));
+                    break;
+                }
             }
         }
         self.shared.transfers.remove(&self.id);