tokio-rs/tokio を読む(1) #

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

struct Hello {
    state: StateHello,
}

enum StateHello {
    Hello,
    World,
}

impl Hello {
    fn new() -> Self {
        Hello {
            state: StateHello::Hello,
        }
    }
}

impl Future for Hello {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        match (*self).state {
            StateHello::Hello => {
                println!("Hello ");
                (*self).state = StateHello::World;
                cx.waker().wake_by_ref();
                Poll::Pending
            }
            StateHello::World => {
                println!("World!");
                Poll::Ready(())
            }
        }
    }
}

fn main() {
    let rt = tokio::runtime::Runtime::new().unwrap();
    rt.block_on(Hello::new());
}

pub fn new() -> std::io::Result<Runtime> {
    Builder::new_multi_thread().enable_all().build()
}

pub struct Runtime {
    /// Task scheduler
    scheduler: Scheduler,
    /// Handle to runtime, also contains driver handles
    handle: Handle,
    /// Blocking pool handle, used to signal shutdown
    blocking_pool: BlockingPool,
}

pub fn new_multi_thread() -> Builder {
    // The number `61` is fairly arbitrary. I believe this value was copied from golang.
    Builder::new(Kind::MultiThread, 61, 61)
}

let core_threads = self.worker_threads.unwrap_or_else(num_cpus);

let (driver, driver_handle) = driver::Driver::new(self.get_cfg())?;

let blocking_pool =
    blocking::create_blocking_pool(self, self.max_blocking_threads + core_threads);
let blocking_spawner = blocking_pool.spawner().clone();

pub(crate) struct BlockingPool {
    spawner: Spawner,
    shutdown_rx: shutdown::Receiver,
}

#[derive(Clone)]
pub(crate) struct Spawner {
    inner: Arc<Inner>,
}

struct Inner {
    /// State shared between worker threads.
    shared: Mutex<Shared>,
    /// Pool threads wait on this.
    condvar: Condvar,
    /// Spawned threads use this name.
    thread_name: ThreadNameFn,
    /// Spawned thread stack size.
    stack_size: Option<usize>,
    /// Call after a thread starts.
    after_start: Option<Callback>,
    /// Call before a thread stops.
    before_stop: Option<Callback>,
    // Maximum number of threads.
    thread_cap: usize,
    // Customizable wait timeout.
    keep_alive: Duration,
    // Metrics about the pool.
    metrics: SpawnerMetrics,
}

struct Shared {
    queue: VecDeque<Task>,
    num_notify: u32,
    shutdown: bool,
    shutdown_tx: Option<shutdown::Sender>,
    /// Prior to shutdown, we clean up JoinHandles by having each timed-out
    /// thread join on the previous timed-out thread. This is not strictly
    /// necessary but helps avoid Valgrind false positives, see
    /// <https://github.com/tokio-rs/tokio/commit/646fbae76535e397ef79dbcaacb945d4c829f666>
    /// for more information.
    last_exiting_thread: Option<thread::JoinHandle<()>>,
    /// This holds the JoinHandles for all running threads; on shutdown, the thread
    /// calling shutdown handles joining on these.
    worker_threads: HashMap<usize, thread::JoinHandle<()>>,
    /// This is a counter used to iterate worker_threads in a consistent order (for loom's
    /// benefit).
    worker_thread_index: usize,
}

pub(crate) struct Task {
    task: task::UnownedTask<BlockingSchedule>,
    mandatory: Mandatory,
}

pub(crate) fn new(builder: &Builder, thread_cap: usize) -> BlockingPool {
    let (shutdown_tx, shutdown_rx) = shutdown::channel();
    let keep_alive = builder.keep_alive.unwrap_or(KEEP_ALIVE);

    BlockingPool {
        spawner: Spawner {
            inner: Arc::new(Inner {
                shared: Mutex::new(Shared {
                    queue: VecDeque::new(),
                    num_notify: 0,
                    shutdown: false,
                    shutdown_tx: Some(shutdown_tx),
                    last_exiting_thread: None,
                    worker_threads: HashMap::new(),
                    worker_thread_index: 0,
                }),
                condvar: Condvar::new(),
                thread_name: builder.thread_name.clone(),
                stack_size: builder.thread_stack_size,
                after_start: builder.after_start.clone(),
                before_stop: builder.before_stop.clone(),
                thread_cap,
                keep_alive,
                metrics: Default::default(),
            }),
        },
        shutdown_rx,
    }
}

