├── .gitignore ├── Cargo.toml ├── LICENSE ├── README.md ├── benches └── executor.rs └── src ├── lib.rs └── thread_pool ├── builder.rs ├── enter.rs ├── handle.rs ├── mod.rs ├── queue.rs ├── random.rs ├── runtime.rs ├── scheduler.rs ├── semaphore.rs ├── task.rs ├── thread.rs └── waker.rs /.gitignore: -------------------------------------------------------------------------------- 1 | .vscode/ 2 | .history/ 3 | 4 | target/ 5 | Cargo.lock 6 | NUL -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "uasync" 3 | version = "0.1.1" 4 | edition = "2021" 5 | license = "MIT" 6 | readme = "README.md" 7 | repository = "https://github.com/kprotty/uasync/" 8 | description = "fast, safe, async executor" 9 | documentation = "https://docs.rs/uasync" 10 | keywords = ["asynchronous", "executor", "spawn"] 11 | categories = ["asynchronous", "concurrency"] 12 | 13 | [[bench]] 14 | name = "executor" 15 | harness = false 16 | 17 | [dev-dependencies] 18 | bencher = "0.1" 19 | num_cpus = "1" 20 | easy-parallel = "3" 21 | futures-lite = "1" 22 | async-channel = "1" 23 | async-executor = "1" 24 | async-task = "4" 25 | tokio = { version = "1", features = ["rt-multi-thread", "sync"] } -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 kprotty 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | uAsync 2 | [![Crates.io](https://img.shields.io/crates/v/uasync.svg)](https://crates.io/crates/uasync) 3 | [![Documentation](https://docs.rs/uasync/badge.svg)](https://docs.rs/uasync/) 4 | ==== 5 | 6 | A fast, `forbid(unsafe_code)`, no dependency, async executor. 7 | 8 | This crate was made primarily to see if such a thing was pheasable. The executor performs surprisingly well against others and I went through the effort of optimizing it and providing a drop-in API with tokio. Benchmarks are provided but results on my machine are explicitly **not** provided because people always end up drawing weird conclusions from them that are annoying to try and dispute. `uasync` for now should be treated as mostly an experiment. As per the MIT license, I'm not liable for whatever you do with it. I will most likely accept PRs for documentation, cleanup, and general improvements. 9 | 10 | ```toml 11 | [dependencies] 12 | uasync = "0.1.1" -------------------------------------------------------------------------------- /benches/executor.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | future::Future, 3 | pin::Pin, 4 | sync::{ 5 | atomic::{AtomicUsize, Ordering}, 6 | Arc, 7 | }, 8 | task::{Context, Poll}, 9 | thread, 10 | }; 11 | 12 | #[macro_use] 13 | extern crate bencher; 14 | use bencher::Bencher; 15 | 16 | struct WaitGroup(Arc<(AtomicUsize, thread::Thread)>); 17 | 18 | impl WaitGroup { 19 | fn new() -> Self { 20 | Self(Arc::new((AtomicUsize::new(0), thread::current()))) 21 | } 22 | 23 | fn add(&self) -> Self { 24 | self.0 .0.fetch_add(1, Ordering::Relaxed); 25 | Self(self.0.clone()) 26 | } 27 | 28 | fn wait(&self) { 29 | while self.0 .0.load(Ordering::Acquire) != 0 { 30 | thread::park(); 31 | } 32 | } 33 | 34 | fn done(&self) { 35 | if self.0 .0.fetch_sub(1, Ordering::Release) == 1 { 36 | self.0 .1.unpark(); 37 | } 38 | } 39 | } 40 | 41 | trait BenchExecutor { 42 | type JoinHandle: Future + Send; 43 | fn block_on>(future: F); 44 | fn spawn + Send + 'static>(future: F) -> Self::JoinHandle; 45 | } 46 | 47 | struct TokioJoinHandle(tokio::task::JoinHandle); 48 | 49 | impl Future for TokioJoinHandle { 50 | type Output = T; 51 | 52 | fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll { 53 | match Pin::new(&mut self.0).poll(ctx) { 54 | Poll::Ready(result) => Poll::Ready(result.unwrap()), 55 | Poll::Pending => Poll::Pending, 56 | } 57 | } 58 | } 59 | 60 | struct TokioExecutor; 61 | 62 | impl BenchExecutor for TokioExecutor { 63 | type JoinHandle = TokioJoinHandle<()>; 64 | 65 | fn block_on>(future: F) { 66 | tokio::runtime::Builder::new_multi_thread() 67 | .build() 68 | .unwrap() 69 | .block_on(future) 70 | } 71 | 72 | fn spawn + Send + 'static>(future: F) -> Self::JoinHandle { 73 | TokioJoinHandle(tokio::spawn(future)) 74 | } 75 | } 76 | 77 | struct AsyncExecutor; 78 | 79 | struct AsyncTask(Option>); 80 | 81 | impl Future for AsyncTask { 82 | type Output = T; 83 | 84 | fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll { 85 | Pin::new((&mut self.0).as_mut().unwrap()).poll(ctx) 86 | } 87 | } 88 | 89 | impl Drop for AsyncTask { 90 | fn drop(&mut self) { 91 | self.0.take().unwrap().detach(); 92 | } 93 | } 94 | 95 | static GLOBAL_ASYNC_EXECUTOR: async_executor::Executor = async_executor::Executor::new(); 96 | 97 | impl BenchExecutor for AsyncExecutor { 98 | type JoinHandle = AsyncTask<()>; 99 | 100 | fn block_on>(future: F) { 101 | let (signal, shutdown) = async_channel::unbounded::<()>(); 102 | easy_parallel::Parallel::new() 103 | .each(0..num_cpus::get(), |_| { 104 | futures_lite::future::block_on(GLOBAL_ASYNC_EXECUTOR.run(shutdown.recv())) 105 | }) 106 | .finish(|| { 107 | futures_lite::future::block_on(async { 108 | future.await; 109 | std::mem::drop(signal); 110 | }) 111 | }); 112 | } 113 | 114 | fn spawn + Send + 'static>(future: F) -> Self::JoinHandle { 115 | AsyncTask(Some(GLOBAL_ASYNC_EXECUTOR.spawn(future))) 116 | } 117 | } 118 | 119 | struct UasyncExecutor; 120 | 121 | struct UasyncJoinHandle(uasync::task::JoinHandle); 122 | 123 | impl Future for UasyncJoinHandle { 124 | type Output = T; 125 | 126 | fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll { 127 | match Pin::new(&mut self.0).poll(ctx) { 128 | Poll::Ready(result) => Poll::Ready(result.unwrap()), 129 | Poll::Pending => Poll::Pending, 130 | } 131 | } 132 | } 133 | 134 | impl BenchExecutor for UasyncExecutor { 135 | type JoinHandle = UasyncJoinHandle<()>; 136 | 137 | fn block_on>(future: F) { 138 | uasync::runtime::Builder::new_multi_thread() 139 | .worker_threads(num_cpus::get()) 140 | .build() 141 | .unwrap() 142 | .block_on(future) 143 | } 144 | 145 | fn spawn + Send + 'static>(future: F) -> Self::JoinHandle { 146 | UasyncJoinHandle(uasync::spawn(future)) 147 | } 148 | } 149 | 150 | fn bench_spawn_inject_uasync(b: &mut Bencher) { 151 | bench_spawn_inject::(b) 152 | } 153 | 154 | fn bench_spawn_inject_tokio(b: &mut Bencher) { 155 | bench_spawn_inject::(b) 156 | } 157 | 158 | fn bench_spawn_inject_async_executor(b: &mut Bencher) { 159 | bench_spawn_inject::(b) 160 | } 161 | 162 | fn bench_spawn_inject(b: &mut Bencher) { 163 | E::block_on(async move { 164 | b.iter(|| { 165 | E::spawn(async move {}); 166 | }); 167 | }); 168 | } 169 | 170 | benchmark_group!( 171 | spawn_inject, 172 | bench_spawn_inject_uasync, 173 | bench_spawn_inject_tokio, 174 | bench_spawn_inject_async_executor 175 | ); 176 | 177 | fn bench_spawn_in_worker_uasync(b: &mut Bencher) { 178 | bench_spawn_in_worker::(b) 179 | } 180 | 181 | fn bench_spawn_in_worker_tokio(b: &mut Bencher) { 182 | bench_spawn_in_worker::(b) 183 | } 184 | 185 | fn bench_spawn_in_worker_async_executor(b: &mut Bencher) { 186 | bench_spawn_in_worker::(b) 187 | } 188 | 189 | fn bench_spawn_in_worker(b: &mut Bencher) { 190 | E::block_on(async move { 191 | let b_ptr = b as *mut Bencher as usize; 192 | E::spawn(async move { 193 | let b = unsafe { &mut *(b_ptr as *mut Bencher) }; 194 | b.iter(|| { 195 | E::spawn(async move {}); 196 | }) 197 | }) 198 | .await; 199 | }); 200 | } 201 | 202 | benchmark_group!( 203 | spawn_in_worker, 204 | bench_spawn_in_worker_uasync, 205 | bench_spawn_in_worker_tokio, 206 | bench_spawn_in_worker_async_executor 207 | ); 208 | 209 | fn bench_multi_spawner_uasync(b: &mut Bencher) { 210 | bench_multi_spawner::(b) 211 | } 212 | 213 | fn bench_multi_spawner_tokio(b: &mut Bencher) { 214 | bench_multi_spawner::(b) 215 | } 216 | 217 | fn bench_multi_spawner_async_executor(b: &mut Bencher) { 218 | bench_multi_spawner::(b) 219 | } 220 | 221 | fn bench_multi_spawner(b: &mut Bencher) { 222 | E::block_on(async move { 223 | let wg = WaitGroup::new(); 224 | b.iter(|| { 225 | for _ in 0..10 { 226 | let wg = wg.add(); 227 | E::spawn(async move { 228 | let handles = (0..1000).map(|_| E::spawn(async move {})); 229 | for handle in handles.collect::>() { 230 | handle.await; 231 | } 232 | wg.done(); 233 | }); 234 | } 235 | wg.wait(); 236 | }); 237 | }); 238 | } 239 | 240 | benchmark_group!( 241 | multi_spawner, 242 | bench_multi_spawner_uasync, 243 | bench_multi_spawner_tokio, 244 | bench_multi_spawner_async_executor, 245 | ); 246 | 247 | fn bench_ping_pong_uasync(b: &mut Bencher) { 248 | bench_ping_pong::(b) 249 | } 250 | 251 | fn bench_ping_pong_tokio(b: &mut Bencher) { 252 | bench_ping_pong::(b) 253 | } 254 | 255 | fn bench_ping_pong_async_executor(b: &mut Bencher) { 256 | bench_ping_pong::(b) 257 | } 258 | 259 | fn bench_ping_pong(b: &mut Bencher) { 260 | E::block_on(async move { 261 | let wg = WaitGroup::new(); 262 | b.iter(|| { 263 | for _ in 0..1000 { 264 | let wg = wg.add(); 265 | E::spawn(async move { 266 | let (tx1, rx1) = tokio::sync::oneshot::channel(); 267 | let (tx2, rx2) = tokio::sync::oneshot::channel(); 268 | 269 | E::spawn(async move { 270 | rx1.await.unwrap(); 271 | tx2.send(()).unwrap(); 272 | }); 273 | 274 | tx1.send(()).unwrap(); 275 | rx2.await.unwrap(); 276 | wg.done(); 277 | }); 278 | } 279 | wg.wait(); 280 | }); 281 | }); 282 | } 283 | 284 | benchmark_group!( 285 | ping_pong, 286 | bench_ping_pong_uasync, 287 | bench_ping_pong_tokio, 288 | bench_ping_pong_async_executor, 289 | ); 290 | 291 | fn bench_chain_uasync(b: &mut Bencher) { 292 | bench_chain::(b) 293 | } 294 | 295 | fn bench_chain_tokio(b: &mut Bencher) { 296 | bench_chain::(b) 297 | } 298 | 299 | fn bench_chain_async_executor(b: &mut Bencher) { 300 | bench_chain::(b) 301 | } 302 | 303 | fn bench_chain(b: &mut Bencher) { 304 | E::block_on(async move { 305 | let wg = WaitGroup::new(); 306 | b.iter(|| { 307 | fn chain_iter(iter: usize, wg: WaitGroup) { 308 | match iter { 309 | 0 => wg.done(), 310 | n => std::mem::drop(E::spawn(async move { chain_iter::(n - 1, wg) })), 311 | } 312 | } 313 | 314 | chain_iter::(1000, wg.add()); 315 | wg.wait(); 316 | }); 317 | }); 318 | } 319 | 320 | benchmark_group!( 321 | chain, 322 | bench_chain_uasync, 323 | bench_chain_tokio, 324 | bench_chain_async_executor, 325 | ); 326 | 327 | fn bench_yield_uasync(b: &mut Bencher) { 328 | bench_yield::(b) 329 | } 330 | 331 | fn bench_yield_tokio(b: &mut Bencher) { 332 | bench_yield::(b) 333 | } 334 | 335 | fn bench_yield_async_executor(b: &mut Bencher) { 336 | bench_yield::(b) 337 | } 338 | 339 | fn bench_yield(b: &mut Bencher) { 340 | E::block_on(async move { 341 | let wg = WaitGroup::new(); 342 | b.iter(|| { 343 | for _ in 0..1000 { 344 | let wg = wg.add(); 345 | E::spawn(async move { 346 | for _ in 0..200 { 347 | tokio::task::yield_now().await; 348 | } 349 | wg.done(); 350 | }); 351 | } 352 | wg.wait(); 353 | }); 354 | }); 355 | } 356 | 357 | benchmark_group!( 358 | yield_now, 359 | bench_yield_uasync, 360 | bench_yield_tokio, 361 | bench_yield_async_executor, 362 | ); 363 | 364 | benchmark_main!( 365 | spawn_inject, 366 | spawn_in_worker, 367 | multi_spawner, 368 | ping_pong, 369 | chain, 370 | yield_now, 371 | ); 372 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | #![forbid(unsafe_code)] 2 | #![warn( 3 | rust_2018_idioms, 4 | unreachable_pub, 5 | // missing_docs, 6 | missing_debug_implementations, 7 | )] 8 | 9 | mod thread_pool; 10 | 11 | pub mod runtime { 12 | pub use crate::thread_pool::{ 13 | builder::Builder, 14 | enter::EnterGuard, 15 | handle::{Handle, TryCurrentError}, 16 | runtime::Runtime, 17 | }; 18 | } 19 | 20 | pub mod task { 21 | pub use crate::thread_pool::task::{spawn, yield_now, JoinError, JoinHandle}; 22 | } 23 | 24 | pub use task::spawn; 25 | -------------------------------------------------------------------------------- /src/thread_pool/builder.rs: -------------------------------------------------------------------------------- 1 | use super::{handle::Handle, runtime::Runtime, scheduler::Scheduler}; 2 | use std::{fmt, io, num::NonZeroUsize, sync::Arc, thread}; 3 | 4 | pub struct Builder { 5 | stack_size: Option, 6 | worker_threads: Option, 7 | on_thread_park: Option>, 8 | on_thread_unpark: Option>, 9 | on_thread_start: Option>, 10 | on_thread_stop: Option>, 11 | on_thread_name: Option String + Send + Sync + 'static>>, 12 | } 13 | 14 | impl fmt::Debug for Builder { 15 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 16 | f.debug_struct("Builder").finish() 17 | } 18 | } 19 | 20 | impl Builder { 21 | pub fn new_multi_thread() -> Self { 22 | Self { 23 | stack_size: None, 24 | worker_threads: None, 25 | on_thread_park: None, 26 | on_thread_unpark: None, 27 | on_thread_start: None, 28 | on_thread_stop: None, 29 | on_thread_name: None, 30 | } 31 | } 32 | 33 | pub fn worker_threads(&mut self, worker_threads: usize) -> &mut Self { 34 | self.worker_threads = NonZeroUsize::new(worker_threads); 35 | self 36 | } 37 | 38 | pub fn thread_stack_size(&mut self, stack_size: usize) -> &mut Self { 39 | self.stack_size = NonZeroUsize::new(stack_size); 40 | self 41 | } 42 | 43 | pub fn thread_name(&mut self, name: impl Into) -> &mut Self { 44 | let name: String = name.into(); 45 | self.thread_name_fn(move || name.clone()) 46 | } 47 | 48 | pub fn thread_name_fn( 49 | &mut self, 50 | callback: impl Fn() -> String + Send + Sync + 'static, 51 | ) -> &mut Self { 52 | self.on_thread_name = Some(Arc::new(callback)); 53 | self 54 | } 55 | 56 | pub fn on_thread_park(&mut self, callback: impl Fn() + Send + Sync + 'static) -> &mut Self { 57 | self.on_thread_park = Some(Arc::new(callback)); 58 | self 59 | } 60 | 61 | pub fn on_thread_unpark(&mut self, callback: impl Fn() + Send + Sync + 'static) -> &mut Self { 62 | self.on_thread_unpark = Some(Arc::new(callback)); 63 | self 64 | } 65 | 66 | pub fn on_thread_start(&mut self, callback: impl Fn() + Send + Sync + 'static) -> &mut Self { 67 | self.on_thread_start = Some(Arc::new(callback)); 68 | self 69 | } 70 | 71 | pub fn on_thread_stop(&mut self, callback: impl Fn() + Send + Sync + 'static) -> &mut Self { 72 | self.on_thread_stop = Some(Arc::new(callback)); 73 | self 74 | } 75 | 76 | pub fn build(&mut self) -> io::Result { 77 | let scheduler = Arc::new(Scheduler::new( 78 | self.worker_threads.or(NonZeroUsize::new(1)).unwrap(), 79 | self.on_thread_park.as_ref().map(Arc::clone), 80 | self.on_thread_unpark.as_ref().map(Arc::clone), 81 | )); 82 | 83 | for queue_index in 0..scheduler.worker_threads().get() { 84 | let mut builder = thread::Builder::new(); 85 | if let Some(on_thread_name) = self.on_thread_name.as_ref() { 86 | builder = builder.name((on_thread_name)()); 87 | } 88 | if let Some(stack_size) = self.stack_size { 89 | builder = builder.stack_size(stack_size.get()); 90 | } 91 | 92 | let spawned = { 93 | let scheduler = scheduler.clone(); 94 | let on_thread_start = self.on_thread_start.as_ref().map(Arc::clone); 95 | let on_thread_stop = self.on_thread_stop.as_ref().map(Arc::clone); 96 | 97 | builder.spawn(move || { 98 | if let Some(on_thread_start) = on_thread_start { 99 | (on_thread_start)(); 100 | } 101 | 102 | scheduler.run_worker(queue_index); 103 | 104 | if let Some(on_thread_stop) = on_thread_stop { 105 | (on_thread_stop)(); 106 | } 107 | }) 108 | }; 109 | 110 | if let Err(error) = spawned { 111 | scheduler.shutdown(); 112 | return Err(error); 113 | } 114 | } 115 | 116 | Ok(Runtime { 117 | handle: Handle { scheduler }, 118 | shutdown_on_drop: true, 119 | }) 120 | } 121 | } 122 | -------------------------------------------------------------------------------- /src/thread_pool/enter.rs: -------------------------------------------------------------------------------- 1 | use super::{scheduler::Scheduler, thread::Thread}; 2 | use std::{fmt, sync::Arc}; 3 | 4 | pub struct EnterGuard<'a> { 5 | _scheduler: &'a Scheduler, 6 | _thread: Thread, 7 | } 8 | 9 | impl<'a> fmt::Debug for EnterGuard<'a> { 10 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 11 | f.debug_struct("EnterGuard").finish() 12 | } 13 | } 14 | 15 | impl<'a> EnterGuard<'a> { 16 | pub(super) fn from(scheduler: &'a Arc) -> Self { 17 | Self { 18 | _scheduler: &*scheduler, 19 | _thread: Thread::enter(scheduler, None), 20 | } 21 | } 22 | } 23 | -------------------------------------------------------------------------------- /src/thread_pool/handle.rs: -------------------------------------------------------------------------------- 1 | use super::{ 2 | enter::EnterGuard, 3 | scheduler::Scheduler, 4 | task::{JoinHandle, Task}, 5 | thread::Thread, 6 | }; 7 | use std::{fmt, future::Future, sync::Arc}; 8 | 9 | pub struct TryCurrentError { 10 | missing: bool, 11 | } 12 | 13 | impl TryCurrentError { 14 | pub fn is_missing_context(&self) -> bool { 15 | self.missing 16 | } 17 | 18 | pub fn is_thread_local_destroyed(&self) -> bool { 19 | !self.missing 20 | } 21 | } 22 | 23 | impl std::error::Error for TryCurrentError {} 24 | 25 | impl fmt::Debug for TryCurrentError { 26 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 27 | match self.missing { 28 | true => write!(f, "NoContext"), 29 | false => write!(f, "ThreadLocalDestroyed"), 30 | } 31 | } 32 | } 33 | 34 | impl fmt::Display for TryCurrentError { 35 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 36 | match self.missing { 37 | true => f.write_str(Thread::CONTEXT_MISSING_ERROR), 38 | false => f.write_str(Thread::CONTEXT_DESTROYED_ERROR), 39 | } 40 | } 41 | } 42 | 43 | pub struct Handle { 44 | pub(super) scheduler: Arc, 45 | } 46 | 47 | impl fmt::Debug for Handle { 48 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 49 | f.debug_struct("Handle").finish() 50 | } 51 | } 52 | 53 | impl Clone for Handle { 54 | fn clone(&self) -> Self { 55 | Self { 56 | scheduler: self.scheduler.clone(), 57 | } 58 | } 59 | } 60 | 61 | impl Handle { 62 | pub fn current() -> Self { 63 | Self { 64 | scheduler: Thread::current().scheduler.clone(), 65 | } 66 | } 67 | 68 | pub fn try_current() -> Result { 69 | match Thread::try_current() { 70 | Ok(Some(thread)) => Ok(Self { 71 | scheduler: thread.scheduler.clone(), 72 | }), 73 | Ok(None) => Err(TryCurrentError { missing: true }), 74 | Err(_) => Err(TryCurrentError { missing: false }), 75 | } 76 | } 77 | 78 | pub fn enter(&self) -> EnterGuard<'_> { 79 | EnterGuard::from(&self.scheduler) 80 | } 81 | 82 | pub fn block_on(&self, future: F) -> F::Output { 83 | let future = Box::pin(future); 84 | Task::block_on(&self.scheduler, future) 85 | } 86 | 87 | pub fn spawn(&self, future: F) -> JoinHandle 88 | where 89 | F: Future + Send + 'static, 90 | F::Output: Send + 'static, 91 | { 92 | let future = Box::pin(future); 93 | let thread = Thread::enter(&self.scheduler, None); 94 | Task::spawn(&self.scheduler, Some(&thread), future) 95 | } 96 | } 97 | -------------------------------------------------------------------------------- /src/thread_pool/mod.rs: -------------------------------------------------------------------------------- 1 | pub(super) mod builder; 2 | pub(super) mod enter; 3 | pub(super) mod handle; 4 | mod queue; 5 | mod random; 6 | pub(super) mod runtime; 7 | mod scheduler; 8 | mod semaphore; 9 | pub(super) mod task; 10 | mod thread; 11 | mod waker; 12 | -------------------------------------------------------------------------------- /src/thread_pool/queue.rs: -------------------------------------------------------------------------------- 1 | use super::task::Runnable; 2 | use std::{ 3 | collections::VecDeque, 4 | sync::atomic::{AtomicBool, Ordering}, 5 | sync::{Arc, Mutex}, 6 | }; 7 | 8 | #[derive(Default)] 9 | pub(super) struct Queue { 10 | pending: AtomicBool, 11 | deque: Mutex>>, 12 | } 13 | 14 | impl Queue { 15 | pub(super) fn is_empty(&self) -> bool { 16 | !self.pending.load(Ordering::Acquire) 17 | } 18 | 19 | pub(super) fn push(&self, runnable: Arc) { 20 | let mut deque = self.deque.lock().unwrap(); 21 | deque.push_back(runnable); 22 | self.pending.store(true, Ordering::Relaxed); 23 | } 24 | 25 | pub(super) fn pop(&self) -> Option> { 26 | if self.is_empty() { 27 | return None; 28 | } 29 | 30 | let mut deque = self.deque.lock().unwrap(); 31 | let runnable = deque.pop_front()?; 32 | 33 | self.pending.store(deque.len() > 0, Ordering::Relaxed); 34 | Some(runnable) 35 | } 36 | 37 | pub(super) fn steal(&self) -> Result>, ()> { 38 | if self.is_empty() { 39 | return Ok(None); 40 | } 41 | 42 | let mut deque = match self.deque.try_lock() { 43 | Ok(guard) => guard, 44 | Err(_) => return Err(()), 45 | }; 46 | 47 | let runnable = match deque.pop_front() { 48 | Some(runnable) => runnable, 49 | None => return Ok(None), 50 | }; 51 | 52 | self.pending.store(deque.len() > 0, Ordering::Relaxed); 53 | Ok(Some(runnable)) 54 | } 55 | } 56 | -------------------------------------------------------------------------------- /src/thread_pool/random.rs: -------------------------------------------------------------------------------- 1 | use std::{mem::replace, num::NonZeroUsize}; 2 | 3 | pub(super) struct Rng { 4 | xorshift: NonZeroUsize, 5 | } 6 | 7 | impl Rng { 8 | pub(super) fn new(seed: usize) -> Self { 9 | Self { 10 | xorshift: NonZeroUsize::new(seed) 11 | .or(NonZeroUsize::new(0xdeadbeef)) 12 | .unwrap(), 13 | } 14 | } 15 | } 16 | 17 | impl Iterator for Rng { 18 | type Item = NonZeroUsize; 19 | 20 | fn next(&mut self) -> Option { 21 | let shifts = match usize::BITS { 22 | 64 => (13, 7, 17), 23 | 32 => (13, 17, 5), 24 | _ => unreachable!(), 25 | }; 26 | 27 | let mut xs = self.xorshift.get(); 28 | xs ^= xs >> shifts.0; 29 | xs ^= xs << shifts.1; 30 | xs ^= xs >> shifts.2; 31 | 32 | self.xorshift = NonZeroUsize::new(xs).unwrap(); 33 | Some(self.xorshift) 34 | } 35 | } 36 | 37 | pub(super) struct RandomSequence { 38 | range: NonZeroUsize, 39 | co_prime: NonZeroUsize, 40 | } 41 | 42 | impl RandomSequence { 43 | pub(super) fn new(range: NonZeroUsize) -> Self { 44 | fn gcd(mut a: usize, mut b: usize) -> usize { 45 | while a != b { 46 | if a > b { 47 | a -= b; 48 | } else { 49 | b -= a; 50 | } 51 | } 52 | a 53 | } 54 | 55 | Self { 56 | range, 57 | co_prime: (range.get() / 2..range.get()) 58 | .find(|&n| gcd(n, range.get()) == 1) 59 | .and_then(NonZeroUsize::new) 60 | .or(NonZeroUsize::new(1)) 61 | .unwrap(), 62 | } 63 | } 64 | 65 | pub(super) fn iter(&self, seed: usize) -> impl Iterator { 66 | let prime = self.co_prime.get(); 67 | let range = self.range.get(); 68 | let mut index = seed % range; 69 | 70 | (0..range).map(move |_| { 71 | let mut new_index = index + prime; 72 | if new_index >= range { 73 | new_index -= range; 74 | } 75 | 76 | assert!(new_index < range); 77 | replace(&mut index, new_index) 78 | }) 79 | } 80 | } 81 | -------------------------------------------------------------------------------- /src/thread_pool/runtime.rs: -------------------------------------------------------------------------------- 1 | use super::{ 2 | builder::Builder, 3 | enter::EnterGuard, 4 | handle::Handle, 5 | task::{JoinHandle, Task}, 6 | }; 7 | use std::{fmt, future::Future, io, mem::replace, time::Duration}; 8 | 9 | pub struct Runtime { 10 | pub(super) handle: Handle, 11 | pub(super) shutdown_on_drop: bool, 12 | } 13 | 14 | impl fmt::Debug for Runtime { 15 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 16 | f.debug_struct("Runtime").finish() 17 | } 18 | } 19 | 20 | impl Runtime { 21 | pub fn new() -> io::Result { 22 | Builder::new_multi_thread().build() 23 | } 24 | 25 | pub fn handle(&self) -> &Handle { 26 | &self.handle 27 | } 28 | 29 | pub fn enter(&self) -> EnterGuard<'_> { 30 | self.handle.enter() 31 | } 32 | 33 | pub fn block_on(&self, future: F) -> F::Output { 34 | let future = Box::pin(future); 35 | Task::block_on(&self.handle.scheduler, future) 36 | } 37 | 38 | pub fn spawn(&self, future: F) -> JoinHandle 39 | where 40 | F: Future + Send + 'static, 41 | F::Output: Send + 'static, 42 | { 43 | let future = Box::pin(future); 44 | Task::spawn(&self.handle.scheduler, None, future) 45 | } 46 | 47 | pub fn shutdown_background(self) { 48 | let mut mut_self = self; 49 | mut_self.shutdown_and_join(None); 50 | } 51 | 52 | pub fn shutdown_timeout(self, timeout: Duration) { 53 | let mut mut_self = self; 54 | mut_self.shutdown_and_join(Some(Some(timeout))); 55 | } 56 | 57 | fn shutdown_and_join(&mut self, join_timeout: Option>) { 58 | if replace(&mut self.shutdown_on_drop, false) { 59 | self.handle.scheduler.shutdown(); 60 | 61 | if let Some(timeout) = join_timeout { 62 | self.handle.scheduler.join(timeout); 63 | } 64 | } 65 | } 66 | } 67 | 68 | impl Drop for Runtime { 69 | fn drop(&mut self) { 70 | self.shutdown_and_join(Some(None)) 71 | } 72 | } 73 | -------------------------------------------------------------------------------- /src/thread_pool/scheduler.rs: -------------------------------------------------------------------------------- 1 | use super::{ 2 | queue::Queue, 3 | random::{RandomSequence, Rng}, 4 | semaphore::Semaphore, 5 | task::Runnable, 6 | thread::Thread, 7 | }; 8 | use std::{ 9 | num::NonZeroUsize, 10 | sync::atomic::{fence, AtomicUsize, Ordering}, 11 | sync::Arc, 12 | time::Duration, 13 | }; 14 | 15 | pub(super) struct Scheduler { 16 | join_semaphore: Semaphore, 17 | idle_semaphore: Semaphore, 18 | rand_seq: RandomSequence, 19 | tasks: AtomicUsize, 20 | state: AtomicUsize, 21 | injector: Queue, 22 | run_queues: Box<[Queue]>, 23 | on_thread_park: Option>, 24 | on_thread_unpark: Option>, 25 | } 26 | 27 | impl Scheduler { 28 | pub(super) fn new( 29 | worker_threads: NonZeroUsize, 30 | on_thread_park: Option>, 31 | on_thread_unpark: Option>, 32 | ) -> Self { 33 | let worker_threads = worker_threads.get().min(Self::STATE_MASK); 34 | let worker_threads = NonZeroUsize::new(worker_threads).unwrap(); 35 | 36 | Self { 37 | join_semaphore: Semaphore::default(), 38 | idle_semaphore: Semaphore::default(), 39 | rand_seq: RandomSequence::new(worker_threads), 40 | tasks: AtomicUsize::new(0), 41 | state: AtomicUsize::new(0), 42 | injector: Queue::default(), 43 | run_queues: (0..worker_threads.get()) 44 | .map(|_| Queue::default()) 45 | .collect(), 46 | on_thread_park, 47 | on_thread_unpark, 48 | } 49 | } 50 | 51 | pub(super) fn worker_threads(&self) -> NonZeroUsize { 52 | NonZeroUsize::new(self.run_queues.len()).unwrap() 53 | } 54 | 55 | pub(super) fn schedule( 56 | &self, 57 | runnable: Arc, 58 | thread: Option<&Thread>, 59 | be_fair: bool, 60 | ) { 61 | if let Some(thread) = thread { 62 | if let Some(queue_index) = thread.queue_index { 63 | let mut runnable = Some(runnable); 64 | if be_fair { 65 | thread.be_fair.set(true); 66 | } else { 67 | runnable = thread.lifo_slot.replace(runnable); 68 | } 69 | 70 | if let Some(runnable) = runnable { 71 | self.run_queues[queue_index].push(runnable); 72 | self.unpark(); 73 | } 74 | 75 | return; 76 | } 77 | } 78 | 79 | self.injector.push(runnable); 80 | fence(Ordering::SeqCst); 81 | self.unpark() 82 | } 83 | 84 | const STATE_BITS: u32 = (usize::BITS - 1) / 2; 85 | const STATE_MASK: usize = (1 << Self::STATE_BITS) - 1; 86 | const SHUTDOWN_MASK: usize = 1 << (usize::BITS - 1); 87 | 88 | #[allow(clippy::erasing_op)] 89 | const IDLE_SHIFT: u32 = Self::STATE_BITS * 0; 90 | 91 | #[allow(clippy::identity_op)] 92 | const SEARCHING_SHIFT: u32 = Self::STATE_BITS * 1; 93 | 94 | fn unpark(&self) { 95 | self.state 96 | .fetch_update(Ordering::Release, Ordering::Relaxed, |state| { 97 | if state & Self::SHUTDOWN_MASK != 0 { 98 | return None; 99 | } 100 | 101 | let searching = (state >> Self::SEARCHING_SHIFT) & Self::STATE_MASK; 102 | assert!(searching <= self.run_queues.len()); 103 | if searching > 0 { 104 | return None; 105 | } 106 | 107 | let mut idle = (state >> Self::IDLE_SHIFT) & Self::STATE_MASK; 108 | assert!(idle <= self.run_queues.len()); 109 | idle = idle.checked_sub(1)?; 110 | 111 | Some((1 << Self::SEARCHING_SHIFT) | (idle << Self::IDLE_SHIFT)) 112 | }) 113 | .map(|_| self.idle_semaphore.post(1)) 114 | .unwrap_or(()) 115 | } 116 | 117 | fn search(&self, was_searching: bool) -> bool { 118 | if was_searching { 119 | return true; 120 | } 121 | 122 | let state = self.state.load(Ordering::Relaxed); 123 | let searching = (state >> Self::SEARCHING_SHIFT) & Self::STATE_MASK; 124 | assert!(searching <= self.run_queues.len()); 125 | if (2 * searching) >= self.run_queues.len() { 126 | return false; 127 | } 128 | 129 | let state = self 130 | .state 131 | .fetch_add(1 << Self::SEARCHING_SHIFT, Ordering::Acquire); 132 | 133 | let searching = (state >> Self::SEARCHING_SHIFT) & Self::STATE_MASK; 134 | assert!(searching < self.run_queues.len()); 135 | true 136 | } 137 | 138 | fn discovered(&self, was_searching: bool) -> bool { 139 | if was_searching { 140 | let state = self 141 | .state 142 | .fetch_sub(1 << Self::SEARCHING_SHIFT, Ordering::Release); 143 | 144 | let searching = (state >> Self::SEARCHING_SHIFT) & Self::STATE_MASK; 145 | assert!(searching <= self.run_queues.len()); 146 | assert!(searching > 0); 147 | 148 | if searching == 1 { 149 | self.unpark(); 150 | } 151 | } 152 | 153 | false 154 | } 155 | 156 | fn park(&self, was_searching: bool) -> Option { 157 | let mut update: usize = 1 << Self::IDLE_SHIFT; 158 | if was_searching { 159 | update = update.wrapping_sub(1 << Self::SEARCHING_SHIFT); 160 | } 161 | 162 | let state = self.state.fetch_add(update, Ordering::SeqCst); 163 | let idle = (state >> Self::IDLE_SHIFT) & Self::STATE_MASK; 164 | assert!(idle < self.run_queues.len()); 165 | 166 | let searching = (state >> Self::SEARCHING_SHIFT) & Self::STATE_MASK; 167 | assert!(searching <= self.run_queues.len()); 168 | assert!(searching >= was_searching as usize); 169 | 170 | if state & Self::SHUTDOWN_MASK != 0 { 171 | let state = self 172 | .state 173 | .fetch_sub(1 << Self::IDLE_SHIFT, Ordering::Relaxed); 174 | 175 | let idle = (state >> Self::IDLE_SHIFT) & Self::STATE_MASK; 176 | assert!(idle <= self.run_queues.len()); 177 | assert!(idle > 0); 178 | return None; 179 | } 180 | 181 | if was_searching && searching == 1 && !self.is_empty() { 182 | self.unpark(); 183 | } 184 | 185 | if let Some(callback) = self.on_thread_park.as_ref() { 186 | (callback)(); 187 | } 188 | 189 | self.idle_semaphore.wait(None); 190 | 191 | if let Some(callback) = self.on_thread_unpark.as_ref() { 192 | (callback)(); 193 | } 194 | 195 | Some(true) 196 | } 197 | 198 | fn is_empty(&self) -> bool { 199 | self.run_queues 200 | .iter() 201 | .map(|queue| queue.is_empty()) 202 | .find(|&is_empty| !is_empty) 203 | .unwrap_or_else(|| self.injector.is_empty()) 204 | } 205 | 206 | pub(super) fn shutdown(&self) { 207 | self.state 208 | .fetch_update(Ordering::AcqRel, Ordering::Relaxed, |mut state| { 209 | let idle = (state >> Self::IDLE_SHIFT) & Self::STATE_MASK; 210 | state -= idle << Self::IDLE_SHIFT; 211 | state += idle << Self::SEARCHING_SHIFT; 212 | 213 | let searching = (state >> Self::SEARCHING_SHIFT) & Self::STATE_MASK; 214 | assert!(searching <= self.run_queues.len()); 215 | 216 | state |= Self::SHUTDOWN_MASK; 217 | Some(state) 218 | }) 219 | .map(|state| { 220 | let idle = (state >> Self::IDLE_SHIFT) & Self::STATE_MASK; 221 | if idle > 0 { 222 | self.idle_semaphore.post(idle); 223 | } 224 | }) 225 | .unwrap_or(()) 226 | } 227 | 228 | pub(super) fn on_task_begin(&self) { 229 | let tasks = self.tasks.fetch_add(1, Ordering::Relaxed); 230 | assert_ne!(tasks, usize::MAX); 231 | } 232 | 233 | pub(super) fn on_task_complete(&self) { 234 | let tasks = self.tasks.fetch_sub(1, Ordering::AcqRel); 235 | assert_ne!(tasks, 0); 236 | 237 | if tasks == 1 { 238 | let state = self.state.load(Ordering::Relaxed); 239 | if state & Self::SHUTDOWN_MASK != 0 { 240 | self.join_semaphore.post(1); 241 | } 242 | } 243 | } 244 | 245 | pub(super) fn join(&self, timeout: Option) { 246 | let mut tasks = self.tasks.load(Ordering::Acquire); 247 | if tasks > 0 { 248 | self.join_semaphore.wait(timeout); 249 | tasks = self.tasks.load(Ordering::Acquire); 250 | } 251 | 252 | if timeout.is_none() { 253 | assert_eq!(tasks, 0); 254 | } 255 | } 256 | 257 | pub(super) fn run_worker(self: &Arc, queue_index: usize) { 258 | let thread = Thread::enter(self, Some(queue_index)); 259 | 260 | let mut tick = queue_index; 261 | let mut is_searching = false; 262 | let mut rng = Rng::new(queue_index); 263 | 264 | let run_queue = &self.run_queues[queue_index]; 265 | loop { 266 | let fairness = match thread.be_fair.take() || (tick % 61 == 0) { 267 | true => self.injector.steal().unwrap_or(None), 268 | false => None, 269 | }; 270 | 271 | let polled = fairness.or_else(|| { 272 | thread.lifo_slot.take().or_else(|| { 273 | run_queue.pop().or_else(|| { 274 | is_searching = self.search(is_searching); 275 | if is_searching { 276 | for _ in 0..32 { 277 | let mut is_empty = false; 278 | let rand_seed = rng.next().unwrap().get(); 279 | 280 | for target_index in self.rand_seq.iter(rand_seed) { 281 | if target_index != queue_index { 282 | match self.run_queues[target_index].steal() { 283 | Ok(Some(runnable)) => return Some(runnable), 284 | Ok(None) => {} 285 | Err(_) => is_empty = false, 286 | } 287 | } 288 | } 289 | 290 | match self.injector.steal() { 291 | Ok(Some(runnable)) => return Some(runnable), 292 | Ok(None) => {} 293 | Err(_) => is_empty = false, 294 | } 295 | 296 | if is_empty { 297 | break; 298 | } 299 | } 300 | } 301 | 302 | None 303 | }) 304 | }) 305 | }); 306 | 307 | if let Some(runnable) = polled { 308 | is_searching = self.discovered(is_searching); 309 | tick = tick.wrapping_add(1); 310 | runnable.run(&thread); 311 | continue; 312 | } 313 | 314 | is_searching = match self.park(is_searching) { 315 | Some(searching) => searching, 316 | None => break, 317 | }; 318 | } 319 | } 320 | } 321 | -------------------------------------------------------------------------------- /src/thread_pool/semaphore.rs: -------------------------------------------------------------------------------- 1 | #![allow(clippy::mutex_atomic)] 2 | 3 | use std::{ 4 | sync::atomic::{AtomicIsize, Ordering}, 5 | sync::{Condvar, Mutex}, 6 | time::Duration, 7 | }; 8 | 9 | #[derive(Default)] 10 | pub(super) struct Semaphore { 11 | value: AtomicIsize, 12 | inner: InnerSemaphore, 13 | } 14 | 15 | impl Semaphore { 16 | pub(super) fn wait(&self, timeout: Option) { 17 | let value = self.value.fetch_sub(1, Ordering::Acquire); 18 | if value > 0 { 19 | return; 20 | } 21 | 22 | if self.inner.wait(timeout) { 23 | return; 24 | } 25 | 26 | let _ = self 27 | .value 28 | .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |state| { 29 | if state < 0 { 30 | Some(state + 1) 31 | } else { 32 | assert!(self.inner.wait(None)); 33 | None 34 | } 35 | }); 36 | } 37 | 38 | pub(super) fn post(&self, n: usize) { 39 | let inc: isize = n.try_into().unwrap(); 40 | let value = self.value.fetch_add(inc, Ordering::Release); 41 | if value >= 0 { 42 | return; 43 | } 44 | 45 | let wake: usize = inc.min(-value).try_into().unwrap(); 46 | self.inner.post(wake) 47 | } 48 | } 49 | 50 | #[derive(Default)] 51 | struct InnerSemaphore { 52 | count: Mutex, 53 | cond: Condvar, 54 | } 55 | 56 | impl InnerSemaphore { 57 | #[cold] 58 | fn wait(&self, timeout: Option) -> bool { 59 | let mut count = self.count.lock().unwrap(); 60 | let cond = |count: &mut usize| *count == 0; 61 | count = match timeout { 62 | None => self.cond.wait_while(count, cond).unwrap(), 63 | Some(dur) => self.cond.wait_timeout_while(count, dur, cond).unwrap().0, 64 | }; 65 | 66 | *count > 0 && { 67 | *count -= 1; 68 | true 69 | } 70 | } 71 | 72 | #[cold] 73 | fn post(&self, wake: usize) { 74 | *self.count.lock().unwrap() += wake; 75 | match wake { 76 | 1 => self.cond.notify_one(), 77 | _ => self.cond.notify_all(), 78 | } 79 | } 80 | } 81 | -------------------------------------------------------------------------------- /src/thread_pool/task.rs: -------------------------------------------------------------------------------- 1 | use super::{scheduler::Scheduler, thread::Thread, waker::AtomicWaker}; 2 | use std::{ 3 | any::Any, 4 | fmt, 5 | future::Future, 6 | io, 7 | mem::{drop, replace}, 8 | panic::{catch_unwind, resume_unwind, AssertUnwindSafe}, 9 | pin::Pin, 10 | sync::atomic::{AtomicBool, AtomicU8, Ordering}, 11 | sync::{Arc, Mutex}, 12 | task::{Context, Poll, Wake, Waker}, 13 | thread, 14 | }; 15 | 16 | pub fn spawn(future: F) -> JoinHandle 17 | where 18 | F: Future + Send + 'static, 19 | F::Output: Send + 'static, 20 | { 21 | let future = Box::pin(future); 22 | let thread = Thread::current(); 23 | Task::spawn(&thread.scheduler, Some(&thread), future) 24 | } 25 | 26 | pub struct JoinHandle { 27 | joinable: Option>>, 28 | } 29 | 30 | impl JoinHandle { 31 | pub fn abort(&self) { 32 | if let Some(joinable) = self.joinable.as_ref() { 33 | if joinable.poll_abort() { 34 | joinable.clone().abort(); 35 | } 36 | } 37 | } 38 | } 39 | 40 | impl fmt::Debug for JoinHandle { 41 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 42 | f.debug_struct("JoinHandle").finish() 43 | } 44 | } 45 | 46 | impl Future for JoinHandle { 47 | type Output = Result; 48 | 49 | fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll { 50 | let joinable = self 51 | .joinable 52 | .take() 53 | .expect("JoinHandle polled after completion"); 54 | 55 | if let Poll::Ready(result) = joinable.poll_join(ctx) { 56 | return Poll::Ready(result); 57 | } 58 | 59 | self.joinable = Some(joinable); 60 | Poll::Pending 61 | } 62 | } 63 | 64 | pub struct JoinError { 65 | panic: Option>, 66 | } 67 | 68 | impl std::error::Error for JoinError {} 69 | 70 | impl From for io::Error { 71 | fn from(this: JoinError) -> io::Error { 72 | io::Error::new( 73 | io::ErrorKind::Other, 74 | match &this.panic { 75 | Some(_) => "task panicked", 76 | None => "task was cancelled", 77 | }, 78 | ) 79 | } 80 | } 81 | 82 | impl fmt::Display for JoinError { 83 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 84 | match &self.panic { 85 | Some(_) => write!(f, "panic"), 86 | None => write!(f, "cancelled"), 87 | } 88 | } 89 | } 90 | 91 | impl fmt::Debug for JoinError { 92 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 93 | match &self.panic { 94 | Some(_) => write!(f, "JoinError::Panic(_)"), 95 | None => write!(f, "JoinError::Cancelled"), 96 | } 97 | } 98 | } 99 | 100 | impl JoinError { 101 | pub fn is_cancelled(&self) -> bool { 102 | self.panic.is_none() 103 | } 104 | 105 | pub fn is_panic(&self) -> bool { 106 | self.panic.is_some() 107 | } 108 | 109 | pub fn try_into_panic(self) -> Result, JoinError> { 110 | match self.panic { 111 | Some(panic) => Ok(panic), 112 | None => Err(self), 113 | } 114 | } 115 | 116 | pub fn into_panic(self) -> Box { 117 | self.try_into_panic().unwrap() 118 | } 119 | } 120 | 121 | pub async fn yield_now() { 122 | struct YieldFuture { 123 | yielded: bool, 124 | } 125 | 126 | impl Future for YieldFuture { 127 | type Output = (); 128 | 129 | fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll { 130 | if self.yielded { 131 | return Poll::Ready(()); 132 | } 133 | 134 | ctx.waker().wake_by_ref(); 135 | self.yielded = true; 136 | Poll::Pending 137 | } 138 | } 139 | 140 | YieldFuture { yielded: false }.await 141 | } 142 | 143 | #[derive(Default)] 144 | struct TaskState { 145 | state: AtomicU8, 146 | } 147 | 148 | impl TaskState { 149 | const IDLE: u8 = 0; 150 | const SCHEDULED: u8 = 1; 151 | const RUNNING: u8 = 2; 152 | const NOTIFIED: u8 = 3; 153 | const ABORTED: u8 = 4; 154 | 155 | fn transition_to_scheduled(&self) -> bool { 156 | self.state 157 | .fetch_update(Ordering::Release, Ordering::Relaxed, |state| match state { 158 | Self::RUNNING => Some(Self::NOTIFIED), 159 | Self::IDLE => Some(Self::SCHEDULED), 160 | _ => None, 161 | }) 162 | .map(|state| state == Self::IDLE) 163 | .unwrap_or(false) 164 | } 165 | 166 | fn transition_to_running(&self) -> bool { 167 | match self.state.swap(Self::RUNNING, Ordering::Acquire) { 168 | Self::SCHEDULED => true, 169 | Self::ABORTED => false, 170 | _ => unreachable!(), 171 | } 172 | } 173 | 174 | fn transition_to_idle(&self) -> bool { 175 | self.state 176 | .fetch_update(Ordering::AcqRel, Ordering::Acquire, |state| match state { 177 | Self::RUNNING => Some(Self::IDLE), 178 | Self::NOTIFIED => Some(Self::SCHEDULED), 179 | Self::ABORTED => None, 180 | _ => unreachable!(), 181 | }) 182 | .map(|state| state == Self::RUNNING) 183 | .unwrap_or(false) 184 | } 185 | 186 | fn transition_to_aborted(&self) -> bool { 187 | self.state.swap(Self::ABORTED, Ordering::Release) == Self::IDLE 188 | } 189 | } 190 | 191 | enum TaskData { 192 | Empty, 193 | Polling(Pin>, Waker), 194 | Ready(F::Output), 195 | Panic(Box), 196 | Aborted, 197 | } 198 | 199 | pub(super) struct Task { 200 | state: TaskState, 201 | waker: AtomicWaker, 202 | data: Mutex>, 203 | scheduler: Arc, 204 | } 205 | 206 | impl Task { 207 | pub(super) fn block_on(scheduler: &Arc, mut future: Pin>) -> F::Output { 208 | struct Parker { 209 | thread: thread::Thread, 210 | unparked: AtomicBool, 211 | } 212 | 213 | impl Parker { 214 | fn new() -> Self { 215 | Self { 216 | thread: thread::current(), 217 | unparked: AtomicBool::new(false), 218 | } 219 | } 220 | 221 | fn park(&self) { 222 | while !self.unparked.swap(false, Ordering::Acquire) { 223 | thread::park(); 224 | } 225 | } 226 | } 227 | 228 | impl Wake for Parker { 229 | fn wake(self: Arc) { 230 | self.wake_by_ref() 231 | } 232 | 233 | fn wake_by_ref(self: &Arc) { 234 | if !self.unparked.swap(true, Ordering::Release) { 235 | self.thread.unpark(); 236 | } 237 | } 238 | } 239 | 240 | let result = { 241 | let parker = Arc::new(Parker::new()); 242 | let waker = Waker::from(parker.clone()); 243 | let mut ctx = Context::from_waker(&waker); 244 | 245 | let thread = Thread::enter(scheduler, None); 246 | if !Arc::ptr_eq(&thread.scheduler, scheduler) { 247 | unreachable!("nested block_on() is not supported"); 248 | } 249 | 250 | scheduler.on_task_begin(); 251 | catch_unwind(AssertUnwindSafe(|| loop { 252 | match future.as_mut().poll(&mut ctx) { 253 | Poll::Ready(output) => break output, 254 | Poll::Pending => parker.park(), 255 | } 256 | })) 257 | }; 258 | 259 | scheduler.on_task_complete(); 260 | match result { 261 | Ok(output) => output, 262 | Err(error) => resume_unwind(error), 263 | } 264 | } 265 | } 266 | 267 | impl Task 268 | where 269 | F: Future + Send + 'static, 270 | F::Output: Send + 'static, 271 | { 272 | pub(super) fn spawn( 273 | scheduler: &Arc, 274 | thread: Option<&Thread>, 275 | future: Pin>, 276 | ) -> JoinHandle { 277 | let task = Arc::new(Task { 278 | state: TaskState::default(), 279 | waker: AtomicWaker::default(), 280 | data: Mutex::new(TaskData::Empty), 281 | scheduler: scheduler.clone(), 282 | }); 283 | 284 | let waker = Waker::from(task.clone()); 285 | *task.data.try_lock().unwrap() = TaskData::Polling(future, waker); 286 | 287 | scheduler.on_task_begin(); 288 | assert!(task.state.transition_to_scheduled()); 289 | scheduler.schedule(task.clone(), thread, false); 290 | 291 | JoinHandle { 292 | joinable: Some(task), 293 | } 294 | } 295 | } 296 | 297 | impl Wake for Task 298 | where 299 | F: Future + Send + 'static, 300 | F::Output: Send + 'static, 301 | { 302 | fn wake(self: Arc) { 303 | if self.state.transition_to_scheduled() { 304 | match Thread::try_current().ok().flatten() { 305 | Some(thread) => thread.scheduler.schedule(self, Some(&thread), false), 306 | None => self.scheduler.schedule(self.clone(), None, false), 307 | } 308 | } 309 | } 310 | 311 | fn wake_by_ref(self: &Arc) { 312 | if self.state.transition_to_scheduled() { 313 | let task = self.clone(); 314 | match Thread::try_current().ok().flatten() { 315 | Some(thread) => thread.scheduler.schedule(task, Some(&thread), false), 316 | None => self.scheduler.schedule(task, None, false), 317 | } 318 | } 319 | } 320 | } 321 | 322 | pub(super) trait Runnable: Send + Sync { 323 | fn run(self: Arc, thread: &Thread); 324 | } 325 | 326 | impl Runnable for Task 327 | where 328 | F: Future + Send + 'static, 329 | F::Output: Send + 'static, 330 | { 331 | fn run(self: Arc, thread: &Thread) { 332 | let mut data = self.data.try_lock().unwrap(); 333 | *data = match self.state.transition_to_running() { 334 | false => TaskData::Aborted, 335 | _ => { 336 | let poll_result = match &mut *data { 337 | TaskData::Polling(future, waker) => { 338 | let future = future.as_mut(); 339 | let mut ctx = Context::from_waker(waker); 340 | catch_unwind(AssertUnwindSafe(|| future.poll(&mut ctx))) 341 | } 342 | _ => unreachable!(), 343 | }; 344 | 345 | match poll_result { 346 | Err(error) => TaskData::Panic(error), 347 | Ok(Poll::Ready(output)) => TaskData::Ready(output), 348 | Ok(Poll::Pending) => { 349 | drop(data); 350 | if self.state.transition_to_idle() { 351 | return; 352 | } 353 | 354 | thread.scheduler.schedule(self, Some(thread), true); 355 | return; 356 | } 357 | } 358 | } 359 | }; 360 | 361 | drop(data); 362 | self.waker.wake(); 363 | thread.scheduler.on_task_complete(); 364 | } 365 | } 366 | 367 | trait Joinable: Send + Sync { 368 | fn poll_join(&self, ctx: &mut Context<'_>) -> Poll>; 369 | fn poll_abort(&self) -> bool; 370 | fn abort(self: Arc); 371 | } 372 | 373 | impl Joinable for Task 374 | where 375 | F: Future + Send + 'static, 376 | F::Output: Send + 'static, 377 | { 378 | fn poll_join(&self, ctx: &mut Context<'_>) -> Poll> { 379 | if self.waker.poll(ctx).is_pending() { 380 | return Poll::Pending; 381 | } 382 | 383 | match replace(&mut *self.data.try_lock().unwrap(), TaskData::Empty) { 384 | TaskData::Ready(output) => Poll::Ready(Ok(output)), 385 | TaskData::Aborted => Poll::Ready(Err(JoinError { panic: None })), 386 | TaskData::Panic(error) => Poll::Ready(Err(JoinError { panic: Some(error) })), 387 | _ => unreachable!(), 388 | } 389 | } 390 | 391 | fn poll_abort(&self) -> bool { 392 | self.state.transition_to_aborted() 393 | } 394 | 395 | fn abort(self: Arc) { 396 | self.wake() 397 | } 398 | } 399 | -------------------------------------------------------------------------------- /src/thread_pool/thread.rs: -------------------------------------------------------------------------------- 1 | use super::{scheduler::Scheduler, task::Runnable}; 2 | use std::{ 3 | cell::{Cell, RefCell}, 4 | ops::Deref, 5 | rc::Rc, 6 | sync::Arc, 7 | }; 8 | 9 | pub(super) struct ThreadContext { 10 | pub(super) scheduler: Arc, 11 | pub(super) queue_index: Option, 12 | pub(super) be_fair: Cell, 13 | pub(super) lifo_slot: Cell>>, 14 | } 15 | 16 | impl ThreadContext { 17 | fn try_with(f: impl FnOnce(&mut Option>) -> F) -> Result { 18 | thread_local!(static TLS: RefCell>> = RefCell::new(None)); 19 | TLS.try_with(|ref_cell| f(&mut *ref_cell.borrow_mut())) 20 | .map_err(|_| ()) 21 | } 22 | } 23 | 24 | pub(super) struct Thread { 25 | context: Rc, 26 | } 27 | 28 | impl Thread { 29 | pub(super) const CONTEXT_MISSING_ERROR: &'static str = 30 | "Thread not running in the context of a Runtime"; 31 | pub(super) const CONTEXT_DESTROYED_ERROR: &'static str = 32 | "ThreadLocal runtime context was destroyed"; 33 | 34 | pub(super) fn try_current() -> Result, ()> { 35 | ThreadContext::try_with(|tls| { 36 | tls.as_ref().map(|context| Self { 37 | context: context.clone(), 38 | }) 39 | }) 40 | } 41 | 42 | pub(super) fn current() -> Self { 43 | match Self::try_current() { 44 | Ok(Some(thread)) => thread, 45 | Ok(None) => unreachable!("{}", Self::CONTEXT_MISSING_ERROR), 46 | Err(_) => unreachable!("{}", Self::CONTEXT_DESTROYED_ERROR), 47 | } 48 | } 49 | 50 | pub(super) fn enter(scheduler: &Arc, queue_index: Option) -> Self { 51 | ThreadContext::try_with(|tls| { 52 | if let Some(context) = tls { 53 | return Self { 54 | context: context.clone(), 55 | }; 56 | } 57 | 58 | let context = Rc::new(ThreadContext { 59 | scheduler: scheduler.clone(), 60 | queue_index, 61 | be_fair: Cell::new(false), 62 | lifo_slot: Cell::new(None), 63 | }); 64 | 65 | *tls = Some(context.clone()); 66 | Self { context } 67 | }) 68 | .unwrap() 69 | } 70 | } 71 | 72 | impl Deref for Thread { 73 | type Target = ThreadContext; 74 | 75 | fn deref(&self) -> &Self::Target { 76 | &*self.context 77 | } 78 | } 79 | 80 | impl Drop for Thread { 81 | fn drop(&mut self) { 82 | if Rc::strong_count(&self.context) == 2 { 83 | let context = ThreadContext::try_with(|tls| tls.take()); 84 | let context = context.unwrap().unwrap(); 85 | assert!(Rc::ptr_eq(&context, &self.context)); 86 | } 87 | } 88 | } 89 | -------------------------------------------------------------------------------- /src/thread_pool/waker.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | mem::replace, 3 | sync::atomic::{AtomicU8, Ordering}, 4 | sync::Mutex, 5 | task::{Context, Poll, Waker}, 6 | }; 7 | 8 | #[derive(Default)] 9 | pub(super) struct AtomicWaker { 10 | state: AtomicU8, 11 | waker: Mutex>, 12 | } 13 | 14 | impl AtomicWaker { 15 | const EMPTY: u8 = 0; 16 | const UPDATING: u8 = 1; 17 | const WAITING: u8 = 2; 18 | const NOTIFIED: u8 = 3; 19 | 20 | pub(super) fn poll(&self, ctx: &mut Context<'_>) -> Poll<()> { 21 | let state = self.state.load(Ordering::Acquire); 22 | match state { 23 | Self::EMPTY | Self::WAITING => {} 24 | Self::NOTIFIED => return Poll::Ready(()), 25 | _ => unreachable!(), 26 | } 27 | 28 | if let Err(state) = 29 | self.state 30 | .compare_exchange(state, Self::UPDATING, Ordering::Acquire, Ordering::Acquire) 31 | { 32 | assert_eq!(state, Self::NOTIFIED); 33 | return Poll::Ready(()); 34 | } 35 | 36 | { 37 | let mut waker = self.waker.try_lock().unwrap(); 38 | let will_wake = waker 39 | .as_ref() 40 | .map(|w| ctx.waker().will_wake(w)) 41 | .unwrap_or(false); 42 | 43 | if !will_wake { 44 | *waker = Some(ctx.waker().clone()); 45 | } 46 | } 47 | 48 | match self.state.compare_exchange( 49 | Self::UPDATING, 50 | Self::WAITING, 51 | Ordering::AcqRel, 52 | Ordering::Acquire, 53 | ) { 54 | Ok(_) => Poll::Pending, 55 | Err(Self::NOTIFIED) => Poll::Ready(()), 56 | Err(_) => unreachable!(), 57 | } 58 | } 59 | 60 | pub(super) fn wake(&self) { 61 | match self.state.swap(Self::NOTIFIED, Ordering::AcqRel) { 62 | Self::EMPTY | Self::UPDATING => return, 63 | Self::WAITING => {} 64 | _ => unreachable!(), 65 | } 66 | 67 | let mut waker = self.waker.try_lock().unwrap(); 68 | replace(&mut *waker, None).unwrap().wake(); 69 | } 70 | } 71 | --------------------------------------------------------------------------------