refactor(scheduler): use Own and RawRef for TaskQueue

Author: Paraworker

compio-runtime/src/runtime/scheduler/mod.rs

@@ -1,16 +1,17 @@
-use std::{cell::RefCell, future::Future, marker::PhantomData, rc::Rc, sync::Arc, task::Waker};
+use std::{cell::RefCell, future::Future, marker::PhantomData, rc::Rc, task::Waker};
 
 use async_task::{Runnable, Task};
 use compio_driver::NotifyHandle;
 use crossbeam_queue::SegQueue;
 use slab::Slab;
 
 use crate::runtime::scheduler::{
-    drop_hook::DropHook, local_queue::LocalQueue, send_wrapper::SendWrapper,
+    drop_hook::DropHook, local_queue::LocalQueue, raw_ref::Own, send_wrapper::SendWrapper,
 };
 
 mod drop_hook;
 mod local_queue;
+mod raw_ref;
 mod send_wrapper;
 
 /// A task queue consisting of a local queue and a synchronized queue.
@@ -99,7 +100,7 @@ impl TaskQueue {
 /// A scheduler for managing and executing tasks.
 pub(crate) struct Scheduler {
     /// Queue for scheduled tasks.
-    task_queue: Arc<TaskQueue>,
+    task_queue: Own<TaskQueue>,
 
     /// `Waker` of active tasks.
     active_tasks: Rc<RefCell<Slab<Waker>>>,
@@ -115,7 +116,7 @@ impl Scheduler {
     /// Creates a new `Scheduler`.
     pub(crate) fn new(event_interval: usize) -> Self {
         Self {
-            task_queue: Arc::new(TaskQueue::new()),
+            task_queue: Own::new(TaskQueue::new()),
             active_tasks: Rc::new(RefCell::new(Slab::new())),
             event_interval,
             _local_marker: PhantomData,
@@ -149,16 +150,15 @@ impl Scheduler {
         };
 
         let schedule = {
-            // The schedule closure is managed by the `Waker` and may be dropped on another
-            // thread, so use `Weak` to ensure the `TaskQueue` is always dropped
-            // on the creator thread.
-            let task_queue = Arc::downgrade(&self.task_queue);
+            let task_queue = self.task_queue.raw_ref();
 
             move |runnable| {
-                // The `upgrade()` never fails because all tasks are dropped when the
-                // `Scheduler` is dropped, if a `Waker` is used after that, the
-                // schedule closure will never be called.
-                task_queue.upgrade().unwrap().push(runnable, &notify);
+                // SAFETY:
+                // `Scheduler` drains and drops all `Runnable`s before it is dropped.
+                // If this closure is still invoked and can push a `Runnable`, the
+                // `Scheduler` is necessarily still alive, so the `task_queue` reference
+                // obtained from `raw_ref()` is valid.
+                unsafe { task_queue.as_ref() }.push(runnable, &notify);
             }
         };
 

compio-runtime/src/runtime/scheduler/raw_ref.rs

@@ -0,0 +1,57 @@
+use std::{ops::Deref, ptr::NonNull};
+
+/// A value with ownership.
+pub(crate) struct Own<T: ?Sized>(Box<T>);
+
+impl<T> Own<T> {
+    /// Creates a new [`Own`].
+    pub(crate) fn new(value: T) -> Self {
+        Own(Box::new(value))
+    }
+}
+
+impl<T: ?Sized> Own<T> {
+    /// Returns a [`RawRef`] to the owned value.
+    pub(crate) fn raw_ref(&self) -> RawRef<T> {
+        RawRef(NonNull::from(&*self.0))
+    }
+}
+
+impl<T: ?Sized> Deref for Own<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+/// A reference to an owned value without lifetime tracking.
+pub(crate) struct RawRef<T: ?Sized>(NonNull<T>);
+
+impl<T: ?Sized> RawRef<T> {
+    /// Returns a shared reference to the value.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure the associated [`Own<T>`] outlives the returned
+    /// reference.
+    pub(crate) const unsafe fn as_ref(&self) -> &T {
+        // SAFETY:
+        // - The `NonNull` is created from a valid reference in `Own::raw_ref()`.
+        // - Only shared reference is returned, so aliasing rules are not violated.
+        // - The validity of the returned reference is ensured by the caller.
+        unsafe { self.0.as_ref() }
+    }
+}
+
+/// `Sync` and `Send` implementations follow `&T`.
+unsafe impl<T: ?Sized + Sync> Sync for RawRef<T> {}
+unsafe impl<T: ?Sized + Sync> Send for RawRef<T> {}
+
+impl<T: ?Sized> Clone for RawRef<T> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+impl<T: ?Sized> Copy for RawRef<T> {}