task: reduce duplication between basic_scheduler and LocalSet

This commit factors out the common task queue logic in the basic
scheduler runtime and the `LocalSet` struct in `tokio::task`. This is
because as more work was done on the `LocalSet`, it has gotten closer
and closer to the basic scheduler in behavior, and factoring out the
shared code reduces the risk of errors caused by `LocalSet` not doing
something that the basic scheduler does.

In particular, I noticed the basic scheduler has a flag that indicates
the remote queue has been closed, which is set when dropping the
scheduler. This prevents tasks from being added after the scheduler has
started shutting down, stopping a potential task leak. Rather than
duplicating this code in `LocalSet`, I thought it was probably better to
factor it out into a shared type.

There are a few cases where there are small differences in behavior,
though, so there is still a need for separate types implemented _using_
the new `Queues` struct. However, it should cover most of the identical
code.

Signed-off-by: Eliza Weisman <[email protected]>

tokio/src/runtime/basic_scheduler.rs

@@ -7,7 +7,7 @@ use std::fmt;
 use std::future::Future;
 use std::mem::ManuallyDrop;
 use std::ptr;
-use std::sync::{Arc, Mutex};
+use std::sync::{Arc, Mutex, MutexGuard};
 use std::task::{RawWaker, RawWakerVTable, Waker};
 use std::time::Duration;
 
@@ -29,15 +29,13 @@ pub(crate) struct Spawner {
     scheduler: Arc<SchedulerPriv>,
 }
 
