├── .gitignore ├── .travis.yml ├── Cargo.toml ├── LICENSE-APACHE ├── LICENSE-MIT ├── README.md ├── benches ├── basic.rs └── depth.rs ├── examples ├── depth.rs └── smoke.rs ├── src ├── deque.rs ├── lib.rs └── task.rs └── tests └── pool.rs /.gitignore: -------------------------------------------------------------------------------- 1 | target 2 | Cargo.lock 3 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | --- 2 | language: rust 3 | sudo: false 4 | 5 | matrix: 6 | include: 7 | # Minimum rustc version 8 | - rust: 1.20.0 9 | - rust: nightly 10 | 11 | script: 12 | - cargo test 13 | - 'if [ "$TRAVIS_RUST_VERSION" == "nightly" ]; then cargo test --benches; fi' 14 | - rustdoc --test README.md -L target/debug/deps 15 | - cargo doc --no-deps 16 | 17 | notifications: 18 | email: 19 | on_success: never 20 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "futures-pool" 3 | version = "0.1.0" 4 | documentation = "https://docs.rs/futures-pool" 5 | repository = "https://github.com/carllerche/futures-pool" 6 | homepage = "https://github.com/carllerche/futures-pool" 7 | license = "MIT/Apache-2.0" 8 | authors = ["Carl Lerche "] 9 | description = """ 10 | A Future aware thread pool based on work stealing. 11 | """ 12 | keywords = ["futures"] 13 | categories = ["concurrency", "asynchronous"] 14 | 15 | [dependencies] 16 | coco = "0.3" 17 | futures = "0.1" 18 | num_cpus = "1.2" 19 | rand = "0.3" 20 | log = "0.3" 21 | 22 | [dev-dependencies] 23 | tokio-timer = "0.1" 24 | env_logger = "0.4" 25 | futures-cpupool = "0.1.7" 26 | -------------------------------------------------------------------------------- /LICENSE-APACHE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright 2017 Carl Lerche 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /LICENSE-MIT: -------------------------------------------------------------------------------- 1 | Copyright (c) 2017 Carl Lerche 2 | 3 | Permission is hereby granted, free of charge, to any 4 | person obtaining a copy of this software and associated 5 | documentation files (the "Software"), to deal in the 6 | Software without restriction, including without 7 | limitation the rights to use, copy, modify, merge, 8 | publish, distribute, sublicense, and/or sell copies of 9 | the Software, and to permit persons to whom the Software 10 | is furnished to do so, subject to the following 11 | conditions: 12 | 13 | The above copyright notice and this permission notice 14 | shall be included in all copies or substantial portions 15 | of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF 18 | ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED 19 | TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A 20 | PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT 21 | SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY 22 | CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 23 | OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR 24 | IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 25 | DEALINGS IN THE SOFTWARE. 26 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Futures pool 2 | 3 | A library for scheduling execution of futures concurrently across a pool of 4 | threads. 5 | 6 | **Note**: This library isn't quite ready for use. 7 | 8 | [![Travis Build Status](https://travis-ci.org/carllerche/futures-pool.svg?branch=master)](https://travis-ci.org/carllerche/futures-pool) 9 | 10 | ### Why not Rayon? 11 | 12 | Rayon is designed to handle parallelizing single computations by breaking them 13 | into smaller chunks. The scheduling for each individual chunk doesn't matter as 14 | long as the root computation completes in a timely fashion. In other words, 15 | Rayon does not provide any guarantees of fairness with regards to how each task 16 | gets scheduled. 17 | 18 | On the other hand, `futures-pool` is a general purpose scheduler and attempts to 19 | schedule each task fairly. This is the ideal behavior when scheduling a set of 20 | unrelated tasks. 21 | 22 | ### Why not futures-cpupool? 23 | 24 | It's 10x slower. 25 | 26 | ## Usage 27 | 28 | To use `futures-pool`, first add this to your `Cargo.toml`: 29 | 30 | ```toml 31 | [dependencies] 32 | futures-pool = { git = "https://github.com/carllerche/futures-pool" } # Soon on crates.io 33 | ``` 34 | 35 | Next, add this to your crate: 36 | 37 | ```rust 38 | extern crate futures_pool; 39 | 40 | fn main() { 41 | // ... 42 | } 43 | ``` 44 | ## Examples 45 | 46 | ```rust 47 | extern crate futures; 48 | extern crate futures_pool; 49 | 50 | use futures::*; 51 | use futures::sync::oneshot; 52 | use futures_pool::*; 53 | 54 | pub fn main() { 55 | let (tx, _pool) = Pool::new(); 56 | 57 | let res = oneshot::spawn(future::lazy(|| { 58 | println!("Running on the pool"); 59 | Ok::<_, ()>("complete") 60 | }), &tx); 61 | 62 | println!("Result: {:?}", res.wait()); 63 | } 64 | ``` 65 | 66 | ## License 67 | 68 | `futures-pool` is primarily distributed under the terms of both the MIT license 69 | and the Apache License (Version 2.0), with portions covered by various BSD-like 70 | licenses. 71 | 72 | See LICENSE-APACHE, and LICENSE-MIT for details. 73 | -------------------------------------------------------------------------------- /benches/basic.rs: -------------------------------------------------------------------------------- 1 | #![feature(test)] 2 | 3 | extern crate futures; 4 | extern crate futures_pool; 5 | extern crate futures_cpupool; 6 | extern crate num_cpus; 7 | extern crate test; 8 | 9 | const NUM_SPAWN: usize = 10_000; 10 | const NUM_YIELD: usize = 1_000; 11 | const TASKS_PER_CPU: usize = 50; 12 | 13 | mod us { 14 | use futures::{task, Async}; 15 | use futures::future::{self, Executor}; 16 | use futures_pool::*; 17 | use num_cpus; 18 | use test; 19 | use std::sync::{mpsc, Arc}; 20 | use std::sync::atomic::AtomicUsize; 21 | use std::sync::atomic::Ordering::SeqCst; 22 | 23 | #[bench] 24 | fn spawn_many(b: &mut test::Bencher) { 25 | let (sched_tx, _scheduler) = Pool::new(); 26 | 27 | let (tx, rx) = mpsc::sync_channel(10); 28 | let rem = Arc::new(AtomicUsize::new(0)); 29 | 30 | b.iter(move || { 31 | rem.store(super::NUM_SPAWN, SeqCst); 32 | 33 | for _ in 0..super::NUM_SPAWN { 34 | let tx = tx.clone(); 35 | let rem = rem.clone(); 36 | 37 | sched_tx.execute(future::lazy(move || { 38 | if 1 == rem.fetch_sub(1, SeqCst) { 39 | tx.send(()).unwrap(); 40 | } 41 | 42 | Ok(()) 43 | })).ok().unwrap(); 44 | } 45 | 46 | let _ = rx.recv().unwrap(); 47 | }); 48 | } 49 | 50 | #[bench] 51 | fn yield_many(b: &mut test::Bencher) { 52 | let (sched_tx, _scheduler) = Pool::new(); 53 | let tasks = super::TASKS_PER_CPU * num_cpus::get(); 54 | 55 | let (tx, rx) = mpsc::sync_channel(tasks); 56 | 57 | b.iter(move || { 58 | for _ in 0..tasks { 59 | let mut rem = super::NUM_YIELD; 60 | let tx = tx.clone(); 61 | 62 | sched_tx.execute(future::poll_fn(move || { 63 | rem -= 1; 64 | 65 | if rem == 0 { 66 | tx.send(()).unwrap(); 67 | Ok(Async::Ready(())) 68 | } else { 69 | // Notify the current task 70 | task::current().notify(); 71 | 72 | // Not ready 73 | Ok(Async::NotReady) 74 | } 75 | })).ok().unwrap(); 76 | } 77 | 78 | for _ in 0..tasks { 79 | let _ = rx.recv().unwrap(); 80 | } 81 | }); 82 | } 83 | } 84 | 85 | // In this case, CPU pool completes the benchmark faster, but this is due to how 86 | // CpuPool currently behaves, starving other futures. This completes the 87 | // benchmark quickly but results in poor runtime characteristics for a thread 88 | // pool. 89 | // 90 | // See alexcrichton/futures-rs#617 91 | // 92 | mod cpupool { 93 | use futures::{task, Async}; 94 | use futures::future::{self, Executor}; 95 | use futures_cpupool::*; 96 | use num_cpus; 97 | use test; 98 | use std::sync::{mpsc, Arc}; 99 | use std::sync::atomic::AtomicUsize; 100 | use std::sync::atomic::Ordering::SeqCst; 101 | 102 | #[bench] 103 | fn spawn_many(b: &mut test::Bencher) { 104 | let pool = CpuPool::new(num_cpus::get()); 105 | 106 | let (tx, rx) = mpsc::sync_channel(10); 107 | let rem = Arc::new(AtomicUsize::new(0)); 108 | 109 | b.iter(move || { 110 | rem.store(super::NUM_SPAWN, SeqCst); 111 | 112 | for _ in 0..super::NUM_SPAWN { 113 | let tx = tx.clone(); 114 | let rem = rem.clone(); 115 | 116 | pool.execute(future::lazy(move || { 117 | if 1 == rem.fetch_sub(1, SeqCst) { 118 | tx.send(()).unwrap(); 119 | } 120 | 121 | Ok(()) 122 | })).ok().unwrap(); 123 | } 124 | 125 | let _ = rx.recv().unwrap(); 126 | }); 127 | } 128 | 129 | #[bench] 130 | fn yield_many(b: &mut test::Bencher) { 131 | let pool = CpuPool::new(num_cpus::get()); 132 | let tasks = super::TASKS_PER_CPU * num_cpus::get(); 133 | 134 | let (tx, rx) = mpsc::sync_channel(tasks); 135 | 136 | b.iter(move || { 137 | for _ in 0..tasks { 138 | let mut rem = super::NUM_YIELD; 139 | let tx = tx.clone(); 140 | 141 | pool.execute(future::poll_fn(move || { 142 | rem -= 1; 143 | 144 | if rem == 0 { 145 | tx.send(()).unwrap(); 146 | Ok(Async::Ready(())) 147 | } else { 148 | // Notify the current task 149 | task::current().notify(); 150 | 151 | // Not ready 152 | Ok(Async::NotReady) 153 | } 154 | })).ok().unwrap(); 155 | } 156 | 157 | for _ in 0..tasks { 158 | let _ = rx.recv().unwrap(); 159 | } 160 | }); 161 | } 162 | } 163 | -------------------------------------------------------------------------------- /benches/depth.rs: -------------------------------------------------------------------------------- 1 | #![feature(test)] 2 | 3 | extern crate futures; 4 | extern crate futures_pool; 5 | extern crate futures_cpupool; 6 | extern crate num_cpus; 7 | extern crate test; 8 | 9 | const ITER: usize = 20_000; 10 | 11 | mod us { 12 | use futures::future::{self, Executor}; 13 | use futures_pool::*; 14 | use test; 15 | use std::sync::mpsc; 16 | 17 | #[bench] 18 | fn chained_spawn(b: &mut test::Bencher) { 19 | let (sched_tx, _scheduler) = Pool::new(); 20 | 21 | fn spawn(sched_tx: Sender, res_tx: mpsc::Sender<()>, n: usize) { 22 | if n == 0 { 23 | res_tx.send(()).unwrap(); 24 | } else { 25 | let sched_tx2 = sched_tx.clone(); 26 | sched_tx.execute(future::lazy(move || { 27 | spawn(sched_tx2, res_tx, n - 1); 28 | Ok(()) 29 | })).ok().unwrap(); 30 | } 31 | } 32 | 33 | b.iter(move || { 34 | let (res_tx, res_rx) = mpsc::channel(); 35 | 36 | spawn(sched_tx.clone(), res_tx, super::ITER); 37 | res_rx.recv().unwrap(); 38 | }); 39 | } 40 | } 41 | 42 | mod cpupool { 43 | use futures::future::{self, Executor}; 44 | use futures_cpupool::*; 45 | use num_cpus; 46 | use test; 47 | use std::sync::mpsc; 48 | 49 | #[bench] 50 | fn chained_spawn(b: &mut test::Bencher) { 51 | let pool = CpuPool::new(num_cpus::get()); 52 | 53 | fn spawn(pool: CpuPool, res_tx: mpsc::Sender<()>, n: usize) { 54 | if n == 0 { 55 | res_tx.send(()).unwrap(); 56 | } else { 57 | let pool2 = pool.clone(); 58 | pool.execute(future::lazy(move || { 59 | spawn(pool2, res_tx, n - 1); 60 | Ok(()) 61 | })).ok().unwrap(); 62 | } 63 | } 64 | 65 | b.iter(move || { 66 | let (res_tx, res_rx) = mpsc::channel(); 67 | 68 | spawn(pool.clone(), res_tx, super::ITER); 69 | res_rx.recv().unwrap(); 70 | }); 71 | } 72 | } 73 | -------------------------------------------------------------------------------- /examples/depth.rs: -------------------------------------------------------------------------------- 1 | extern crate futures; 2 | extern crate futures_pool; 3 | extern crate env_logger; 4 | 5 | use futures::future::{self, Executor}; 6 | use futures_pool::*; 7 | 8 | use std::sync::mpsc; 9 | 10 | const ITER: usize = 2_000_000; 11 | // const ITER: usize = 30; 12 | 13 | fn chained_spawn() { 14 | let (sched_tx, _scheduler) = Pool::new(); 15 | 16 | fn spawn(sched_tx: Sender, res_tx: mpsc::Sender<()>, n: usize) { 17 | if n == 0 { 18 | res_tx.send(()).unwrap(); 19 | } else { 20 | let sched_tx2 = sched_tx.clone(); 21 | sched_tx.execute(future::lazy(move || { 22 | spawn(sched_tx2, res_tx, n - 1); 23 | Ok(()) 24 | })).ok().unwrap(); 25 | } 26 | } 27 | 28 | loop { 29 | println!("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"); 30 | let (res_tx, res_rx) = mpsc::channel(); 31 | 32 | for _ in 0..10 { 33 | spawn(sched_tx.clone(), res_tx.clone(), ITER); 34 | } 35 | 36 | for _ in 0..10 { 37 | res_rx.recv().unwrap(); 38 | } 39 | } 40 | } 41 | 42 | pub fn main() { 43 | let _ = ::env_logger::init(); 44 | chained_spawn(); 45 | } 46 | -------------------------------------------------------------------------------- /examples/smoke.rs: -------------------------------------------------------------------------------- 1 | extern crate futures; 2 | extern crate futures_pool; 3 | extern crate tokio_timer; 4 | extern crate env_logger; 5 | 6 | use futures::*; 7 | use futures::sync::oneshot::spawn; 8 | use futures_pool::*; 9 | 10 | use tokio_timer::Timer; 11 | 12 | use std::thread; 13 | use std::time::Duration; 14 | 15 | pub fn main() { 16 | let _ = ::env_logger::init(); 17 | 18 | let timer = Timer::default(); 19 | { 20 | let (tx, _scheduler) = Pool::new(); 21 | 22 | let fut = timer.interval(Duration::from_millis(300)) 23 | .for_each(|_| { 24 | println!("~~~~~ Hello ~~~"); 25 | Ok(()) 26 | }) 27 | .map_err(|_| unimplemented!()); 28 | 29 | spawn(fut, &tx).wait().unwrap(); 30 | } 31 | 32 | thread::sleep(Duration::from_millis(100)); 33 | } 34 | -------------------------------------------------------------------------------- /src/deque.rs: -------------------------------------------------------------------------------- 1 | //! A (mostly) lock-free concurrent work-stealing deque 2 | //! 3 | //! This module contains an implementation of the Chase-Lev work stealing deque 4 | //! described in "Dynamic Circular Work-Stealing Deque". The implementation is 5 | //! heavily based on the implementation using C11 atomics in "Correct and 6 | //! Efficient Work Stealing for Weak Memory Models". 7 | //! 8 | //! The only potentially lock-synchronized portion of this deque is the 9 | //! occasional call to the memory allocator when growing the deque. Otherwise 10 | //! all operations are lock-free. 11 | 12 | pub use self::Poll::*; 13 | 14 | use std::mem::forget; 15 | use std::ptr; 16 | 17 | use std::sync::atomic::{AtomicIsize, AtomicPtr, fence}; 18 | use std::sync::atomic::Ordering::{SeqCst, Acquire, Release, Relaxed}; 19 | 20 | // Initial size for a buffer. 21 | // TODO: Much too small... 22 | static MIN_SIZE: usize = 32; 23 | 24 | #[derive(Debug)] 25 | pub struct Deque { 26 | bottom: AtomicIsize, 27 | top: AtomicIsize, 28 | array: AtomicPtr>, 29 | } 30 | 31 | /// When stealing some data, this is an enumeration of the possible outcomes. 32 | #[derive(Debug, PartialEq)] 33 | pub enum Poll { 34 | /// The deque was empty at the time of stealing 35 | Empty, 36 | /// The stealer lost the race for stealing data, and a retry may return more 37 | /// data. 38 | Inconsistent, 39 | /// The stealer has successfully stolen some data. 40 | Data(T), 41 | } 42 | 43 | /// An internal buffer used by the chase-lev deque. This structure is actually 44 | /// implemented as a circular buffer, and is used as the intermediate storage of 45 | /// the data in the deque. 46 | /// 47 | /// This type is implemented with *T instead of Vec for two reasons: 48 | /// 49 | /// 1. There is nothing safe about using this buffer. This easily allows the 50 | /// same value to be read twice in to rust, and there is nothing to 51 | /// prevent this. The usage by the deque must ensure that one of the 52 | /// values is forgotten. Furthermore, we only ever want to manually run 53 | /// destructors for values in this buffer (on drop) because the bounds 54 | /// are defined by the deque it's owned by. 55 | /// 56 | /// 2. We can certainly avoid bounds checks using *T instead of Vec, although 57 | /// LLVM is probably pretty good at doing this already. 58 | /// 59 | /// Note that we keep old buffers around after growing because stealers may still 60 | /// be concurrently accessing them. The buffers are kept in a linked list, with 61 | /// each buffer pointing to the previous, smaller buffer. This doesn't leak any 62 | /// memory because all buffers in the list are freed when the deque is dropped. 63 | #[derive(Debug)] 64 | struct Buffer { 65 | storage: *mut T, 66 | size: usize, 67 | prev: Option>>, 68 | } 69 | 70 | impl Deque { 71 | pub fn new() -> Deque { 72 | let buf = Box::new(unsafe { Buffer::new(MIN_SIZE) }); 73 | Deque { 74 | bottom: AtomicIsize::new(0), 75 | top: AtomicIsize::new(0), 76 | array: AtomicPtr::new(Box::into_raw(buf)), 77 | } 78 | } 79 | 80 | #[inline] 81 | pub unsafe fn push(&self, data: T) { 82 | let b = self.bottom.load(Relaxed); 83 | let t = self.top.load(Acquire); 84 | let mut a = self.array.load(Relaxed); 85 | 86 | // Grow the buffer if it is full. 87 | let size = b.wrapping_sub(t); 88 | if size == (*a).size() { 89 | a = Box::into_raw(Box::from_raw(a).grow(b, t)); 90 | self.array.store(a, Release); 91 | } 92 | 93 | (*a).put(b, data); 94 | fence(Release); 95 | self.bottom.store(b.wrapping_add(1), Relaxed); 96 | } 97 | 98 | #[inline] 99 | pub fn poll(&self) -> Poll { 100 | unsafe { 101 | // Make sure top is read before bottom. 102 | let t = self.top.load(Acquire); 103 | fence(SeqCst); 104 | let b = self.bottom.load(Acquire); 105 | 106 | // Exit if the queue is empty. 107 | let size = b.wrapping_sub(t); 108 | if size <= 0 { 109 | return Empty; 110 | } 111 | 112 | // Fetch the element from the queue. 113 | let a = self.array.load(Acquire); 114 | let data = (*a).get(t); 115 | 116 | // Attempt to increment top. 117 | if self.top.compare_and_swap(t, t.wrapping_add(1), SeqCst) == t { 118 | Data(data) 119 | } else { 120 | forget(data); // Someone else stole this value 121 | Inconsistent 122 | } 123 | } 124 | } 125 | } 126 | 127 | impl Drop for Deque { 128 | fn drop(&mut self) { 129 | let t = self.top.load(Relaxed); 130 | let b = self.bottom.load(Relaxed); 131 | let a = self.array.load(Relaxed); 132 | 133 | // Free whatever is leftover in the deque, and then free the buffer. 134 | // This will also free all linked buffers. 135 | let mut i = t; 136 | while i != b { 137 | unsafe { (*a).get(i) }; 138 | i = i.wrapping_add(1); 139 | } 140 | unsafe { Box::from_raw(a) }; 141 | } 142 | } 143 | 144 | #[inline] 145 | unsafe fn take_ptr_from_vec(mut buf: Vec) -> *mut T { 146 | let ptr = buf.as_mut_ptr(); 147 | forget(buf); 148 | ptr 149 | } 150 | 151 | #[inline] 152 | unsafe fn allocate(number: usize) -> *mut T { 153 | let v = Vec::with_capacity(number); 154 | take_ptr_from_vec(v) 155 | } 156 | 157 | #[inline] 158 | unsafe fn deallocate(ptr: *mut T, number: usize) { 159 | Vec::from_raw_parts(ptr, 0, number); 160 | } 161 | 162 | impl Buffer { 163 | unsafe fn new(size: usize) -> Buffer { 164 | Buffer { 165 | storage: allocate(size), 166 | size: size, 167 | prev: None, 168 | } 169 | } 170 | 171 | fn size(&self) -> isize { self.size as isize } 172 | 173 | fn mask(&self) -> isize { self.size as isize - 1 } 174 | 175 | unsafe fn elem(&self, i: isize) -> *mut T { 176 | self.storage.offset(i & self.mask()) 177 | } 178 | 179 | // This does not protect against loading duplicate values of the same cell, 180 | // nor does this clear out the contents contained within. Hence, this is a 181 | // very unsafe method which the caller needs to treat specially in case a 182 | // race is lost. 183 | unsafe fn get(&self, i: isize) -> T { 184 | ptr::read(self.elem(i)) 185 | } 186 | 187 | // Unsafe because this unsafely overwrites possibly uninitialized or 188 | // initialized data. 189 | unsafe fn put(&self, i: isize, t: T) { 190 | ptr::write(self.elem(i), t); 191 | } 192 | 193 | // Again, unsafe because this has incredibly dubious ownership violations. 194 | // It is assumed that this buffer is immediately dropped. 195 | unsafe fn grow(self: Box>, b: isize, t: isize) -> Box> { 196 | let mut buf = Box::new(Buffer::new(self.size * 2)); 197 | let mut i = t; 198 | while i != b { 199 | buf.put(i, self.get(i)); 200 | i = i.wrapping_add(1); 201 | } 202 | buf.prev = Some(self); 203 | return buf; 204 | } 205 | } 206 | 207 | impl Drop for Buffer { 208 | fn drop(&mut self) { 209 | // It is assumed that all buffers are empty on drop. 210 | unsafe { deallocate(self.storage, self.size) } 211 | } 212 | } 213 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | //! A work-stealing based thread pool for executing futures. 2 | 3 | #![deny(warnings, missing_docs, missing_debug_implementations)] 4 | 5 | extern crate coco; 6 | extern crate futures; 7 | extern crate num_cpus; 8 | extern crate rand; 9 | 10 | #[macro_use] 11 | extern crate log; 12 | 13 | mod task; 14 | 15 | use coco::deque; 16 | use task::Task; 17 | 18 | use futures::{Future, Poll, Async}; 19 | use futures::executor::Notify; 20 | use futures::future::{Executor, ExecuteError}; 21 | use futures::task::AtomicTask; 22 | 23 | use rand::{Rng, SeedableRng, XorShiftRng}; 24 | 25 | use std::{fmt, mem, thread, usize}; 26 | use std::cell::{Cell, UnsafeCell}; 27 | use std::sync::{Arc, Weak, Mutex, Condvar}; 28 | use std::sync::atomic::AtomicUsize; 29 | use std::sync::atomic::Ordering::{AcqRel, Acquire, Release, Relaxed}; 30 | use std::time::{Instant, Duration}; 31 | 32 | /// Thread pool 33 | #[derive(Debug)] 34 | pub struct Pool { 35 | inner: Arc, 36 | } 37 | 38 | /// Sends work to the pool 39 | #[derive(Debug)] 40 | pub struct Sender { 41 | inner: Arc, 42 | } 43 | 44 | /// Pool configuration 45 | #[derive(Debug)] 46 | pub struct Builder { 47 | /// Thread pool specific configuration values 48 | config: Config, 49 | 50 | /// Number of workers to spawn 51 | pool_size: usize, 52 | } 53 | 54 | /// Thread pool specific configuration values 55 | #[derive(Debug, Clone)] 56 | struct Config { 57 | keep_alive: Option, 58 | // Used to configure a worker thread 59 | name_prefix: Option, 60 | stack_size: Option, 61 | after_start: Option, 62 | before_stop: Option, 63 | } 64 | 65 | #[derive(Debug)] 66 | struct Inner { 67 | // Pool state 68 | state: AtomicUsize, 69 | 70 | // Stack tracking sleeping workers. 71 | sleep_stack: AtomicUsize, 72 | 73 | // Number of workers who haven't reached the final state of shutdown 74 | // 75 | // This is only used to know when to single `shutdown_task` once the 76 | // shutdown process has completed. 77 | num_workers: AtomicUsize, 78 | 79 | // Used to generate a thread local RNG seed 80 | next_thread_id: AtomicUsize, 81 | 82 | // Storage for workers 83 | // 84 | // This will *usually* be a small number 85 | workers: Box<[WorkerEntry]>, 86 | 87 | // Task notified when the worker shuts down 88 | shutdown_task: AtomicTask, 89 | 90 | // Configuration 91 | config: Config, 92 | } 93 | 94 | #[derive(Clone)] 95 | struct Callback { 96 | f: Arc, 97 | } 98 | 99 | /// Implements the future `Notify` API. 100 | /// 101 | /// This is how external events are able to signal the task, informing it to try 102 | /// to poll the future again. 103 | #[derive(Debug)] 104 | struct Notifier { 105 | inner: Weak, 106 | } 107 | 108 | /// Pool state. 109 | /// 110 | /// The least significant bit is a flag indicating if the pool is shutting down 111 | /// (0 for active, 1 for shutting down). The remaining bits represent the number 112 | /// of futures that still need to complete. 113 | #[derive(Eq, PartialEq, Clone, Copy)] 114 | struct State(usize); 115 | 116 | /// Flag used to track if the pool is running 117 | const SHUTDOWN: usize = 1; 118 | 119 | /// Max number of futures the pool can handle. 120 | const MAX_FUTURES: usize = usize::MAX >> 1; 121 | 122 | /// State related to the stack of sleeping workers. 123 | /// 124 | /// - Parked head 16 bits 125 | /// - Sequence remaining 126 | /// 127 | /// The parked head value has a couple of special values: 128 | /// 129 | /// - EMPTY: No sleepers 130 | /// - TERMINATED: Don't spawn more threads 131 | #[derive(Eq, PartialEq, Clone, Copy)] 132 | struct SleepStack(usize); 133 | 134 | /// Extracts the head of the worker stack from the scheduler state 135 | const STACK_MASK: usize = ((1 << 16) - 1); 136 | 137 | /// Max number of workers that can be part of a pool. This is the most that can 138 | /// fit in the scheduler state. Note, that this is the max number of **active** 139 | /// threads. There can be more standby threads. 140 | const MAX_WORKERS: usize = 1 << 15; 141 | 142 | /// Used to mark the stack as empty 143 | const EMPTY: usize = MAX_WORKERS; 144 | 145 | /// Used to mark the stack as terminated 146 | const TERMINATED: usize = EMPTY + 1; 147 | 148 | /// How many bits the treiber ABA guard is offset by 149 | const ABA_GUARD_SHIFT: usize = 16; 150 | 151 | #[cfg(target_pointer_width = "64")] 152 | const ABA_GUARD_MASK: usize = (1 << (64 - ABA_GUARD_SHIFT)) - 1; 153 | 154 | #[cfg(target_pointer_width = "32")] 155 | const ABA_GUARD_MASK: usize = (1 << (32 - ABA_GUARD_SHIFT)) - 1; 156 | 157 | // Some constants used to work with State 158 | // const A: usize: 0; 159 | 160 | // TODO: This should be split up between what is accessed by each thread and 161 | // what is concurrent. The bits accessed by each thread should be sized to 162 | // exactly one cache line. 163 | #[derive(Debug)] 164 | struct WorkerEntry { 165 | // Worker state. This is mutated when notifying the worker. 166 | state: AtomicUsize, 167 | 168 | // Next entry in the parked Trieber stack 169 | next_sleeper: UnsafeCell, 170 | 171 | // Worker half of deque 172 | deque: deque::Worker, 173 | 174 | // Stealer half of deque 175 | steal: deque::Stealer, 176 | 177 | // Park mutex 178 | park_mutex: Mutex<()>, 179 | 180 | // Park condvar 181 | park_condvar: Condvar, 182 | 183 | // MPSC queue of jobs submitted to the worker from an external source. 184 | inbound: task::Queue, 185 | } 186 | 187 | /// Tracks worker state 188 | #[derive(Clone, Copy, Eq, PartialEq)] 189 | struct WorkerState(usize); 190 | 191 | /// Set when the worker is pushed onto the scheduler's stack of sleeping 192 | /// threads. 193 | const PUSHED_MASK: usize = 0b001; 194 | 195 | /// Manages the worker lifecycle part of the state 196 | const WORKER_LIFECYCLE_MASK: usize = 0b1110; 197 | const WORKER_LIFECYCLE_SHIFT: usize = 1; 198 | 199 | /// The worker does not currently have an associated thread. 200 | const WORKER_SHUTDOWN: usize = 0; 201 | 202 | /// The worker is currently processing its task. 203 | const WORKER_RUNNING: usize = 1; 204 | 205 | /// The worker is currently asleep in the condvar 206 | const WORKER_SLEEPING: usize = 2; 207 | 208 | /// The worker has been notified it should process more work. 209 | const WORKER_NOTIFIED: usize = 3; 210 | 211 | /// A stronger form of notification. In this case, the worker is expected to 212 | /// wakeup and try to acquire more work... if it enters this state while already 213 | /// busy with other work, it is expected to signal another worker. 214 | const WORKER_SIGNALED: usize = 4; 215 | 216 | struct Worker { 217 | // Shared scheduler data 218 | inner: Arc, 219 | 220 | // WorkerEntry index 221 | idx: usize, 222 | } 223 | 224 | // Pointer to the current worker info 225 | thread_local!(static CURRENT_WORKER: Cell<*const Worker> = Cell::new(0 as *const _)); 226 | 227 | // ===== impl Builder ===== 228 | 229 | impl Builder { 230 | /// Returns a builder with default values 231 | fn new() -> Builder { 232 | let num_cpus = num_cpus::get(); 233 | 234 | Builder { 235 | pool_size: num_cpus, 236 | config: Config { 237 | keep_alive: None, 238 | name_prefix: None, 239 | stack_size: None, 240 | after_start: None, 241 | before_stop: None, 242 | }, 243 | } 244 | } 245 | 246 | /// Set the scheduler's pool size 247 | pub fn pool_size(&mut self, val: usize) -> &mut Self { 248 | assert!(val >= 1, "at least one thread required"); 249 | 250 | self.pool_size = val; 251 | self 252 | } 253 | 254 | /// Set the thread keep alive duration 255 | pub fn keep_alive(&mut self, val: Duration) -> &mut Self { 256 | self.config.keep_alive = Some(val); 257 | self 258 | } 259 | 260 | /// Set name prefix of threads spawned by the scheduler 261 | /// 262 | /// Thread name prefix is used for generating thread names. For example, if 263 | /// prefix is `my-pool-`, then threads in the pool will get names like 264 | /// `my-pool-1` etc. 265 | pub fn name_prefix>(&mut self, val: S) -> &mut Self { 266 | self.config.name_prefix = Some(val.into()); 267 | self 268 | } 269 | 270 | /// Set the stack size of threads spawned by the scheduler 271 | pub fn stack_size(&mut self, val: usize) -> &mut Self { 272 | self.config.stack_size = Some(val); 273 | self 274 | } 275 | 276 | /// Execute function `f` right after each thread is started but before 277 | /// running any tasks on it 278 | /// 279 | /// This is initially intended for bookkeeping and monitoring uses 280 | pub fn after_start(&mut self, f: F) -> &mut Self 281 | where F: Fn() + Send + Sync + 'static 282 | { 283 | self.config.after_start = Some(Callback::new(f)); 284 | self 285 | } 286 | 287 | /// Execute function `f` before each worker thread stops 288 | /// 289 | /// This is initially intended for bookkeeping and monitoring uses 290 | pub fn before_stop(&mut self, f: F) -> &mut Self 291 | where F: Fn() + Send + Sync + 'static 292 | { 293 | self.config.before_stop = Some(Callback::new(f)); 294 | self 295 | } 296 | 297 | /// Build and return the configured thread pool 298 | pub fn build(&self) -> (Sender, Pool) { 299 | let mut workers = vec![]; 300 | 301 | trace!("build; num-workers={}", self.pool_size); 302 | 303 | for _ in 0..self.pool_size { 304 | workers.push(WorkerEntry::new()); 305 | } 306 | 307 | let inner = Arc::new(Inner { 308 | state: AtomicUsize::new(State::new().into()), 309 | sleep_stack: AtomicUsize::new(SleepStack::new().into()), 310 | num_workers: AtomicUsize::new(self.pool_size), 311 | next_thread_id: AtomicUsize::new(0), 312 | workers: workers.into_boxed_slice(), 313 | shutdown_task: AtomicTask::new(), 314 | config: self.config.clone(), 315 | }); 316 | 317 | // Now, we prime the sleeper stack 318 | for i in 0..self.pool_size { 319 | inner.push_sleeper(i).unwrap(); 320 | } 321 | 322 | let sender = Sender { 323 | inner: inner.clone(), 324 | }; 325 | 326 | let scheduler = Pool { 327 | inner: inner, 328 | }; 329 | 330 | (sender, scheduler) 331 | } 332 | } 333 | 334 | // ===== impl Pool ===== 335 | 336 | impl Pool { 337 | /// Create a new Pool with default configuration. 338 | pub fn new() -> (Sender, Pool) { 339 | Builder::new().build() 340 | } 341 | 342 | /// Build a Pool with custom settings 343 | pub fn builder() -> Builder { 344 | Builder::new() 345 | } 346 | 347 | /// Start a core thread, causing it to idly wait for work. 348 | /// 349 | /// This overrides the default policy of starting core threads only when new 350 | /// tasks are executed. This function will return `false` if all core 351 | /// threads have already been started. 352 | pub fn prestart_core_thread(&self) -> bool { 353 | unimplemented!(); 354 | } 355 | 356 | /// Start all core threads, causing them to idly wait for work. 357 | /// 358 | /// This overrides the default policy of starting core threads only when new 359 | /// tasks are executed. 360 | pub fn prestart_core_threads(&self) { 361 | while self.prestart_core_thread() {} 362 | } 363 | 364 | /// Returns true if the pool is currently running. 365 | pub fn is_running(&self) -> bool { 366 | let state: State = self.inner.state.load(Relaxed).into(); 367 | state.is_running() 368 | } 369 | 370 | /// Shutdown the pool 371 | /// 372 | /// Perform an orderly shutdown. 373 | pub fn shutdown(&mut self) { 374 | self.inner.shutdown(false); 375 | } 376 | 377 | /// Shutdown the pool immediately 378 | /// 379 | /// All queued futures are dropped. 380 | pub fn force_shutdown(&mut self) { 381 | self.inner.shutdown(true); 382 | } 383 | } 384 | 385 | impl Future for Pool { 386 | type Item = (); 387 | type Error = (); 388 | 389 | fn poll(&mut self) -> Poll<(), ()> { 390 | trace!("Pool::poll"); 391 | 392 | // Implicitly shutdown... 393 | self.shutdown(); 394 | 395 | self.inner.shutdown_task.register(); 396 | 397 | if 0 != self.inner.num_workers.load(Acquire) { 398 | return Ok(Async::NotReady); 399 | } 400 | 401 | Ok(().into()) 402 | } 403 | } 404 | 405 | // ===== impl Sender ====== 406 | 407 | impl Sender { 408 | } 409 | 410 | impl Executor for Sender 411 | where T: Future + Send + 'static, 412 | { 413 | fn execute(&self, f: T) -> Result<(), ExecuteError> { 414 | use futures::future::ExecuteErrorKind::*; 415 | 416 | let mut state: State = self.inner.state.load(Acquire).into(); 417 | 418 | // Increment the number of futures spawned on the pool as well as 419 | // validate that the pool is still running/ 420 | loop { 421 | let mut next = state; 422 | 423 | if next.num_futures() == MAX_FUTURES { 424 | // No capacity 425 | return Err(ExecuteError::new(NoCapacity, f)); 426 | } 427 | 428 | if !next.is_running() { 429 | // Cannot execute the future 430 | return Err(ExecuteError::new(Shutdown, f)); 431 | } 432 | 433 | next.inc_num_futures(); 434 | 435 | let actual = self.inner.state.compare_and_swap( 436 | state.into(), next.into(), AcqRel).into(); 437 | 438 | if actual == state { 439 | trace!("execute; count={:?}", next.num_futures()); 440 | break; 441 | } 442 | 443 | state = actual; 444 | } 445 | 446 | // At this point, the pool has accepted the future, so schedule it for 447 | // execution. 448 | 449 | // Create a new task for the future 450 | let task = Task::new(f); 451 | 452 | // TODO: handle the error 453 | let _ = self.inner.submit(task, &self.inner); 454 | 455 | Ok(()) 456 | } 457 | } 458 | 459 | impl Clone for Sender { 460 | #[inline] 461 | fn clone(&self) -> Sender { 462 | let inner = self.inner.clone(); 463 | Sender { inner } 464 | } 465 | } 466 | 467 | // ===== impl Inner ===== 468 | 469 | macro_rules! worker_loop { 470 | ($probe_var:ident = $init:expr; $len:expr; $body:expr) => {{ 471 | let mut $probe_var = $init; 472 | let start = $probe_var; 473 | let len = $len; 474 | 475 | loop { 476 | if $probe_var < len { 477 | $body 478 | $probe_var += 1; 479 | } else { 480 | $probe_var = 0; 481 | } 482 | 483 | if $probe_var == start { 484 | break; 485 | } 486 | } 487 | }}; 488 | } 489 | 490 | impl Inner { 491 | /// Start shutting down the pool. This means that no new futures will be 492 | /// accepted. 493 | fn shutdown(&self, force: bool) { 494 | let mut state: State = self.state.load(Acquire).into(); 495 | 496 | trace!("shutdown; state={:?}", state); 497 | 498 | // Start by setting the SHUTDOWN flag 499 | loop { 500 | let mut next = state; 501 | 502 | if next.is_shutdown() { 503 | // Already transitioned to shutting down state 504 | return; 505 | } 506 | 507 | next.set_shutdown(); 508 | 509 | if force { 510 | next.clear_num_futures(); 511 | } 512 | 513 | let actual = self.state.compare_and_swap( 514 | state.into(), next.into(), AcqRel).into(); 515 | 516 | if state == actual { 517 | state = next; 518 | break; 519 | } 520 | 521 | state = actual; 522 | } 523 | 524 | trace!(" -> transitioned to shutdown"); 525 | 526 | // Only transition to terminate if there are no futures currently on the 527 | // pool 528 | if state.num_futures() != 0 { 529 | return; 530 | } 531 | 532 | self.terminate_sleeping_workers(); 533 | } 534 | 535 | fn terminate_sleeping_workers(&self) { 536 | trace!(" -> shutting down workers"); 537 | // Wakeup all sleeping workers. They will wake up, see the state 538 | // transition, and terminate. 539 | while let Some((idx, worker_state)) = self.pop_sleeper(WORKER_SIGNALED, TERMINATED) { 540 | trace!(" -> shutdown worker; idx={:?}; state={:?}", idx, worker_state); 541 | self.signal_stop(idx, worker_state); 542 | } 543 | } 544 | 545 | /// Signals to the worker that it should stop 546 | fn signal_stop(&self, idx: usize, mut state: WorkerState) { 547 | let worker = &self.workers[idx]; 548 | 549 | // Transition the worker state to signaled 550 | loop { 551 | let mut next = state; 552 | 553 | match state.lifecycle() { 554 | WORKER_SHUTDOWN => { 555 | trace!("signal_stop -- WORKER_SHUTDOWN; idx={}", idx); 556 | // If the worker is in the shutdown state, then it will never be 557 | // started again. 558 | self.worker_terminated(); 559 | 560 | return; 561 | } 562 | WORKER_RUNNING | WORKER_SLEEPING => {} 563 | _ => { 564 | trace!("signal_stop -- skipping; idx={}; state={:?}", idx, state); 565 | // All other states will naturally converge to a state of 566 | // shutdown. 567 | return; 568 | } 569 | } 570 | 571 | next.set_lifecycle(WORKER_SIGNALED); 572 | 573 | let actual = worker.state.compare_and_swap( 574 | state.into(), next.into(), AcqRel).into(); 575 | 576 | if actual == state { 577 | break; 578 | } 579 | 580 | state = actual; 581 | } 582 | 583 | // Wakeup the worker 584 | worker.wakeup(); 585 | } 586 | 587 | fn worker_terminated(&self) { 588 | let prev = self.num_workers.fetch_sub(1, AcqRel); 589 | 590 | trace!("worker_terminated; num_workers={}", prev - 1); 591 | 592 | if 1 == prev { 593 | trace!("notifying shutdown task"); 594 | self.shutdown_task.notify(); 595 | } 596 | } 597 | 598 | /// Submit a task to the scheduler. 599 | /// 600 | /// Called from either inside or outside of the scheduler. If currently on 601 | /// the scheduler, then a fast path is taken. 602 | fn submit(&self, task: Task, inner: &Arc) -> Result<(), Task> { 603 | Worker::with_current(|worker| { 604 | match worker { 605 | Some(worker) => { 606 | let idx = worker.idx; 607 | 608 | trace!(" -> submit internal; idx={}", idx); 609 | 610 | worker.inner.workers[idx].submit_internal(task); 611 | worker.inner.signal_work(inner); 612 | } 613 | None => { 614 | self.submit_external(task, inner); 615 | } 616 | } 617 | }); 618 | 619 | Ok(()) 620 | } 621 | 622 | /// Submit a task to the scheduler from off worker 623 | /// 624 | /// Called from outside of the scheduler, this function is how new tasks 625 | /// enter the system. 626 | fn submit_external(&self, task: Task, inner: &Arc) { 627 | // First try to get a handle to a sleeping worker. This ensures that 628 | // sleeping tasks get woken up 629 | if let Some((idx, state)) = self.pop_sleeper(WORKER_NOTIFIED, EMPTY) { 630 | trace!("submit to existing worker; idx={}; state={:?}", idx, state); 631 | self.submit_to_external(idx, task, state, inner); 632 | return; 633 | } 634 | 635 | // All workers are active, so pick a random worker and submit the 636 | // task to it. 637 | let len = self.workers.len(); 638 | let idx = self.rand_usize() % len; 639 | 640 | trace!(" -> submitting to random; idx={}", idx); 641 | 642 | let state: WorkerState = self.workers[idx].state.load(Acquire).into(); 643 | self.submit_to_external(idx, task, state, inner); 644 | } 645 | 646 | fn submit_to_external(&self, 647 | idx: usize, 648 | task: Task, 649 | state: WorkerState, 650 | inner: &Arc) 651 | { 652 | let entry = &self.workers[idx]; 653 | 654 | if !entry.submit_external(task, state) { 655 | Worker::spawn(idx, inner); 656 | } 657 | } 658 | 659 | /// If there are any other workers currently relaxing, signal them that work 660 | /// is available so that they can try to find more work to process. 661 | fn signal_work(&self, inner: &Arc) { 662 | if let Some((idx, mut state)) = self.pop_sleeper(WORKER_SIGNALED, EMPTY) { 663 | let entry = &self.workers[idx]; 664 | 665 | // Transition the worker state to signaled 666 | loop { 667 | let mut next = state; 668 | 669 | // pop_sleeper should skip these 670 | debug_assert!(state.lifecycle() != WORKER_SIGNALED); 671 | next.set_lifecycle(WORKER_SIGNALED); 672 | 673 | let actual = entry.state.compare_and_swap( 674 | state.into(), next.into(), AcqRel).into(); 675 | 676 | if actual == state { 677 | break; 678 | } 679 | 680 | state = actual; 681 | } 682 | 683 | // The state has been transitioned to signal, now we need to wake up 684 | // the worker if necessary. 685 | match state.lifecycle() { 686 | WORKER_SLEEPING => { 687 | trace!("signal_work -- wakeup; idx={}", idx); 688 | self.workers[idx].wakeup(); 689 | } 690 | WORKER_SHUTDOWN => { 691 | trace!("signal_work -- spawn; idx={}", idx); 692 | Worker::spawn(idx, inner); 693 | } 694 | _ => {} 695 | } 696 | } 697 | } 698 | 699 | /// Push a worker on the sleep stack 700 | /// 701 | /// Returns `Err` if the pool has been terminated 702 | fn push_sleeper(&self, idx: usize) -> Result<(), ()> { 703 | let mut state: SleepStack = self.sleep_stack.load(Acquire).into(); 704 | 705 | debug_assert!(WorkerState::from(self.workers[idx].state.load(Relaxed)).is_pushed()); 706 | 707 | loop { 708 | let mut next = state; 709 | 710 | let head = state.head(); 711 | 712 | if head == TERMINATED { 713 | // The pool is terminated, cannot push the sleeper. 714 | return Err(()); 715 | } 716 | 717 | self.workers[idx].set_next_sleeper(head); 718 | next.set_head(idx); 719 | 720 | let actual = self.sleep_stack.compare_and_swap( 721 | state.into(), next.into(), AcqRel).into(); 722 | 723 | if state == actual { 724 | return Ok(()); 725 | } 726 | 727 | state = actual; 728 | } 729 | } 730 | 731 | /// Pop a worker from the sleep stack 732 | fn pop_sleeper(&self, max_lifecycle: usize, terminal: usize) 733 | -> Option<(usize, WorkerState)> 734 | { 735 | debug_assert!(terminal == EMPTY || terminal == TERMINATED); 736 | 737 | let mut state: SleepStack = self.sleep_stack.load(Acquire).into(); 738 | 739 | loop { 740 | let head = state.head(); 741 | 742 | if head == EMPTY { 743 | let mut next = state; 744 | next.set_head(terminal); 745 | 746 | if next == state { 747 | debug_assert!(terminal == EMPTY); 748 | return None; 749 | } 750 | 751 | let actual = self.sleep_stack.compare_and_swap( 752 | state.into(), next.into(), AcqRel).into(); 753 | 754 | if actual != state { 755 | state = actual; 756 | continue; 757 | } 758 | 759 | return None; 760 | } else if head == TERMINATED { 761 | return None; 762 | } 763 | 764 | debug_assert!(head < MAX_WORKERS); 765 | 766 | let mut next = state; 767 | 768 | let next_head = self.workers[head].next_sleeper(); 769 | 770 | // TERMINATED can never be set as the "next pointer" on a worker. 771 | debug_assert!(next_head != TERMINATED); 772 | 773 | if next_head == EMPTY { 774 | next.set_head(terminal); 775 | } else { 776 | next.set_head(next_head); 777 | } 778 | 779 | let actual = self.sleep_stack.compare_and_swap( 780 | state.into(), next.into(), AcqRel).into(); 781 | 782 | if actual == state { 783 | // The worker has been removed from the stack, so the pushed bit 784 | // can be unset. Release ordering is used to ensure that this 785 | // operation happens after actually popping the task. 786 | debug_assert_eq!(1, PUSHED_MASK); 787 | 788 | // Unset the PUSHED flag and get the current state. 789 | let state: WorkerState = self.workers[head].state 790 | .fetch_sub(PUSHED_MASK, Release).into(); 791 | 792 | if state.lifecycle() >= max_lifecycle { 793 | // If the worker has already been notified, then it is 794 | // warming up to do more work. In this case, try to pop 795 | // another thread that might be in a relaxed state. 796 | continue; 797 | } 798 | 799 | return Some((head, state)); 800 | } 801 | 802 | state = actual; 803 | } 804 | } 805 | 806 | /// Generates a random number 807 | /// 808 | /// Uses a thread-local seeded XorShift. 809 | fn rand_usize(&self) -> usize { 810 | // Use a thread-local random number generator. If the thread does not 811 | // have one yet, then seed a new one 812 | thread_local!(static THREAD_RNG_KEY: UnsafeCell> = UnsafeCell::new(None)); 813 | 814 | THREAD_RNG_KEY.with(|t| { 815 | #[cfg(target_pointer_width = "32")] 816 | fn new_rng(thread_id: usize) -> XorShiftRng { 817 | XorShiftRng::from_seed([ 818 | thread_id as u32, 819 | 0x00000000, 820 | 0xa8a7d469, 821 | 0x97830e05]) 822 | } 823 | 824 | #[cfg(target_pointer_width = "64")] 825 | fn new_rng(thread_id: usize) -> XorShiftRng { 826 | XorShiftRng::from_seed([ 827 | thread_id as u32, 828 | (thread_id >> 32) as u32, 829 | 0xa8a7d469, 830 | 0x97830e05]) 831 | } 832 | 833 | let thread_id = self.next_thread_id.fetch_add(1, Relaxed); 834 | let rng = unsafe { &mut *t.get() }; 835 | 836 | if rng.is_none() { 837 | *rng = Some(new_rng(thread_id)); 838 | } 839 | 840 | rng.as_mut().unwrap().next_u32() as usize 841 | }) 842 | } 843 | } 844 | 845 | impl Notify for Notifier { 846 | fn notify(&self, id: usize) { 847 | trace!("Notifier::notify; id=0x{:x}", id); 848 | 849 | let id = id as usize; 850 | let task = unsafe { Task::from_notify_id_ref(&id) }; 851 | 852 | if !task.schedule() { 853 | trace!(" -> task already scheduled"); 854 | // task is already scheduled, there is nothing more to do 855 | return; 856 | } 857 | 858 | // TODO: Check if the pool is still running 859 | 860 | // Bump the ref count 861 | let task = task.clone(); 862 | 863 | if let Some(inner) = self.inner.upgrade() { 864 | let _ = inner.submit(task, &inner); 865 | } 866 | } 867 | 868 | fn clone_id(&self, id: usize) -> usize { 869 | unsafe { 870 | let handle = Task::from_notify_id_ref(&id); 871 | mem::forget(handle.clone()); 872 | } 873 | 874 | id 875 | } 876 | 877 | fn drop_id(&self, id: usize) { 878 | unsafe { 879 | let _ = Task::from_notify_id(id); 880 | } 881 | } 882 | } 883 | 884 | unsafe impl Send for Inner {} 885 | unsafe impl Sync for Inner {} 886 | 887 | // ===== impl Worker ===== 888 | 889 | impl Worker { 890 | fn spawn(idx: usize, inner: &Arc) { 891 | trace!("spawning new worker thread; idx={}", idx); 892 | 893 | let worker = Worker { 894 | inner: inner.clone(), 895 | idx: idx, 896 | }; 897 | 898 | let mut th = thread::Builder::new(); 899 | 900 | if let Some(ref prefix) = inner.config.name_prefix { 901 | th = th.name(format!("{}{}", prefix, idx)); 902 | } 903 | 904 | if let Some(stack) = inner.config.stack_size { 905 | th = th.stack_size(stack); 906 | } 907 | 908 | th.spawn(move || { 909 | worker.run(); 910 | }).unwrap(); 911 | } 912 | 913 | fn with_current) -> R, R>(f: F) -> R { 914 | CURRENT_WORKER.with(move |c| { 915 | let ptr = c.get(); 916 | 917 | if ptr.is_null() { 918 | f(None) 919 | } else { 920 | f(Some(unsafe { &*ptr })) 921 | } 922 | }) 923 | } 924 | 925 | fn run(&self) { 926 | CURRENT_WORKER.with(move |c| { 927 | c.set(self as *const _); 928 | self.run2(); 929 | }); 930 | } 931 | 932 | fn run2(&self) { 933 | if let Some(ref f) = self.inner.config.after_start { 934 | f.call(); 935 | } 936 | 937 | // Get the notifier. 938 | let notify = Arc::new(Notifier { 939 | inner: Arc::downgrade(&self.inner), 940 | }); 941 | 942 | let mut first = true; 943 | let mut spin_cnt = 0; 944 | 945 | while self.check_run_state(first) { 946 | first = false; 947 | 948 | // Poll inbound until empty, transfering all tasks to the internal 949 | // queue. 950 | let consistent = self.drain_inbound(); 951 | 952 | // Run the next available task 953 | if self.try_run_task(¬ify) { 954 | spin_cnt = 0; 955 | // As long as there is work, keep looping. 956 | continue; 957 | } 958 | 959 | // No work in this worker's queue, it is time to try stealing. 960 | if self.try_steal_task(¬ify) { 961 | spin_cnt = 0; 962 | continue; 963 | } 964 | 965 | if !consistent { 966 | spin_cnt = 0; 967 | continue; 968 | } 969 | 970 | // Starting to get sleeeeepy 971 | if spin_cnt < 32 { 972 | spin_cnt += 1; 973 | 974 | // Don't do anything further 975 | } else if spin_cnt < 256 { 976 | spin_cnt += 1; 977 | 978 | // Yield the thread 979 | thread::yield_now(); 980 | } else { 981 | if !self.sleep() { 982 | // The sleep expired and now the worker should shutdown due 983 | // to being idle. 984 | if let Some(ref f) = self.inner.config.before_stop { 985 | f.call(); 986 | } 987 | 988 | return; 989 | } 990 | } 991 | 992 | // If there still isn't any work to do, shutdown the worker? 993 | } 994 | 995 | trace!("shutting down thread; idx={}", self.idx); 996 | 997 | // Drain all work 998 | self.drain_inbound(); 999 | 1000 | while let Some(_) = self.entry().deque.pop() { 1001 | } 1002 | 1003 | if let Some(ref f) = self.inner.config.before_stop { 1004 | f.call(); 1005 | } 1006 | 1007 | // TODO: Drain the work queue... 1008 | self.inner.worker_terminated(); 1009 | } 1010 | 1011 | /// Checks the worker's current state, updating it as needed. 1012 | /// 1013 | /// Returns `true` if the worker should run. 1014 | #[inline] 1015 | fn check_run_state(&self, first: bool) -> bool { 1016 | let mut state: WorkerState = self.entry().state.load(Acquire).into(); 1017 | 1018 | loop { 1019 | let pool_state: State = self.inner.state.load(Acquire).into(); 1020 | 1021 | if !pool_state.is_running() && pool_state.num_futures() == 0 { 1022 | return false; 1023 | } 1024 | 1025 | let mut next = state; 1026 | 1027 | match state.lifecycle() { 1028 | WORKER_RUNNING => break, 1029 | WORKER_NOTIFIED | WORKER_SIGNALED => { 1030 | // transition back to running 1031 | next.set_lifecycle(WORKER_RUNNING); 1032 | } 1033 | lifecycle => panic!("unexpected worker state; lifecycle={}", lifecycle), 1034 | } 1035 | 1036 | let actual = self.entry().state.compare_and_swap( 1037 | state.into(), next.into(), AcqRel).into(); 1038 | 1039 | if actual == state { 1040 | break; 1041 | } 1042 | 1043 | state = actual; 1044 | } 1045 | 1046 | // If this is the first iteration of the worker loop, then the state can 1047 | // be signaled. 1048 | if !first && state.is_signaled() { 1049 | trace!("Worker::check_run_state; delegate signal"); 1050 | // This worker is not ready to be signaled, so delegate the signal 1051 | // to another worker. 1052 | self.inner.signal_work(&self.inner); 1053 | } 1054 | 1055 | true 1056 | } 1057 | 1058 | /// Runs the next task on this worker's queue. 1059 | /// 1060 | /// Returns `true` if work was found. 1061 | #[inline] 1062 | fn try_run_task(&self, notify: &Arc) -> bool { 1063 | use coco::deque::Steal::*; 1064 | 1065 | // Poll the internal queue for a task to run 1066 | match self.entry().deque.steal_weak() { 1067 | Data(task) => { 1068 | self.run_task(task, notify); 1069 | true 1070 | } 1071 | Empty => false, 1072 | Inconsistent => true, 1073 | } 1074 | } 1075 | 1076 | /// Tries to steal a task from another worker. 1077 | /// 1078 | /// Returns `true` if work was found 1079 | #[inline] 1080 | fn try_steal_task(&self, notify: &Arc) -> bool { 1081 | use coco::deque::Steal::*; 1082 | 1083 | let len = self.inner.workers.len(); 1084 | 1085 | let mut found_work = false; 1086 | 1087 | worker_loop!(idx = len; len; { 1088 | match self.inner.workers[idx].steal.steal_weak() { 1089 | Data(task) => { 1090 | trace!("stole task"); 1091 | 1092 | self.run_task(task, notify); 1093 | 1094 | trace!("try_steal_task -- signal_work; self={}; from={}", 1095 | self.idx, idx); 1096 | 1097 | // Signal other workers that work is available 1098 | self.inner.signal_work(&self.inner); 1099 | 1100 | return true; 1101 | } 1102 | Empty => {} 1103 | Inconsistent => found_work = true, 1104 | } 1105 | }); 1106 | 1107 | found_work 1108 | } 1109 | 1110 | fn run_task(&self, task: Task, notify: &Arc) { 1111 | use task::Run::*; 1112 | 1113 | match task.run(notify) { 1114 | Idle => {} 1115 | Schedule => { 1116 | self.entry().push_internal(task); 1117 | } 1118 | Complete => { 1119 | let mut state: State = self.inner.state.load(Acquire).into(); 1120 | 1121 | loop { 1122 | let mut next = state; 1123 | next.dec_num_futures(); 1124 | 1125 | let actual = self.inner.state.compare_and_swap( 1126 | state.into(), next.into(), AcqRel).into(); 1127 | 1128 | if actual == state { 1129 | trace!("task complete; state={:?}", next); 1130 | 1131 | if !next.is_shutdown() { 1132 | return; 1133 | } 1134 | 1135 | if state.num_futures() == 1 && next.num_futures() == 0 { 1136 | self.inner.terminate_sleeping_workers(); 1137 | } 1138 | 1139 | // The worker's run loop will detect the shutdown state 1140 | // next iteration. 1141 | return; 1142 | } 1143 | 1144 | state = actual; 1145 | } 1146 | } 1147 | } 1148 | } 1149 | 1150 | /// Drains all tasks on the extern queue and pushes them onto the internal 1151 | /// queue. 1152 | /// 1153 | /// Returns `true` if the operation was able to complete in a consistent 1154 | /// state. 1155 | #[inline] 1156 | fn drain_inbound(&self) -> bool { 1157 | use task::Poll::*; 1158 | 1159 | let mut found_work = false; 1160 | 1161 | loop { 1162 | let task = unsafe { self.entry().inbound.poll() }; 1163 | 1164 | match task { 1165 | Empty => { 1166 | if found_work { 1167 | trace!("found work while draining; signal_work"); 1168 | self.inner.signal_work(&self.inner); 1169 | } 1170 | 1171 | return true; 1172 | } 1173 | Inconsistent => { 1174 | if found_work { 1175 | trace!("found work while draining; signal_work"); 1176 | self.inner.signal_work(&self.inner); 1177 | } 1178 | 1179 | return false; 1180 | } 1181 | Data(task) => { 1182 | found_work = true; 1183 | self.entry().push_internal(task); 1184 | } 1185 | } 1186 | } 1187 | } 1188 | 1189 | /// Put the worker to sleep 1190 | /// 1191 | /// Returns `true` if woken up due to new work arriving. 1192 | #[inline] 1193 | fn sleep(&self) -> bool { 1194 | trace!("Worker::sleep; idx={}", self.idx); 1195 | 1196 | let mut state: WorkerState = self.entry().state.load(Acquire).into(); 1197 | 1198 | // The first part of the sleep process is to transition the worker state 1199 | // to "pushed". Now, it may be that the worker is already pushed on the 1200 | // sleeper stack, in which case, we don't push again. However, part of 1201 | // this process is also to do some final state checks to avoid entering 1202 | // the mutex if at all possible. 1203 | 1204 | loop { 1205 | let mut next = state; 1206 | 1207 | match state.lifecycle() { 1208 | WORKER_RUNNING => { 1209 | // Try setting the pushed state 1210 | next.set_pushed(); 1211 | } 1212 | WORKER_NOTIFIED | WORKER_SIGNALED => { 1213 | // No need to sleep, transition back to running and move on. 1214 | next.set_lifecycle(WORKER_RUNNING); 1215 | } 1216 | actual => panic!("unexpected worker state; {}", actual), 1217 | } 1218 | 1219 | let actual = self.entry().state.compare_and_swap( 1220 | state.into(), next.into(), AcqRel).into(); 1221 | 1222 | if actual == state { 1223 | if state.is_notified() { 1224 | // The previous state was notified, so we don't need to 1225 | // sleep. 1226 | return true; 1227 | } 1228 | 1229 | if !state.is_pushed() { 1230 | debug_assert!(next.is_pushed()); 1231 | 1232 | trace!(" sleeping -- push to stack; idx={}", self.idx); 1233 | 1234 | // We obtained permission to push the worker into the 1235 | // sleeper queue. 1236 | if let Err(_) = self.inner.push_sleeper(self.idx) { 1237 | trace!(" sleeping -- push to stack failed; idx={}", self.idx); 1238 | // The push failed due to the pool being terminated. 1239 | // 1240 | // This is true because the "work" being woken up for is 1241 | // shutting down. 1242 | return true; 1243 | } 1244 | } 1245 | 1246 | break; 1247 | } 1248 | 1249 | state = actual; 1250 | } 1251 | 1252 | // Acquire the sleep mutex, the state is transitioned to sleeping within 1253 | // the mutex in order to avoid losing wakeup notifications. 1254 | let mut lock = self.entry().park_mutex.lock().unwrap(); 1255 | 1256 | // Transition the state to sleeping, a CAS is still needed as other 1257 | // state transitions could happen unrelated to the sleep / wakeup 1258 | // process. We also have to redo the lifecycle check done above as 1259 | // the state could have been transitioned before entering the mutex. 1260 | loop { 1261 | let mut next = state; 1262 | 1263 | match state.lifecycle() { 1264 | WORKER_RUNNING => {} 1265 | WORKER_NOTIFIED | WORKER_SIGNALED => { 1266 | // Release the lock, sleep will not happen this call. 1267 | drop(lock); 1268 | 1269 | // Transition back to running 1270 | loop { 1271 | let mut next = state; 1272 | next.set_lifecycle(WORKER_RUNNING); 1273 | 1274 | let actual = self.entry().state.compare_and_swap( 1275 | state.into(), next.into(), AcqRel).into(); 1276 | 1277 | if actual == state { 1278 | return true; 1279 | } 1280 | 1281 | state = actual; 1282 | } 1283 | } 1284 | _ => unreachable!(), 1285 | } 1286 | 1287 | trace!(" sleeping -- set WORKER_SLEEPING; idx={}", self.idx); 1288 | 1289 | next.set_lifecycle(WORKER_SLEEPING); 1290 | 1291 | let actual = self.entry().state.compare_and_swap( 1292 | state.into(), next.into(), AcqRel).into(); 1293 | 1294 | if actual == state { 1295 | break; 1296 | } 1297 | 1298 | state = actual; 1299 | } 1300 | 1301 | trace!(" -> starting to sleep; idx={}", self.idx); 1302 | 1303 | let sleep_until = self.inner.config.keep_alive 1304 | .map(|dur| Instant::now() + dur); 1305 | 1306 | // The state has been transitioned to sleeping, we can now wait on the 1307 | // condvar. This is done in a loop as condvars can wakeup spuriously. 1308 | loop { 1309 | let mut drop_thread = false; 1310 | 1311 | lock = match sleep_until { 1312 | Some(when) => { 1313 | let now = Instant::now(); 1314 | 1315 | if when >= now { 1316 | drop_thread = true; 1317 | } 1318 | 1319 | let dur = when - now; 1320 | 1321 | self.entry().park_condvar 1322 | .wait_timeout(lock, dur) 1323 | .unwrap().0 1324 | } 1325 | None => { 1326 | self.entry().park_condvar.wait(lock).unwrap() 1327 | } 1328 | }; 1329 | 1330 | trace!(" -> wakeup; idx={}", self.idx); 1331 | 1332 | // Reload the state 1333 | state = self.entry().state.load(Acquire).into(); 1334 | 1335 | loop { 1336 | match state.lifecycle() { 1337 | WORKER_SLEEPING => {} 1338 | WORKER_NOTIFIED | WORKER_SIGNALED => { 1339 | // Release the lock, done sleeping 1340 | drop(lock); 1341 | 1342 | // Transition back to running 1343 | loop { 1344 | let mut next = state; 1345 | next.set_lifecycle(WORKER_RUNNING); 1346 | 1347 | let actual = self.entry().state.compare_and_swap( 1348 | state.into(), next.into(), AcqRel).into(); 1349 | 1350 | if actual == state { 1351 | return true; 1352 | } 1353 | 1354 | state = actual; 1355 | } 1356 | } 1357 | _ => unreachable!(), 1358 | } 1359 | 1360 | if !drop_thread { 1361 | break; 1362 | } 1363 | 1364 | let mut next = state; 1365 | next.set_lifecycle(WORKER_SHUTDOWN); 1366 | 1367 | let actual = self.entry().state.compare_and_swap( 1368 | state.into(), next.into(), AcqRel).into(); 1369 | 1370 | if actual == state { 1371 | // Transitioned to a shutdown state 1372 | return false; 1373 | } 1374 | 1375 | state = actual; 1376 | } 1377 | 1378 | // The worker hasn't been notified, go back to sleep 1379 | } 1380 | } 1381 | 1382 | fn entry(&self) -> &WorkerEntry { 1383 | &self.inner.workers[self.idx] 1384 | } 1385 | } 1386 | 1387 | // ===== impl State ===== 1388 | 1389 | impl State { 1390 | #[inline] 1391 | fn new() -> State { 1392 | State(0) 1393 | } 1394 | 1395 | /// Returns the number of futures still pending completion. 1396 | fn num_futures(&self) -> usize { 1397 | self.0 >> 1 1398 | } 1399 | 1400 | /// Increment the number of futures pending completion. 1401 | /// 1402 | /// Returns false on failure. 1403 | fn inc_num_futures(&mut self) { 1404 | debug_assert!(self.num_futures() < MAX_FUTURES); 1405 | debug_assert!(self.is_running()); 1406 | 1407 | self.0 += 2; 1408 | } 1409 | 1410 | /// Decrement the number of futures pending completion. 1411 | fn dec_num_futures(&mut self) { 1412 | if self.0 & !SHUTDOWN == 0 { 1413 | // Already zero 1414 | debug_assert_eq!(self.num_futures(), 0); 1415 | return; 1416 | } 1417 | 1418 | self.0 -= 2; 1419 | } 1420 | 1421 | /// Set the number of futures pending completion to zero 1422 | fn clear_num_futures(&mut self) { 1423 | self.0 = self.0 & SHUTDOWN 1424 | } 1425 | 1426 | /// Returns true if the shutdown flag is set 1427 | fn is_shutdown(&self) -> bool { 1428 | self.0 & SHUTDOWN == SHUTDOWN 1429 | } 1430 | 1431 | /// Returns true if not shutdown 1432 | fn is_running(&self) -> bool { 1433 | !self.is_shutdown() 1434 | } 1435 | 1436 | /// Set the shutdown flag 1437 | fn set_shutdown(&mut self) { 1438 | debug_assert!(!self.is_shutdown()); 1439 | self.0 |= SHUTDOWN; 1440 | } 1441 | } 1442 | 1443 | impl From for State { 1444 | fn from(src: usize) -> Self { 1445 | State(src) 1446 | } 1447 | } 1448 | 1449 | impl From for usize { 1450 | fn from(src: State) -> Self { 1451 | src.0 1452 | } 1453 | } 1454 | 1455 | impl fmt::Debug for State { 1456 | fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { 1457 | fmt.debug_struct("State") 1458 | .field("is_shutdown", &self.is_shutdown()) 1459 | .field("num_futures", &self.num_futures()) 1460 | .finish() 1461 | } 1462 | } 1463 | 1464 | // ===== impl SleepStack ===== 1465 | 1466 | impl SleepStack { 1467 | #[inline] 1468 | fn new() -> SleepStack { 1469 | SleepStack(EMPTY) 1470 | } 1471 | 1472 | #[inline] 1473 | fn head(&self) -> usize { 1474 | self.0 & STACK_MASK 1475 | } 1476 | 1477 | #[inline] 1478 | fn set_head(&mut self, val: usize) { 1479 | // The ABA guard protects against the ABA problem w/ treiber stacks 1480 | let aba_guard = ((self.0 >> ABA_GUARD_SHIFT) + 1) & ABA_GUARD_MASK; 1481 | 1482 | self.0 = (aba_guard << ABA_GUARD_SHIFT) | val; 1483 | } 1484 | } 1485 | 1486 | impl From for SleepStack { 1487 | fn from(src: usize) -> Self { 1488 | SleepStack(src) 1489 | } 1490 | } 1491 | 1492 | impl From for usize { 1493 | fn from(src: SleepStack) -> Self { 1494 | src.0 1495 | } 1496 | } 1497 | 1498 | impl fmt::Debug for SleepStack { 1499 | fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { 1500 | let head = self.head(); 1501 | 1502 | let mut fmt = fmt.debug_struct("SleepStack"); 1503 | 1504 | if head < MAX_WORKERS { 1505 | fmt.field("head", &head); 1506 | } else if head == EMPTY { 1507 | fmt.field("head", &"EMPTY"); 1508 | } else if head == TERMINATED { 1509 | fmt.field("head", &"TERMINATED"); 1510 | } 1511 | 1512 | fmt.finish() 1513 | } 1514 | } 1515 | 1516 | // ===== impl WorkerEntry ===== 1517 | 1518 | impl WorkerEntry { 1519 | fn new() -> Self { 1520 | let (w, s) = deque::new(); 1521 | 1522 | WorkerEntry { 1523 | state: AtomicUsize::new(WorkerState::default().into()), 1524 | next_sleeper: UnsafeCell::new(0), 1525 | deque: w, 1526 | steal: s, 1527 | inbound: task::Queue::new(), 1528 | park_mutex: Mutex::new(()), 1529 | park_condvar: Condvar::new(), 1530 | } 1531 | } 1532 | 1533 | #[inline] 1534 | fn submit_internal(&self, task: Task) { 1535 | self.push_internal(task); 1536 | } 1537 | 1538 | /// Submits a task to the worker. This assumes that the caller is external 1539 | /// to the worker. Internal submissions go through another path. 1540 | /// 1541 | /// Returns `false` if the worker needs to be spawned. 1542 | fn submit_external(&self, task: Task, mut state: WorkerState) -> bool { 1543 | // Push the task onto the external queue 1544 | self.push_external(task); 1545 | 1546 | loop { 1547 | let mut next = state; 1548 | next.notify(); 1549 | 1550 | let actual = self.state.compare_and_swap( 1551 | state.into(), next.into(), 1552 | AcqRel).into(); 1553 | 1554 | if state == actual { 1555 | break; 1556 | } 1557 | 1558 | state = actual; 1559 | } 1560 | 1561 | match state.lifecycle() { 1562 | WORKER_SLEEPING => { 1563 | // The worker is currently sleeping, the condition variable must 1564 | // be signaled 1565 | self.wakeup(); 1566 | true 1567 | } 1568 | WORKER_SHUTDOWN => false, 1569 | _ => true, 1570 | } 1571 | } 1572 | 1573 | #[inline] 1574 | fn push_external(&self, task: Task) { 1575 | self.inbound.push(task); 1576 | } 1577 | 1578 | #[inline] 1579 | fn push_internal(&self, task: Task) { 1580 | self.deque.push(task); 1581 | } 1582 | 1583 | #[inline] 1584 | fn wakeup(&self) { 1585 | let _lock = self.park_mutex.lock().unwrap(); 1586 | self.park_condvar.notify_one(); 1587 | } 1588 | 1589 | #[inline] 1590 | fn next_sleeper(&self) -> usize { 1591 | unsafe { *self.next_sleeper.get() } 1592 | } 1593 | 1594 | #[inline] 1595 | fn set_next_sleeper(&self, val: usize) { 1596 | unsafe { *self.next_sleeper.get() = val; } 1597 | } 1598 | } 1599 | 1600 | // ===== impl WorkerState ===== 1601 | 1602 | impl WorkerState { 1603 | /// Returns true if the worker entry is pushed in the sleeper stack 1604 | fn is_pushed(&self) -> bool { 1605 | self.0 & PUSHED_MASK == PUSHED_MASK 1606 | } 1607 | 1608 | fn set_pushed(&mut self) { 1609 | self.0 |= PUSHED_MASK 1610 | } 1611 | 1612 | fn is_notified(&self) -> bool { 1613 | match self.lifecycle() { 1614 | WORKER_NOTIFIED | WORKER_SIGNALED => true, 1615 | _ => false, 1616 | } 1617 | } 1618 | 1619 | fn lifecycle(&self) -> usize { 1620 | (self.0 & WORKER_LIFECYCLE_MASK) >> WORKER_LIFECYCLE_SHIFT 1621 | } 1622 | 1623 | fn set_lifecycle(&mut self, val: usize) { 1624 | self.0 = (self.0 & !WORKER_LIFECYCLE_MASK) | 1625 | (val << WORKER_LIFECYCLE_SHIFT) 1626 | } 1627 | 1628 | fn is_signaled(&self) -> bool { 1629 | self.lifecycle() == WORKER_SIGNALED 1630 | } 1631 | 1632 | fn notify(&mut self) { 1633 | if self.lifecycle() != WORKER_SIGNALED { 1634 | self.set_lifecycle(WORKER_NOTIFIED) 1635 | } 1636 | } 1637 | } 1638 | 1639 | impl Default for WorkerState { 1640 | fn default() -> WorkerState { 1641 | // All workers will start pushed in the sleeping stack 1642 | WorkerState(PUSHED_MASK) 1643 | } 1644 | } 1645 | 1646 | impl From for WorkerState { 1647 | fn from(src: usize) -> Self { 1648 | WorkerState(src) 1649 | } 1650 | } 1651 | 1652 | impl From for usize { 1653 | fn from(src: WorkerState) -> Self { 1654 | src.0 1655 | } 1656 | } 1657 | 1658 | impl fmt::Debug for WorkerState { 1659 | fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { 1660 | fmt.debug_struct("WorkerState") 1661 | .field("lifecycle", &match self.lifecycle() { 1662 | WORKER_SHUTDOWN => "WORKER_SHUTDOWN", 1663 | WORKER_RUNNING => "WORKER_RUNNING", 1664 | WORKER_SLEEPING => "WORKER_SLEEPING", 1665 | WORKER_NOTIFIED => "WORKER_NOTIFIED", 1666 | WORKER_SIGNALED => "WORKER_SIGNALED", 1667 | _ => unreachable!(), 1668 | }) 1669 | .field("is_pushed", &self.is_pushed()) 1670 | .finish() 1671 | } 1672 | } 1673 | 1674 | // ===== impl Callback ===== 1675 | 1676 | impl Callback { 1677 | fn new(f: F) -> Self 1678 | where F: Fn() + Send + Sync + 'static 1679 | { 1680 | Callback { f: Arc::new(f) } 1681 | } 1682 | 1683 | pub fn call(&self) { 1684 | (self.f)() 1685 | } 1686 | } 1687 | 1688 | impl fmt::Debug for Callback { 1689 | fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { 1690 | write!(fmt, "Fn") 1691 | } 1692 | } 1693 | -------------------------------------------------------------------------------- /src/task.rs: -------------------------------------------------------------------------------- 1 | use Notifier; 2 | 3 | use futures::{future, Future, Async}; 4 | use futures::executor::{self, Spawn}; 5 | 6 | use std::{fmt, mem, ptr}; 7 | use std::cell::Cell; 8 | use std::sync::Arc; 9 | use std::sync::atomic::{self, AtomicUsize, AtomicPtr}; 10 | use std::sync::atomic::Ordering::{AcqRel, Acquire, Release, Relaxed}; 11 | 12 | pub(crate) struct Task { 13 | ptr: *mut Inner, 14 | } 15 | 16 | #[derive(Debug)] 17 | pub(crate) struct Queue { 18 | head: AtomicPtr, 19 | tail: Cell<*mut Inner>, 20 | stub: Box, 21 | } 22 | 23 | #[derive(Debug)] 24 | pub(crate) enum Poll { 25 | Empty, 26 | Inconsistent, 27 | Data(Task), 28 | } 29 | 30 | #[derive(Debug)] 31 | pub(crate) enum Run { 32 | Idle, 33 | Schedule, 34 | Complete, 35 | } 36 | 37 | struct Inner { 38 | // Next pointer in the queue that submits tasks to a worker. 39 | next: AtomicPtr, 40 | 41 | // Task state 42 | state: AtomicUsize, 43 | 44 | // Number of outstanding references to the task 45 | ref_count: AtomicUsize, 46 | 47 | // Store the future at the head of the struct 48 | // 49 | // The future is dropped immediately when it transitions to Complete 50 | future: Option>, 51 | } 52 | 53 | #[derive(Debug, Clone, Copy, Eq, PartialEq)] 54 | enum State { 55 | /// Task is currently idle 56 | Idle, 57 | /// Task is currently running 58 | Running, 59 | /// Task is currently running, but has been notified that it must run again. 60 | Notified, 61 | /// Task has been scheduled 62 | Scheduled, 63 | /// Task is complete 64 | Complete, 65 | } 66 | 67 | type BoxFuture = Box + Send + 'static>; 68 | 69 | 70 | 71 | // ===== impl Task ===== 72 | 73 | impl Task { 74 | /// Create a new task handle 75 | pub fn new + Send + 'static>(f: T) -> Task { 76 | let inner = Box::new(Inner { 77 | next: AtomicPtr::new(ptr::null_mut()), 78 | state: AtomicUsize::new(State::new().into()), 79 | ref_count: AtomicUsize::new(1), 80 | future: Some(executor::spawn(Box::new(f))), 81 | }); 82 | 83 | Task { ptr: Box::into_raw(inner) } 84 | } 85 | 86 | /// Transmute a u64 to a Task 87 | pub unsafe fn from_notify_id(unpark_id: usize) -> Task { 88 | mem::transmute(unpark_id) 89 | } 90 | 91 | /// Transmute a u64 to a task ref 92 | pub unsafe fn from_notify_id_ref<'a>(unpark_id: &'a usize) -> &'a Task { 93 | mem::transmute(unpark_id) 94 | } 95 | 96 | /// Execute the task returning `true` if the task needs to be scheduled again. 97 | pub fn run(&self, unpark: &Arc) -> Run { 98 | use self::State::*; 99 | 100 | // Transition task to running state. At this point, the task must be 101 | // scheduled. 102 | let actual: State = self.inner().state.compare_and_swap( 103 | Scheduled.into(), Running.into(), AcqRel).into(); 104 | 105 | trace!("running; state={:?}", actual); 106 | 107 | match actual { 108 | Scheduled => {}, 109 | _ => panic!("unexpected task state; {:?}", actual), 110 | } 111 | 112 | trace!("Task::run; state={:?}", State::from(self.inner().state.load(Relaxed))); 113 | 114 | let res = self.inner_mut().future.as_mut().unwrap() 115 | .poll_future_notify(unpark, self.ptr as usize); 116 | 117 | match res { 118 | Ok(Async::Ready(_)) | Err(_) => { 119 | trace!(" -> task complete"); 120 | 121 | // Drop the future 122 | self.inner_mut().drop_future(); 123 | 124 | // Transition to the completed state 125 | self.inner().state.store(State::Complete.into(), Release); 126 | 127 | Run::Complete 128 | } 129 | _ => { 130 | trace!(" -> not ready"); 131 | 132 | // Attempt to transition from Running -> Idle, if successful, 133 | // then the task does not need to be scheduled again. If the CAS 134 | // fails, then the task has been unparked concurrent to running, 135 | // in which case it transitions immediately back to scheduled 136 | // and we return `true`. 137 | let prev: State = self.inner().state.compare_and_swap( 138 | Running.into(), Idle.into(), AcqRel).into(); 139 | 140 | match prev { 141 | Running => Run::Idle, 142 | Notified => { 143 | self.inner().state.store(Scheduled.into(), Release); 144 | Run::Schedule 145 | } 146 | _ => unreachable!(), 147 | } 148 | } 149 | } 150 | } 151 | 152 | /// Transition the task state to scheduled. 153 | /// 154 | /// Returns `true` if the caller is permitted to schedule the task. 155 | pub fn schedule(&self) -> bool { 156 | use self::State::*; 157 | 158 | loop { 159 | let actual = self.inner().state.compare_and_swap( 160 | Idle.into(), 161 | Scheduled.into(), 162 | Relaxed).into(); 163 | 164 | match actual { 165 | Idle => return true, 166 | Running => { 167 | let actual = self.inner().state.compare_and_swap( 168 | Running.into(), Notified.into(), Relaxed).into(); 169 | 170 | match actual { 171 | Idle => continue, 172 | _ => return false, 173 | } 174 | } 175 | Complete | Notified | Scheduled => return false, 176 | } 177 | } 178 | } 179 | 180 | #[inline] 181 | fn inner(&self) -> &Inner { 182 | unsafe { &*self.ptr } 183 | } 184 | 185 | #[inline] 186 | fn inner_mut(&self) -> &mut Inner { 187 | unsafe { &mut *self.ptr } 188 | } 189 | } 190 | 191 | impl fmt::Debug for Task { 192 | fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { 193 | fmt.debug_struct("Task") 194 | .field("inner", self.inner()) 195 | .finish() 196 | } 197 | } 198 | 199 | impl Clone for Task { 200 | fn clone(&self) -> Task { 201 | use std::isize; 202 | 203 | const MAX_REFCOUNT: usize = (isize::MAX) as usize; 204 | // Using a relaxed ordering is alright here, as knowledge of the 205 | // original reference prevents other threads from erroneously deleting 206 | // the object. 207 | // 208 | // As explained in the [Boost documentation][1], Increasing the 209 | // reference counter can always be done with memory_order_relaxed: New 210 | // references to an object can only be formed from an existing 211 | // reference, and passing an existing reference from one thread to 212 | // another must already provide any required synchronization. 213 | // 214 | // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html) 215 | let old_size = self.inner().ref_count.fetch_add(1, Relaxed); 216 | 217 | // However we need to guard against massive refcounts in case someone 218 | // is `mem::forget`ing Arcs. If we don't do this the count can overflow 219 | // and users will use-after free. We racily saturate to `isize::MAX` on 220 | // the assumption that there aren't ~2 billion threads incrementing 221 | // the reference count at once. This branch will never be taken in 222 | // any realistic program. 223 | // 224 | // We abort because such a program is incredibly degenerate, and we 225 | // don't care to support it. 226 | if old_size > MAX_REFCOUNT { 227 | // TODO: abort 228 | panic!(); 229 | } 230 | 231 | Task { ptr: self.ptr } 232 | } 233 | } 234 | 235 | impl Drop for Task { 236 | fn drop(&mut self) { 237 | // Because `fetch_sub` is already atomic, we do not need to synchronize 238 | // with other threads unless we are going to delete the object. This 239 | // same logic applies to the below `fetch_sub` to the `weak` count. 240 | if self.inner().ref_count.fetch_sub(1, Release) != 1 { 241 | return; 242 | } 243 | 244 | // This fence is needed to prevent reordering of use of the data and 245 | // deletion of the data. Because it is marked `Release`, the decreasing 246 | // of the reference count synchronizes with this `Acquire` fence. This 247 | // means that use of the data happens before decreasing the reference 248 | // count, which happens before this fence, which happens before the 249 | // deletion of the data. 250 | // 251 | // As explained in the [Boost documentation][1], 252 | // 253 | // > It is important to enforce any possible access to the object in one 254 | // > thread (through an existing reference) to *happen before* deleting 255 | // > the object in a different thread. This is achieved by a "release" 256 | // > operation after dropping a reference (any access to the object 257 | // > through this reference must obviously happened before), and an 258 | // > "acquire" operation before deleting the object. 259 | // 260 | // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html) 261 | atomic::fence(Acquire); 262 | 263 | unsafe { 264 | let _ = Box::from_raw(self.ptr); 265 | } 266 | } 267 | } 268 | 269 | unsafe impl Send for Task {} 270 | 271 | // ===== impl Inner ===== 272 | 273 | impl Inner { 274 | fn stub() -> Inner { 275 | Inner { 276 | next: AtomicPtr::new(ptr::null_mut()), 277 | state: AtomicUsize::new(State::stub().into()), 278 | ref_count: AtomicUsize::new(0), 279 | future: Some(executor::spawn(Box::new(future::empty()))), 280 | } 281 | } 282 | 283 | fn drop_future(&mut self) { 284 | let _ = self.future.take(); 285 | } 286 | } 287 | 288 | impl Drop for Inner { 289 | fn drop(&mut self) { 290 | self.drop_future(); 291 | } 292 | } 293 | 294 | impl fmt::Debug for Inner { 295 | fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { 296 | fmt.debug_struct("Inner") 297 | .field("next", &self.next) 298 | .field("state", &self.state) 299 | .field("ref_count", &self.ref_count) 300 | .field("future", &"Spawn") 301 | .finish() 302 | } 303 | } 304 | 305 | // ===== impl Queue ===== 306 | 307 | impl Queue { 308 | pub fn new() -> Queue { 309 | let stub = Box::new(Inner::stub()); 310 | let ptr = &*stub as *const _ as *mut _; 311 | 312 | Queue { 313 | head: AtomicPtr::new(ptr), 314 | tail: Cell::new(ptr), 315 | stub: stub, 316 | } 317 | } 318 | 319 | pub fn push(&self, handle: Task) { 320 | unsafe { 321 | self.push2(handle.ptr); 322 | 323 | // Forgetting the handle is necessary to avoid the ref dec 324 | mem::forget(handle); 325 | } 326 | } 327 | 328 | unsafe fn push2(&self, handle: *mut Inner) { 329 | // Set the next pointer. This does not require an atomic operation as 330 | // this node is not accessible. The write will be flushed with the next 331 | // operation 332 | (*handle).next = AtomicPtr::new(ptr::null_mut()); 333 | 334 | // Update the head to point to the new node. We need to see the previous 335 | // node in order to update the next pointer as well as release `handle` 336 | // to any other threads calling `push`. 337 | let prev = self.head.swap(handle, AcqRel); 338 | 339 | // Release `handle` to the consume end. 340 | (*prev).next.store(handle, Release); 341 | } 342 | 343 | pub unsafe fn poll(&self) -> Poll { 344 | let mut tail = self.tail.get(); 345 | let mut next = (*tail).next.load(Acquire); 346 | let stub = &*self.stub as *const _ as *mut _; 347 | 348 | if tail == stub { 349 | if next.is_null() { 350 | return Poll::Empty; 351 | } 352 | 353 | self.tail.set(next); 354 | tail = next; 355 | next = (*next).next.load(Acquire); 356 | } 357 | 358 | if !next.is_null() { 359 | self.tail.set(next); 360 | 361 | // No ref_count inc is necessary here as this poll is paired 362 | // with a `push` which "forgets" the handle. 363 | return Poll::Data(Task { 364 | ptr: tail, 365 | }); 366 | } 367 | 368 | if self.head.load(Acquire) != tail { 369 | return Poll::Inconsistent; 370 | } 371 | 372 | self.push2(stub); 373 | 374 | next = (*tail).next.load(Acquire); 375 | 376 | if !next.is_null() { 377 | self.tail.set(next); 378 | return Poll::Data(Task { 379 | ptr: tail, 380 | }); 381 | } 382 | 383 | Poll::Inconsistent 384 | } 385 | } 386 | 387 | // ===== impl State ===== 388 | 389 | impl State { 390 | /// Returns the initial task state. 391 | /// 392 | /// Tasks start in the scheduled state as they are immediately scheduled on 393 | /// creation. 394 | fn new() -> State { 395 | State::Scheduled 396 | } 397 | 398 | fn stub() -> State { 399 | State::Idle 400 | } 401 | } 402 | 403 | impl From for State { 404 | fn from(src: usize) -> Self { 405 | use self::State::*; 406 | 407 | match src { 408 | 0 => Idle, 409 | 1 => Running, 410 | 2 => Notified, 411 | 3 => Scheduled, 412 | 4 => Complete, 413 | _ => unreachable!(), 414 | } 415 | } 416 | } 417 | 418 | impl From for usize { 419 | fn from(src: State) -> Self { 420 | use self::State::*; 421 | 422 | match src { 423 | Idle => 0, 424 | Running => 1, 425 | Notified => 2, 426 | Scheduled => 3, 427 | Complete => 4, 428 | } 429 | } 430 | } 431 | -------------------------------------------------------------------------------- /tests/pool.rs: -------------------------------------------------------------------------------- 1 | extern crate env_logger; 2 | extern crate futures; 3 | extern crate futures_pool; 4 | 5 | use futures::{Poll, Sink, Stream, Async}; 6 | use futures::future::{Future, Executor, lazy}; 7 | use futures_pool::*; 8 | 9 | use std::cell::Cell; 10 | use std::sync::{mpsc, Arc}; 11 | use std::sync::atomic::{AtomicUsize, ATOMIC_USIZE_INIT}; 12 | use std::sync::atomic::Ordering::Relaxed; 13 | use std::time::Duration; 14 | 15 | thread_local!(static FOO: Cell = Cell::new(0)); 16 | 17 | #[test] 18 | fn natural_shutdown_simple_futures() { 19 | let _ = ::env_logger::init(); 20 | 21 | for _ in 0..1_000 { 22 | static NUM_INC: AtomicUsize = ATOMIC_USIZE_INIT; 23 | static NUM_DEC: AtomicUsize = ATOMIC_USIZE_INIT; 24 | 25 | FOO.with(|f| { 26 | f.set(1); 27 | 28 | let (tx, pool) = Pool::builder() 29 | .after_start(|| { 30 | NUM_INC.fetch_add(1, Relaxed); 31 | }) 32 | .before_stop(|| { 33 | NUM_DEC.fetch_add(1, Relaxed); 34 | }) 35 | .build(); 36 | 37 | let a = { 38 | let (t, rx) = mpsc::channel(); 39 | tx.execute(lazy(move || { 40 | // Makes sure this runs on a worker thread 41 | FOO.with(|f| assert_eq!(f.get(), 0)); 42 | 43 | t.send("one").unwrap(); 44 | Ok(()) 45 | })).unwrap(); 46 | rx 47 | }; 48 | 49 | let b = { 50 | let (t, rx) = mpsc::channel(); 51 | tx.execute(lazy(move || { 52 | // Makes sure this runs on a worker thread 53 | FOO.with(|f| assert_eq!(f.get(), 0)); 54 | 55 | t.send("two").unwrap(); 56 | Ok(()) 57 | })).unwrap(); 58 | rx 59 | }; 60 | 61 | drop(tx); 62 | 63 | assert_eq!("one", a.recv().unwrap()); 64 | assert_eq!("two", b.recv().unwrap()); 65 | 66 | // Wait for the pool to shutdown 67 | pool.wait().unwrap(); 68 | 69 | // Assert that at least one thread started 70 | let num_inc = NUM_INC.load(Relaxed); 71 | assert!(num_inc > 0); 72 | 73 | // Assert that all threads shutdown 74 | let num_dec = NUM_DEC.load(Relaxed); 75 | assert_eq!(num_inc, num_dec); 76 | }); 77 | } 78 | } 79 | 80 | #[test] 81 | fn force_shutdown_drops_futures() { 82 | let _ = ::env_logger::init(); 83 | 84 | for _ in 0..1_000 { 85 | let num_inc = Arc::new(AtomicUsize::new(0)); 86 | let num_dec = Arc::new(AtomicUsize::new(0)); 87 | let num_drop = Arc::new(AtomicUsize::new(0)); 88 | 89 | struct Never(Arc); 90 | 91 | impl Future for Never { 92 | type Item = (); 93 | type Error = (); 94 | 95 | fn poll(&mut self) -> Poll<(), ()> { 96 | Ok(Async::NotReady) 97 | } 98 | } 99 | 100 | impl Drop for Never { 101 | fn drop(&mut self) { 102 | self.0.fetch_add(1, Relaxed); 103 | } 104 | } 105 | 106 | let a = num_inc.clone(); 107 | let b = num_dec.clone(); 108 | 109 | let (tx, mut pool) = Pool::builder() 110 | .after_start(move || { 111 | a.fetch_add(1, Relaxed); 112 | }) 113 | .before_stop(move || { 114 | b.fetch_add(1, Relaxed); 115 | }) 116 | .build(); 117 | 118 | tx.execute(Never(num_drop.clone())).unwrap(); 119 | 120 | pool.force_shutdown(); 121 | 122 | // Wait for the pool to shutdown 123 | pool.wait().unwrap(); 124 | 125 | // Assert that only a single thread was spawned. 126 | let a = num_inc.load(Relaxed); 127 | assert!(a >= 1); 128 | 129 | // Assert that all threads shutdown 130 | let b = num_dec.load(Relaxed); 131 | assert_eq!(a, b); 132 | 133 | // Assert that the future was dropped 134 | let c = num_drop.load(Relaxed); 135 | assert_eq!(c, 1); 136 | } 137 | } 138 | 139 | #[test] 140 | fn thread_shutdown_timeout() { 141 | use std::sync::Mutex; 142 | 143 | let _ = ::env_logger::init(); 144 | 145 | let (shutdown_tx, shutdown_rx) = mpsc::channel(); 146 | let (complete_tx, complete_rx) = mpsc::channel(); 147 | 148 | let t = Mutex::new(shutdown_tx); 149 | 150 | let (tx, pool) = Pool::builder() 151 | .keep_alive(Duration::from_millis(200)) 152 | .before_stop(move || { 153 | // There could be multiple threads here 154 | let _ = t.lock().unwrap().send(()); 155 | }) 156 | .build(); 157 | 158 | let t = complete_tx.clone(); 159 | tx.execute(lazy(move || { 160 | t.send(()).unwrap(); 161 | Ok(()) 162 | })).unwrap(); 163 | 164 | // The future completes 165 | complete_rx.recv().unwrap(); 166 | 167 | // The thread shuts down eventually 168 | shutdown_rx.recv().unwrap(); 169 | 170 | // Futures can still be run 171 | tx.execute(lazy(move || { 172 | complete_tx.send(()).unwrap(); 173 | Ok(()) 174 | })).unwrap(); 175 | 176 | complete_rx.recv().unwrap(); 177 | 178 | drop(pool); 179 | } 180 | 181 | #[test] 182 | fn many_oneshot_futures() { 183 | const NUM: usize = 10_000; 184 | 185 | let _ = ::env_logger::init(); 186 | 187 | for _ in 0..50 { 188 | let (tx, pool) = Pool::new(); 189 | let cnt = Arc::new(AtomicUsize::new(0)); 190 | 191 | for _ in 0..NUM { 192 | let cnt = cnt.clone(); 193 | tx.execute(lazy(move || { 194 | cnt.fetch_add(1, Relaxed); 195 | Ok(()) 196 | })).unwrap(); 197 | } 198 | 199 | // Wait for the pool to shutdown 200 | pool.wait().unwrap(); 201 | 202 | let num = cnt.load(Relaxed); 203 | assert_eq!(num, NUM); 204 | } 205 | } 206 | 207 | #[test] 208 | fn many_multishot_futures() { 209 | use futures::sync::mpsc; 210 | 211 | const CHAIN: usize = 200; 212 | const CYCLES: usize = 5; 213 | const TRACKS: usize = 50; 214 | 215 | let _ = ::env_logger::init(); 216 | 217 | for _ in 0..50 { 218 | let (pool_tx, pool) = Pool::new(); 219 | 220 | let mut start_txs = Vec::with_capacity(TRACKS); 221 | let mut final_rxs = Vec::with_capacity(TRACKS); 222 | 223 | for _ in 0..TRACKS { 224 | let (start_tx, mut chain_rx) = mpsc::channel(10); 225 | 226 | for _ in 0..CHAIN { 227 | let (next_tx, next_rx) = mpsc::channel(10); 228 | 229 | let rx = chain_rx 230 | .map_err(|e| panic!("{:?}", e)); 231 | 232 | // Forward all the messages 233 | pool_tx.execute(next_tx 234 | .send_all(rx) 235 | .map(|_| ()) 236 | .map_err(|e| panic!("{:?}", e)) 237 | ).unwrap(); 238 | 239 | chain_rx = next_rx; 240 | } 241 | 242 | // This final task cycles if needed 243 | let (final_tx, final_rx) = mpsc::channel(10); 244 | let cycle_tx = start_tx.clone(); 245 | let mut rem = CYCLES; 246 | 247 | pool_tx.execute(chain_rx.take(CYCLES as u64).for_each(move |msg| { 248 | rem -= 1; 249 | let send = if rem == 0 { 250 | final_tx.clone().send(msg) 251 | } else { 252 | cycle_tx.clone().send(msg) 253 | }; 254 | 255 | send.then(|res| { 256 | res.unwrap(); 257 | Ok(()) 258 | }) 259 | })).unwrap(); 260 | 261 | start_txs.push(start_tx); 262 | final_rxs.push(final_rx); 263 | } 264 | 265 | for start_tx in start_txs { 266 | start_tx.send("ping").wait().unwrap(); 267 | } 268 | 269 | for final_rx in final_rxs { 270 | final_rx.wait().next().unwrap().unwrap(); 271 | } 272 | 273 | // Shutdown the pool 274 | pool.wait().unwrap(); 275 | } 276 | } 277 | --------------------------------------------------------------------------------