Add sanitizer builds.

Author: sugarraysam

ringolo/src/context/core.rs

@@ -122,25 +122,6 @@ impl Core {
         self.pending_ios.fetch_sub(1, Ordering::Release);
     }
 
-    pub(crate) fn increment_task_pending_ios(&self, waker: &Waker) {
-        self.modify_task_pending_ios(waker, 1);
-    }
-
-    pub(crate) fn decrement_task_pending_ios(&self, waker: &Waker) {
-        self.modify_task_pending_ios(waker, -1);
-    }
-
-    fn modify_task_pending_ios(&self, waker: &Waker, delta: i32) {
-        // Don't track owned_resources for the root_future. The reason is because it
-        // *is not backed* by a regular task and is never stealable by other threads.
-        if !self.is_polling_root() {
-            unsafe {
-                let ptr = NonNull::new_unchecked(waker.data() as *mut Header);
-                Header::modify_pending_ios(ptr, delta);
-            }
-        }
-    }
-
     pub(crate) fn increment_task_owned_resources(&self, waker: &Waker) -> Option<NonNull<Header>> {
         // Don't track owned_resources for the root_future. The reason is because it
         // *is not backed* by a regular task and is never stealable by other threads.
@@ -159,11 +140,8 @@ impl Core {
 #[cfg(test)]
 mod tests {
     use crate as ringolo;
-    use crate::context::{self, PendingIoOp, init_stealing_context};
-    use crate::runtime::runtime::Kind;
-    use crate::test_utils::{init_local_runtime_and_context, init_stealing_runtime_and_context};
+    use crate::context::{self, PendingIoOp};
     use anyhow::Result;
-    use rstest::rstest;
 
     #[ringolo::test]
     async fn test_pending_io_on_core_and_shared() -> Result<()> {

ringolo/src/context/maintenance/cleanup.rs

@@ -238,7 +238,6 @@ impl CleanupHandler {
     }
 }
 
-#[doc(hidden)]
 fn cvt(ops: &[CleanupOp]) -> Vec<AnyOp<'_>> {
     ops.iter().map(|op| op.into()).collect()
 }

ringolo/src/context/mod.rs

@@ -144,6 +144,30 @@ where
     })
 }
 
+#[inline(always)]
+pub(crate) fn with_context<F, R>(f: F) -> R
+where
+    F: FnOnce(&Context) -> R,
+{
+    CONTEXT.with(|ctx| {
+        let root = ctx.get().expect("Context not initialized").borrow();
+        f(&root.context)
+    })
+}
+
+#[inline(always)]
+pub(crate) fn try_with_context<F, R>(f: F) -> Option<R>
+where
+    F: FnOnce(&Context) -> R,
+{
+    CONTEXT
+        .try_with(|ctx| {
+            let root = ctx.get().expect("Context not initialized").borrow();
+            f(&root.context)
+        })
+        .ok()
+}
+
 #[inline(always)]
 pub(crate) fn with_core<F, R>(f: F) -> R
 where
@@ -155,6 +179,17 @@ where
     })
 }
 
+#[inline(always)]
+pub(crate) fn try_with_core<F, R>(f: F) -> Option<R>
+where
+    F: FnOnce(&Core) -> R,
+{
+    try_with_context(|outer| match outer {
+        Context::Local(c) => c.with_core(f),
+        Context::Stealing(c) => c.with_core(f),
+    })
+}
+
 #[inline(always)]
 pub(crate) fn with_shared<F, R>(f: F) -> R
 where
@@ -233,18 +268,6 @@ pub(crate) fn set_current_task(task: Option<Arc<TaskNode>>) -> Option<Arc<TaskNo
     with_core(|core| core.current_task.replace(task))
 }
 
-// Private helpers.
-#[inline(always)]
-pub(crate) fn with_context<F, R>(f: F) -> R
-where
-    F: FnOnce(&Context) -> R,
-{
-    CONTEXT.with(|ctx| {
-        let root = ctx.get().expect("Context not initialized").borrow();
-        f(&root.context)
-    })
-}
-
 #[inline(always)]
 pub(crate) fn with_scheduler<F, R>(f: F) -> R
 where