let (scheduler, handle, launch) = MultiThread::new(
    core_threads,
    driver,
    driver_handle,
    blocking_spawner,
    seed_generator_2,
    Config {
        before_park: self.before_park.clone(),
        after_unpark: self.after_unpark.clone(),
        global_queue_interval: self.global_queue_interval,
        event_interval: self.event_interval,
        #[cfg(tokio_unstable)]
        unhandled_panic: self.unhandled_panic.clone(),
        disable_lifo_slot: self.disable_lifo_slot,
        seed_generator: seed_generator_1,
    },
);

let mut cores = Vec::with_capacity(size);
let mut remotes = Vec::with_capacity(size);
let mut worker_metrics = Vec::with_capacity(size);

for _ in 0..size {
    let (steal, run_queue) = queue::local();
    // ...
}

pub(crate) fn local<T: 'static>() -> (Steal<T>, Local<T>) {
    let mut buffer = Vec::with_capacity(LOCAL_QUEUE_CAPACITY);

    for _ in 0..LOCAL_QUEUE_CAPACITY {
        buffer.push(UnsafeCell::new(MaybeUninit::uninit()));
    }

    let inner = Arc::new(Inner {
        head: AtomicUnsignedLong::new(0),
        tail: AtomicUnsignedShort::new(0),
        buffer: make_fixed_size(buffer.into_boxed_slice()),
    });

    let local = Local {
        inner: inner.clone(),
    };

    let remote = Steal(inner);

    (remote, local)
}

pub(crate) struct Inner<T: 'static> {
    /// Concurrently updated by many threads.
    ///
    /// Contains two `UnsignedShort` values. The LSB byte is the "real" head of
    /// the queue. The `UnsignedShort` in the MSB is set by a stealer in process
    /// of stealing values. It represents the first value being stolen in the
    /// batch. The `UnsignedShort` indices are intentionally wider than strictly
    /// required for buffer indexing in order to provide ABA mitigation and make
    /// it possible to distinguish between full and empty buffers.
    ///
    /// When both `UnsignedShort` values are the same, there is no active
    /// stealer.
    ///
    /// Tracking an in-progress stealer prevents a wrapping scenario.
    head: AtomicUnsignedLong,

    /// Only updated by producer thread but read by many threads.
    tail: AtomicUnsignedShort,

    /// Elements
    buffer: Box<[UnsafeCell<MaybeUninit<task::Notified<T>>>; LOCAL_QUEUE_CAPACITY]>,
}

for _ in 0..size {
    let (steal, run_queue) = queue::local();
    let park = park.clone();
    let unpark = park.unpark();
    cores.push(Box::new(Core {
        tick: 0,
        lifo_slot: None,
        run_queue,
        is_searching: false,
        is_shutdown: false,
        park: Some(park),
        metrics: MetricsBatch::new(),
        rand: FastRand::new(config.seed_generator.next_seed()),
    }));
    remotes.push(Remote { steal, unpark });
    worker_metrics.push(WorkerMetrics::new());
}

let handle = Arc::new(Handle {
    shared: Shared {
        remotes: remotes.into_boxed_slice(),
        inject: Inject::new(),
        idle: Idle::new(size),
        owned: OwnedTasks::new(),
        shutdown_cores: Mutex::new(vec![]),
        config,
        scheduler_metrics: SchedulerMetrics::new(),
        worker_metrics: worker_metrics.into_boxed_slice(),
    },
    driver: driver_handle,
    blocking_spawner,
    seed_generator,
});

    let mut launch = Launch(vec![]);
    for (index, core) in cores.drain(..).enumerate() {
        launch.0.push(Arc::new(Worker {
            handle: handle.clone(),
            index,
            core: AtomicCell::new(Some(core)),
        }));
    }
    (handle, launch)