-/// The scheduler component.
+pub(crate) struct Queues<S: 'static> {
-pub(super) struct SchedulerPriv {
     /// List of all active tasks spawned onto this executor.
     ///
     /// # Safety
     ///
     /// Must only be accessed from the primary thread
-    owned_tasks: UnsafeCell<task::OwnedList<Self>>,
+    owned_tasks: UnsafeCell<task::OwnedList<S>>,
-
     /// Local run queue.
     ///
     /// Tasks notified from the current thread are pushed into this queue.
@@ -46,16 +44,20 @@ pub(super) struct SchedulerPriv {
     ///
     /// References should not be handed out. Only call `push` / `pop` functions.
     /// Only call from the owning thread.
-    local_queue: UnsafeCell<VecDeque<Task<SchedulerPriv>>>,
+    local_queue: UnsafeCell<VecDeque<Task<S>>>,
 
     /// Remote run queue.
     ///
     /// Tasks notified from another thread are pushed into this queue.
-    remote_queue: Mutex<RemoteQueue>,
+    remote_queue: Mutex<RemoteQueue<S>>,
 
     /// Tasks pending drop
-    pending_drop: task::TransferStack<Self>,
+    pending_drop: task::TransferStack<S>,
+}
 
+/// The scheduler component.
+pub(super) struct SchedulerPriv {
+    queues: Queues<Self>,
     /// Unpark the blocked thread
     unpark: Box<dyn Unpark>,
 }
@@ -73,10 +75,9 @@ struct LocalState<P> {
     park: P,
 }
 
-#[derive(Debug)]
+pub(crate) struct RemoteQueue<S: 'static> {
-struct RemoteQueue {
     /// FIFO list of tasks
-    queue: VecDeque<Task<SchedulerPriv>>,
+    queue: VecDeque<Task<S>>,
 
     /// `true` when a task can be pushed into the queue, false otherwise.
     open: bool,
@@ -85,9 +86,6 @@ struct RemoteQueue {
 /// Max number of tasks to poll per tick.
 const MAX_TASKS_PER_TICK: usize = 61;
 
-/// How often to check the remote queue first
-const CHECK_REMOTE_INTERVAL: u8 = 13;
-
 thread_local! {
     static ACTIVE: Cell<*const SchedulerPriv> = Cell::new(ptr::null())
 }
@@ -101,13 +99,7 @@ where
 
         BasicScheduler {
             scheduler: Arc::new(SchedulerPriv {
-                owned_tasks: UnsafeCell::new(task::OwnedList::new()),
+                queues: Queues::new(),
-                local_queue: UnsafeCell::new(VecDeque::with_capacity(64)),
-                remote_queue: Mutex::new(RemoteQueue {
-                    queue: VecDeque::with_capacity(64),
-                    open: true,
-                }),
-                pending_drop: task::TransferStack::new(),
                 unpark: Box::new(unpark),
             }),
             local: LocalState { tick: 0, park },
@@ -155,9 +147,7 @@ where
 
         // Track the current scheduler
         let _guard = ACTIVE.with(|cell| {
-            let guard = Guard {
+            let guard = Guard { old: cell.get() };
-                old: cell.get(),
-            };
 
             cell.set(scheduler as *const SchedulerPriv);
 
@@ -188,7 +178,9 @@ where
                 scheduler.tick(local);
 
                 // Maintenance work
-                scheduler.drain_pending_drop();
+                unsafe {
+                    scheduler.queues.drain_pending_drop();
+                }
             }
         })
     }
@@ -216,6 +208,221 @@ impl Spawner {
     }
 }
 
+impl<S> Queues<S>
+where
+    S: Schedule + 'static,
+{
+    pub(crate) const INITIAL_CAPACITY: usize = 64;
+
+    /// How often to check the remote queue first
+    pub(crate) const CHECK_REMOTE_INTERVAL: u8 = 13;
+
+    pub(crate) fn new() -> Self {
+        Self {
+            owned_tasks: UnsafeCell::new(task::OwnedList::new()),
+            local_queue: UnsafeCell::new(VecDeque::with_capacity(Self::INITIAL_CAPACITY)),
+            pending_drop: task::TransferStack::new(),
+            remote_queue: Mutex::new(RemoteQueue {
+                queue: VecDeque::with_capacity(Self::INITIAL_CAPACITY),
+                open: true,
+            }),
+        }
+    }
+
+    /// Add a new task to the scheduler.
+    ///
+    /// # Safety
+    ///
+    /// This *must* be called only from the thread that owns the scheduler.
+    pub(crate) unsafe fn add_task(&self, task: &Task<S>) {
+        (*self.owned_tasks.get()).insert(task);
+    }
+
+    /// Push a task to the local queue.
+    ///
+    /// # Safety
+    ///
+    /// This *must* be called only from the thread that owns the scheduler.
+    pub(crate) unsafe fn push_local(&self, task: Task<S>) {
+        (*self.local_queue.get()).push_back(task);
+    }
+
+    /// Push a task to the local queue.
+    ///
+    /// # Safety
+    ///
+    /// This *must* be called only from the thread that owns the scheduler.
+    pub(crate) unsafe fn release_local(&self, task: &Task<S>) {
+        (*self.owned_tasks.get()).remove(task);
+    }
+
+    /// Lock the remote queue, returning a `MutexGuard`.
+    ///
+    /// This can be used to push to the remote queue and perform other
+    /// operations while holding the lock.
+    ///
+    /// # Panics
+    ///
+    /// If the remote queue mutex is poisoned.
+    pub(crate) fn remote(&self) -> MutexGuard<'_, RemoteQueue<S>> {
+        self.remote_queue
+            .lock()
+            .expect("failed to lock remote queue")
+    }
+
+    /// Release a task from outside of the thread that owns the scheduler.
+    ///
+    /// This simply pushes the task to the pending drop queue.
+    pub(crate) fn release_remote(&self, task: Task<S>) {
+        self.pending_drop.push(task);
+    }
+
+    /// Returns the next task from the remote *or* local queue.
+    ///
+    /// Typically, this checks the local queue before the remote queue, and only
+    /// checks the remote queue if the local queue is empty. However, to avoid
+    /// starving the remote queue, it is checked first every
+    /// `CHECK_REMOTE_INTERVAL` ticks.
+    ///
+    /// # Safety
+    ///
+    /// This *must* be called only from the thread that owns the scheduler.
+    pub(crate) unsafe fn next_task(&self, tick: u8) -> Option<Task<S>> {
+        if 0 == tick % Self::CHECK_REMOTE_INTERVAL {
+            self.next_remote_task().or_else(|| self.next_local_task())
+        } else {
+            self.next_local_task().or_else(|| self.next_remote_task())
+        }
+    }
+
+    /// Returns the next task from the local queue.
+    ///
+    /// # Safety
+    ///
+    /// This *must* be called only from the thread that owns the scheduler.
+    pub(crate) unsafe fn next_local_task(&self) -> Option<Task<S>> {
+        (*self.local_queue.get()).pop_front()
+    }
+
+    /// Returns the next task from the remote queue.
+    ///
+    /// # Panics
+    /// - If the mutex around the remote queue is poisoned _and_ the current
+    ///   thread is not already panicking. This is safe to call in a `Drop` impl.
+    pub(crate) fn next_remote_task(&self) -> Option<Task<S>> {
+        // there is no semantic information in the `PoisonError`, and it
+        // doesn't implement `Debug`, but clippy thinks that it's bad to
+        // match all errors here...
+        #[allow(clippy::match_wild_err_arm)]
+        let mut lock = match self.remote_queue.lock() {
+            // If the lock is poisoned, but the thread is already panicking,
+            // avoid a double panic. This is necessary since `next_task` (which
+            // calls `next_remote_task`) can be called in the `Drop` impl.
+            Err(_) if std::thread::panicking() => return None,
+            Err(_) => panic!("mutex poisoned"),
+            Ok(lock) => lock,
+        };
+        lock.queue.pop_front()
+    }
+
+    /// Returns true if any owned tasks are still bound to this scheduler.
+    ///
+    /// # Safety
+    ///
+    /// This *must* be called only from the thread that owns the scheduler.
+    pub(crate) unsafe fn has_tasks_remaining(&self) -> bool {
+        !(*self.owned_tasks.get()).is_empty()
+    }
+
+    /// Drain any tasks that have previously been released from other threads.
+    ///
+    /// # Safety
+    ///
+    /// This *must* be called only from the thread that owns the scheduler.
+    pub(crate) unsafe fn drain_pending_drop(&self) {
+        for task in self.pending_drop.drain() {
+            (*self.owned_tasks.get()).remove(&task);
+            drop(task);
+        }
+    }
+
+    /// Shut down the scheduler's owned task list.
+    ///
+    /// # Safety
+    ///
+    /// This *must* be called only from the thread that owns the scheduler.
+    pub(crate) unsafe fn shutdown(&self) {
+        (*self.owned_tasks.get()).shutdown();
+    }
+
+    /// Drain the remote queue, and shut down its tasks.
+    ///
+    /// This closes the remote queue. Any additional tasks added to it will be
+    /// shut down instead.
+    ///
+    /// # Panics
+    /// - If the mutex around the remote queue is poisoned _and_ the current
+    ///   thread is not already panicking. This is safe to call in a `Drop` impl.
+    pub(crate) fn close_remote(&self) {
+        #[allow(clippy::match_wild_err_arm)]
+        let mut lock = match self.remote_queue.lock() {
+            // If the lock is poisoned, but the thread is already panicking,
+            // avoid a double panic. This is necessary since this fn can be
+            // called in a drop impl.
+            Err(_) if std::thread::panicking() => return,
+            Err(_) => panic!("mutex poisoned"),
+            Ok(lock) => lock,
+        };
+        lock.open = false;
+
+        while let Some(task) = lock.queue.pop_front() {
+            task.shutdown();
+        }
+    }
+}
+
+impl<S> fmt::Debug for Queues<S> {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt.debug_struct("Queues")
+            .field("owned_tasks", &self.owned_tasks)
+            .field("remote_queue", &self.remote_queue)
+            .field("local_queue", &self.local_queue)
+            .finish()
+    }
+}
+
+// === impl RemoteQueue ===
+
+impl<S> RemoteQueue<S>
+where
+    S: Schedule,
+{
+    /// Schedule a remote task.
+    ///
+    /// If the queue is open to accept new tasks, the task is pushed to the back
+    /// of the queue. Otherwise, if the queue is closed (the scheduler is
+    /// shutting down), the new task
+    /// will be shut down immediately.
+    pub(crate) fn schedule(&mut self, task: Task<S>) {
+        if self.open {
+            self.queue.push_back(task);
+        } else {
+            task.shutdown();
+        }
+    }
+}
+
+impl<S> fmt::Debug for RemoteQueue<S> {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt.debug_struct("RemoteQueue")
+            .field("queue", &self.queue)
+            .field("open", &self.open)
+            .finish()
+    }
+}
+
+// === impl SchedulerPriv ===
+
 impl SchedulerPriv {
     fn tick(&self, local: &mut LocalState<impl Park>) {
         for _ in 0..MAX_TASKS_PER_TICK {
@@ -225,7 +432,7 @@ impl SchedulerPriv {
             // Increment the tick
             local.tick = tick.wrapping_add(1);
 
-            let task = match self.next_task(tick) {
+            let task = match unsafe { self.queues.next_task(tick) } {
                 Some(task) => task,
                 None => {
                     local.park.park().ok().expect("failed to park");
@@ -235,7 +442,7 @@ impl SchedulerPriv {
 
             if let Some(task) = task.run(&mut || Some(self.into())) {
                 unsafe {
-                    self.schedule_local(task);
+                    self.queues.push_local(task);
                 }
             }
         }
@@ -247,15 +454,6 @@ impl SchedulerPriv {
             .expect("failed to park");
     }
 
-    fn drain_pending_drop(&self) {
-        for task in self.pending_drop.drain() {
-            unsafe {
-                (*self.owned_tasks.get()).remove(&task);
-            }
-            drop(task);
-        }
-    }
-
     /// # Safety
     ///
     /// Must be called from the same thread that holds the `BasicScheduler`
@@ -266,63 +464,36 @@ impl SchedulerPriv {
         F::Output: Send + 'static,
     {
         let (task, handle) = task::joinable(future);
-        self.schedule_local(task);
+        self.queues.push_local(task);
         handle
     }
-
-    unsafe fn schedule_local(&self, task: Task<Self>) {
-        (*self.local_queue.get()).push_back(task);
-    }
-
-    fn next_task(&self, tick: u8) -> Option<Task<Self>> {
-        if 0 == tick % CHECK_REMOTE_INTERVAL {
-            self.next_remote_task().or_else(|| self.next_local_task())
-        } else {
-            self.next_local_task().or_else(|| self.next_remote_task())
-        }
-    }
-
-    fn next_local_task(&self) -> Option<Task<Self>> {
-        unsafe { (*self.local_queue.get()).pop_front() }
-    }
-
-    fn next_remote_task(&self) -> Option<Task<Self>> {
-        self.remote_queue.lock().unwrap().queue.pop_front()
-    }
 }
 
 impl Schedule for SchedulerPriv {
     fn bind(&self, task: &Task<Self>) {
         unsafe {
-            (*self.owned_tasks.get()).insert(task);
+            self.queues.add_task(task);
         }
     }
 
     fn release(&self, task: Task<Self>) {
-        self.pending_drop.push(task);
+        self.queues.release_remote(task);
     }
 
     fn release_local(&self, task: &Task<Self>) {
         unsafe {
-            (*self.owned_tasks.get()).remove(task);
+            self.queues.release_local(task);
         }
     }
 
     fn schedule(&self, task: Task<Self>) {
-        let is_current = ACTIVE.with(|cell| {
+        let is_current = ACTIVE.with(|cell| cell.get() == self as *const SchedulerPriv);
-            cell.get() == self as *const SchedulerPriv
-        });
 
         if is_current {
-            unsafe { self.schedule_local(task) };
+            unsafe { self.queues.push_local(task) };
         } else {
-            let mut lock = self.remote_queue.lock().unwrap();
+            let mut lock = self.queues.remote();
-
+            lock.schedule(task);
-            if lock.open {
-                lock.queue.push_back(task);
-            } else {
-                task.shutdown();
-            }
 
             // while locked, call unpark
             self.unpark.unpark();
@@ -340,38 +511,37 @@ where
 {
     fn drop(&mut self) {
         // Close the remote queue
-        let mut lock = self.scheduler.remote_queue.lock().unwrap();
+        self.scheduler.queues.close_remote();
-        lock.open = false;
-
-        while let Some(task) = lock.queue.pop_front() {
-            task.shutdown();
-        }
-
-        drop(lock);
 
         // Drain all local tasks
-        while let Some(task) = self.scheduler.next_local_task() {
+        while let Some(task) = unsafe { self.scheduler.queues.next_local_task() } {
             task.shutdown();
         }
 
         // Release owned tasks
         unsafe {
-            (*self.scheduler.owned_tasks.get()).shutdown();
+            self.scheduler.queues.shutdown();
         }
 
-        self.scheduler.drain_pending_drop();
+        unsafe {
+            self.scheduler.queues.drain_pending_drop();
+        }
 
         // Wait until all tasks have been released.
-        while unsafe { !(*self.scheduler.owned_tasks.get()).is_empty() } {
+        while unsafe { self.scheduler.queues.has_tasks_remaining() } {
             self.local.park.park().ok().expect("park failed");
-            self.scheduler.drain_pending_drop();
+            unsafe {
+                self.scheduler.queues.drain_pending_drop();
+            }
         }
     }
 }
 
 impl fmt::Debug for SchedulerPriv {
     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
-        fmt.debug_struct("Scheduler").finish()
+        fmt.debug_struct("Scheduler")
+            .field("queues", &self.queues)
+            .finish()
     }
 }
 

tokio/src/runtime/mod.rs

@@ -183,7 +183,7 @@
 mod tests;
 
 cfg_rt_core! {
-    mod basic_scheduler;
+    pub(crate) mod basic_scheduler;
     use basic_scheduler::BasicScheduler;
 }
 

tokio/src/task/local.rs

@@ -1,15 +1,13 @@
 //! Runs `!Send` futures on the current thread.
+use crate::runtime::basic_scheduler::Queues;
 use crate::sync::AtomicWaker;
-use crate::task::{self, JoinHandle, Schedule, Task, TransferStack};
+use crate::task::{self, JoinHandle, Schedule, Task};
 
-use std::cell::{Cell, UnsafeCell};
+use std::cell::Cell;
-use std::collections::VecDeque;
-use std::fmt;
 use std::future::Future;
 use std::pin::Pin;
 use std::ptr::{self, NonNull};
 use std::rc::Rc;
-use std::sync::Mutex;
 use std::task::{Context, Poll};
 
 use pin_project_lite::pin_project;
@@ -82,33 +80,12 @@ cfg_rt_util! {
         scheduler: Rc<Scheduler>,
     }
 }
+
+#[derive(Debug)]
 struct Scheduler {
-    /// List of all active tasks spawned onto this executor.
-    ///
-    /// # Safety
-    ///
-    /// Must only be accessed from the primary thread
-    tasks: UnsafeCell<task::OwnedList<Scheduler>>,
-
-    /// Local run local_queue.
-    ///
-    /// Tasks notified from the current thread are pushed into this queue.
-    ///
-    /// # Safety
-    ///
-    /// References should not be handed out. Only call `push` / `pop` functions.
-    /// Only call from the owning thread.
-    local_queue: UnsafeCell<VecDeque<Task<Scheduler>>>,
-
     tick: Cell<u8>,
 
-    /// Remote run queue.
+    queues: Queues<Self>,
-    ///
-    /// Tasks notified from another thread are pushed into this queue.
-    remote_queue: Mutex<VecDeque<Task<Scheduler>>>,
-
-    /// Tasks pending drop
-    pending_drop: TransferStack<Self>,
 
     /// Used to notify the `LocalFuture` when a task in the local task set is
     /// notified.
@@ -167,11 +144,12 @@ cfg_rt_util! {
             let current = current
                 .get()
                 .expect("`spawn_local` called from outside of a local::LocalSet!");
+            let (task, handle) = task::joinable_local(future);
             unsafe {
-                let (task, handle) = task::joinable_local(future);
+                current.as_ref().queues.push_local(task);
-                current.as_ref().schedule_local(task);
-                handle
             }
+
+            handle
         })
     }
 }
@@ -179,9 +157,6 @@ cfg_rt_util! {
 /// Max number of tasks to poll per tick.
 const MAX_TASKS_PER_TICK: usize = 61;
 
-/// How often to check the remote queue first
-const CHECK_REMOTE_INTERVAL: u8 = 13;
-
 impl LocalSet {
     /// Returns a new local task set.
     pub fn new() -> Self {
@@ -234,7 +209,7 @@ impl LocalSet {
         unsafe {
             // This is safe: since `LocalSet` is not Send or Sync, this is
             // always being called from the local thread.
-            self.scheduler.schedule_local(task);
+            self.scheduler.queues.push_local(task);
         }
         handle
     }
@@ -341,31 +316,32 @@ impl Schedule for Scheduler {
     fn bind(&self, task: &Task<Self>) {
         assert!(self.is_current());
         unsafe {
-            (*self.tasks.get()).insert(task);
+            self.queues.add_task(task);
         }
     }
 
     fn release(&self, task: Task<Self>) {
         // This will be called when dropping the local runtime.
-        self.pending_drop.push(task);
+        self.queues.release_remote(task);
     }
 
     fn release_local(&self, task: &Task<Self>) {
         debug_assert!(self.is_current());
         unsafe {
-            (*self.tasks.get()).remove(task);
+            self.queues.release_local(task);
         }
     }
 
     fn schedule(&self, task: Task<Self>) {
         if self.is_current() {
-            unsafe {
+            unsafe { self.queues.push_local(task) };
-                self.schedule_local(task);
-            }
         } else {
-            self.remote_queue.lock().unwrap().push_back(task);
+            let mut lock = self.queues.remote();
+            lock.schedule(task);
 
             self.waker.wake();
+
+            drop(lock);
         }
     }
 }
@@ -373,11 +349,8 @@ impl Schedule for Scheduler {
 impl Scheduler {
     fn new() -> Self {
         Self {
-            tasks: UnsafeCell::new(task::OwnedList::new()),
-            local_queue: UnsafeCell::new(VecDeque::with_capacity(64)),
             tick: Cell::new(0),
-            pending_drop: TransferStack::new(),
+            queues: Queues::new(),
-            remote_queue: Mutex::new(VecDeque::with_capacity(64)),
             waker: AtomicWaker::new(),
         }
     }
@@ -401,10 +374,6 @@ impl Scheduler {
         })
     }
 
-    unsafe fn schedule_local(&self, task: Task<Self>) {
-        (*self.local_queue.get()).push_back(task);
-    }
-
     fn is_current(&self) -> bool {
         CURRENT_TASK_SET
             .try_with(|current| {
@@ -416,40 +385,12 @@ impl Scheduler {
             .unwrap_or(false)
     }
 
-    fn next_task(&self, tick: u8) -> Option<Task<Self>> {
-        if 0 == tick % CHECK_REMOTE_INTERVAL {
-            self.next_remote_task().or_else(|| self.next_local_task())
-        } else {
-            self.next_local_task().or_else(|| self.next_remote_task())
-        }
-    }
-
-    fn next_local_task(&self) -> Option<Task<Self>> {
-        unsafe { (*self.local_queue.get()).pop_front() }
-    }
-
-    fn next_remote_task(&self) -> Option<Task<Self>> {
-        // there is no semantic information in the `PoisonError`, and it
-        // doesn't implement `Debug`, but clippy thinks that it's bad to
-        // match all errors here...
-        #[allow(clippy::match_wild_err_arm)]
-        let mut lock = match self.remote_queue.lock() {
-            // If the lock is poisoned, but the thread is already panicking,
-            // avoid a double panic. This is necessary since `next_task` (which
-            // calls `next_remote_task`) can be called in the `Drop` impl.
-            Err(_) if std::thread::panicking() => return None,
-            Err(_) => panic!("mutex poisoned"),
-            Ok(lock) => lock,
-        };
-        lock.pop_front()
-    }
-
     fn tick(&self) {
         assert!(self.is_current());
         for _ in 0..MAX_TASKS_PER_TICK {
             let tick = self.tick.get().wrapping_add(1);
             self.tick.set(tick);
-            let task = match self.next_task(tick) {
+            let task = match unsafe { self.queues.next_task(tick) } {
                 Some(task) => task,
                 None => return,
             };
@@ -457,56 +398,39 @@ impl Scheduler {
             if let Some(task) = task.run(&mut || Some(self.into())) {
                 unsafe {
                     // we are on the local thread, so this is okay.
-                    self.schedule_local(task);
+                    self.queues.push_local(task);
                 }
             }
         }
     }
-
-    fn drain_pending_drop(&self) {
-        for task in self.pending_drop.drain() {
-            unsafe {
-                (*self.tasks.get()).remove(&task);
-            }
-            drop(task);
-        }
-    }
-}
-
-impl fmt::Debug for Scheduler {
-    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
-        fmt.debug_struct("Scheduler { .. }").finish()
-    }
 }
 
 impl Drop for Scheduler {
     fn drop(&mut self) {
         // Close the remote queue
-        let mut lock = self.remote_queue.lock().unwrap();
+        self.queues.close_remote();
-
-        while let Some(task) = lock.pop_front() {
-            task.shutdown();
-        }
-
-        drop(lock);
 
         // Drain all local tasks
-        while let Some(task) = self.next_local_task() {
+        while let Some(task) = unsafe { self.queues.next_local_task() } {
             task.shutdown();
         }
 
         // Release owned tasks
         unsafe {
-            (*self.tasks.get()).shutdown();
+            self.queues.shutdown();
         }
 
-        self.drain_pending_drop();
+        unsafe {
+            self.queues.drain_pending_drop();
+        }
 
         // Wait until all tasks have been released.
         // XXX: this is a busy loop, but we don't really have any way to park
         // the thread here?
-        while unsafe { !(*self.tasks.get()).is_empty() } {
+        unsafe {
-            self.drain_pending_drop();
+            while self.queues.has_tasks_remaining() {
+                self.queues.drain_pending_drop();
+            }
         }
     }
 }