path: root/tokio/src/runtime/thread_pool/worker.rs

use crate::loom::cell::CausalCell;
use crate::loom::sync::Arc;
use crate::park::Park;
use crate::runtime::park::Parker;
use crate::runtime::thread_pool::{current, slice, Owned, Shared, Spawner};
use crate::runtime::{self, blocking};
use crate::task::Task;

use std::cell::Cell;
use std::marker::PhantomData;
use std::sync::atomic::Ordering::Relaxed;
use std::time::Duration;

thread_local! {
    /// Used to handle block_in_place
    static ON_BLOCK: Cell<Option<*const dyn Fn()>> = Cell::new(None)
}

cfg_blocking! {
    pub(crate) fn block_in_place<F, R>(f: F) -> R
    where
        F: FnOnce() -> R,
    {
        // Make the current worker give away its Worker to another thread so that we can safely block
        // this one without preventing progress on other futures the worker owns.
        ON_BLOCK.with(|ob| {
            let allow_blocking = ob
                .get()
                .expect("can only call blocking when on Tokio runtime");

            // This is safe, because ON_BLOCK was set from an &mut dyn FnMut in the worker that wraps
            // the worker's operation, and is unset just prior to when the FnMut is dropped.
            let allow_blocking = unsafe { &*allow_blocking };

            allow_blocking();
            f()
        })
    }
}

pub(crate) struct Worker {
    /// Parks the thread. Requires the calling worker to have obtained unique
    /// access via the generation synchronization action.
    inner: Arc<Inner>,

    /// Scheduler slices
    slices: Arc<slice::Set>,

    /// Slice assigned to this worker
    index: usize,

    /// Worker generation. This is used to synchronize access to the internal
    /// data.
    generation: usize,

    /// To indicate that the Worker has been given away and should no longer be used
    gone: Cell<bool>,
}

/// Internal worker state. This may be referenced from multiple threads, but the
/// generation guard protects unsafe access
struct Inner {
    /// Used to park the thread
    park: CausalCell<Parker>,
}

unsafe impl Send for Worker {}

/// Used to ensure the invariants are respected
struct GenerationGuard<'a> {
    /// Worker reference
    worker: &'a Worker,

    /// Prevent `Sync` access
    _p: PhantomData<Cell<()>>,
}

struct WorkerGone;

// TODO: Move into slices
pub(super) fn create_set(pool_size: usize, parker: Parker) -> (Arc<slice::Set>, Vec<Worker>) {
    // Create the parks...
    let parkers: Vec<_> = (0..pool_size).map(|_| parker.clone()).collect();

    let mut slices = Arc::new(slice::Set::new(&parkers));

    // Establish the circular link between the individual worker state
    // structure and the container.
    Arc::get_mut(&mut slices).unwrap().set_ptr();

    // This will contain each worker.
    let workers = parkers
        .into_iter()
        .enumerate()
        .map(|(index, parker)| Worker::new(slices.clone(), index, parker))
        .collect();

    (slices, workers)
}

/// After how many ticks is the global queue polled. This helps to ensure
/// fairness.
///
/// The number is fairly arbitrary. I believe this value was copied from golang.
const GLOBAL_POLL_INTERVAL: u16 = 61;

impl Worker {
    // Safe as aquiring a lock is required before doing anything potentially
    // dangerous.
    pub(super) fn new(slices: Arc<slice::Set>, index: usize, park: Parker) -> Self {
        Worker {
            inner: Arc::new(Inner {
                park: CausalCell::new(park),
            }),
            slices,
            index,
            generation: 0,
            gone: Cell::new(false),
        }
    }