/// Work-stealing based thread pool for executing futures.
pub(crate) struct MultiThread;

pub(crate) fn block_on<F>(&self, handle: &scheduler::Handle, future: F) -> F::Output
where
    F: Future,
{
    let mut enter = crate::runtime::context::enter_runtime(handle, true);
    enter
        .blocking
        .block_on(future)
        .expect("failed to park thread")
}

pub(crate) fn shutdown(&mut self, handle: &scheduler::Handle) {
    match handle {
        scheduler::Handle::MultiThread(handle) => handle.shutdown(),
        _ => panic!("expected MultiThread scheduler"),
    }
}

let handle = Handle { inner: scheduler::Handle::MultiThread(handle) };
// Spawn the thread pool workers
let _enter = handle.enter();
launch.launch();
Ok(Runtime::from_parts(Scheduler::MultiThread(scheduler), handle, blocking_pool))

pub(crate) fn try_set_current(handle: &scheduler::Handle) -> Option<SetCurrentGuard> {
    CONTEXT.try_with(|ctx| ctx.set_current(handle)).ok()
}

static CONTEXT: Context = {
    Context {
        #[cfg(feature = "rt")]
        thread_id: Cell::new(None),
        /// Tracks the current runtime handle to use when spawning,
        /// accessing drivers, etc...
        #[cfg(feature = "rt")]
        handle: RefCell::new(None),
        #[cfg(feature = "rt")]
        current_task_id: Cell::new(None),
        // ....
    }
}

pub(crate) fn launch(mut self) {
    for worker in self.0.drain(..) {
        runtime::spawn_blocking(move || run(worker));
    }
}

pub(crate) fn spawn_blocking<F, R>(func: F) -> JoinHandle<R>
where
    F: FnOnce() -> R + Send + 'static,
    R: Send + 'static,
{
    let rt = Handle::current();
    rt.spawn_blocking(func)
}

pub(crate) fn spawn_blocking_inner<F, R>(
    &self,
    func: F,
    is_mandatory: Mandatory,
    name: Option<&str>,
    rt: &Handle,
) -> (JoinHandle<R>, Result<(), SpawnError>)
where
    F: FnOnce() -> R + Send + 'static,
    R: Send + 'static,
{
    let fut = BlockingTask::new(func);
    let id = task::Id::next();
    let (task, handle) = task::unowned(fut, BlockingSchedule::new(rt), id);
    let spawned = self.spawn_task(Task::new(task, is_mandatory), rt);
    (handle, spawned)
}

impl<T, R> Future for BlockingTask<T>
where
    T: FnOnce() -> R + Send + 'static,
    R: Send + 'static,
{
    type Output = R;

    fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<R> {
        let me = &mut *self;
        let func = me
            .func
            .take()
            .expect("[internal exception] blocking task ran twice.");

        // This is a little subtle:
        // For convenience, we'd like _every_ call tokio ever makes to Task::poll() to be budgeted
        // using coop. However, the way things are currently modeled, even running a blocking task
        // currently goes through Task::poll(), and so is subject to budgeting. That isn't really
        // what we want; a blocking task may itself want to run tasks (it might be a Worker!), so
        // we want it to start without any budgeting.
        crate::runtime::coop::stop();

        Poll::Ready(func())
    }
}

pub(crate) fn unowned<T, S>(task: T, scheduler: S, id: Id) -> (UnownedTask<S>, JoinHandle<T::Output>)
where
    S: Schedule,
    T: Send + Future + 'static,
    T::Output: Send + 'static,
{
    let (task, notified, join) = new_task(task, scheduler, id);
    // This transfers the ref-count of task and notified into an UnownedTask.
    // This is valid because an UnownedTask holds two ref-counts.
    let unowned = UnownedTask {
        raw: task.raw,
        _p: PhantomData,
    };
    std::mem::forget(task);
    std::mem::forget(notified);
    (unowned, join)
}

