From 0f605b51cadc3e6759209abaeb355ed48f925548 Mon Sep 17 00:00:00 2001 From: Carl Lerche Date: Thu, 8 Jun 2023 14:07:07 -0700 Subject: [PATCH] wip --- tokio/src/runtime/builder.rs | 10 +- tokio/src/runtime/scheduler/mod.rs | 9 + .../runtime/scheduler/multi_thread/handle.rs | 6 +- .../runtime/scheduler/multi_thread/idle.rs | 3 +- .../src/runtime/scheduler/multi_thread/mod.rs | 16 +- .../runtime/scheduler/multi_thread/queue.rs | 5 + .../runtime/scheduler/multi_thread/worker.rs | 391 +++++++++++++----- 7 files changed, 313 insertions(+), 127 deletions(-) diff --git a/tokio/src/runtime/builder.rs b/tokio/src/runtime/builder.rs index af9e0e172..a700342b8 100644 --- a/tokio/src/runtime/builder.rs +++ b/tokio/src/runtime/builder.rs @@ -1172,7 +1172,7 @@ cfg_rt_multi_thread! { fn build_threaded_runtime(&mut self) -> io::Result { use crate::loom::sys::num_cpus; use crate::runtime::{Config, runtime::Scheduler}; - use crate::runtime::scheduler::{self, MultiThread}; + use crate::runtime::scheduler::MultiThread; let core_threads = self.worker_threads.unwrap_or_else(num_cpus); @@ -1187,7 +1187,7 @@ cfg_rt_multi_thread! { let seed_generator_1 = self.seed_generator.next_generator(); let seed_generator_2 = self.seed_generator.next_generator(); - let (scheduler, handle, launch) = MultiThread::new( + let (scheduler, handle) = MultiThread::new( core_threads, driver, driver_handle, @@ -1206,12 +1206,6 @@ cfg_rt_multi_thread! { }, ); - 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)) } } diff --git a/tokio/src/runtime/scheduler/mod.rs b/tokio/src/runtime/scheduler/mod.rs index 3e3151711..81274e9a9 100644 --- a/tokio/src/runtime/scheduler/mod.rs +++ b/tokio/src/runtime/scheduler/mod.rs @@ -123,6 +123,15 @@ cfg_rt! { _ => panic!("not a CurrentThread handle"), } } + + cfg_rt_multi_thread! { + pub(crate) fn expect_multi_thread(&self) -> &Arc { + match self { + Handle::MultiThread(handle) => handle, + _ => panic!("not a `MultiThread` handle"), + } + } + } } cfg_metrics! { diff --git a/tokio/src/runtime/scheduler/multi_thread/handle.rs b/tokio/src/runtime/scheduler/multi_thread/handle.rs index e43c92aba..405200890 100644 --- a/tokio/src/runtime/scheduler/multi_thread/handle.rs +++ b/tokio/src/runtime/scheduler/multi_thread/handle.rs @@ -43,8 +43,7 @@ impl Handle { } pub(crate) fn shutdown(&self) { - // self.close(); - todo!() + self.shared.close(); } pub(super) fn bind_new_task(me: &Arc, future: T, id: task::Id) -> JoinHandle @@ -55,8 +54,7 @@ impl Handle { let (handle, notified) = me.shared.owned.bind(future, me.clone(), id); if let Some(notified) = notified { - // me.schedule_task(notified, false); - todo!() + me.shared.schedule_task(notified, false); } handle diff --git a/tokio/src/runtime/scheduler/multi_thread/idle.rs b/tokio/src/runtime/scheduler/multi_thread/idle.rs index 3c579f613..eaa3aae8e 100644 --- a/tokio/src/runtime/scheduler/multi_thread/idle.rs +++ b/tokio/src/runtime/scheduler/multi_thread/idle.rs @@ -84,8 +84,7 @@ impl Idle { /// work. pub(super) fn transition_worker_to_parked( &self, - shared: &Shared, - worker: usize, + synced: &mut Synced, is_searching: bool, ) -> bool { /* diff --git a/tokio/src/runtime/scheduler/multi_thread/mod.rs b/tokio/src/runtime/scheduler/multi_thread/mod.rs index d85a0ae0a..2599a80aa 100644 --- a/tokio/src/runtime/scheduler/multi_thread/mod.rs +++ b/tokio/src/runtime/scheduler/multi_thread/mod.rs @@ -15,13 +15,10 @@ use self::idle::Idle; mod stats; pub(crate) use stats::Stats; -mod park; -pub(crate) use park::{Parker, Unparker}; - pub(crate) mod queue; mod worker; -pub(crate) use worker::{Context, Launch, Shared}; +pub(crate) use worker::{Context, Shared}; cfg_taskdump! { mod trace; @@ -37,8 +34,8 @@ cfg_not_taskdump! { pub(crate) use worker::block_in_place; -use crate::loom::sync::Arc; use crate::runtime::{ + self, blocking, driver::{self, Driver}, scheduler, Config, @@ -61,18 +58,17 @@ impl MultiThread { blocking_spawner: blocking::Spawner, seed_generator: RngSeedGenerator, config: Config, - ) -> (MultiThread, Arc, Launch) { - let parker = Parker::new(driver); - let (handle, launch) = worker::create( + ) -> (MultiThread, runtime::Handle) { + let handle = worker::create( size, - parker, + driver, driver_handle, blocking_spawner, seed_generator, config, ); - (MultiThread, handle, launch) + (MultiThread, handle) } /// Blocks the current thread waiting for the future to complete. diff --git a/tokio/src/runtime/scheduler/multi_thread/queue.rs b/tokio/src/runtime/scheduler/multi_thread/queue.rs index 6444df88b..03406f94e 100644 --- a/tokio/src/runtime/scheduler/multi_thread/queue.rs +++ b/tokio/src/runtime/scheduler/multi_thread/queue.rs @@ -119,6 +119,11 @@ impl Local { LOCAL_QUEUE_CAPACITY } + /// Returns `true` if there are no entries in the queue + pub(crate) fn is_empty(&self) -> bool { + self.inner.is_empty() + } + /// Returns false if there are any entries in the queue /// /// Separate to is_stealable so that refactors of is_stealable to "protect" diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index b265e0eae..2949c72fa 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -56,16 +56,16 @@ //! the inject queue indefinitely. This would be a ref-count cycle and a memory //! leak. -use crate::loom::sync::{Arc, Condvar, Mutex}; +use crate::loom::sync::{Arc, Condvar, Mutex, MutexGuard}; use crate::runtime; use crate::runtime::context; use crate::runtime::scheduler::multi_thread::{ - idle, queue, Counters, Handle, Idle, Overflow, Parker, Stats, TraceStatus, Unparker, + idle, queue, Counters, Handle, Idle, Overflow, Stats, TraceStatus, }; -use crate::runtime::scheduler::{inject, Defer, Lock}; +use crate::runtime::scheduler::{self, inject, Defer, Lock}; use crate::runtime::task::OwnedTasks; use crate::runtime::{ - blocking, coop, driver, scheduler, task, Config, Driver, SchedulerMetrics, WorkerMetrics, + blocking, coop, driver, task, Config, Driver, SchedulerMetrics, WorkerMetrics, }; use crate::util::atomic_cell::AtomicCell; use crate::util::rand::{FastRand, RngSeedGenerator}; @@ -157,17 +157,6 @@ pub(crate) struct Shared { /// Condition variable used to unblock waiting workers condvar: Condvar, - /// Power's Tokio's I/O, timers, etc... the responsibility of polling the - /// driver is shared across workers. - driver: AtomicCell, - - /// Cores that have observed the shutdown signal - /// - /// The core is **not** placed back in the worker to avoid it from being - /// stolen by a thread that was spawned as part of `block_in_place`. - #[allow(clippy::vec_box)] // we're moving an already-boxed value - shutdown_cores: Mutex>>, - /// The number of cores that have observed the trace signal. pub(super) trace_status: TraceStatus, @@ -191,20 +180,27 @@ pub(crate) struct Synced { /// Cores not currently assigned to workers cores: Vec>, + /// Cores that have observed the shutdown signal + /// + /// The core is **not** placed back in the worker to avoid it from being + /// stolen by a thread that was spawned as part of `block_in_place`. + shutdown_cores: Vec>, + /// Synchronized state for `Idle`. pub(super) idle: idle::Synced, /// Synchronized state for `Inject`. pub(crate) inject: inject::Synced, + + /// Power's Tokio's I/O, timers, etc... the responsibility of polling the + /// driver is shared across workers. + driver: Option>, } /// Used to communicate with a worker from other threads. struct Remote { /// Steals tasks from this worker. pub(super) steal: queue::Steal>, - - /// Unparks the associated worker thread - unpark: Unparker, } /// Thread-local context @@ -219,9 +215,6 @@ pub(crate) struct Context { pub(crate) defer: Defer, } -/// Starts the workers -pub(crate) struct Launch(Vec>); - /// Running a task may consume the core. If the core is still available when /// running the task completes, it is returned. Otherwise, the worker will need /// to stop processing. @@ -241,24 +234,20 @@ const MAX_LIFO_POLLS_PER_TICK: usize = 3; pub(super) fn create( size: usize, - park: Parker, + driver: Driver, driver_handle: driver::Handle, blocking_spawner: blocking::Spawner, seed_generator: RngSeedGenerator, config: Config, -) -> (Arc, Launch) { - /* +) -> runtime::Handle { let mut cores = Vec::with_capacity(size); let mut remotes = Vec::with_capacity(size); - let mut parkers = Vec::with_capacity(size); let mut worker_metrics = Vec::with_capacity(size); // Create the local queues for i in 0..size { let (steal, run_queue) = queue::local(); - let park = park.clone(); - let unpark = park.unpark(); let metrics = WorkerMetrics::from_config(&config); let stats = Stats::new(&metrics); @@ -276,8 +265,7 @@ pub(super) fn create( rand: FastRand::from_seed(config.seed_generator.next_seed()), })); - parkers.push(park); - remotes.push(Remote { steal, unpark }); + remotes.push(Remote { steal }); worker_metrics.push(metrics); } @@ -292,10 +280,13 @@ pub(super) fn create( idle, owned: OwnedTasks::new(), synced: Mutex::new(Synced { + cores, + shutdown_cores: Vec::with_capacity(size), idle: idle_synced, inject: inject_synced, + driver: Some(Box::new(driver)), }), - shutdown_cores: Mutex::new(vec![]), + condvar: Condvar::new(), trace_status: TraceStatus::new(remotes_len), config, scheduler_metrics: SchedulerMetrics::new(), @@ -307,19 +298,23 @@ pub(super) fn create( seed_generator, }); - let mut launch = Launch(vec![]); + let rt_handle = runtime::Handle { + inner: scheduler::Handle::MultiThread(handle), + }; - for (core, park) in cores.drain(..).zip(parkers.drain(..)) { - launch.0.push(Arc::new(Worker { + // Eagerly start worker threads + for _ in 0..size { + let handle = rt_handle.inner.expect_multi_thread(); + let worker = Worker { handle: handle.clone(), - park, - core: AtomicCell::new(Some(core)), - })); + }; + + handle + .blocking_spawner + .spawn_blocking(&rt_handle, move || run(worker)); } - (handle, launch) - */ - todo!() + rt_handle } #[track_caller] @@ -436,17 +431,6 @@ where todo!() } -impl Launch { - pub(crate) fn launch(mut self) { - /* - for worker in self.0.drain(..) { - runtime::spawn_blocking(move || run(worker)); - } - */ - todo!(); - } -} - fn run(mut worker: Worker) { struct AbortOnPanic; @@ -488,12 +472,6 @@ fn run(mut worker: Worker) { // Run the worker worker.run(&cx); - /* - // This should always be an error. It only returns a `Result` to support - // using `?` to short circuit. - assert!(cx.run(core).is_err()); - */ - // Check if there are any deferred tasks to notify. This can happen when // the worker core is lost due to `block_in_place()` being called from // within the task. @@ -506,11 +484,13 @@ impl Worker { fn run(&mut self, cx: &Context) { // First, acquire a core. If no cores are available, the thread will // block until one becomes available. - let mut core = self.acquire_core(); - - // Start as "processing" tasks as polling tasks from the local queue - // will be one of the first things we do. - core.stats.start_processing_scheduled_tasks(); + // + // Acquiring a core will also pull tasks from the injection queue and + // run one, if found. + let mut core = match self.acquire_core(cx, self.shared().synced.lock()) { + Ok(core) => core, + Err(_) => return, + }; while !core.is_shutdown { self.assert_lifo_enabled_is_correct(&core); @@ -549,15 +529,16 @@ impl Worker { Err(_) => return, }; } else { - /* // Wait for work - core = if !self.defer.is_empty() { - self.park_timeout(core, Some(Duration::from_millis(0))) + core = if !cx.defer.is_empty() { + // Just run maintenance + self.park_yield(core) } else { - self.park(core) + match self.park(cx, core) { + Ok(core) => core, + Err(_) => return, + } }; - */ - todo!() } } @@ -569,13 +550,49 @@ impl Worker { todo!() } - fn acquire_core(&mut self) -> Box { - let mut core = todo!(); + fn acquire_core(&self, cx: &Context, mut synced: MutexGuard) -> RunResult { + // Wait until a core is available, then exit the loop. + let mut core = loop { + if let Some(core) = synced.cores.pop() { + break core; + } + + // TODO: not always the case + assert!(cx.defer.is_empty()); + + synced = self.shared().condvar.wait(synced).unwrap(); + }; // Reset `lifo_enabled` here in case the core was previously stolen from // a task that had the LIFO slot disabled. self.reset_lifo_enabled(&mut core); - todo!() + + // At this point, the local queue should be empty + debug_assert!(core.run_queue.is_empty()); + + // Update shutdown state while locked + core.is_shutdown = self.shared().inject.is_closed(&synced.inject); + + if core.is_shutdown { + // Currently shutting down, don't do any more work + return Ok(core); + } + + // TODO: don't hardcode 128 + let n = core.run_queue.max_capacity() / 2; + let maybe_task = self.next_remote_task_batch(&mut synced, &mut core, n); + + drop(synced); + + // Start as "processing" tasks as polling tasks from the local queue + // will be one of the first things we do. + core.stats.start_processing_scheduled_tasks(); + + if let Some(task) = maybe_task { + self.run_task(cx, core, task) + } else { + Ok(core) + } } fn next_task(&self, core: &mut Core) -> Option { @@ -605,25 +622,8 @@ impl Worker { core.run_queue.max_capacity() / 2, ); - // The worker is currently idle, pull a batch of work from the - // injection queue. We don't want to pull *all* the work so other - // workers can also get some. - let n = usize::min( - self.shared().inject.len() / self.shared().remotes.len() + 1, - cap, - ); - let mut synced = self.shared().synced.lock(); - // safety: passing in the correct `inject::Synced`. - let mut tasks = unsafe { self.shared().inject.pop_n(&mut synced.inject, n) }; - - // Pop the first task to return immedietly - let ret = tasks.next(); - - // Push the rest of the on the run queue - core.run_queue.push_back(tasks); - - ret + self.next_remote_task_batch(&mut synced, core, cap) } } @@ -633,10 +633,40 @@ impl Worker { } let mut synced = self.shared().synced.lock(); - // safety: passing in correct `idle::Synced` + self.next_remote_task_synced(&mut synced) + } + + fn next_remote_task_synced(&self, synced: &mut Synced) -> Option { + // safety: we only have access to a valid `Synced` in this file. unsafe { self.shared().inject.pop(&mut synced.inject) } } + fn next_remote_task_batch( + &self, + synced: &mut Synced, + core: &mut Core, + max: usize, + ) -> Option { + // The worker is currently idle, pull a batch of work from the + // injection queue. We don't want to pull *all* the work so other + // workers can also get some. + let n = usize::min( + self.shared().inject.len() / self.shared().remotes.len() + 1, + max, + ); + + // safety: passing in the correct `inject::Synced`. + let mut tasks = unsafe { self.shared().inject.pop_n(&mut synced.inject, n) }; + + // Pop the first task to return immedietly + let ret = tasks.next(); + + // Push the rest of the on the run queue + core.run_queue.push_back(tasks); + + ret + } + fn next_local_task(&self, core: &mut Core) -> Option { core.lifo_slot.take().or_else(|| core.run_queue.pop()) } @@ -768,13 +798,6 @@ impl Worker { core.stats.end_processing_scheduled_tasks(); - /* - // Call `park` with a 0 timeout. This enables the I/O driver, timer, ... - // to run without actually putting the thread to sleep. - core = self.park_timeout(core, Some(Duration::from_millis(0))); - */ - todo!(); - // Run regularly scheduled maintenance self.maintenance(&mut core); @@ -786,6 +809,15 @@ impl Worker { /// Runs maintenance work such as checking the pool's state. fn maintenance(&self, core: &mut Core) { + /* + // Call `park` with a 0 timeout. This enables the I/O driver, timer, ... + // to run without actually putting the thread to sleep. + core = self.park_timeout(core, Some(Duration::from_millis(0))); + */ + if true { + todo!(); + } + core.stats.submit(&self.shared().worker_metrics[core.index]); if !core.is_shutdown { @@ -800,6 +832,57 @@ impl Worker { } } + fn park_yield(&self, mut core: Box) -> Box { + /* + // Call `park` with a 0 timeout. This enables the I/O driver, timer, ... + // to run without actually putting the thread to sleep. + core = self.park_timeout(core, Some(Duration::from_millis(0))); + */ + if true { + todo!(); + } + + self.maintenance(&mut core); + core + } + + fn park(&self, cx: &Context, mut core: Box) -> RunResult { + if let Some(f) = &self.shared().config.before_park { + f(); + } + + if self.transition_to_parked(&mut core) { + debug_assert!(!core.is_shutdown); + debug_assert!(!core.is_traced); + + core.stats.about_to_park(); + core = self.do_park(cx, core)?; + } else { + // Just run maintenance and carry on + core = self.park_yield(core); + } + + if let Some(f) = &self.shared().config.after_unpark { + f(); + } + + Ok(core) + } + + fn do_park(&self, cx: &Context, mut core: Box) -> RunResult { + let mut synced = self.shared().synced.lock(); + + // Return `core` to shared + synced.cores.push(core); + + if let Some(mut driver) = synced.driver.take() { + todo!() + } else { + synced = self.shared().condvar.wait(synced).unwrap(); + self.acquire_core(cx, synced) + } + } + fn transition_to_searching(&self, core: &mut Core) -> bool { if !core.is_searching { core.is_searching = self.shared().idle.transition_worker_to_searching(); @@ -822,6 +905,39 @@ impl Worker { } } + /// Prepares the worker state for parking. + /// + /// Returns true if the transition happened, false if there is work to do first. + fn transition_to_parked(&self, core: &mut Core) -> bool { + // Workers should not park if they have work to do + if core.lifo_slot.is_some() || core.run_queue.has_tasks() || core.is_traced { + return false; + } + + // When the final worker transitions **out** of searching to parked, it + // must check all the queues one last time in case work materialized + // between the last work scan and transitioning out of searching. + let is_last_searcher = self + .shared() + .idle + .transition_worker_to_parked(todo!(), core.is_searching); + + // The worker is no longer searching. Setting this is the local cache + // only. + core.is_searching = false; + + if is_last_searcher { + // worker.handle.notify_if_work_pending(); + todo!() + } + + true + } + + fn transition_from_parked(&self, core: &mut Core) -> bool { + todo!() + } + /// Signals all tasks to shut down, and waits for them to complete. Must run /// before we enter the single-threaded phase of shutdown processing. fn pre_shutdown(&self, core: &mut Core) { @@ -836,28 +952,31 @@ impl Worker { /// /// If all workers have reached this point, the final cleanup is performed. fn shutdown_core(&self, core: Box) { - let mut cores = self.shared().shutdown_cores.lock(); - cores.push(core); + let mut synced = self.shared().synced.lock(); + synced.cores.push(core); - if cores.len() != self.shared().remotes.len() { + if synced.cores.len() != self.shared().remotes.len() { return; } debug_assert!(self.shared().owned.is_empty()); - for mut core in cores.drain(..) { + for mut core in synced.cores.drain(..) { // Drain tasks from the local queue while self.next_local_task(&mut core).is_some() {} } // Shutdown the driver - let mut driver = self.shared().driver.take().expect("driver missing"); + let mut driver = synced.driver.take().expect("driver missing"); driver.shutdown(&self.handle.driver); // Drain the injection queue // - // We already shut down every task, so we can simply drop the tasks. - while let Some(task) = self.next_remote_task() { + // We already shut down every task, so we can simply drop the tasks. We + // cannot call `next_remote_task()` because we already hold the lock. + // + // safety: passing in correct `idle::Synced` + while let Some(task) = self.next_remote_task_synced(&mut synced) { drop(task); } } @@ -898,6 +1017,65 @@ impl Context { } impl Shared { + pub(super) fn schedule_task(&self, task: Notified, is_yield: bool) { + /* + with_current(|maybe_cx| { + if let Some(cx) = maybe_cx { + // Make sure the task is part of the **current** scheduler. + if self.ptr_eq(&cx.worker.handle) { + // And the current thread still holds a core + if let Some(core) = cx.core.borrow_mut().as_mut() { + self.schedule_local(core, task, is_yield); + return; + } + } + } + + // Otherwise, use the inject queue. + self.push_remote_task(task); + self.notify_parked_remote(); + }) + */ + todo!() + } + + fn schedule_local(&self, core: &mut Core, task: Notified, is_yield: bool) { + /* + core.stats.inc_local_schedule_count(); + + // Spawning from the worker thread. If scheduling a "yield" then the + // task must always be pushed to the back of the queue, enabling other + // tasks to be executed. If **not** a yield, then there is more + // flexibility and the task may go to the front of the queue. + let should_notify = if is_yield || !core.lifo_enabled { + core.run_queue + .push_back_or_overflow(task, self, &mut core.stats); + true + } else { + // Push to the LIFO slot + let prev = core.lifo_slot.take(); + let ret = prev.is_some(); + + if let Some(prev) = prev { + core.run_queue + .push_back_or_overflow(prev, self, &mut core.stats); + } + + core.lifo_slot = Some(task); + + ret + }; + + // Only notify if not currently parked. If `park` is `None`, then the + // scheduling is from a resource driver. As notifications often come in + // batches, the notification is delayed until the park is complete. + if should_notify && core.park.is_some() { + self.notify_parked_local(); + } + */ + todo!() + } + fn notify_parked_local(&self) { /* super::counters::inc_num_inc_notify_local(); @@ -928,6 +1106,13 @@ impl Shared { self.inject.push(&mut synced.inject, task); } } + + pub(super) fn close(&self) { + if self.inject.close(&mut self.synced.lock().inject) { + // self.notify_all(); + todo!() + } + } } impl Overflow> for Shared {