[DRAFT] perf(io): improve the performance of copy using linked splice batches

compio-driver/src/lib.rs

@@ -272,6 +272,47 @@ impl Proactor {
         }
     }
 
+    /// Submit single-shot operations as one linked io_uring chain.
+    ///
+    /// All operations must have the same type and produce native, single-shot
+    /// entries. Multishot operations are not supported. Each member except
+    /// the last is hard-linked when its `hardlinks` entry is true, otherwise
+    /// soft-linked; the last entry is ignored. Operations must not set their
+    /// own link or skip-completion flags.
+    ///
+    /// The complete chain is published together, without splitting it across
+    /// submissions. This does not make execution transactional: completed I/O
+    /// is not rolled back when a later member fails.
+    ///
+    /// Returns [`io::ErrorKind::Unsupported`] on non-io_uring backends, for
+    /// unsupported native entries, or if the chain exceeds the submission
+    /// queue's capacity. No blocking fallback is used. On rejection, no member
+    /// is submitted and all operations are returned.
+    pub fn push_linked<T: OpCode + 'static, const N: usize>(
+        &mut self,
+        ops: [T; N],
+        hardlinks: [bool; N],
+        extras: [Extra; N],
+    ) -> Result<[Key<T>; N], (io::Error, [T; N])> {
+        #[cfg(io_uring)]
+        if let Some(driver) = self.driver.as_iour_mut() {
+            let mut extras = extras.into_iter();
+            let keys = ops.map(|op| Key::new(op, extras.next().unwrap(), DriverType::IoUring));
+            match driver.push_linked(std::array::from_fn(|i| keys[i].clone().erase()), hardlinks) {
+                Ok(()) => return Ok(keys),
+                Err(error) => {
+                    let ops = keys.map(|key| {
+                        key.set_result(Err(error.kind().into()));
+                        key.take_result().1
+                    });
+                    return Err((error, ops));
+                }
+            }
+        }
+        let _ = (hardlinks, extras);
+        Err((io::ErrorKind::Unsupported.into(), ops))
+    }
+
     /// Flush the pushed operations to the kernel. It is roughly equivalent to
     /// calling [`poll`](Proactor::poll) with `Some(Duration::ZERO)` on
     /// io-uring, but is only needed if you're waiting on the driver fd with an

compio-driver/src/sys/driver/iour/mod.rs

@@ -1,8 +1,18 @@
 use std::{
-    marker::PhantomData, mem::ManuallyDrop, panic::AssertUnwindSafe, sync::Arc, time::Duration,
+    collections::{HashMap, VecDeque},
+    marker::PhantomData,
+    mem::ManuallyDrop,
+    panic::AssertUnwindSafe,
+    sync::Arc,
+    time::Duration,
 };
 
-use crate::sys::{extra::IourExtra, prelude::*};
+use smallvec::SmallVec;
+
+use crate::sys::{
+    extra::{IourExtra, Linked},
+    prelude::*,
+};
 
 mod_use![op, notify];
 
@@ -35,7 +45,7 @@ use slotmap::{DefaultKey, SlotMap};
 
 use crate::{
     AsyncifyPool, DriverType, Entry, ProactorBuilder,
-    key::{BorrowedKey, ErasedKey},
+    key::{BorrowedKey, ErasedKey, WeakKey},
     panic::catch_unwind_io,
 };
 