impl<S: 'static> Drop for UnownedTask<S> {
    fn drop(&mut self) {
        // Decrement the ref count
        if self.raw.header().state.ref_dec_twice() {
            // Deallocate if this is the final ref count
            self.raw.dealloc();
        }
    }
}

    fn new_task<T, S>(
        task: T,
        scheduler: S,
        id: Id,
    ) -> (Task<S>, Notified<S>, JoinHandle<T::Output>)
    where
        S: Schedule,
        T: Future + 'static,
        T::Output: 'static,
    {
        let raw = RawTask::new::<T, S>(task, scheduler, id);
        let task = Task {
            raw,
            _p: PhantomData,
        };
        let notified = Notified(Task {
            raw,
            _p: PhantomData,
        });
        let join = JoinHandle::new(raw);

        (task, notified, join)
    }

pub(super) fn new<T, S>(task: T, scheduler: S, id: Id) -> RawTask
where
    T: Future,
    S: Schedule,
{
    let ptr = Box::into_raw(Cell::<_, S>::new(task, scheduler, State::new(), id));
    let ptr = unsafe { NonNull::new_unchecked(ptr as *mut Header) };
    RawTask { ptr }
}

fn spawn_task(&self, task: Task, rt: &Handle) -> Result<(), SpawnError> {
    let mut shared = self.inner.shared.lock();
    // ...
    shared.queue.push_back(task);
    self.inner.metrics.inc_queue_depth();
    if self.inner.metrics.num_idle_threads() == 0 {
        // No threads are able to process the task.
        if self.inner.metrics.num_threads() == self.inner.thread_cap {
            // At max number of threads
        } else {
            // ...
                let id = shared.worker_thread_index;
                match self.spawn_thread(shutdown_tx, rt, id) {
                    Ok(handle) => {
                        self.inner.metrics.inc_num_threads();
                        shared.worker_thread_index += 1;
                        shared.worker_threads.insert(id, handle);
                    }
                    Err(e) => {
                        return Err(/*...*/);
                    }
                }
            //
        }
    } else {
        // ....
    }
    Ok(())
}

fn spawn_thread(
    &self,
    shutdown_tx: shutdown::Sender,
    rt: &Handle,
    id: usize,
) -> std::io::Result<thread::JoinHandle<()>> {
    let mut builder = thread::Builder::new().name((self.inner.thread_name)());
    if let Some(stack_size) = self.inner.stack_size {
        builder = builder.stack_size(stack_size);
    }
    let rt = rt.clone();
    builder.spawn(move || {
        // Only the reference should be moved into the closure
        let _enter = rt.enter();
        rt.inner.blocking_spawner().inner.run(id);
        drop(shutdown_tx);
    })
}

'main: loop {
    // BUSY
    while let Some(task) = shared.queue.pop_front() {
        self.metrics.dec_queue_depth();
        drop(shared);
        task.run();
        shared = self.shared.lock();
    }

    // IDLE
    self.metrics.inc_num_idle_threads();
    while !shared.shutdown {
        let lock_result = self.condvar.wait_timeout(shared, self.keep_alive).unwrap();
        shared = lock_result.0;
        let timeout_result = lock_result.1;
        if shared.num_notify != 0 {
            // We have received a legitimate wakeup,
            // acknowledge it by decrementing the counter
            // and transition to the BUSY state.
            shared.num_notify -= 1;
            break;
        }
        // Even if the condvar "timed out", if the pool is entering the
        // shutdown phase, we want to perform the cleanup logic.
        if !shared.shutdown && timeout_result.timed_out() {
            // We'll join the prior timed-out thread's JoinHandle after dropping the lock.
            // This isn't done when shutting down, because the thread calling shutdown will
            // handle joining everything.
            let my_handle = shared.worker_threads.remove(&worker_thread_id);
            join_on_thread = std::mem::replace(&mut shared.last_exiting_thread, my_handle);
            break 'main;
        }
        // Spurious wakeup detected, go back to sleep.
    }
    // ....
}