    pub(super) fn run(self, blocking_pool: blocking::Spawner) {
        // First, acquire a lock on the worker.
        let guard = match self.acquire_lock() {
            Some(guard) => guard,
            None => return,
        };

        let spawner = Spawner::new(self.slices.clone());

        // Track the current worker
        current::set(&self.slices, self.index, || {
            // Enter a runtime context
            let _enter = crate::runtime::enter();

            crate::runtime::global::with_thread_pool(&spawner, || {
                blocking_pool.enter(|| {
                    ON_BLOCK.with(|ob| {
                        // Ensure that the ON_BLOCK is removed from the thread-local context
                        // when leaving the scope. This handles cases that involve panicking.
                        struct Reset<'a>(&'a Cell<Option<*const dyn Fn()>>);

                        impl<'a> Drop for Reset<'a> {
                            fn drop(&mut self) {
                                self.0.set(None);
                            }
                        }

                        let _reset = Reset(ob);

                        let allow_blocking: &dyn Fn() = &|| self.block_in_place(&blocking_pool);

                        ob.set(Some(unsafe {
                            // NOTE: We cannot use a safe cast to raw pointer here, since we are
                            // _also_ erasing the lifetime of these pointers. That is safe here,
                            // because we know that ob will set back to None before allow_blocking
                            // is dropped.
                            #[allow(clippy::useless_transmute)]
                            std::mem::transmute::<_, *const dyn Fn()>(allow_blocking)
                        }));

                        let _ = guard.run();

                        // Ensure that we reset ob before allow_blocking is dropped.
                        drop(_reset);
                    });
                })
            })
        });

        if self.gone.get() {
            // Synchronize with the pool for load(Acquire) in is_closed to get
            // up-to-date value.
            self.slices.wait_for_unlocked();

            if self.slices.is_closed() {
                // If the pool is shutting down, some other thread may be
                // waiting to clean up after the task that we were holding on
                // to. If we completed that task, we did nothing (because
                // task.run() returned None), and so crucially we did not wait
                // up any such thread.
                //
                // So, we have to do that here.
                self.slices.notify_all();
            }
        }
    }

    /// Acquire the lock
    fn acquire_lock(&self) -> Option<GenerationGuard<'_>> {
        // Safety: Only getting `&self` access to access atomic field
        let owned = unsafe { &*self.slices.owned()[self.index].get() };

        // The lock is only to establish mutual exclusion. Other synchronization
        // handles memory orderings
        let prev = owned.generation.compare_and_swap(
            self.generation,
            self.generation.wrapping_add(1),
            Relaxed,
        );

        if prev == self.generation {
            Some(GenerationGuard {
                worker: self,
                _p: PhantomData,
            })
        } else {
            None
        }
    }

    /// Enter an in-place blocking section
    fn block_in_place(&self, blocking_pool: &blocking::Spawner) {
        // If our Worker has already been given away, then blocking is fine!
        if self.gone.get() {
            return;
        }

        // make sure no subsequent code thinks that it is on a worker
        current::clear();

        // Track that the worker is gone
        self.gone.set(true);

        // If this method is called, we need to move the entire worker onto a
        // separate (blocking) thread before returning. Once we return, the
        // caller is going to execute some blocking code which would otherwise
        // block our reactor from making progress. Since we are _in the middle_
        // of running a task, this isn't trivial, as the Worker is "active".
        // We do have the luxury of knowing that we are on the worker thread,
        // so we can assert exclusive access to any Worker-specific state.
        //
        // More specifically, the caller is _currently_ "stuck" in
        // Entry::run_task at:
        //
        //   if let Some(task) = task.run(self.shared().into()) {
        //
        // And _we_ get to decide when it continues (specifically, by choosing
        // when we return from the second callback (i.e., after the FnOnce
        // passed to blocking has returned).
        //
        // Here's what we'll have to do:
        //
        //  - Reconstruct our `Worker` struct
        //  - Spawn the reconstructed `Worker` on another blocking thread
        //  - Clear any state indicating what worker we are on, since at this
        //    point we are effectively no longer "on" that worker.
        //  - Allow the caller of `blocking` to continue.
        //
        // Once the caller completes the blocking operations, we need to ensure
        // that async code can continue running in that context. Luckily, since
        // `Arc<slice::Set>` has a fallback for when
        // curre