├── .gitignore ├── apps ├── dummy.rs ├── Cargo.toml └── qsort.rs ├── Cargo.toml ├── .github └── workflows │ ├── post.yaml │ └── pre.yaml ├── LICENSE ├── src ├── util.rs ├── arc.rs └── lib.rs ├── CHANGELOG.md ├── benches └── ballpark.rs ├── README.md └── tests └── basic.rs /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | Cargo.lock 3 | 4 | .DS_Store 5 | -------------------------------------------------------------------------------- /apps/dummy.rs: -------------------------------------------------------------------------------- 1 | fn main() { 2 | let choir = choir::Choir::new(); 3 | let _workers = (0..2) 4 | .map(|i| choir.add_worker(&format!("worker-{}", i))) 5 | .collect::>(); 6 | let mut end = choir.spawn("end").init_dummy(); 7 | for _ in 0..1_000_000 { 8 | let task = choir.spawn("").init_dummy(); 9 | end.depend_on(&task); 10 | } 11 | end.run_attached(); 12 | } 13 | -------------------------------------------------------------------------------- /apps/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "choir-apps" 3 | version = "0.1.0" 4 | license = "MPL-2.0" 5 | edition = "2021" 6 | publish = false 7 | 8 | [features] 9 | profiler-tracy = ["profiling/profile-with-tracy"] 10 | 11 | [dependencies] 12 | choir = { path = ".." } 13 | env_logger = "0.9" 14 | num_cpus = "1" 15 | profiling = "1" 16 | rand = "0.8" 17 | 18 | [[bin]] 19 | name = "qsort" 20 | path = "qsort.rs" 21 | 22 | [[bin]] 23 | name = "dummy" 24 | path = "dummy.rs" 25 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "choir" 3 | version = "0.7.0" 4 | edition = "2021" 5 | description = "Task Orchestration Framework" 6 | license = "MIT" 7 | documentation = "https://docs.rs/choir" 8 | repository = "https://github.com/kvark/choir" 9 | keywords = ["job-system", "tasks"] 10 | categories = ["concurrency", "game-development"] 11 | 12 | [workspace] 13 | members = [ 14 | "apps", 15 | ] 16 | 17 | [[bench]] 18 | name = "ballpark" 19 | harness = false 20 | 21 | [dependencies] 22 | crossbeam-deque = "0.8" 23 | log = "0.4" 24 | profiling = "1" 25 | 26 | [dev-dependencies] 27 | criterion = "0.5" 28 | env_logger = "0.9" 29 | num_cpus = "1" 30 | -------------------------------------------------------------------------------- /.github/workflows/post.yaml: -------------------------------------------------------------------------------- 1 | # Post-landing jobs 2 | name: Post-check 3 | on: 4 | push: 5 | branches: [main] 6 | 7 | jobs: 8 | coverage: 9 | name: Coverage 10 | runs-on: ubuntu-latest 11 | steps: 12 | - uses: actions/checkout@v2 13 | - uses: actions-rs/toolchain@v1 14 | with: 15 | toolchain: stable 16 | - name: Generate report 17 | uses: actions-rs/tarpaulin@v0.1 18 | with: 19 | version: 0.22.0 20 | args: '--tests --workspace' 21 | - name: Upload to codecov.io 22 | uses: codecov/codecov-action@v1 23 | with: 24 | token: ${{ secrets.CODECOV_TOKEN }} 25 | - name: Archive code coverage results 26 | uses: actions/upload-artifact@v1 27 | with: 28 | name: code-coverage-report 29 | path: cobertura.xml 30 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Dzmitry Malyshau 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. 22 | -------------------------------------------------------------------------------- /src/util.rs: -------------------------------------------------------------------------------- 1 | use super::SubIndex; 2 | use std::{cell::UnsafeCell, iter::FromIterator}; 3 | 4 | /// Helper data structure for holding per-task data. 5 | /// Each element is expected to only be accessed zero or one time. 6 | pub struct PerTaskData { 7 | data: Box<[UnsafeCell>]>, 8 | } 9 | 10 | unsafe impl Sync for PerTaskData {} 11 | 12 | impl PerTaskData { 13 | /// Take an element at a given index. 14 | /// 15 | /// # Safety 16 | /// Can't be executed more than once for any given index. 17 | pub unsafe fn take(&self, index: SubIndex) -> T { 18 | (*self.data[index as usize].get()).take().unwrap() 19 | } 20 | 21 | /// Return the length of the data. 22 | pub fn len(&self) -> SubIndex { 23 | self.data.len() as _ 24 | } 25 | } 26 | 27 | impl FromIterator for PerTaskData { 28 | fn from_iter>(iter: I) -> Self { 29 | Self { 30 | data: iter.into_iter().map(Some).map(UnsafeCell::new).collect(), 31 | } 32 | } 33 | } 34 | 35 | #[test] 36 | fn smoke() { 37 | let choir = super::Choir::new(); 38 | let data: PerTaskData = (0..10).collect(); 39 | choir 40 | .spawn("") 41 | .init_multi(data.len(), move |_, i| { 42 | let v = unsafe { data.take(i) }; 43 | println!("v = {}", v); 44 | }) 45 | .run_attached(); 46 | } 47 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Change Log 2 | 3 | ## v0.7 (07-11-2023) 4 | - make `join()` to work with `&self` 5 | - reimplement `join()` in a robust way 6 | - remove non-static tasks 7 | - fix multiple race conditions 8 | - allow multiple parents of a task 9 | 10 | ## v0.6 (30-05-2023) 11 | - redesign fork semantics that's more robust and supports `join()` 12 | - add an ability to add an existing task as a fork 13 | - support non-`'static` function bodies via `run_attached()` 14 | - truly joining the choir when waiting on a task via `join_active()` 15 | - always use `Arc` 16 | - propagate panics from the tasks/workers 17 | - expose `choir` in the tasks and execution contexts 18 | - `Condvar`-based thread blocking 19 | - MSRV bumped to 1.60 20 | 21 | ## v0.5 (15-08-2022) 22 | - all functors accept an argument of `ExecutionContext` 23 | - ability to fork tasks from a functor body 24 | - `RunningTask::join()` instead of `Choir::wait_all` 25 | - intermediate `ProtoTask` type 26 | - `impl Clone for RunningTask` 27 | - everything implements `Debug` 28 | - new `Linearc` type for linearized `Arc` 29 | - no more spontaneous blocking - the task synchronization is fixed 30 | 31 | ### v0.4.2 (07-06-2022) 32 | - dummy tasks support 33 | 34 | ## v0.4 (01-06-2022) 35 | - iterator tasks support 36 | - auto-schedule idle tasks on drop 37 | - replace `run_xx` and `idle_xx` calls by just `add_xx` 38 | 39 | ## v0.3 (23-04-2022) 40 | - no `Sync` bound 41 | - multi-tasks 42 | - profiling integration 43 | - benchmarks 44 | 45 | ## v0.2 (10-04-2022) 46 | - task dependencies 47 | - proper names for things 48 | - variable worker count 49 | 50 | ## v0.1 (07-04-2022) 51 | - basic task execution 52 | -------------------------------------------------------------------------------- /benches/ballpark.rs: -------------------------------------------------------------------------------- 1 | use criterion::{criterion_group, criterion_main, Criterion}; 2 | 3 | fn work(n: choir::SubIndex) { 4 | let mut vec = vec![1, 1]; 5 | for i in 2..n as usize { 6 | vec.push(vec[i - 1] + vec[i - 2]); 7 | } 8 | } 9 | 10 | fn many_tasks(c: &mut Criterion) { 11 | const TASK_COUNT: choir::SubIndex = 1_000; 12 | //let _ = profiling::tracy_client::Client::start(); 13 | let choir = choir::Choir::new(); 14 | if true { 15 | let _worker = choir.add_worker("main"); 16 | c.bench_function("individual tasks: single worker", |b| { 17 | b.iter(|| { 18 | let mut parent = choir.spawn("parent").init_dummy(); 19 | for i in 0..TASK_COUNT { 20 | let task = choir.spawn("").init(move |_| work(i)); 21 | parent.depend_on(&task); 22 | } 23 | parent.run().join(); 24 | }); 25 | }); 26 | c.bench_function("multi-task: single worker", |b| { 27 | b.iter(|| { 28 | choir 29 | .spawn("") 30 | .init_multi(TASK_COUNT, move |_, i| work(i)) 31 | .run() 32 | .join(); 33 | }); 34 | }); 35 | } 36 | 37 | if true { 38 | let num_cores = num_cpus::get_physical(); 39 | let _workers = (0..num_cores) 40 | .map(|i| choir.add_worker(&format!("worker-{}", i))) 41 | .collect::>(); 42 | c.bench_function(&format!("individual tasks: {} workers", num_cores), |b| { 43 | b.iter(|| { 44 | let mut parent = choir.spawn("parent").init_dummy(); 45 | for i in 0..TASK_COUNT { 46 | let task = choir.spawn("").init(move |_| work(i)); 47 | parent.depend_on(&task); 48 | } 49 | parent.run().join(); 50 | }); 51 | }); 52 | c.bench_function(&format!("multi-task: {} workers", num_cores), |b| { 53 | b.iter(|| { 54 | choir 55 | .spawn("") 56 | .init_multi(TASK_COUNT, move |_, i| work(i)) 57 | .run() 58 | .join(); 59 | }); 60 | }); 61 | } 62 | } 63 | 64 | criterion_group!(benches, many_tasks); 65 | criterion_main!(benches); 66 | -------------------------------------------------------------------------------- /.github/workflows/pre.yaml: -------------------------------------------------------------------------------- 1 | # Regular testing. 2 | name: Check 3 | on: [pull_request] 4 | 5 | jobs: 6 | test-msrv: 7 | name: Test MSRV 8 | runs-on: ubuntu-latest 9 | steps: 10 | - uses: actions/checkout@v2 11 | - uses: actions-rs/toolchain@v1 12 | with: 13 | profile: minimal 14 | toolchain: "1.60.0" 15 | override: true 16 | - uses: actions-rs/cargo@v1 17 | name: Main check 18 | with: 19 | command: check 20 | args: --workspace 21 | - uses: actions-rs/cargo@v1 22 | name: Bench compile check 23 | with: 24 | command: bench 25 | args: --no-run 26 | test: 27 | name: Test Nightly 28 | runs-on: ubuntu-latest 29 | steps: 30 | - uses: actions/checkout@v2 31 | - uses: actions-rs/toolchain@v1 32 | with: 33 | profile: minimal 34 | toolchain: nightly 35 | override: true 36 | - uses: actions-rs/cargo@v1 37 | name: Default test 38 | with: 39 | command: test 40 | - uses: actions-rs/cargo@v1 41 | name: Test all 42 | with: 43 | command: test 44 | args: --workspace 45 | clippy: 46 | name: Clippy 47 | runs-on: ubuntu-latest 48 | steps: 49 | - uses: actions/checkout@v2 50 | - uses: actions-rs/toolchain@v1 51 | with: 52 | profile: minimal 53 | toolchain: stable 54 | override: true 55 | - run: rustup component add clippy 56 | - uses: actions-rs/cargo@v1 57 | with: 58 | command: clippy 59 | args: --workspace --all-features -- -D warnings 60 | documentation: 61 | name: Documentation 62 | runs-on: ubuntu-latest 63 | env: 64 | RUSTDOCFLAGS: -Dwarnings 65 | steps: 66 | - uses: actions/checkout@v2 67 | - uses: actions-rs/toolchain@v1 68 | with: 69 | profile: minimal 70 | toolchain: stable 71 | override: true 72 | - run: rustup component add clippy 73 | - uses: actions-rs/cargo@v1 74 | with: 75 | command: doc 76 | fmt: 77 | name: Format 78 | runs-on: ubuntu-latest 79 | steps: 80 | - uses: actions/checkout@v2 81 | - uses: actions-rs/toolchain@v1 82 | with: 83 | profile: minimal 84 | toolchain: stable 85 | override: true 86 | components: rustfmt 87 | - name: run rustfmt 88 | run: | 89 | cargo fmt -- --check 90 | -------------------------------------------------------------------------------- /apps/qsort.rs: -------------------------------------------------------------------------------- 1 | use rand::Rng as _; 2 | use std::{mem, ptr, slice}; 3 | 4 | type Value = i64; 5 | 6 | fn insertion_sort(data: &mut [Value]) { 7 | for i in 1..data.len() { 8 | let v = data[i]; 9 | let mut j = i; 10 | while j != 0 { 11 | let w = data[j - 1]; 12 | if w <= v { 13 | data[j] = v; 14 | break; 15 | } else { 16 | data[j] = w; 17 | j -= 1; 18 | } 19 | } 20 | } 21 | } 22 | 23 | fn split(data: &mut [Value]) -> (&mut [Value], &mut [Value]) { 24 | let mid = data[0]; 25 | let mut i = 1; 26 | let mut j = data.len() - 1; 27 | 'outer: while i < j { 28 | while data[i] <= mid { 29 | i += 1; 30 | if i == j { 31 | break 'outer; 32 | } 33 | } 34 | while data[j] > mid { 35 | j -= 1; 36 | if i == j { 37 | break 'outer; 38 | } 39 | } 40 | unsafe { 41 | ptr::swap(&mut data[i], &mut data[j]); 42 | } 43 | i += 1; 44 | j -= 1; 45 | } 46 | 47 | let (left, right) = data.split_at_mut(j); 48 | // guarantee that the element in the middle can be excluded 49 | if right[0] <= mid { 50 | mem::swap(&mut left[0], &mut right[0]); 51 | (left, right.split_first_mut().unwrap().1) 52 | } else { 53 | (left, right) 54 | } 55 | } 56 | 57 | fn qsort(data: *mut Value, size: usize, context: choir::ExecutionContext) { 58 | if size > 5 { 59 | let (left, right) = split(unsafe { slice::from_raw_parts_mut(data, size) }); 60 | context 61 | .fork("left") 62 | .init(move |ec| qsort(left.as_mut_ptr(), left.len(), ec)); 63 | context 64 | .fork("right") 65 | .init(move |ec| qsort(right.as_mut_ptr(), right.len(), ec)); 66 | } else if size > 1 { 67 | insertion_sort(unsafe { slice::from_raw_parts_mut(data, size) }) 68 | } 69 | } 70 | 71 | fn main() { 72 | const USE_TASKS: bool = true; 73 | const COUNT: usize = 10000000; 74 | env_logger::init(); 75 | 76 | let mut data = { 77 | let mut rng = rand::thread_rng(); 78 | (0..COUNT).map(|_| rng.gen()).collect::>() 79 | }; 80 | 81 | if USE_TASKS { 82 | let choir = choir::Choir::new(); 83 | let _worker1 = choir.add_worker("worker1"); 84 | let _worker2 = choir.add_worker("worker2"); 85 | let data_ptr = data.as_mut_ptr() as usize; 86 | choir 87 | .spawn("main") 88 | .init(move |ec| qsort(data_ptr as *mut Value, COUNT, ec)) 89 | .run_attached(); 90 | } else { 91 | insertion_sort(&mut data); 92 | } 93 | 94 | // check if sorted 95 | if let Some(position) = data.iter().zip(data[1..].iter()).position(|(a, b)| a > b) { 96 | panic!( 97 | "position {}: {} > {}", 98 | position, 99 | data[position], 100 | data[position + 1] 101 | ); 102 | } 103 | } 104 | -------------------------------------------------------------------------------- /src/arc.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | fmt, mem, ops, 3 | ptr::NonNull, 4 | sync::atomic::{AtomicUsize, Ordering}, 5 | }; 6 | 7 | /// Note: `pub(super)` is only needed for a hack 8 | #[derive(Debug)] 9 | pub(super) struct LinearcInner { 10 | pub(super) ref_count: AtomicUsize, 11 | pub(super) data: T, 12 | } 13 | 14 | /// Linear atomically referenced pointer. 15 | /// Analogous to `Arc` but without `Weak` functionality, 16 | /// and with an extra powers to treat the type as linear. 17 | /// In particular, it supports atomic destruction with the extraction of data. 18 | /// 19 | /// See 20 | /// And 21 | pub struct Linearc { 22 | ptr: NonNull>, 23 | } 24 | 25 | unsafe impl Send for Linearc {} 26 | unsafe impl Sync for Linearc {} 27 | 28 | impl Linearc { 29 | #[inline] 30 | pub(super) fn from_inner(inner: Box>) -> Self { 31 | Self { 32 | ptr: unsafe { NonNull::new_unchecked(Box::into_raw(inner)) }, 33 | } 34 | } 35 | 36 | /// Clone a given pointer. 37 | #[inline] 38 | pub fn clone(arc: &Self) -> Self { 39 | arc.clone() 40 | } 41 | 42 | fn into_box(arc: Self) -> Option>> { 43 | let count = unsafe { arc.ptr.as_ref() } 44 | .ref_count 45 | .fetch_sub(1, Ordering::AcqRel); 46 | if count == 1 { 47 | let inner = unsafe { Box::from_raw(arc.ptr.as_ptr()) }; 48 | mem::forget(arc); 49 | Some(inner) 50 | } else { 51 | mem::forget(arc); 52 | None 53 | } 54 | } 55 | 56 | /// Drop a pointer and return true if this was the last instance. 57 | #[inline] 58 | pub fn drop_last(arc: Self) -> bool { 59 | Linearc::into_box(arc).is_some() 60 | } 61 | } 62 | 63 | impl Linearc { 64 | /// Create a new pointer from data. 65 | #[inline] 66 | pub fn new(data: T) -> Self { 67 | Self::from_inner(Box::new(LinearcInner { 68 | ref_count: AtomicUsize::new(1), 69 | data, 70 | })) 71 | } 72 | 73 | /// Move out the value of this pointer if it's the last instance. 74 | pub fn into_inner(arc: Self) -> Option { 75 | Linearc::into_box(arc).map(|inner| inner.data) 76 | } 77 | } 78 | 79 | impl Clone for Linearc { 80 | fn clone(&self) -> Self { 81 | unsafe { self.ptr.as_ref() } 82 | .ref_count 83 | .fetch_add(1, Ordering::Release); 84 | Self { ptr: self.ptr } 85 | } 86 | } 87 | 88 | impl Drop for Linearc { 89 | fn drop(&mut self) { 90 | let count = unsafe { self.ptr.as_ref() } 91 | .ref_count 92 | .fetch_sub(1, Ordering::AcqRel); 93 | if count == 1 { 94 | let _ = unsafe { Box::from_raw(self.ptr.as_ptr()) }; 95 | } 96 | } 97 | } 98 | 99 | impl ops::Deref for Linearc { 100 | type Target = T; 101 | 102 | #[inline] 103 | fn deref(&self) -> &T { 104 | &unsafe { self.ptr.as_ref() }.data 105 | } 106 | } 107 | 108 | impl fmt::Debug for Linearc { 109 | fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { 110 | unsafe { self.ptr.as_ref() }.fmt(formatter) 111 | } 112 | } 113 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Choir 2 | 3 | [![Crates.io](https://img.shields.io/crates/v/choir.svg?label=choir)](https://crates.io/crates/choir) 4 | [![Docs.rs](https://docs.rs/choir/badge.svg)](https://docs.rs/choir) 5 | [![Build Status](https://github.com/kvark/choir/workflows/Check/badge.svg)](https://github.com/kvark/choir/actions) 6 | ![MSRV](https://img.shields.io/badge/rustc-1.56+-blue.svg) 7 | [![codecov.io](https://codecov.io/gh/kvark/choir/branch/main/graph/badge.svg)](https://codecov.io/gh/kvark/choir) 8 | 9 | Choir is a task orchestration framework. It helps you to organize all the CPU workflow in terms of tasks. 10 | 11 | ### Example: 12 | ```rust 13 | let mut choir = choir::Choir::new(); 14 | let _worker = choir.add_worker("worker"); 15 | let task1 = choir.spawn("foo").init_dummy().run(); 16 | let mut task2 = choir.spawn("bar").init(|_| { println!("bar"); }); 17 | task2.depend_on(&task1); 18 | task2.run().join(); 19 | ``` 20 | 21 | ### Selling Pitch 22 | 23 | What makes Choir _elegant_? Generally when we need to encode the semantics of "wait for dependencies", we think of some sort of a counter. Maybe an atomic, for the dependency number. When it reaches zero (or one), we schedule a task for execution. In _Choir_, the internal data for a task (i.e. the functor itself!) is placed in an `Arc`. Whenever we are able to extract it from the `Arc` (which means there are no other dependencies), we move it to a scheduling queue. I think Rust type system shows its best here. 24 | 25 | Note: it turns out `Arc` doesn't fully support such a "linear" usage as required here, and it's impossible to control where the last reference gets destructed (without logic in `drop()`). For this reason, we introduce our own `Linearc` to be used internally. 26 | 27 | You can also add or remove workers at any time to balance the system load, which may be running other applications at the same time. 28 | 29 | ## API 30 | 31 | General workflow is about creating tasks and setting up dependencies between them. There is a few different kinds of tasks: 32 | - single-run tasks, initialized with `init()` and represented as `FnOnce()` 33 | - dummy tasks, initialized with `init_dummy()`, and having no function body 34 | - multi-run tasks, executed for every index in a range, represented as `Fn(SubIndex)`, and initialized with `init_multi()` 35 | - iteration tasks, executed for every item produced by an iterator, represented as `Fn(T)`, and initialized with `init_iter()` 36 | 37 | Just calling `run()` is done automatically on `IdleTask::drop()` if not called explicitly. 38 | This object also allows adding dependencies before scheduling the task. The running task can be also used as a dependency for others. 39 | 40 | Note that all tasks are pre-empted at the `Fn()` execution boundary. Thus, for example, a long-running multi task will be pre-empted by any incoming single-run tasks. 41 | 42 | ## Users 43 | 44 | [Blade](https://github.com/kvark/blade) heavily relies on Choir for parallelizing the asset loading. See [blade-asset talk](https://youtu.be/1DiA3OYqvqU) at Rust Gamedev meetup for details. 45 | 46 | ### TODO: 47 | - detect when dependencies aren't set up correctly 48 | - test with [Loom](https://github.com/tokio-rs/loom): blocked by https://github.com/crossbeam-rs/crossbeam/pull/849 49 | 50 | ## Overhead 51 | 52 | Machine: MBP 2016, 3.3 GHz Dual-Core Intel Core i7 53 | 54 | - functions `spawn()+init()` (optimized): 237ns 55 | - "steal" task: 61ns 56 | - empty "execute": 37ns 57 | - dummy "unblock": 78ns 58 | 59 | Executing 100k empty tasks: 60 | - individually: 28ms 61 | - as a multi-task: 6ms 62 | 63 | ### Profiling workflow example 64 | 65 | With Tracy: 66 | Add this line to the start of the benchmark: 67 | ```rust 68 | let _ = profiling::tracy_client::Client::start(); 69 | ``` 70 | Then run in command prompt: 71 | ```bash 72 | cargo bench --features "profiling/profile-with-tracy" 73 | ``` 74 | -------------------------------------------------------------------------------- /tests/basic.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | sync::{ 3 | atomic::{AtomicBool, AtomicUsize, Ordering}, 4 | Arc, Mutex, 5 | }, 6 | thread, 7 | time::Duration, 8 | }; 9 | 10 | #[test] 11 | fn parallel() { 12 | let _ = env_logger::try_init(); 13 | let choir = choir::Choir::new(); 14 | let _worker1 = choir.add_worker("P1"); 15 | let _worker2 = choir.add_worker("P2"); 16 | 17 | let value = Arc::new(AtomicUsize::new(0)); 18 | let n = 100; 19 | // Launch N independent tasks, each bumping 20 | // the value. Expect all of them to work. 21 | let mut last = choir.spawn("last").init_dummy(); 22 | for _ in 0..n { 23 | let v = Arc::clone(&value); 24 | let child = choir 25 | .spawn("") 26 | .init(move |_| { 27 | v.fetch_add(1, Ordering::AcqRel); 28 | }) 29 | .run(); 30 | last.depend_on(&child); 31 | } 32 | 33 | last.run().join(); 34 | assert_eq!(value.load(Ordering::Acquire), n); 35 | } 36 | 37 | #[test] 38 | fn sequential() { 39 | let _ = env_logger::try_init(); 40 | let choir = choir::Choir::new(); 41 | let _worker = choir.add_worker("S"); 42 | 43 | let value = Arc::new(Mutex::new(0)); 44 | let mut base = choir.spawn("base").init_dummy(); 45 | let n = 100; 46 | // Launch N tasks, each depending on the previous one 47 | // and each setting a value. 48 | // If they were running in parallel, the resulting 49 | // value would be undetermined. But sequentially, 50 | // it has to be N. 51 | for i in 0..n { 52 | let v = Arc::clone(&value); 53 | let mut next = choir.spawn("").init(move |_| { 54 | *v.lock().unwrap() = i + 1; 55 | }); 56 | next.depend_on(&base); 57 | base = next; 58 | } 59 | base.run().join(); 60 | assert_eq!(*value.lock().unwrap(), n); 61 | } 62 | 63 | #[test] 64 | fn zero_count() { 65 | let choir = choir::Choir::new(); 66 | let _worker1 = choir.add_worker("A"); 67 | choir.spawn("").init_multi(0, |_, _| {}).run().join(); 68 | } 69 | 70 | #[test] 71 | fn multi_sum() { 72 | let _ = env_logger::try_init(); 73 | let choir = choir::Choir::new(); 74 | let _worker1 = choir.add_worker("A"); 75 | let _worker2 = choir.add_worker("B"); 76 | 77 | let value = Arc::new(AtomicUsize::new(0)); 78 | let value_other = Arc::clone(&value); 79 | let n = 100; 80 | choir 81 | .spawn("") 82 | .init_multi(n, move |_, i| { 83 | value_other.fetch_add(i as usize, Ordering::SeqCst); 84 | }) 85 | .run() 86 | .join(); 87 | assert_eq!(value.load(Ordering::Acquire) as u32, (n - 1) * n / 2); 88 | } 89 | 90 | #[test] 91 | fn iter_xor() { 92 | let _ = env_logger::try_init(); 93 | let choir = choir::Choir::new(); 94 | let _worker1 = choir.add_worker("A"); 95 | let _worker2 = choir.add_worker("B"); 96 | 97 | let value = Arc::new(AtomicUsize::new(0)); 98 | let value_other = Arc::clone(&value); 99 | let n = 50; 100 | 101 | choir 102 | .spawn("") 103 | .init_iter(0..n, move |_, item| { 104 | value_other.fetch_xor(item, Ordering::SeqCst); 105 | }) 106 | .run() 107 | .join(); 108 | assert_eq!(value.load(Ordering::Acquire), 1); 109 | } 110 | 111 | #[test] 112 | fn proxy() { 113 | let _ = env_logger::try_init(); 114 | let choir = choir::Choir::new(); 115 | let _worker1 = choir.add_worker("A"); 116 | let _worker2 = choir.add_worker("B"); 117 | 118 | let value = Arc::new(AtomicUsize::new(0)); 119 | let value_other = Arc::clone(&value); 120 | let n = 50; 121 | choir 122 | .spawn("parent") 123 | .init_multi(n, move |ec, i| { 124 | println!("base[{}]", i); 125 | let value_other2 = Arc::clone(&value_other); 126 | value_other.fetch_or(1 << i, Ordering::SeqCst); 127 | ec.fork("proxy").init(move |_| { 128 | println!("proxy[{}]", i); 129 | value_other2.fetch_xor(1 << i, Ordering::SeqCst); 130 | }); 131 | }) 132 | .run() 133 | .join(); 134 | 135 | assert_eq!(value.load(Ordering::Acquire), 0); 136 | } 137 | 138 | #[test] 139 | fn fork_in_flight() { 140 | let _ = env_logger::try_init(); 141 | let choir = choir::Choir::new(); 142 | let _worker1 = choir.add_worker("A"); 143 | let value = Arc::new(AtomicUsize::new(0)); 144 | let value1 = Arc::clone(&value); 145 | // This task deliberately waits, so that we know for sure 146 | // if it's being waited on. 147 | let t1 = choir 148 | .spawn("child") 149 | .init(move |_| { 150 | thread::sleep(Duration::from_millis(10)); 151 | value1.fetch_add(1, Ordering::AcqRel); 152 | }) 153 | .run(); 154 | let value2 = Arc::clone(&value); 155 | // This task decides to add `t1` as a fork, which is already in flight. 156 | choir 157 | .spawn("parent") 158 | .init(move |ec| { 159 | value2.fetch_add(1, Ordering::AcqRel); 160 | ec.add_fork(&t1); 161 | }) 162 | .run() 163 | .join(); 164 | assert_eq!(value.load(Ordering::Acquire), 2); 165 | } 166 | 167 | #[test] 168 | fn fork_interdep() { 169 | let _ = env_logger::try_init(); 170 | let choir = choir::Choir::new(); 171 | let _worker1 = choir.add_worker("A"); 172 | let value = Arc::new(AtomicUsize::new(0)); 173 | let value1 = Arc::clone(&value); 174 | let value2 = Arc::clone(&value); 175 | let value3 = Arc::clone(&value); 176 | choir 177 | .spawn("parent") 178 | .init(move |ec| { 179 | assert_eq!(0, value1.fetch_add(1, Ordering::AcqRel)); 180 | let t1 = ec.choir().spawn("child").init(move |_| { 181 | thread::sleep(Duration::from_millis(20)); 182 | assert_eq!(1, value2.fetch_add(1, Ordering::AcqRel)); 183 | }); 184 | let mut t2 = ec.fork("fork").init(move |_| { 185 | thread::sleep(Duration::from_millis(10)); 186 | assert_eq!(2, value3.fetch_add(1, Ordering::AcqRel)); 187 | }); 188 | t2.depend_on(&t1); 189 | }) 190 | .run() 191 | .join(); 192 | assert_eq!(value.load(Ordering::Acquire), 3); 193 | } 194 | 195 | #[test] 196 | fn unhelpful() { 197 | let choir = choir::Choir::new(); 198 | let done = Arc::new(AtomicBool::new(false)); 199 | let done_final = Arc::clone(&done); 200 | choir 201 | .spawn("task") 202 | .init(move |_| { 203 | done.store(true, Ordering::Release); 204 | }) 205 | .run_attached(); 206 | assert!(done_final.load(Ordering::Acquire)); 207 | } 208 | 209 | #[test] 210 | fn multi_thread_join() { 211 | let _ = env_logger::try_init(); 212 | let choir = choir::Choir::new(); 213 | let _w = choir.add_worker("main"); 214 | let running = choir 215 | .spawn("task") 216 | .init(|_| { 217 | thread::sleep(Duration::from_millis(100)); 218 | }) 219 | .run(); 220 | let r1 = running.clone(); 221 | let t1 = thread::spawn(move || r1.join()); 222 | let r2 = running.clone(); 223 | let t2 = thread::spawn(move || r2.join()); 224 | t1.join().unwrap(); 225 | t2.join().unwrap(); 226 | } 227 | 228 | #[test] 229 | #[should_panic] 230 | fn join_timeout() { 231 | let choir = choir::Choir::new(); 232 | let task = choir.spawn("test").init_dummy(); 233 | task.run().join_debug(Default::default()); 234 | } 235 | 236 | #[test] 237 | #[should_panic] 238 | fn task_panic() { 239 | let choir = choir::Choir::new(); 240 | let _w = choir.add_worker("main"); 241 | choir.spawn("task").init(|_| panic!("Oops!")).run().join(); 242 | } 243 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | /*! Task Orchestration Framework. 2 | 3 | This framework helps to organize the execution of a program 4 | into a live task graph. In this model, all the work is happening 5 | inside tasks, which are scheduled to run by the `Choir`. 6 | 7 | Lifetime of a Task: 8 | 1. Idle: task is just created. 9 | 2. Initialized: function body is assigned. 10 | 3. Scheduled: no more dependencies can be added to the task. 11 | 4. Executing: task was dispatched from the queue by one of the workers. 12 | 5. Done: task is retired. 13 | !*/ 14 | 15 | #![allow( 16 | renamed_and_removed_lints, 17 | clippy::new_without_default, 18 | clippy::unneeded_field_pattern, 19 | clippy::match_like_matches_macro, 20 | clippy::manual_strip, 21 | clippy::if_same_then_else, 22 | clippy::unknown_clippy_lints, 23 | clippy::len_without_is_empty, 24 | clippy::should_implement_trait 25 | )] 26 | #![warn( 27 | missing_docs, 28 | trivial_casts, 29 | trivial_numeric_casts, 30 | unused_extern_crates, 31 | unused_qualifications, 32 | clippy::pattern_type_mismatch 33 | )] 34 | //#![forbid(unsafe_code)] 35 | 36 | /// Better shared pointer. 37 | pub mod arc; 38 | /// Additional utilities. 39 | pub mod util; 40 | 41 | use self::arc::Linearc; 42 | use crossbeam_deque::{Injector, Steal}; 43 | use std::{ 44 | borrow::Cow, 45 | fmt, mem, ops, 46 | sync::{ 47 | atomic::{AtomicBool, AtomicIsize, AtomicUsize, Ordering}, 48 | Arc, Condvar, Mutex, RwLock, 49 | }, 50 | thread, time, 51 | }; 52 | 53 | const BITS_PER_BYTE: usize = 8; 54 | const MAX_WORKERS: usize = mem::size_of::() * BITS_PER_BYTE; 55 | 56 | /// Name to be associated with a task. 57 | pub type Name = Cow<'static, str>; 58 | 59 | #[derive(Debug)] 60 | struct WaitingThread { 61 | // The conditional variable lives on the stack of the 62 | // thread that is currently waiting, during the lifetime 63 | // of this pointer, so it's safe. 64 | condvar: *const Condvar, 65 | } 66 | unsafe impl Send for WaitingThread {} 67 | unsafe impl Sync for WaitingThread {} 68 | 69 | #[derive(Debug, Default)] 70 | struct Continuation { 71 | parents: Vec>, 72 | forks: usize, 73 | dependents: Vec>, 74 | waiting_threads: Vec, 75 | } 76 | 77 | impl Continuation { 78 | fn unpark_waiting(&mut self) { 79 | for wt in self.waiting_threads.drain(..) { 80 | log::trace!("\tresolving a join"); 81 | unsafe { 82 | (*wt.condvar).notify_all(); 83 | } 84 | } 85 | } 86 | } 87 | 88 | /// An object responsible to notify follow-up tasks. 89 | #[derive(Debug)] 90 | pub struct Notifier { 91 | name: Name, 92 | continuation: Mutex>, 93 | } 94 | 95 | impl fmt::Display for Notifier { 96 | fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { 97 | write!(formatter, "'{}'", self.name) 98 | } 99 | } 100 | 101 | /// Context of a task execution body. 102 | pub struct ExecutionContext<'a> { 103 | choir: &'a Arc, 104 | notifier: &'a Arc, 105 | worker_index: isize, 106 | } 107 | 108 | impl<'a> ExecutionContext<'a> { 109 | /// Get the running task handle of the current task. 110 | pub fn self_task(&self) -> RunningTask { 111 | RunningTask { 112 | choir: Arc::clone(self.choir), 113 | notifier: Arc::clone(self.notifier), 114 | } 115 | } 116 | 117 | /// Return the main choir. 118 | pub fn choir(&self) -> &'a Arc { 119 | self.choir 120 | } 121 | 122 | /// Fork the current task. 123 | /// 124 | /// The new task will block all the same dependents as the currently executing task. 125 | /// 126 | /// This is useful because it allows creating tasks on the fly from within 127 | /// other tasks. Generally, making a task in flight depend on anything is impossible. 128 | /// But one can fork a dependency instead. 129 | pub fn fork>(&self, name: N) -> ProtoTask { 130 | let name = name.into(); 131 | log::trace!("Forking task {} as '{}'", self.notifier, name); 132 | ProtoTask { 133 | choir: self.choir, 134 | name, 135 | parent: Some(Arc::clone(self.notifier)), 136 | dependents: Vec::new(), 137 | } 138 | } 139 | 140 | /// Register an existing task as a fork. 141 | /// 142 | /// It will block all the dependents of the currently executing task. 143 | /// 144 | /// Will panic if the other task is already a fork of something. 145 | pub fn add_fork>(&self, other: &D) { 146 | if let Some(ref mut continuation) = *other.as_ref().continuation.lock().unwrap() { 147 | continuation.parents.push(Arc::clone(self.notifier)); 148 | self.notifier 149 | .continuation 150 | .lock() 151 | .unwrap() 152 | .as_mut() 153 | .unwrap() 154 | .forks += 1; 155 | } 156 | } 157 | } 158 | 159 | impl Drop for ExecutionContext<'_> { 160 | fn drop(&mut self) { 161 | if thread::panicking() { 162 | self.choir.issue_panic(self.worker_index); 163 | let mut guard = self.notifier.continuation.lock().unwrap(); 164 | if let Some(mut cont) = guard.take() { 165 | cont.unpark_waiting(); 166 | } 167 | } 168 | } 169 | } 170 | 171 | /// Index of a sub-task inside a multi-task. 172 | pub type SubIndex = u32; 173 | 174 | enum Functor { 175 | Dummy, 176 | Once(Box), 177 | Multi( 178 | ops::Range, 179 | Linearc, 180 | ), 181 | } 182 | 183 | impl fmt::Debug for Functor { 184 | fn fmt(&self, serializer: &mut fmt::Formatter) -> fmt::Result { 185 | match *self { 186 | Self::Dummy => serializer.debug_struct("Dummy").finish(), 187 | Self::Once(_) => serializer.debug_struct("Once").finish(), 188 | Self::Multi(ref range, _) => serializer 189 | .debug_struct("Multi") 190 | .field("range", range) 191 | .finish(), 192 | } 193 | } 194 | } 195 | 196 | // This is totally safe. See: 197 | // https://internals.rust-lang.org/t/dyn-fnonce-should-always-be-sync/16470 198 | unsafe impl Sync for Functor {} 199 | 200 | #[derive(Debug)] 201 | struct Task { 202 | /// Body of the task. 203 | functor: Functor, 204 | /// This notifier is only really shared with `RunningTask` 205 | notifier: Arc, 206 | } 207 | 208 | struct Worker { 209 | name: String, 210 | alive: AtomicBool, 211 | } 212 | 213 | struct WorkerContext {} 214 | 215 | struct WorkerPool { 216 | contexts: [Option; MAX_WORKERS], 217 | } 218 | 219 | /// Return type for `join` functions that have to detect panics. 220 | /// Does actual panic on drop if any of the tasks have panicked. 221 | /// Note: an idiomatic `Result` is not used because it's not actionable. 222 | pub struct MaybePanic { 223 | worker_index: isize, 224 | } 225 | impl Drop for MaybePanic { 226 | fn drop(&mut self) { 227 | assert_eq!( 228 | self.worker_index, -1, 229 | "Panic occurred on worker {}", 230 | self.worker_index, 231 | ); 232 | } 233 | } 234 | 235 | /// Main structure for managing tasks. 236 | pub struct Choir { 237 | injector: Injector, 238 | condvar: Condvar, 239 | parked_mask_mutex: Mutex, 240 | workers: RwLock, 241 | panic_worker: AtomicIsize, 242 | } 243 | 244 | impl Default for Choir { 245 | fn default() -> Self { 246 | const NO_WORKER: Option = None; 247 | let injector = Injector::new(); 248 | Self { 249 | injector, 250 | condvar: Condvar::default(), 251 | parked_mask_mutex: Mutex::new(0), 252 | workers: RwLock::new(WorkerPool { 253 | contexts: [NO_WORKER; MAX_WORKERS], 254 | }), 255 | panic_worker: AtomicIsize::new(-1), 256 | } 257 | } 258 | } 259 | 260 | impl Choir { 261 | /// Create a new task system. 262 | pub fn new() -> Arc { 263 | Arc::new(Self::default()) 264 | } 265 | 266 | /// Add a new worker thread. 267 | /// 268 | /// Note: A system can't have more than `MAX_WORKERS` workers 269 | /// enabled at any time. 270 | pub fn add_worker(self: &Arc, name: &str) -> WorkerHandle { 271 | let worker = Arc::new(Worker { 272 | name: name.to_string(), 273 | alive: AtomicBool::new(true), 274 | }); 275 | let worker_clone = Arc::clone(&worker); 276 | let choir = Arc::clone(self); 277 | 278 | let join_handle = thread::Builder::new() 279 | .name(name.to_string()) 280 | .spawn(move || choir.work_loop(&worker_clone)) 281 | .unwrap(); 282 | 283 | WorkerHandle { 284 | worker, 285 | join_handle: Some(join_handle), 286 | choir: Arc::clone(self), 287 | } 288 | } 289 | 290 | /// Spawn a new task. 291 | pub fn spawn<'a, N: Into>(self: &'a Arc, name: N) -> ProtoTask<'a> { 292 | let name = name.into(); 293 | log::trace!("Creating task '{}'", name); 294 | ProtoTask { 295 | choir: self, 296 | name, 297 | parent: None, 298 | dependents: Vec::new(), 299 | } 300 | } 301 | 302 | fn schedule(&self, task: Task) { 303 | log::trace!("Task {} is scheduled", task.notifier); 304 | self.injector.push(task); 305 | self.condvar.notify_one(); 306 | } 307 | 308 | fn execute(self: &Arc, task: Task, worker_index: isize) { 309 | let execontext = ExecutionContext { 310 | choir: self, 311 | notifier: &task.notifier, 312 | worker_index, 313 | }; 314 | match task.functor { 315 | Functor::Dummy => { 316 | log::debug!( 317 | "Task {} (dummy) runs on thread[{}]", 318 | task.notifier, 319 | worker_index 320 | ); 321 | drop(execontext); 322 | } 323 | Functor::Once(fun) => { 324 | log::debug!("Task {} runs on thread[{}]", task.notifier, worker_index); 325 | profiling::scope!(task.notifier.name.as_ref()); 326 | (fun)(execontext); 327 | } 328 | Functor::Multi(mut sub_range, fun) => { 329 | log::debug!( 330 | "Task {} ({}) runs on thread[{}]", 331 | task.notifier, 332 | sub_range.start, 333 | worker_index, 334 | ); 335 | debug_assert!(sub_range.start < sub_range.end); 336 | let middle = (sub_range.end + sub_range.start) >> 1; 337 | // split the task if needed 338 | if middle != sub_range.start && *self.parked_mask_mutex.lock().unwrap() != 0 { 339 | self.injector.push(Task { 340 | functor: Functor::Multi(middle..sub_range.end, Linearc::clone(&fun)), 341 | notifier: Arc::clone(&task.notifier), 342 | }); 343 | log::trace!("\tsplit out {:?}", middle..sub_range.end); 344 | sub_range.end = middle; 345 | self.condvar.notify_one(); 346 | } 347 | // fun the functor 348 | { 349 | profiling::scope!(task.notifier.name.as_ref()); 350 | (fun)(execontext, sub_range.start); 351 | } 352 | // are we done yet? 353 | sub_range.start += 1; 354 | if sub_range.start == sub_range.end { 355 | if !Linearc::drop_last(fun) { 356 | return; 357 | } 358 | } else { 359 | // Put it back to the queue, with the next sub-index. 360 | // Note: we aren't calling `schedule` because we know at least this very thread 361 | // will be able to pick it up, so no need to wake up anybody. 362 | self.injector.push(Task { 363 | functor: Functor::Multi(sub_range, fun), 364 | notifier: task.notifier, 365 | }); 366 | return; 367 | } 368 | } 369 | } 370 | 371 | let mut notifiers = self.finish(&task.notifier); 372 | while let Some(notifier) = notifiers.pop() { 373 | notifiers.extend(self.finish(¬ifier)); 374 | } 375 | } 376 | 377 | #[profiling::function] 378 | fn finish(&self, notifier: &Notifier) -> Vec> { 379 | // mark the task as done 380 | log::trace!("Finishing task {}", notifier); 381 | 382 | let continuation = { 383 | let mut guard = notifier.continuation.lock().unwrap(); 384 | if let Some(ref mut cont) = *guard { 385 | if cont.forks != 0 { 386 | log::trace!("\t{} forks are still alive", cont.forks); 387 | cont.forks -= 1; 388 | return Vec::new(); 389 | } 390 | } 391 | let mut cont = guard.take().unwrap(); 392 | //Note: this is important to do within the lock, 393 | // so that anything waiting for the unpark signal 394 | // is guaranteed to be notified. 395 | cont.unpark_waiting(); 396 | cont 397 | }; 398 | 399 | // unblock dependencies if needed 400 | for dependent in continuation.dependents { 401 | if let Some(ready) = Linearc::into_inner(dependent) { 402 | self.schedule(ready); 403 | } 404 | } 405 | 406 | continuation.parents 407 | } 408 | 409 | fn register(&self) -> Option { 410 | let mut pool = self.workers.write().unwrap(); 411 | let index = pool.contexts.iter_mut().position(|c| c.is_none())?; 412 | pool.contexts[index] = Some(WorkerContext {}); 413 | Some(index) 414 | } 415 | 416 | fn unregister(&self, index: usize) { 417 | self.workers.write().unwrap().contexts[index] = None; 418 | // Avoid a situation where choir is expecting this thread 419 | // to help with more tasks. 420 | self.condvar.notify_one(); 421 | } 422 | 423 | fn work_loop(self: &Arc, worker: &Worker) { 424 | profiling::register_thread!(); 425 | let index = self.register().unwrap(); 426 | log::info!("Thread[{}] = '{}' started", index, worker.name); 427 | 428 | while worker.alive.load(Ordering::Acquire) { 429 | match self.injector.steal() { 430 | Steal::Empty => { 431 | log::trace!("Thread[{}] sleeps", index); 432 | let mask = 1 << index; 433 | let mut parked_mask = self.parked_mask_mutex.lock().unwrap(); 434 | // Note: the check for `injector.is_empty()` here ensures that 435 | // we handle a race condition between something pushing a task, 436 | // and the thread going on the way to sleep. 437 | *parked_mask |= mask; 438 | parked_mask = self 439 | .condvar 440 | .wait_while(parked_mask, |_| { 441 | worker.alive.load(Ordering::Acquire) && self.injector.is_empty() 442 | }) 443 | .unwrap(); 444 | *parked_mask &= !mask; 445 | } 446 | Steal::Success(task) => { 447 | self.execute(task, index as isize); 448 | } 449 | Steal::Retry => {} 450 | } 451 | } 452 | 453 | log::info!("Thread '{}' dies", worker.name); 454 | self.unregister(index); 455 | } 456 | 457 | fn flush_queue(&self) { 458 | let mut num_tasks = 0; 459 | loop { 460 | match self.injector.steal() { 461 | Steal::Empty => { 462 | break; 463 | } 464 | Steal::Success(task) => { 465 | num_tasks += 1; 466 | let mut guard = task.notifier.continuation.lock().unwrap(); 467 | if let Some(mut cont) = guard.take() { 468 | cont.unpark_waiting(); 469 | } 470 | } 471 | Steal::Retry => {} 472 | } 473 | } 474 | log::trace!("\tflushed {} tasks down the drain", num_tasks); 475 | } 476 | 477 | fn issue_panic(&self, worker_index: isize) { 478 | log::debug!("panic on worker {}", worker_index); 479 | self.panic_worker.store(worker_index, Ordering::Release); 480 | self.flush_queue(); 481 | } 482 | 483 | /// Check if any of the workers terminated with panic. 484 | pub fn check_panic(&self) -> MaybePanic { 485 | let worker_index = self.panic_worker.load(Ordering::Acquire); 486 | MaybePanic { worker_index } 487 | } 488 | } 489 | 490 | /// Handle object holding a worker thread alive. 491 | pub struct WorkerHandle { 492 | worker: Arc, 493 | join_handle: Option>, 494 | choir: Arc, 495 | } 496 | 497 | enum MaybeArc { 498 | Unique(T), 499 | Shared(Linearc), 500 | Null, 501 | } 502 | 503 | impl MaybeArc { 504 | const NULL_ERROR: &'static str = "Value is gone!"; 505 | 506 | fn new(value: T) -> Self { 507 | Self::Unique(value) 508 | } 509 | 510 | fn share(&mut self) -> Linearc { 511 | let arc = match mem::replace(self, Self::Null) { 512 | Self::Unique(value) => Linearc::new(value), 513 | Self::Shared(arc) => arc, 514 | Self::Null => panic!("{}", Self::NULL_ERROR), 515 | }; 516 | *self = Self::Shared(Linearc::clone(&arc)); 517 | arc 518 | } 519 | 520 | fn as_ref(&self) -> &T { 521 | match *self { 522 | Self::Unique(ref value) => value, 523 | Self::Shared(ref arc) => arc, 524 | Self::Null => panic!("{}", Self::NULL_ERROR), 525 | } 526 | } 527 | 528 | fn extract(&mut self) -> Option { 529 | match mem::replace(self, Self::Null) { 530 | // No dependencies, can be launched now. 531 | Self::Unique(value) => Some(value), 532 | // There are dependencies, potentially all resolved now. 533 | Self::Shared(arc) => Linearc::into_inner(arc), 534 | // Already been extracted. 535 | Self::Null => None, 536 | } 537 | } 538 | } 539 | 540 | /// Task construct without any functional logic. 541 | pub struct ProtoTask<'c> { 542 | choir: &'c Arc, 543 | name: Name, 544 | parent: Option>, 545 | dependents: Vec>, 546 | } 547 | 548 | /// Task that is created but not running yet. 549 | /// It will be scheduled on `run()` or on drop. 550 | /// The 'a lifetime is responsible for the data 551 | /// in the closure of the task function. 552 | pub struct IdleTask { 553 | choir: Arc, 554 | task: MaybeArc, 555 | } 556 | 557 | impl AsRef for IdleTask { 558 | fn as_ref(&self) -> &Notifier { 559 | &self.task.as_ref().notifier 560 | } 561 | } 562 | 563 | //HACK: turns out, it's impossible to re-implement `Arc` in Rust stable userspace, 564 | // Because `Arc` relies on `trait Unsized` and `std::ops::CoerceUnsized`, which 565 | // are still nightly-only behind a feature. 566 | impl<'a> Linearc { 567 | fn new_unsized(fun: impl Fn(ExecutionContext, SubIndex) + Send + Sync + 'a) -> Self { 568 | Self::from_inner(Box::new(arc::LinearcInner { 569 | ref_count: AtomicUsize::new(1), 570 | data: fun, 571 | })) 572 | } 573 | } 574 | 575 | impl Drop for ProtoTask<'_> { 576 | fn drop(&mut self) { 577 | for dependent in self.dependents.drain(..) { 578 | if let Some(ready) = Linearc::into_inner(dependent) { 579 | self.choir.schedule(ready); 580 | } 581 | } 582 | } 583 | } 584 | 585 | impl ProtoTask<'_> { 586 | /// Block another task from starting until this one is finished. 587 | /// This is the reverse of `depend_on` but could be done faster because 588 | /// the `self` task is not yet shared with anything. 589 | pub fn as_blocker_for(mut self, other: &mut IdleTask) -> Self { 590 | self.dependents.push(other.task.share()); 591 | self 592 | } 593 | 594 | fn fill(mut self, functor: Functor) -> IdleTask { 595 | // Only register the fork here, so that nothing happens if a `ProtoTask` is dropped. 596 | if let Some(ref parent_notifier) = self.parent { 597 | parent_notifier 598 | .continuation 599 | .lock() 600 | .unwrap() 601 | .as_mut() 602 | .unwrap() 603 | .forks += 1; 604 | } 605 | IdleTask { 606 | choir: Arc::clone(self.choir), 607 | task: MaybeArc::new(Task { 608 | functor, 609 | notifier: Arc::new(Notifier { 610 | name: mem::take(&mut self.name), 611 | continuation: Mutex::new(Some(Continuation { 612 | parents: self.parent.take().into_iter().collect(), 613 | forks: 0, 614 | dependents: mem::take(&mut self.dependents), 615 | waiting_threads: Vec::new(), 616 | })), 617 | }), 618 | }), 619 | } 620 | } 621 | 622 | /// Init task with no function. 623 | /// Can be useful to aggregate dependencies, for example 624 | /// if a function returns a task handle, and it launches 625 | /// multiple sub-tasks in parallel. 626 | pub fn init_dummy(self) -> IdleTask { 627 | self.fill(Functor::Dummy) 628 | } 629 | 630 | /// Init task to execute a standalone function. 631 | /// The function body will be executed once the task is scheduled, 632 | /// and all of its dependencies are fulfilled. 633 | pub fn init(self, fun: F) -> IdleTask { 634 | let b: Box = Box::new(fun); 635 | // Transmute is for the lifetime bound only: it's stored as `'static`, 636 | // but the only way to run it is `run_attached`, which would be blocking. 637 | self.fill(Functor::Once(unsafe { mem::transmute(b) })) 638 | } 639 | 640 | /// Init task to execute a function multiple times. 641 | /// Every invocation is given an index in 0..count 642 | /// There are no ordering guarantees between the indices. 643 | pub fn init_multi( 644 | self, 645 | count: SubIndex, 646 | fun: F, 647 | ) -> IdleTask { 648 | self.fill(if count == 0 { 649 | Functor::Dummy 650 | } else { 651 | let arc: Linearc = 652 | Linearc::new_unsized(fun); 653 | Functor::Multi(0..count, unsafe { mem::transmute(arc) }) 654 | }) 655 | } 656 | 657 | /// Init task to execute a function on each element of a finite iterator. 658 | /// Similarly to `init_multi`, each invocation is executed 659 | /// indepdently and can be out of order. 660 | pub fn init_iter(self, iter: I, fun: F) -> IdleTask 661 | where 662 | I: Iterator, 663 | I::Item: Send + 'static, 664 | F: Fn(ExecutionContext, I::Item) + Send + Sync + 'static, 665 | { 666 | let task_data = iter.collect::>(); 667 | self.init_multi(task_data.len(), move |exe_context, index| unsafe { 668 | fun(exe_context, task_data.take(index)) 669 | }) 670 | } 671 | } 672 | 673 | /// Task that is already scheduled for running. 674 | #[derive(Clone)] 675 | pub struct RunningTask { 676 | choir: Arc, 677 | notifier: Arc, 678 | } 679 | 680 | impl fmt::Debug for RunningTask { 681 | fn fmt(&self, serializer: &mut fmt::Formatter) -> fmt::Result { 682 | self.notifier.fmt(serializer) 683 | } 684 | } 685 | 686 | impl AsRef for RunningTask { 687 | fn as_ref(&self) -> &Notifier { 688 | &self.notifier 689 | } 690 | } 691 | 692 | impl RunningTask { 693 | /// Return true if this task is done. 694 | pub fn is_done(&self) -> bool { 695 | self.notifier.continuation.lock().unwrap().is_none() 696 | } 697 | 698 | /// Block until the task has finished executing. 699 | #[profiling::function] 700 | pub fn join(&self) -> MaybePanic { 701 | log::debug!("Joining {}", self.notifier); 702 | let mut guard = self.notifier.continuation.lock().unwrap(); 703 | // This code is a bit magical, 704 | // and it's amazing to see it compiling and working. 705 | if let Some(ref mut cont) = *guard { 706 | let condvar = Condvar::new(); 707 | cont.waiting_threads 708 | .push(WaitingThread { condvar: &condvar }); 709 | let _ = condvar.wait_while(guard, |cont| cont.is_some()); 710 | } 711 | self.choir.check_panic() 712 | } 713 | 714 | /// Block until the task has finished executing. 715 | /// Also, use the current thread to help in the meantime. 716 | #[profiling::function] 717 | pub fn join_active(&self) -> MaybePanic { 718 | let condvar; 719 | match *self.notifier.continuation.lock().unwrap() { 720 | Some(ref mut cont) => { 721 | condvar = Condvar::new(); 722 | cont.waiting_threads 723 | .push(WaitingThread { condvar: &condvar }); 724 | } 725 | None => return self.choir.check_panic(), 726 | } 727 | let index = self.choir.register().unwrap(); 728 | log::info!("Join thread[{}] started", index); 729 | 730 | loop { 731 | let is_done = match self.choir.injector.steal() { 732 | Steal::Empty => { 733 | log::trace!("Thread[{}] sleeps", index); 734 | let guard = self.notifier.continuation.lock().unwrap(); 735 | if guard.is_some() { 736 | let guard = condvar.wait(guard).unwrap(); 737 | guard.is_none() 738 | } else { 739 | false 740 | } 741 | } 742 | Steal::Success(task) => { 743 | self.choir.execute(task, index as isize); 744 | self.is_done() 745 | } 746 | Steal::Retry => false, 747 | }; 748 | if is_done { 749 | break; 750 | } 751 | } 752 | 753 | log::info!("Thread[{}] is released", index); 754 | self.choir.unregister(index); 755 | self.choir.check_panic() 756 | } 757 | 758 | /// Block until the task has finished executing, with timeout. 759 | /// Panics and prints helpful info if the timeout is reached. 760 | pub fn join_debug(&self, timeout: time::Duration) -> MaybePanic { 761 | log::debug!("Joining {}", self.notifier); 762 | let mut guard = self.notifier.continuation.lock().unwrap(); 763 | if let Some(ref mut cont) = *guard { 764 | let condvar = Condvar::new(); 765 | cont.waiting_threads 766 | .push(WaitingThread { condvar: &condvar }); 767 | let (guard, wait_result) = condvar 768 | .wait_timeout_while(guard, timeout, |cont| cont.is_some()) 769 | .unwrap(); 770 | if wait_result.timed_out() { 771 | println!("Join timeout reached for {}", self.notifier); 772 | println!("Continuation: {:?}", guard); 773 | panic!(""); 774 | } 775 | } 776 | self.choir.check_panic() 777 | } 778 | } 779 | 780 | impl Drop for WorkerHandle { 781 | fn drop(&mut self) { 782 | self.worker.alive.store(false, Ordering::Release); 783 | let handle = self.join_handle.take().unwrap(); 784 | // make sure it wakes up and checks if it's still alive 785 | // Locking the mutex is required to guarantee that the worker loop 786 | // actually receives the notification. 787 | if let Ok(_guard) = self.choir.parked_mask_mutex.lock() { 788 | self.choir.condvar.notify_all(); 789 | } 790 | let _ = handle.join(); 791 | } 792 | } 793 | 794 | impl IdleTask { 795 | /// Schedule this task for running. 796 | /// 797 | /// It will only be executed once the dependencies are fulfilled. 798 | pub fn run(self) -> RunningTask { 799 | let task = self.task.as_ref(); 800 | RunningTask { 801 | choir: Arc::clone(&self.choir), 802 | notifier: Arc::clone(&task.notifier), 803 | } 804 | } 805 | 806 | /// Run the task now and block until it's executed. 807 | /// Use the current thread to help the choir in the meantime. 808 | pub fn run_attached(mut self) { 809 | let task = self.task.as_ref(); 810 | let notifier = Arc::clone(&task.notifier); 811 | 812 | if let Some(ready) = self.task.extract() { 813 | // Task has no dependencies. No need to join the pool, 814 | // just execute it right here instead. 815 | self.choir.execute(ready, -1); 816 | } else { 817 | RunningTask { 818 | choir: Arc::clone(&self.choir), 819 | notifier, 820 | } 821 | .join_active(); 822 | } 823 | } 824 | 825 | /// Add a dependency on another task, which is possibly running. 826 | #[profiling::function] 827 | pub fn depend_on>(&mut self, dependency: &D) { 828 | if let Some(ref mut cont) = *dependency.as_ref().continuation.lock().unwrap() { 829 | cont.dependents.push(self.task.share()); 830 | } 831 | } 832 | } 833 | 834 | impl Drop for IdleTask { 835 | fn drop(&mut self) { 836 | if let Some(ready) = self.task.extract() { 837 | self.choir.schedule(ready); 838 | } 839 | } 840 | } 841 | --------------------------------------------------------------------------------