while let Some(task) = shared.queue.pop_front() {
    self.metrics.dec_queue_depth();
    drop(shared);
    task.run();
    shared = self.shared.lock();
}

pub(crate) fn run(self) {
    let raw = self.raw;
    mem::forget(self);
    // Transfer one ref-count to a Task object.
    let task = Task::<S> {
        raw,
        _p: PhantomData,
    };
    // Use the other ref-count to poll the task.
    raw.poll();
    // Decrement our extra ref-count
    drop(task);
}

pub(super) fn poll(self) {
    let vtable = self.header().vtable;
    unsafe { (vtable.poll)(self.ptr) }
}

unsafe fn poll<T: Future, S: Schedule>(ptr: NonNull<Header>) {
    let harness = Harness::<T, S>::from_raw(ptr);
    harness.poll();
}

pub(super) fn poll(&self, mut cx: Context<'_>) -> Poll<T::Output> { 
    let res = {
        self.stage.stage.with_mut(|ptr| {
            // Safety: The caller ensures mutual exclusion to the field.
            let future = match unsafe { &mut *ptr } {
                Stage::Running(future) => future,
                _ => unreachable!("unexpected stage"),
            };
            // Safety: The caller ensures the future is pinned.
            let future = unsafe { Pin::new_unchecked(future) };
            let _guard = TaskIdGuard::enter(self.task_id);
            future.poll(&mut cx) 
        })
    };
    if res.is_ready() {
        self.drop_future_or_output();
    }
    res
}

pub(super) fn new(future: T, scheduler: S, state: State, task_id: Id) -> Box<Cell<T, S>> {
    let result = Box::new(Cell {
        header: Header {
            state,
            queue_next: UnsafeCell::new(None),
            vtable: raw::vtable::<T, S>(),
            owner_id: UnsafeCell::new(0),
        },
        core: Core {
            scheduler,
            stage: CoreStage {
                stage: UnsafeCell::new(Stage::Running(future)),
            },
            task_id,
        },
        trailer: Trailer {
            waker: UnsafeCell::new(None),
            owned: linked_list::Pointers::new(),
        },
    });
}

pub(super) fn vtable<T: Future, S: Schedule>() -> &'static Vtable {
    &Vtable {
        poll: poll::<T, S>,
        schedule: schedule::<S>,
        dealloc: dealloc::<T, S>,
        try_read_output: try_read_output::<T, S>,
        drop_join_handle_slow: drop_join_handle_slow::<T, S>,
        drop_abort_handle: drop_abort_handle::<T, S>,
        shutdown: shutdown::<T, S>,
        trailer_offset: OffsetHelper::<T, S>::TRAILER_OFFSET,
        scheduler_offset: OffsetHelper::<T, S>::SCHEDULER_OFFSET,
        id_offset: OffsetHelper::<T, S>::ID_OFFSET,
    }
}

unsafe fn poll<T: Future, S: Schedule>(ptr: NonNull<Header>) {
    let harness = Harness::<T, S>::from_raw(ptr);
    harness.poll();
}

// IDLE
self.metrics.inc_num_idle_threads();
while !shared.shutdown {
    let lock_result = self.condvar.wait_timeout(shared, self.keep_alive).unwrap();
    shared = lock_result.0;
    let timeout_result = lock_result.1;
    if shared.num_notify != 0 {
        // We have received a legitimate wakeup,
        // acknowledge it by decrementing the counter
        // and transition to the BUSY state.
        shared.num_notify -= 1;
        break;
    }
    // Even if the condvar "timed out", if the pool is entering the
    // shutdown phase, we want to perform the cleanup logic.
    if !shared.shutdown && timeout_result.timed_out() {
        // We'll join the prior timed-out thread's JoinHandle after dropping the lock.
        // This isn't done when shutting down, because the thread calling shutdown will
        // handle joining everything.
        let my_handle = shared.worker_threads.remove(&worker_thread_id);
        join_on_thread = std::mem::replace(&mut shared.last_exiting_thread, my_handle);
        break 'main;
    }
    // Spurious wakeup detected, go back to sleep.
}