@@ -70,11 +80,15 @@ pub(crate) struct Driver {
     flags: DriverFlags,
     /// Keys leaked via `into_raw()` into io_uring user_data, freed on drop.
     in_flight: SlotMap<DefaultKey, usize>,
+    // Weak references keep cancellation addresses from being recycled without
+    // preventing completed operations from being extracted by their futures.
+    cancellations: HashMap<usize, (WeakKey, bool)>,
+    pending_cancels: VecDeque<usize>,
+    deferred_cancels: HashMap<DefaultKey, WeakKey>,
     _p: PhantomData<ErasedKey>,
 }
 
 impl Driver {
-    const CANCEL: u64 = u64::MAX;
     const NOTIFY: u64 = u64::MAX - 1;
 
     pub fn new(builder: &ProactorBuilder) -> io::Result<Self> {
@@ -141,6 +155,9 @@ impl Driver {
             pool: builder.create_or_get_thread_pool(),
             flags,
             in_flight: SlotMap::new(),
+            cancellations: HashMap::new(),
+            pending_cancels: VecDeque::new(),
+            deferred_cancels: HashMap::new(),
             _p: PhantomData,
         })
     }
@@ -278,9 +295,9 @@ impl Driver {
     fn poll_entries(&mut self) -> bool {
         let cqueue = self.inner.completion();
         let has_entry = !cqueue.is_empty();
+        let mut cancellations: SmallVec<[ErasedKey; LINKED_BATCH_INLINE]> = SmallVec::new();
         for entry in cqueue {
             match entry.user_data() {
-                Self::CANCEL => {}
                 Self::NOTIFY => {
                     let flags = entry.flags();
                     if !more(flags) {
@@ -290,6 +307,21 @@ impl Driver {
                         error!("failed to clear notifier: {e:?}");
                     }
                 }
+                // Operation addresses are aligned, so their low bit is zero.
+                // We set that bit in AsyncCancel's user_data to distinguish
+                // its completion from the target operation's own completion.
+                // Clear the tag bit to recover the target's cancellation key.
+                tag if tag & 1 != 0 => {
+                    let cancelled = self.cancellations.remove(&((tag & !1) as usize));
+                    if entry.result() == -libc::ENOENT
+                        && let Some((key, true)) = cancelled
+                        && let Some(key) = key.upgrade()
+                    {
+                        // The predecessor completed, but its successor may
+                        // not yet have reached the kernel's cancellation
+                        cancellations.push(key);
+                    }
+                }
                 key => {
                     let flags = entry.flags();
                     if more(flags) {
@@ -304,12 +336,22 @@ impl Driver {
                         key.wake_by_ref();
                     } else {
                         let entry = create_entry(entry);
-                        Self::remove_in_flight(&mut self.in_flight, &entry.key);
+                        if let Some(slot) = Self::remove_in_flight(&mut self.in_flight, &entry.key)
+                            && let Some(key) = self
+                                .deferred_cancels
+                                .remove(&slot)
+                                .and_then(|key| key.upgrade())
+                        {
+                            cancellations.push(key);
+                        }
                         entry.notify()
                     }
                 }
             }
         }
+        for key in cancellations {
+            self.cancel(key);
+        }
         has_entry
     }
 
@@ -323,22 +365,62 @@ impl Driver {
 
     pub fn cancel(&mut self, key: ErasedKey) {
         instrument!(compio_log::Level::TRACE, "cancel", ?key);
-        trace!("cancel RawOp");
-        unsafe {
+        if key.has_result() {
+            return;
+        }
+        let linked = key.borrow().extra().as_iour().linked;
+        if let Some(Linked::After(previous)) = linked
+            && self.in_flight.contains_key(previous)
+        {
+            // Hard-linked successors are not necessarily visible to
+            // AsyncCancel yet. Submit their cancellation only after the
+            // predecessor's CQE, even if cancellation tokens arrive in
+            // arbitrary order.
+            self.deferred_cancels.insert(previous, key.downgrade());
+            return;
+        }
+        let address = key.as_raw();
+        if self.cancellations.contains_key(&address) {
+            return;
+        }
+        self.cancellations
+            .insert(address, (key.downgrade(), linked.is_some()));
+        self.pending_cancels.push_back(address);
+    }
+
+    /// Stage at most one SQ's worth of cancellation work without submitting,
+    /// polling, or waiting. In particular, dropping a future never flushes a
+    /// full queue or recursively dispatches cancellation retries.
+    fn stage_cancellations(&mut self) {
+        if self.pending_cancels.is_empty() {
+            return;
+        }
+        let limit = self
+            .pending_cancels
+            .len()
+            .min(self.inner.submission().capacity());
+        for _ in 0..limit {
+            let address = *self.pending_cancels.front().unwrap();
+            let key = self.cancellations[&address].0.upgrade();
+            if key.as_ref().is_none_or(|key| key.has_result()) {
+                self.pending_cancels.pop_front();
+                self.cancellations.remove(&address);
+                continue;
+            }
+            // Key allocations are pointer-aligned. The weak reference remains
+            // until the cancellation CQE, preventing address reuse meanwhile.
+            debug_assert_eq!(address & 1, 0);
             #[allow(clippy::useless_conversion)]
-            if self
-                .inner
-                .submission()
-                .push(
-                    &AsyncCancel::new(key.as_raw() as _)
-                        .build()
-                        .user_data(Self::CANCEL)
-                        .into(),
-                )
-                .is_err()
-            {
-                warn!("could not push AsyncCancel entry");
+            let entry = AsyncCancel::new(address as _)
+                .build()
+                .user_data((address as u64) | 1)
+                .into();
+            // SAFETY: AsyncCancel has no buffers; its target address remains
+            // reserved by the weak reference in `cancellations`.
+            if unsafe { self.inner.submission().push(&entry) }.is_err() {
+                break;
             }
+            self.pending_cancels.pop_front();
         }
     }
 
@@ -387,6 +469,79 @@ impl Driver {
         }
     }
 
+    /// Validate and publish an entire linked group without splitting it across
+    /// kernel submissions, including when older operations fill the SQ.
+    pub(crate) fn push_linked<const N: usize>(
+        &mut self,
+        keys: [ErasedKey; N],
+        hardlinks: [bool; N],
+    ) -> io::Result<()> {
+        if N > self.inner.submission().capacity() {
+            return Err(unsupported_linked());
+        }
+        let mut entries: SmallVec<[SEntry; LINKED_BATCH_INLINE]> = SmallVec::with_capacity(N);
+        for (i, key) in keys.iter().enumerate() {
+            let mut entry: SEntry = match key.borrow().create_entry::<false>() {
+                #[allow(clippy::useless_conversion)]
+                OpEntry::Submission(entry) => entry.into(),
+                #[cfg(feature = "io-uring-sqe128")]
+                OpEntry::Submission128(entry) => entry,
+                OpEntry::Blocking => return Err(unsupported_linked()),
+            };
+            if !is_op_supported(entry.get_opcode() as _) {
+                return Err(unsupported_linked());
+            }
+            if i + 1 < N {
+                entry = entry.flags(link_flags(hardlinks[i]));
+            }
+            entries.push(entry.user_data(key.as_raw() as _));
+        }
+
+        while {
+            let queue = self.inner.submission();
+            queue.capacity() - queue.len() < N
+        } {
+            match self.submit_auto(Some(Duration::ZERO), false) {
+                Ok(()) => {}
+                Err(e)
+                    if matches!(
+                        e.kind(),
+                        io::ErrorKind::TimedOut | io::ErrorKind::Interrupted
+                    ) => {}
+                Err(e) => return Err(e),
+            }
+            if self.poll_entries() {
+                self.notifier.waker().wake();
+            }
+        }
+
+        // Do all allocation and bookkeeping before publishing, since SQPOLL
+        // can start executing immediately. Only the kernel can free SQ slots
+        // while we hold the driver, so the capacity check remains valid.
+        self.in_flight.reserve(N);
+        let slots = std::array::from_fn::<_, N, _>(|i| self.in_flight.insert(keys[i].as_raw()));
+        for (i, key) in keys.iter().enumerate() {
+            let mut key = key.borrow();
+            let extra = key.extra_mut().as_iour_mut();
+            extra.set_in_flight(slots[i]);
+            extra.linked = Some(if i == 0 {
+                Linked::Head
+            } else {
+                Linked::After(slots[i - 1])
+            });
+        }
+        // SAFETY: the keys retain every operation's parameters. The queue
+        // publishes its tail only after the entire group has been written.
+        unsafe { self.inner.submission().push_multiple(&entries) }
+            .expect("reserved space for the whole linked group");
+        // No allocation or fallible work after publication. The driver alone
+        // reclaims these references when processing CQEs or shutting down.
+        for key in keys {
+            key.into_raw();
+        }
+        Ok(())
+    }
+
     pub fn push(&mut self, key: ErasedKey) -> Poll<io::Result<usize>> {
         instrument!(compio_log::Level::TRACE, "push", ?key);
         let mut op_entry = key.borrow().create_entry::<false>();
@@ -443,10 +598,11 @@ impl Driver {
     }
 
     pub fn flush(&mut self) -> bool {
+        self.stage_cancellations();
         let succeed = self.submit_auto(Some(Duration::ZERO), false).is_ok();
         // If submission failed, return true to let the driver wake up
         // immediately.
-        !succeed | self.notifier.reset()
+        !succeed | self.notifier.reset() | !self.pending_cancels.is_empty()
     }
 
     pub fn poll(&mut self, timeout: Option<Duration>) -> io::Result<()> {
@@ -472,7 +628,11 @@ impl Driver {
             self.flags.remove(DriverFlags::NEED_PUSH_NOTIFIER);
         }
 
-        self.submit_auto(timeout, need_wait)?;
+        self.stage_cancellations();
+        // A full SQ may have prevented staging a cancellation needed to wake
+        // the very operation we would wait for. Flush once without waiting;
+        // the next driver pass can stage it after the kernel consumes the SQ.
+        self.submit_auto(timeout, need_wait && self.pending_cancels.is_empty())?;
 
         self.notifier.set_awake();
         self.poll_entries();
@@ -494,12 +654,17 @@ impl Driver {
 
     /// Take an op out of the `in_flight` map. The slot lives in the op's
     /// extra, and is only set while the op is in flight.
-    fn remove_in_flight(in_flight: &mut SlotMap<DefaultKey, usize>, key: &ErasedKey) {
-        let Some(slot) = key.borrow().extra_mut().as_iour_mut().take_in_flight() else {
-            return;
-        };
+    fn remove_in_flight(
+        in_flight: &mut SlotMap<DefaultKey, usize>,
+        key: &ErasedKey,
+    ) -> Option<DefaultKey> {
+        let mut op = key.borrow();
+        let extra = op.extra_mut().as_iour_mut();
+        let slot = extra.take_in_flight()?;
+        extra.linked = None;
         let removed = in_flight.remove(slot);
         debug_assert_eq!(removed, Some(key.as_raw()));
+        Some(slot)
     }
 }
 
@@ -557,8 +722,27 @@ fn create_result(result: i32) -> io::Result<usize> {
     }
 }
 
+fn unsupported_linked() -> io::Error {
+    io::Error::new(
+        io::ErrorKind::Unsupported,
+        "linked submission is not supported for this driver, opcode or queue capacity",
+    )
+}
+
+fn link_flags(hardlink: bool) -> io_uring::squeue::Flags {
+    if hardlink {
+        io_uring::squeue::Flags::IO_HARDLINK
+    } else {
+        io_uring::squeue::Flags::IO_LINK
+    }
+}
+
 fn timespec(duration: std::time::Duration) -> Timespec {
     Timespec::new()
         .sec(duration.as_secs())
         .nsec(duration.subsec_nanos())
 }
+
+/// Inline capacity for the linked-batch staging buffers, sized for eight
+/// fill/drain splice pairs (16 members) without heap allocation.
+const LINKED_BATCH_INLINE: usize = 16;