feat(runtime): faster SendWrapper with cached thread id

Author: Berrysoft

compio-runtime/Cargo.toml

@@ -40,7 +40,6 @@ crossbeam-queue = { workspace = true }
 futures-util = { workspace = true }
 once_cell = { workspace = true }
 scoped-tls = "1.0.1"
-send_wrapper = "0.6.0"
 slab = { workspace = true, optional = true }
 smallvec = "1.11.1"
 socket2 = { workspace = true }

compio-runtime/src/runtime/mod.rs

@@ -20,13 +20,15 @@ use compio_driver::{
 use compio_log::{debug, instrument};
 use crossbeam_queue::SegQueue;
 use futures_util::{future::Either, FutureExt};
-use send_wrapper::SendWrapper;
 use smallvec::SmallVec;
 
 pub(crate) mod op;
 #[cfg(feature = "time")]
 pub(crate) mod time;
 
+mod send_wrapper;
+use send_wrapper::SendWrapper;
+
 #[cfg(feature = "time")]
 use crate::runtime::time::{TimerFuture, TimerRuntime};
 use crate::{runtime::op::OpFuture, BufResult};
@@ -120,8 +122,8 @@ impl Runtime {
             .handle()
             .expect("cannot create notify handle of the proactor");
         let schedule = move |runnable| {
-            if local_runnables.valid() {
-                local_runnables.borrow_mut().push_back(runnable);
+            if let Some(runnables) = local_runnables.get() {
+                runnables.borrow_mut().push_back(runnable);
             } else {
                 sync_runnables.push(runnable);
                 handle.notify().ok();
@@ -136,9 +138,8 @@ impl Runtime {
     ///
     /// Run the scheduled tasks.
     pub fn run(&self) {
-        use std::ops::Deref;
-
-        let local_runnables = self.local_runnables.deref().deref();
+        // SAFETY: self is !Send + !Sync.
+        let local_runnables = unsafe { self.local_runnables.get_unchecked() };
         loop {
             let next_task = local_runnables.borrow_mut().pop_front();
             let has_local_task = next_task.is_some();

compio-runtime/src/runtime/send_wrapper.rs

@@ -0,0 +1,99 @@
+use std::{
+    cell::Cell,
+    mem::{self, ManuallyDrop},
+    thread::{self, ThreadId},
+};
+
+thread_local! {
+    static THREAD_ID: Cell<ThreadId> = Cell::new(thread::current().id());
+}
+
+/// A wrapper that copied from `send_wrapper` crate, with our own optimizations.
+pub struct SendWrapper<T> {
+    data: ManuallyDrop<T>,
+    thread_id: ThreadId,
+}
+
+impl<T> SendWrapper<T> {
+    /// Create a `SendWrapper<T>` wrapper around a value of type `T`.
+    /// The wrapper takes ownership of the value.
+    pub fn new(data: T) -> SendWrapper<T> {
+        SendWrapper {
+            data: ManuallyDrop::new(data),
+            thread_id: THREAD_ID.get(),
+        }
+    }
+
+    /// Returns `true` if the value can be safely accessed from within the
+    /// current thread.
+    pub fn valid(&self) -> bool {
+        self.thread_id == THREAD_ID.get()
+    }
+
+    /// Returns a reference to the contained value.
+    ///
+    /// # Safety
+    ///
+    /// The caller should be in the same thread as the creator.
+    pub unsafe fn get_unchecked(&self) -> &T {
+        &self.data
+    }
+
+    /// Returns a reference to the contained value, if valid.
+    pub fn get(&self) -> Option<&T> {
+        if self.valid() { Some(&self.data) } else { None }
+    }
+}
+
+unsafe impl<T> Send for SendWrapper<T> {}
+unsafe impl<T> Sync for SendWrapper<T> {}
+
+impl<T> Drop for SendWrapper<T> {
+    /// Drops the contained value.
+    ///
+    /// # Panics
+    ///
+    /// Dropping panics if it is done from a different thread than the one the
+    /// `SendWrapper<T>` instance has been created with.
+    ///
+    /// Exceptions:
+    /// - There is no extra panic if the thread is already panicking/unwinding.
+    ///   This is because otherwise there would be double panics (usually
+    ///   resulting in an abort) when dereferencing from a wrong thread.
+    /// - If `T` has a trivial drop ([`needs_drop::<T>()`] is false) then this
+    ///   method never panics.
+    ///
+    /// [`needs_drop::<T>()`]: std::mem::needs_drop
+    #[track_caller]
+    fn drop(&mut self) {
+        // If the drop is trivial (`needs_drop` = false), then dropping `T` can't access
+        // it and so it can be safely dropped on any thread.
+        if !mem::needs_drop::<T>() || self.valid() {
+            unsafe {
+                // Drop the inner value
+                //
+                // Safety:
+                // - We've just checked that it's valid to drop `T` on this thread
+                // - We only move out from `self.data` here and in drop, so `self.data` is
+                //   present
+                ManuallyDrop::drop(&mut self.data);
+            }
+        } else {
+            invalid_drop()
+        }
+    }
+}
+
+#[cold]
+#[inline(never)]
+#[track_caller]
+fn invalid_drop() {
+    const DROP_ERROR: &str = "Dropped SendWrapper<T> variable from a thread different to the one \
+                              it has been created with.";
+
+    if !thread::panicking() {
+        // panic because of dropping from wrong thread
+        // only do this while not unwinding (could be caused by deref from wrong thread)
+        panic!("{}", DROP_ERROR)
+    }
+}