Skip to content

Commit

Permalink
Nip some potentially unsound behavior in the bud
Browse files Browse the repository at this point in the history
  • Loading branch information
notgull committed Jul 2, 2023
1 parent e93561c commit e436d2b
Showing 1 changed file with 17 additions and 2 deletions.
19 changes: 17 additions & 2 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -231,7 +231,7 @@ impl<'a> Executor<'a> {
let mut rng = fastrand::Rng::new();

// Set the local queue while we're running.
LocalQueue::set(&runner.local, {
LocalQueue::set(self.state(), &runner.local, {
let runner = &runner;
async move {
// A future that runs tasks forever.
Expand Down Expand Up @@ -262,6 +262,11 @@ impl<'a> Executor<'a> {

// Try to push into the local queue.
LocalQueue::with(|local_queue| {
// Make sure that we don't accidentally push to an executor that isn't ours.
if !std::ptr::eq(local_queue.state, &*state) {
return;
}

if let Err(e) = local_queue.queue.push(runnable.take().unwrap()) {
runnable = Some(e.into_inner());
return;
Expand Down Expand Up @@ -844,6 +849,11 @@ impl Drop for Runner<'_> {

/// The state of the currently running local queue.
struct LocalQueue {
/// The pointer to the state of the executor.
///
/// Used to make sure we don't push runnables to the wrong executor.
state: *const State,

/// The concurrent queue.
queue: Arc<ConcurrentQueue<Runnable>>,

Expand All @@ -861,14 +871,19 @@ impl LocalQueue {

impl LocalQueue {
/// Run a function with a set local queue.
async fn set<F>(queue: &Arc<ConcurrentQueue<Runnable>>, fut: F) -> F::Output
async fn set<F>(
state: &State,
queue: &Arc<ConcurrentQueue<Runnable>>,
fut: F,
) -> F::Output
where
F: Future,
{
// Store the local queue and the current waker.
let mut old = with_waker(|waker| {
LOCAL_QUEUE.with(move |slot| {
slot.borrow_mut().replace(LocalQueue {
state: state as *const State,
queue: queue.clone(),
waker: waker.clone(),
})
Expand Down

0 comments on commit e436d2b

Please sign in to comment.