ringolo/src/context/ring.rs

@@ -146,38 +146,48 @@ impl SingleIssuerRing {
         self.sq().taskrun()
     }
 
-    // Will busy loop until `num_to_complete` has been achieved. It is the caller's
+    // Will busy loop until `num_to_visit` CQE have been visited. It is the caller's
     // responsibility to make sure the CQ will see that many completions, otherwise
     // this will result in an infinite loop.
     pub(crate) fn process_cqes(
         &mut self,
         slab: &mut RawSqeSlab,
-        num_to_complete: Option<usize>,
+        num_to_visit: Option<usize>,
     ) -> Result<usize> {
+        // We need to split `num_visited`/`num_completed` counters because if we
+        // cancel an SQE inflight, the entry might be dropped from the slab while
+        // still producing a CQE. In that case, the loop would only increment the
+        // `num_visited` counter.
+        let mut num_visited = 0;
         let mut num_completed = 0;
-        let mut should_sync = false;
 
-        let to_complete = num_to_complete.unwrap_or(self.cq().len());
+        let mut cq = self.cq();
+        let to_visit = num_to_visit.unwrap_or(cq.len());
 
-        while num_completed < to_complete {
-            let mut cq = self.cq();
+        let mut should_sync = false;
 
-            // Avoid syncing on first pass
+        while num_visited < to_visit {
+            // Avoid syncing on first pass.
             if should_sync {
+                // If we get here, we are spinning waiting for completions so
+                // let's give the CPU a tiny breather.
+                std::thread::yield_now();
                 cq.sync();
             }
 
-            for cqe in cq {
+            for cqe in &mut cq {
+                num_visited += 1;
+
                 let raw_sqe = match slab.get_mut(cqe.user_data() as usize) {
                     Err(e) => {
                         eprintln!("CQE user data not found in RawSqeSlab: {:?}", e);
                         continue;
                     }
                     Ok(sqe) => {
-                        if !matches!(sqe.get_state(), RawSqeState::Pending | RawSqeState::Ready) {
+                        if !matches!(sqe.state, RawSqeState::Pending | RawSqeState::Ready) {
                             // Ignore unknown CQEs which might have valid index in
                             // the Slab. Can this even happen?
-                            eprintln!("SQE in unexpected state: {:?}", sqe.get_state());
+                            eprintln!("SQE in unexpected state: {:?}", sqe.state);
                             continue;
                         }
                         sqe
@@ -189,7 +199,12 @@ impl SingleIssuerRing {
                 if let Some(CompletionEffect::WakeHead { head }) =
                     raw_sqe.on_completion(cqe.result(), cqe.flags().into())?
                 {
-                    slab.get_mut(head)?.wake()?;
+                    slab.get_mut(head)?.wake_by_ref()?;
+                }
+
+                // Exit early if user provided an override for `to_visit`.
+                if num_visited >= to_visit {
+                    break;
                 }
             }
 
@@ -273,6 +288,8 @@ mod tests {
         let builder = Builder::new_local().sq_ring_size(sq_ring_size);
         let (_runtime, _scheduler) = init_local_runtime_and_context(Some(builder))?;
 
+        let (waker, _) = mock_waker();
+
         context::with_slab_and_ring_mut(|slab, ring| -> Result<()> {
             {
                 let sq = ring.sq();
@@ -289,7 +306,7 @@ mod tests {
 
                 (0..n).for_each(|i| {
                     nops.push(nop().user_data(indices[i] as u64));
-                    raws.push(RawSqe::new(CompletionHandler::new_single()));
+                    raws.push(RawSqe::new(&waker, CompletionHandler::new_single()));
                 });
 
                 let _ = batch.commit(raws)?;
@@ -342,15 +359,9 @@ mod tests {
                 (0..n).for_each(|i| {
                     let nop = nop().user_data(indices[i] as u64);
 
-                    let mut raw = RawSqe::new(CompletionHandler::new_single());
-                    raw.set_waker(&waker);
-
+                    let raw = RawSqe::new(&waker, CompletionHandler::new_single());
                     nops.push(nop);
                     raws.push(raw);
-
-                    // We use single SQE even though we use batch API on slab,
-                    // anyways a bit messed up but need to count N pending ios.
-                    context::with_core(|core| core.increment_pending_ios());
                 });
 
                 let _ = batch.commit(raws)?;

ringolo/src/context/slab.rs

@@ -12,6 +12,7 @@ use std::ops::{Deref, DerefMut};
 /// This structure manages the lifecycle of raw pointers used in I/O requests.
 /// It utilizes a "Reserve-Commit" pattern to ensure that entries are only
 /// permanently occupied if the I/O submission logic succeeds.
+#[derive(Debug)]
 pub(crate) struct RawSqeSlab {
     slab: slab::Slab<RawSqe>,
 }
@@ -191,14 +192,15 @@ mod tests {
     #[test]
     fn test_insert_and_len() -> Result<()> {
         let (_runtime, _scheduler) = init_local_runtime_and_context(None)?;
+        let (waker, _) = mock_waker();
 
         let n_sqes = 3;
         let mut slab = RawSqeSlab::new(n_sqes);
 
         for i in 1..=n_sqes {
             let reserved = slab.reserve_entry()?;
             let idx = reserved.key();
-            reserved.commit(RawSqe::new(CompletionHandler::new_single()));
+            reserved.commit(RawSqe::new(&waker, CompletionHandler::new_single()));
 
             assert!(slab.get(idx).is_ok());
             assert_eq!(slab.len(), i);
@@ -210,13 +212,14 @@ mod tests {
     #[test]
     fn test_slab_full() -> Result<()> {
         let (_runtime, _scheduler) = init_local_runtime_and_context(None)?;
+        let (waker, _) = mock_waker();
 
         let n_sqes = 3;
         let mut slab = RawSqeSlab::new(n_sqes - 1);
 
         for i in 1..=n_sqes {
             let res = slab.reserve_entry().map(|reserved| {
-                reserved.commit(RawSqe::new(CompletionHandler::new_single()));
+                reserved.commit(RawSqe::new(&waker, CompletionHandler::new_single()));
             });
 
             if i == n_sqes {
@@ -232,6 +235,7 @@ mod tests {
     #[test]
     fn test_reserve_batch() -> Result<()> {
         let (_runtime, _scheduler) = init_local_runtime_and_context(None)?;
+        let (waker, _) = mock_waker();
 
         context::with_slab_mut(|slab| -> Result<()> {
             let n_sqes = 16;
@@ -256,7 +260,7 @@ mod tests {
 
                 // Commit the batch
                 let entries = (0..n_sqes)
-                    .map(|_| RawSqe::new(CompletionHandler::new_single()))
+                    .map(|_| RawSqe::new(&waker, CompletionHandler::new_single()))
                     .collect::<SmallVec<[_; SPILL_TO_HEAP_THRESHOLD]>>();
 
                 batch.commit(entries)?;

ringolo/src/runtime/cancel.rs

@@ -174,12 +174,13 @@ mod tests {
             let stats = ringolo::recursive_cancel_all_metadata(&TaskMetadata::from(["dead"]))?;
             assert_eq!(stats.cancelled, 1);
             assert_eq!(stats.visited, n + 1);
+            drop(cancelled);
 
-            // Check orphanage - only happens after we drop the join_handle
+            // Next poll on the cancelled task creates orphans
             let orphan_root = get_orphan_root();
-            assert_eq!(orphan_root.num_children(), 0);
-            drop(cancelled);
-            assert_eq!(orphan_root.num_children(), n);
+            while orphan_root.num_children() < n {
+                assert!(YieldNow::new(Some(AddMode::Fifo)).await.is_ok());
+            }
 
             // Now cancel remaining orphans
             let stats = ringolo::recursive_cancel_all_orphans()?;

ringolo/src/runtime/local/worker.rs

@@ -7,7 +7,7 @@ use crate::runtime::{AddMode, EventLoop};
 use crate::runtime::{Ticker, TickerData, TickerEvents};
 use crate::task::TaskNodeGuard;
 use crate::task::ThreadId;
-use anyhow::{Result, anyhow};
+use anyhow::{anyhow, Result};
 use std::cell::RefCell;
 use std::collections::VecDeque;
 use std::pin::pin;
@@ -56,7 +56,7 @@ impl EventLoop for Worker {
         // the end corresponds to the front.
         match mode {
             AddMode::Fifo => q.push_front(task),
-            AddMode::Lifo => q.push_back(task),
+            AddMode::Lifo | AddMode::Cancel => q.push_back(task),
         }
     }
 
@@ -105,7 +105,9 @@ impl Worker {
                 })?;
             } else {
                 match root_result.is_some() {
-                    true => return Ok(root_result),
+                    true => {
+                        return Ok(root_result);
+                    }
                     false => scheduler.set_root_woken(),
                 }
             }

ringolo/src/runtime/mod.rs

@@ -153,6 +153,12 @@ pub enum AddMode {
     /// polled in the order they were yielded, but may be less cache-efficient
     /// than `Lifo`.
     Fifo,
+
+    /// Schedules a task to be cancelled.
+    ///
+    /// We leverage the `RawTask::remote_abort` API to make sure we cancel the
+    /// task on the owning thread so we can free resources in the right place.
+    Cancel,
 }
 
 #[doc(hidden)]

ringolo/src/runtime/registry.rs

@@ -229,7 +229,7 @@ impl<S: Schedule> OwnedTasks<S> {
 /// Expose functionality of the TaskRegistry common to any scheduler.
 pub(crate) trait TaskRegistry: Send + Sync + std::fmt::Debug {
     /// Allow shutting down a specific task.
-    fn shutdown(&self, id: &Id);
+    fn remote_abort(&self, id: &Id);
 
     /// Returns true if the registry was closed as part of runtime shutdown.
     fn is_closed(&self) -> bool;
@@ -239,9 +239,9 @@ pub(crate) trait TaskRegistry: Send + Sync + std::fmt::Debug {
 }
 
 impl<S: Schedule> TaskRegistry for OwnedTasks<S> {
-    fn shutdown(&self, id: &Id) {
-        if let Some((_, task)) = self.tasks.remove(id) {
-            task.shutdown()
+    fn remote_abort(&self, id: &Id) {
+        if let Some(task) = self.tasks.get(id) {
+            task.remote_abort();
         }
     }
 

ringolo/src/runtime/stealing/root_worker.rs

@@ -56,10 +56,11 @@ impl RootWorker {
 impl EventLoop for RootWorker {
     type Task = StealableTask;
 
-    fn add_task(&self, task: Self::Task, _mode: AddMode) {
+    fn add_task(&self, task: Self::Task, mode: AddMode) {
         debug_assert!(
-            task.is_maintenance_task(),
-            "Only maintenance task can be scheduled on root_worker."
+            task.is_maintenance_task() || matches!(mode, AddMode::Cancel),
+            "Can only schedule maintenance task or unclaimed tasks
+                        that needs cancellation on root_worker"
         );
         task.set_owner_id(self.thread_id);
         self.pollable.borrow_mut().push_back(task);
@@ -101,10 +102,6 @@ impl RootWorker {
 
         loop {
             if let Some(task) = self.find_task() {
-                debug_assert!(
-                    task.is_maintenance_task(),
-                    "Only maintenance task should be running on root_worker"
-                );
                 task.run();
             } else if ctx.with_core(|core| core.get_pending_ios() > 0) {
                 // Block thread waiting for next completion.