fn run(worker: Arc<Worker>) {
    // ...
    let core = match worker.core.take() {
        Some(core) => core,
        None => return,
    };

    let handle = scheduler::Handle::MultiThread(worker.handle.clone());
    let _enter = crate::runtime::context::enter_runtime(&handle, true);

    // Set the worker context.
    let cx = Context {
        worker,
        core: RefCell::new(None),
    };

    CURRENT.set(&cx, || {
        assert!(cx.run(core).is_err());
        wake_deferred_tasks();
    });
}

// Tracks thread-local state
scoped_thread_local!(static CURRENT: Context);

/// Sets a reference as a thread-local.
macro_rules! scoped_thread_local {
    ($(#[$attrs:meta])* $vis:vis static $name:ident: $ty:ty) => (
        $(#[$attrs])*
        $vis static $name: $crate::macros::scoped_tls::ScopedKey<$ty>
            = $crate::macros::scoped_tls::ScopedKey {
                inner: {
                    tokio_thread_local!(static FOO: ::std::cell::Cell<*const ()> = const {
                        std::cell::Cell::new(::std::ptr::null())
                    });
                    &FOO
                },
                _marker: ::std::marker::PhantomData,
            };
    )
}

fn run(&self, mut core: Box<Core>) -> RunResult {
    while !core.is_shutdown {
        core.tick();
        core = self.maintenance(core);
        // First, check work available to the current worker.
        if let Some(task) = core.next_task(&self.worker) {
            core = self.run_task(task, core)?;
            continue;
        }
        // There is no more **local** work to process, try to steal work
        // from other workers.
        if let Some(task) = core.steal_work(&self.worker) {
            core = self.run_task(task, core)?;
        } else {
            // Wait for work
            core = if did_defer_tasks() {
                self.park_timeout(core, Some(Duration::from_millis(0)))
            } else {
                self.park(core)
            };
        }
    }
    // ...
}

pub fn block_on<F: Future>(&self, future: F) -> F::Output {
    let _enter = self.enter();
    match &self.scheduler {
        Scheduler::CurrentThread(exec) => exec.block_on(&self.handle.inner, future),
        #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
        Scheduler::MultiThread(exec) => exec.block_on(&self.handle.inner, future),
    }
}

pub(crate) fn block_on<F>(&self, handle: &scheduler::Handle, future: F) -> F::Output
where
    F: Future,
{
    let mut enter = crate::runtime::context::enter_runtime(handle, true);
    enter
        .blocking
        .block_on(future)
        .expect("failed to park thread")
}

pub(crate) fn block_on<F: Future>(&mut self, f: F) -> Result<F::Output, AccessError> {
    use std::task::Context;
    use std::task::Poll::Ready;
    // `get_unpark()` should not return a Result
    let waker = self.waker()?;
    let mut cx = Context::from_waker(&waker);
    pin!(f);
    loop {
        if let Ready(v) = crate::runtime::coop::budget(|| f.as_mut().poll(&mut cx)) {
            return Ok(v);
        }
        // Wake any yielded tasks before parking in order to avoid blocking.
        crate::runtime::context::with_defer(|defer| defer.wake());
        self.park();
    }
}

pub(crate) fn waker(&self) -> Result<Waker, AccessError> {
    self.unpark().map(|unpark| unpark.into_waker())
}

fn unpark(&self) -> Result<UnparkThread, AccessError> {
    self.with_current(|park_thread| park_thread.unpark())
}

fn with_current<F, R>(&self, f: F) -> Result<R, AccessError>
where
    F: FnOnce(&ParkThread) -> R,
{
    CURRENT_PARKER.try_with(|inner| f(inner))
}

tokio_thread_local! {
    static CURRENT_PARKER: ParkThread = ParkThread::new();
}

impl ParkThread {
    pub(crate) fn new() -> Self {
        Self {
            inner: Arc::new(Inner {
                state: AtomicUsize::new(EMPTY),
                mutex: Mutex::new(()),
                condvar: Condvar::new(),
            }),
        }
    }
}