├── .gitignore ├── .github ├── dependabot.yml └── workflows │ ├── release.yml │ └── ci.yml ├── benches └── spawn.rs ├── LICENSE-MIT ├── Cargo.toml ├── examples ├── spawn.rs ├── spawn-on-thread.rs ├── spawn-local.rs └── with-metadata.rs ├── tests ├── metadata.rs ├── cancel.rs ├── ready.rs ├── panic.rs ├── waker_ready.rs ├── basic.rs ├── waker_panic.rs ├── waker_pending.rs └── join.rs ├── README.md ├── CHANGELOG.md ├── src ├── state.rs ├── utils.rs ├── lib.rs ├── header.rs ├── task.rs ├── raw.rs └── runnable.rs └── LICENSE-APACHE /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | **/*.rs.bk 3 | Cargo.lock 4 | -------------------------------------------------------------------------------- /.github/dependabot.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | updates: 3 | - package-ecosystem: cargo 4 | directory: / 5 | schedule: 6 | interval: weekly 7 | commit-message: 8 | prefix: '' 9 | labels: [] 10 | -------------------------------------------------------------------------------- /benches/spawn.rs: -------------------------------------------------------------------------------- 1 | #![feature(test)] 2 | 3 | extern crate test; 4 | 5 | use smol::future; 6 | use test::Bencher; 7 | 8 | #[bench] 9 | fn task_create(b: &mut Bencher) { 10 | b.iter(|| { 11 | let _ = async_task::spawn(async {}, drop); 12 | }); 13 | } 14 | 15 | #[bench] 16 | fn task_run(b: &mut Bencher) { 17 | b.iter(|| { 18 | let (runnable, task) = async_task::spawn(async {}, drop); 19 | runnable.run(); 20 | future::block_on(task); 21 | }); 22 | } 23 | -------------------------------------------------------------------------------- /.github/workflows/release.yml: -------------------------------------------------------------------------------- 1 | name: Release 2 | 3 | permissions: 4 | contents: write 5 | 6 | on: 7 | push: 8 | tags: 9 | - v[0-9]+.* 10 | 11 | jobs: 12 | create-release: 13 | if: github.repository_owner == 'smol-rs' 14 | runs-on: ubuntu-latest 15 | steps: 16 | - uses: taiki-e/checkout-action@v1 17 | - uses: taiki-e/create-gh-release-action@v1 18 | with: 19 | changelog: CHANGELOG.md 20 | branch: master 21 | env: 22 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 23 | -------------------------------------------------------------------------------- /LICENSE-MIT: -------------------------------------------------------------------------------- 1 | Permission is hereby granted, free of charge, to any 2 | person obtaining a copy of this software and associated 3 | documentation files (the "Software"), to deal in the 4 | Software without restriction, including without 5 | limitation the rights to use, copy, modify, merge, 6 | publish, distribute, sublicense, and/or sell copies of 7 | the Software, and to permit persons to whom the Software 8 | is furnished to do so, subject to the following 9 | conditions: 10 | 11 | The above copyright notice and this permission notice 12 | shall be included in all copies or substantial portions 13 | of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF 16 | ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED 17 | TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A 18 | PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT 19 | SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY 20 | CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 21 | OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR 22 | IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 23 | DEALINGS IN THE SOFTWARE. 24 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "async-task" 3 | # When publishing a new version: 4 | # - Update CHANGELOG.md 5 | # - Create "v4.x.y" git tag 6 | version = "4.7.1" 7 | authors = ["Stjepan Glavina "] 8 | edition = "2021" 9 | rust-version = "1.57" 10 | license = "Apache-2.0 OR MIT" 11 | repository = "https://github.com/smol-rs/async-task" 12 | description = "Task abstraction for building executors" 13 | keywords = ["futures", "task", "executor", "spawn"] 14 | categories = ["asynchronous", "concurrency", "no-std"] 15 | exclude = ["/.*"] 16 | 17 | [features] 18 | default = ["std"] 19 | std = [] 20 | 21 | [dependencies] 22 | # Uses portable-atomic polyfill atomics on targets without them 23 | portable-atomic = { version = "1", optional = true, default-features = false } 24 | 25 | [dev-dependencies] 26 | atomic-waker = "1" 27 | easy-parallel = "3" 28 | flaky_test = "0.2" 29 | flume = { version = "0.12", default-features = false } 30 | futures-lite = "2.0.0" 31 | once_cell = "1" 32 | pin-project-lite = "0.2.10" 33 | smol = "2" 34 | 35 | # rewrite dependencies to use this version of async-task when running tests 36 | [patch.crates-io] 37 | async-task = { path = "." } 38 | -------------------------------------------------------------------------------- /examples/spawn.rs: -------------------------------------------------------------------------------- 1 | //! A simple single-threaded executor. 2 | 3 | use std::future::Future; 4 | use std::panic::catch_unwind; 5 | use std::thread; 6 | 7 | use async_task::{Runnable, Task}; 8 | use once_cell::sync::Lazy; 9 | use smol::future; 10 | 11 | /// Spawns a future on the executor. 12 | fn spawn(future: F) -> Task 13 | where 14 | F: Future + Send + 'static, 15 | T: Send + 'static, 16 | { 17 | // A queue that holds scheduled tasks. 18 | static QUEUE: Lazy> = Lazy::new(|| { 19 | let (sender, receiver) = flume::unbounded::(); 20 | 21 | // Start the executor thread. 22 | thread::spawn(|| { 23 | for runnable in receiver { 24 | // Ignore panics inside futures. 25 | let _ignore_panic = catch_unwind(|| runnable.run()); 26 | } 27 | }); 28 | 29 | sender 30 | }); 31 | 32 | // Create a task that is scheduled by pushing it into the queue. 33 | let schedule = |runnable| QUEUE.send(runnable).unwrap(); 34 | let (runnable, task) = async_task::spawn(future, schedule); 35 | 36 | // Schedule the task by pushing it into the queue. 37 | runnable.schedule(); 38 | 39 | task 40 | } 41 | 42 | fn main() { 43 | // Spawn a future and await its result. 44 | let task = spawn(async { 45 | println!("Hello, world!"); 46 | }); 47 | future::block_on(task); 48 | } 49 | -------------------------------------------------------------------------------- /examples/spawn-on-thread.rs: -------------------------------------------------------------------------------- 1 | //! A function that runs a future to completion on a dedicated thread. 2 | 3 | use std::future::Future; 4 | use std::sync::Arc; 5 | use std::thread; 6 | 7 | use async_task::Task; 8 | use smol::future; 9 | 10 | /// Spawns a future on a new dedicated thread. 11 | /// 12 | /// The returned task can be used to await the output of the future. 13 | fn spawn_on_thread(future: F) -> Task 14 | where 15 | F: Future + Send + 'static, 16 | T: Send + 'static, 17 | { 18 | // Create a channel that holds the task when it is scheduled for running. 19 | let (sender, receiver) = flume::unbounded(); 20 | let sender = Arc::new(sender); 21 | let s = Arc::downgrade(&sender); 22 | 23 | // Wrap the future into one that disconnects the channel on completion. 24 | let future = async move { 25 | // When the inner future completes, the sender gets dropped and disconnects the channel. 26 | let _sender = sender; 27 | future.await 28 | }; 29 | 30 | // Create a task that is scheduled by sending it into the channel. 31 | let schedule = move |runnable| s.upgrade().unwrap().send(runnable).unwrap(); 32 | let (runnable, task) = async_task::spawn(future, schedule); 33 | 34 | // Schedule the task by sending it into the channel. 35 | runnable.schedule(); 36 | 37 | // Spawn a thread running the task to completion. 38 | thread::spawn(move || { 39 | // Keep taking the task from the channel and running it until completion. 40 | for runnable in receiver { 41 | runnable.run(); 42 | } 43 | }); 44 | 45 | task 46 | } 47 | 48 | fn main() { 49 | // Spawn a future on a dedicated thread. 50 | future::block_on(spawn_on_thread(async { 51 | println!("Hello, world!"); 52 | })); 53 | } 54 | -------------------------------------------------------------------------------- /tests/metadata.rs: -------------------------------------------------------------------------------- 1 | use async_task::{Builder, Runnable}; 2 | use flume::unbounded; 3 | use smol::future; 4 | 5 | use std::sync::atomic::{AtomicUsize, Ordering}; 6 | 7 | #[test] 8 | fn metadata_use_case() { 9 | // Each future has a counter that is incremented every time it is scheduled. 10 | let (sender, receiver) = unbounded::>(); 11 | let schedule = move |runnable: Runnable| { 12 | runnable.metadata().fetch_add(1, Ordering::SeqCst); 13 | sender.send(runnable).ok(); 14 | }; 15 | 16 | async fn my_future(counter: &AtomicUsize) { 17 | loop { 18 | // Loop until we've been scheduled five times. 19 | let count = counter.load(Ordering::SeqCst); 20 | if count < 5 { 21 | // Make sure that we are immediately scheduled again. 22 | future::yield_now().await; 23 | continue; 24 | } 25 | 26 | // We've been scheduled five times, so we're done. 27 | break; 28 | } 29 | } 30 | 31 | let make_task = || { 32 | // SAFETY: We are spawning a non-'static future, so we need to use the unsafe API. 33 | // The borrowed variables, in this case the metadata, are guaranteed to outlive the runnable. 34 | let (runnable, task) = unsafe { 35 | Builder::new() 36 | .metadata(AtomicUsize::new(0)) 37 | .spawn_unchecked(my_future, schedule.clone()) 38 | }; 39 | 40 | runnable.schedule(); 41 | task 42 | }; 43 | 44 | // Make tasks. 45 | let t1 = make_task(); 46 | let t2 = make_task(); 47 | 48 | // Run the tasks. 49 | while let Ok(runnable) = receiver.try_recv() { 50 | runnable.run(); 51 | } 52 | 53 | // Unwrap the tasks. 54 | smol::future::block_on(async move { 55 | t1.await; 56 | t2.await; 57 | }); 58 | } 59 | -------------------------------------------------------------------------------- /examples/spawn-local.rs: -------------------------------------------------------------------------------- 1 | //! A simple single-threaded executor that can spawn non-`Send` futures. 2 | 3 | use std::cell::Cell; 4 | use std::future::Future; 5 | use std::rc::Rc; 6 | 7 | use async_task::{Runnable, Task}; 8 | 9 | thread_local! { 10 | // A queue that holds scheduled tasks. 11 | static QUEUE: (flume::Sender, flume::Receiver) = flume::unbounded(); 12 | } 13 | 14 | /// Spawns a future on the executor. 15 | fn spawn(future: F) -> Task 16 | where 17 | F: Future + 'static, 18 | T: 'static, 19 | { 20 | // Create a task that is scheduled by pushing itself into the queue. 21 | let schedule = |runnable| QUEUE.with(|(s, _)| s.send(runnable).unwrap()); 22 | let (runnable, task) = async_task::spawn_local(future, schedule); 23 | 24 | // Schedule the task by pushing it into the queue. 25 | runnable.schedule(); 26 | 27 | task 28 | } 29 | 30 | /// Runs a future to completion. 31 | fn run(future: F) -> T 32 | where 33 | F: Future + 'static, 34 | T: 'static, 35 | { 36 | // Spawn a task that sends its result through a channel. 37 | let (s, r) = flume::unbounded(); 38 | spawn(async move { drop(s.send(future.await)) }).detach(); 39 | 40 | loop { 41 | // If the original task has completed, return its result. 42 | if let Ok(val) = r.try_recv() { 43 | return val; 44 | } 45 | 46 | // Otherwise, take a task from the queue and run it. 47 | QUEUE.with(|(_, r)| r.recv().unwrap().run()); 48 | } 49 | } 50 | 51 | fn main() { 52 | let val = Rc::new(Cell::new(0)); 53 | 54 | // Run a future that increments a non-`Send` value. 55 | run({ 56 | let val = val.clone(); 57 | async move { 58 | // Spawn a future that increments the value. 59 | let task = spawn({ 60 | let val = val.clone(); 61 | async move { 62 | val.set(dbg!(val.get()) + 1); 63 | } 64 | }); 65 | 66 | val.set(dbg!(val.get()) + 1); 67 | task.await; 68 | } 69 | }); 70 | 71 | // The value should be 2 at the end of the program. 72 | dbg!(val.get()); 73 | } 74 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # async-task 2 | 3 | [![Build](https://github.com/smol-rs/async-task/actions/workflows/ci.yml/badge.svg)]( 4 | https://github.com/smol-rs/async-task/actions) 5 | [![License](https://img.shields.io/badge/license-Apache--2.0_OR_MIT-blue.svg)]( 6 | https://github.com/smol-rs/async-task) 7 | [![Cargo](https://img.shields.io/crates/v/async-task.svg)]( 8 | https://crates.io/crates/async-task) 9 | [![Documentation](https://docs.rs/async-task/badge.svg)]( 10 | https://docs.rs/async-task) 11 | 12 | Task abstraction for building executors. 13 | 14 | To spawn a future onto an executor, we first need to allocate it on the heap and keep some 15 | state attached to it. The state indicates whether the future is ready for polling, waiting to 16 | be woken up, or completed. Such a stateful future is called a *task*. 17 | 18 | All executors have a queue that holds scheduled tasks: 19 | 20 | ```rust 21 | let (sender, receiver) = flume::unbounded(); 22 | ``` 23 | 24 | A task is created using either `spawn()`, `spawn_local()`, or `spawn_unchecked()` which 25 | returns a `Runnable` and a `Task`: 26 | 27 | ```rust 28 | // A future that will be spawned. 29 | let future = async { 1 + 2 }; 30 | 31 | // A function that schedules the task when it gets woken up. 32 | let schedule = move |runnable| sender.send(runnable).unwrap(); 33 | 34 | // Construct a task. 35 | let (runnable, task) = async_task::spawn(future, schedule); 36 | 37 | // Push the task into the queue by invoking its schedule function. 38 | runnable.schedule(); 39 | ``` 40 | 41 | The `Runnable` is used to poll the task's future, and the `Task` is used to await its 42 | output. 43 | 44 | Finally, we need a loop that takes scheduled tasks from the queue and runs them: 45 | 46 | ```rust 47 | for runnable in receiver { 48 | runnable.run(); 49 | } 50 | ``` 51 | 52 | Method `run()` polls the task's future once. Then, the `Runnable` 53 | vanishes and only reappears when its `Waker` wakes the task, thus 54 | scheduling it to be run again. 55 | 56 | ## License 57 | 58 | Licensed under either of 59 | 60 | * Apache License, Version 2.0 ([LICENSE-APACHE](LICENSE-APACHE) or http://www.apache.org/licenses/LICENSE-2.0) 61 | * MIT license ([LICENSE-MIT](LICENSE-MIT) or http://opensource.org/licenses/MIT) 62 | 63 | at your option. 64 | 65 | #### Contribution 66 | 67 | Unless you explicitly state otherwise, any contribution intentionally submitted 68 | for inclusion in the work by you, as defined in the Apache-2.0 license, shall be 69 | dual licensed as above, without any additional terms or conditions. 70 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Version 4.7.1 2 | 3 | - Improve the panic message for when a task is polled after completion. (#73) 4 | 5 | # Version 4.7.0 6 | 7 | - Add `from_raw` and `into_raw` functions for `Runnable` to ease passing it 8 | across an FFI boundary. (#65) 9 | 10 | # Version 4.6.0 11 | 12 | - Bump MSRV to 1.57. (#63) 13 | - Task layout computation failures are now a compile-time error instead of a 14 | runtime abort. (#63) 15 | 16 | # Version 4.5.0 17 | 18 | - Add a `portable-atomic` feature that enables the usage of fallback primitives for CPUs without atomics. (#58) 19 | 20 | # Version 4.4.1 21 | 22 | - Clarify safety documentation for `spawn_unchecked`. (#49) 23 | 24 | # Version 4.4.0 25 | 26 | - Ensure that the allocation doesn't exceed `isize::MAX` (#32) 27 | - Add `FallibleTask::is_finished()` (#34) 28 | - Add a metadata generic parameter to tasks (#33) 29 | - Add panic propagation to tasks (#37) 30 | - Add a way to tell if the task was woken while running from the schedule function (#42) 31 | 32 | # Version 4.3.0 33 | 34 | - Bump MSRV to Rust 1.47. (#30) 35 | - Evaluate the layouts for the tasks at compile time. (#30) 36 | - Add layout_info field to TaskVTable so that debuggers can decode raw tasks. (#29) 37 | 38 | # Version 4.2.0 39 | 40 | - Add `Task::is_finished`. (#19) 41 | 42 | # Version 4.1.0 43 | 44 | - Add `FallibleTask`. (#21) 45 | 46 | # Version 4.0.3 47 | 48 | - Document the return value of `Runnable::run()` better. 49 | 50 | # Version 4.0.2 51 | 52 | - Nits in the docs. 53 | 54 | # Version 4.0.1 55 | 56 | - Nits in the docs. 57 | 58 | # Version 4.0.0 59 | 60 | - Rename `Task` to `Runnable`. 61 | - Rename `JoinHandle` to `Task`. 62 | - Cancel `Task` on drop. 63 | - Add `Task::detach()` and `Task::cancel()`. 64 | - Add `spawn_unchecked()`. 65 | 66 | # Version 3.0.0 67 | 68 | - Use `ThreadId` in `spawn_local` because OS-provided IDs can get recycled. 69 | - Add `std` feature to `Cargo.toml`. 70 | 71 | # Version 2.1.1 72 | 73 | - Allocate large futures on the heap. 74 | 75 | # Version 2.1.0 76 | 77 | - `JoinHandle` now only evaluates after the task's future has been dropped. 78 | 79 | # Version 2.0.0 80 | 81 | - Return `true` in `Task::run()`. 82 | 83 | # Version 1.3.1 84 | 85 | - Make `spawn_local` available only on unix and windows. 86 | 87 | # Version 1.3.0 88 | 89 | - Add `waker_fn`. 90 | 91 | # Version 1.2.1 92 | 93 | - Add the `no-std` category to the package. 94 | 95 | # Version 1.2.0 96 | 97 | - The crate is now marked with `#![no_std]`. 98 | - Add `Task::waker` and `JoinHandle::waker`. 99 | - Add `Task::into_raw` and `Task::from_raw`. 100 | 101 | # Version 1.1.1 102 | 103 | - Fix a use-after-free bug where the schedule function is dropped while running. 104 | 105 | # Version 1.1.0 106 | 107 | - If a task is dropped or canceled outside the `run` method, it gets re-scheduled. 108 | - Add `spawn_local` constructor. 109 | 110 | # Version 1.0.0 111 | 112 | - Initial release 113 | -------------------------------------------------------------------------------- /src/state.rs: -------------------------------------------------------------------------------- 1 | /// Set if the task is scheduled for running. 2 | /// 3 | /// A task is considered to be scheduled whenever its `Runnable` exists. 4 | /// 5 | /// This flag can't be set when the task is completed. However, it can be set while the task is 6 | /// running, in which case it will be rescheduled as soon as polling finishes. 7 | pub(crate) const SCHEDULED: usize = 1 << 0; 8 | 9 | /// Set if the task is running. 10 | /// 11 | /// A task is in running state while its future is being polled. 12 | /// 13 | /// This flag can't be set when the task is completed. However, it can be in scheduled state while 14 | /// it is running, in which case it will be rescheduled as soon as polling finishes. 15 | pub(crate) const RUNNING: usize = 1 << 1; 16 | 17 | /// Set if the task has been completed. 18 | /// 19 | /// This flag is set when polling returns `Poll::Ready`. The output of the future is then stored 20 | /// inside the task until it becomes closed. In fact, `Task` picks up the output by marking 21 | /// the task as closed. 22 | /// 23 | /// This flag can't be set when the task is scheduled or running. 24 | pub(crate) const COMPLETED: usize = 1 << 2; 25 | 26 | /// Set if the task is closed. 27 | /// 28 | /// If a task is closed, that means it's either canceled or its output has been consumed by the 29 | /// `Task`. A task becomes closed in the following cases: 30 | /// 31 | /// 1. It gets canceled by `Runnable::drop()`, `Task::drop()`, or `Task::cancel()`. 32 | /// 2. Its output gets awaited by the `Task`. 33 | /// 3. It panics while polling the future. 34 | /// 4. It is completed and the `Task` gets dropped. 35 | pub(crate) const CLOSED: usize = 1 << 3; 36 | 37 | /// Set if the `Task` still exists. 38 | /// 39 | /// The `Task` is a special case in that it is only tracked by this flag, while all other 40 | /// task references (`Runnable` and `Waker`s) are tracked by the reference count. 41 | pub(crate) const TASK: usize = 1 << 4; 42 | 43 | /// Set if the `Task` is awaiting the output. 44 | /// 45 | /// This flag is set while there is a registered awaiter of type `Waker` inside the task. When the 46 | /// task gets closed or completed, we need to wake the awaiter. This flag can be used as a fast 47 | /// check that tells us if we need to wake anyone. 48 | pub(crate) const AWAITER: usize = 1 << 5; 49 | 50 | /// Set if an awaiter is being registered. 51 | /// 52 | /// This flag is set when `Task` is polled and we are registering a new awaiter. 53 | pub(crate) const REGISTERING: usize = 1 << 6; 54 | 55 | /// Set if the awaiter is being notified. 56 | /// 57 | /// This flag is set when notifying the awaiter. If an awaiter is concurrently registered and 58 | /// notified, whichever side came first will take over the responsibility of resolving the race. 59 | pub(crate) const NOTIFYING: usize = 1 << 7; 60 | 61 | /// A single reference. 62 | /// 63 | /// The lower bits in the state contain various flags representing the task state, while the upper 64 | /// bits contain the reference count. The value of `REFERENCE` represents a single reference in the 65 | /// total reference count. 66 | /// 67 | /// Note that the reference counter only tracks the `Runnable` and `Waker`s. The `Task` is 68 | /// tracked separately by the `TASK` flag. 69 | pub(crate) const REFERENCE: usize = 1 << 8; 70 | -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | 3 | permissions: 4 | contents: read 5 | 6 | on: 7 | pull_request: 8 | push: 9 | branches: 10 | - master 11 | schedule: 12 | - cron: '0 2 * * 0' 13 | 14 | env: 15 | CARGO_INCREMENTAL: 0 16 | CARGO_NET_GIT_FETCH_WITH_CLI: true 17 | CARGO_NET_RETRY: 10 18 | CARGO_TERM_COLOR: always 19 | RUST_BACKTRACE: 1 20 | RUSTFLAGS: -D warnings 21 | RUSTDOCFLAGS: -D warnings 22 | RUSTUP_MAX_RETRIES: 10 23 | 24 | defaults: 25 | run: 26 | shell: bash 27 | 28 | jobs: 29 | fmt: 30 | uses: smol-rs/.github/.github/workflows/fmt.yml@main 31 | clippy: 32 | uses: smol-rs/.github/.github/workflows/clippy.yml@main 33 | with: 34 | bench: false # We run clippy using stable rustc, but our benchmarks use #![feature(test)]. 35 | security_audit: 36 | uses: smol-rs/.github/.github/workflows/security_audit.yml@main 37 | permissions: 38 | checks: write 39 | contents: read 40 | issues: write 41 | secrets: inherit 42 | 43 | test: 44 | runs-on: ${{ matrix.os }} 45 | strategy: 46 | fail-fast: false 47 | matrix: 48 | os: [ubuntu-latest] 49 | rust: [nightly, beta, stable] 50 | steps: 51 | - uses: taiki-e/checkout-action@v1 52 | - name: Install Rust 53 | run: rustup update ${{ matrix.rust }} && rustup default ${{ matrix.rust }} 54 | - run: rustup target add thumbv7m-none-eabi 55 | - name: Install cargo-hack 56 | uses: taiki-e/install-action@cargo-hack 57 | - name: Install valgrind 58 | uses: taiki-e/install-action@valgrind 59 | - run: cargo build --all --all-features --all-targets 60 | if: startsWith(matrix.rust, 'nightly') 61 | - run: cargo hack build --feature-powerset --no-dev-deps 62 | - run: cargo hack build --feature-powerset --no-dev-deps --target thumbv7m-none-eabi --skip std,default 63 | - run: cargo test 64 | - name: Run cargo test (with valgrind) 65 | run: cargo test -- --test-threads=1 66 | env: 67 | # TODO: ignore possible and reachable leaks due to upstream issue (https://github.com/rust-lang/rust/issues/135608) 68 | CARGO_TARGET_X86_64_UNKNOWN_LINUX_GNU_RUNNER: valgrind -v --error-exitcode=1 --error-limit=no --leak-check=full --show-leak-kinds=definite,indirect --errors-for-leak-kinds=definite,indirect --track-origins=yes --fair-sched=yes --gen-suppressions=all 69 | - name: Run cargo test (with portable-atomic enabled) 70 | run: cargo test --features portable-atomic 71 | - name: Clone async-executor 72 | run: git clone https://github.com/smol-rs/async-executor.git 73 | - name: Add patch section 74 | run: | 75 | echo '[patch.crates-io]' >> async-executor/Cargo.toml 76 | echo 'async-task = { path = ".." }' >> async-executor/Cargo.toml 77 | - name: Test async-executor 78 | run: cargo test --manifest-path async-executor/Cargo.toml 79 | 80 | msrv: 81 | runs-on: ubuntu-latest 82 | steps: 83 | - uses: taiki-e/checkout-action@v1 84 | - name: Install cargo-hack 85 | uses: taiki-e/install-action@cargo-hack 86 | - run: cargo hack build --feature-powerset --no-dev-deps --rust-version 87 | 88 | miri: 89 | runs-on: ubuntu-latest 90 | steps: 91 | - uses: taiki-e/checkout-action@v1 92 | - name: Install Rust 93 | run: rustup toolchain install nightly --component miri && rustup default nightly 94 | - run: cargo miri test 95 | env: 96 | # -Zmiri-ignore-leaks is needed because we use detached threads in doctests: https://github.com/rust-lang/miri/issues/1371 97 | # disable preemption due to https://github.com/rust-lang/rust/issues/55005 98 | MIRIFLAGS: -Zmiri-strict-provenance -Zmiri-symbolic-alignment-check -Zmiri-disable-isolation -Zmiri-ignore-leaks -Zmiri-preemption-rate=0 99 | RUSTFLAGS: ${{ env.RUSTFLAGS }} -Z randomize-layout 100 | -------------------------------------------------------------------------------- /src/utils.rs: -------------------------------------------------------------------------------- 1 | use core::alloc::Layout as StdLayout; 2 | use core::mem; 3 | 4 | /// Aborts the process. 5 | /// 6 | /// To abort, this function simply panics while panicking. 7 | pub(crate) fn abort() -> ! { 8 | struct Panic; 9 | 10 | impl Drop for Panic { 11 | fn drop(&mut self) { 12 | panic!("aborting the process"); 13 | } 14 | } 15 | 16 | let _panic = Panic; 17 | panic!("aborting the process"); 18 | } 19 | 20 | pub(crate) struct Bomb; 21 | 22 | impl Drop for Bomb { 23 | fn drop(&mut self) { 24 | abort(); 25 | } 26 | } 27 | 28 | /// Calls a function and aborts if it panics. 29 | /// 30 | /// This is useful in unsafe code where we can't recover from panics. 31 | #[inline] 32 | pub(crate) fn abort_on_panic(f: impl FnOnce() -> T) -> T { 33 | let bomb = Bomb; 34 | let t = f(); 35 | mem::forget(bomb); 36 | t 37 | } 38 | 39 | /// A version of `alloc::alloc::Layout` that can be used in the const 40 | /// position. 41 | #[derive(Clone, Copy, Debug)] 42 | pub(crate) struct Layout { 43 | size: usize, 44 | align: usize, 45 | } 46 | 47 | impl Layout { 48 | /// Creates a new `Layout` with the given size and alignment. 49 | #[inline] 50 | pub(crate) const fn from_size_align(size: usize, align: usize) -> Self { 51 | Self { size, align } 52 | } 53 | 54 | /// Creates a new `Layout` for the given sized type. 55 | #[inline] 56 | pub(crate) const fn new() -> Self { 57 | Self::from_size_align(mem::size_of::(), mem::align_of::()) 58 | } 59 | 60 | /// Convert this into the standard library's layout type. 61 | /// 62 | /// # Safety 63 | /// 64 | /// - `align` must be non-zero and a power of two. 65 | /// - When rounded up to the nearest multiple of `align`, the size 66 | /// must not overflow. 67 | #[inline] 68 | pub(crate) const unsafe fn into_std(self) -> StdLayout { 69 | StdLayout::from_size_align_unchecked(self.size, self.align) 70 | } 71 | 72 | /// Get the alignment of this layout. 73 | #[inline] 74 | pub(crate) const fn align(&self) -> usize { 75 | self.align 76 | } 77 | 78 | /// Get the size of this layout. 79 | #[inline] 80 | pub(crate) const fn size(&self) -> usize { 81 | self.size 82 | } 83 | 84 | /// Returns the layout for `a` followed by `b` and the offset of `b`. 85 | /// 86 | /// This function was adapted from the `Layout::extend()`: 87 | /// https://doc.rust-lang.org/nightly/std/alloc/struct.Layout.html#method.extend 88 | #[inline] 89 | pub(crate) const fn extend(self, other: Layout) -> Option<(Layout, usize)> { 90 | let new_align = max(self.align(), other.align()); 91 | let pad = self.padding_needed_for(other.align()); 92 | 93 | let offset = leap!(self.size().checked_add(pad)); 94 | let new_size = leap!(offset.checked_add(other.size())); 95 | 96 | // return None if any of the following are true: 97 | // - align is 0 (implied false by is_power_of_two()) 98 | // - align is not a power of 2 99 | // - size rounded up to align overflows 100 | if !new_align.is_power_of_two() || new_size > isize::MAX as usize - (new_align - 1) { 101 | return None; 102 | } 103 | 104 | let layout = Layout::from_size_align(new_size, new_align); 105 | Some((layout, offset)) 106 | } 107 | 108 | /// Returns the padding after `layout` that aligns the following address to `align`. 109 | /// 110 | /// This function was adapted from the `Layout::padding_needed_for()`: 111 | /// https://doc.rust-lang.org/nightly/std/alloc/struct.Layout.html#method.padding_needed_for 112 | #[inline] 113 | const fn padding_needed_for(self, align: usize) -> usize { 114 | let len = self.size(); 115 | let len_rounded_up = len.wrapping_add(align).wrapping_sub(1) & !align.wrapping_sub(1); 116 | len_rounded_up.wrapping_sub(len) 117 | } 118 | } 119 | 120 | #[inline] 121 | pub(crate) const fn max(left: usize, right: usize) -> usize { 122 | if left > right { 123 | left 124 | } else { 125 | right 126 | } 127 | } 128 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | //! Task abstraction for building executors. 2 | //! 3 | //! To spawn a future onto an executor, we first need to allocate it on the heap and keep some 4 | //! state attached to it. The state indicates whether the future is ready for polling, waiting to 5 | //! be woken up, or completed. Such a stateful future is called a *task*. 6 | //! 7 | //! All executors have a queue that holds scheduled tasks: 8 | //! 9 | //! ``` 10 | //! let (sender, receiver) = flume::unbounded(); 11 | //! # 12 | //! # // A future that will get spawned. 13 | //! # let future = async { 1 + 2 }; 14 | //! # 15 | //! # // A function that schedules the task when it gets woken up. 16 | //! # let schedule = move |runnable| sender.send(runnable).unwrap(); 17 | //! # 18 | //! # // Create a task. 19 | //! # let (runnable, task) = async_task::spawn(future, schedule); 20 | //! ``` 21 | //! 22 | //! A task is created using either [`spawn()`], [`spawn_local()`], or [`spawn_unchecked()`] which 23 | //! returns a [`Runnable`] and a [`Task`]: 24 | //! 25 | //! ``` 26 | //! # let (sender, receiver) = flume::unbounded(); 27 | //! # 28 | //! // A future that will be spawned. 29 | //! let future = async { 1 + 2 }; 30 | //! 31 | //! // A function that schedules the task when it gets woken up. 32 | //! let schedule = move |runnable| sender.send(runnable).unwrap(); 33 | //! 34 | //! // Construct a task. 35 | //! let (runnable, task) = async_task::spawn(future, schedule); 36 | //! 37 | //! // Push the task into the queue by invoking its schedule function. 38 | //! runnable.schedule(); 39 | //! # let handle = std::thread::spawn(move || { for runnable in receiver { runnable.run(); }}); 40 | //! # smol::future::block_on(task); 41 | //! # handle.join().unwrap(); 42 | //! ``` 43 | //! 44 | //! The [`Runnable`] is used to poll the task's future, and the [`Task`] is used to await its 45 | //! output. 46 | //! 47 | //! Finally, we need a loop that takes scheduled tasks from the queue and runs them: 48 | //! 49 | //! ```no_run 50 | //! # let (sender, receiver) = flume::unbounded(); 51 | //! # 52 | //! # // A future that will get spawned. 53 | //! # let future = async { 1 + 2 }; 54 | //! # 55 | //! # // A function that schedules the task when it gets woken up. 56 | //! # let schedule = move |runnable| sender.send(runnable).unwrap(); 57 | //! # 58 | //! # // Create a task. 59 | //! # let (runnable, task) = async_task::spawn(future, schedule); 60 | //! # 61 | //! # // Push the task into the queue by invoking its schedule function. 62 | //! # runnable.schedule(); 63 | //! # 64 | //! for runnable in receiver { 65 | //! runnable.run(); 66 | //! } 67 | //! ``` 68 | //! 69 | //! Method [`run()`][`Runnable::run()`] polls the task's future once. Then, the [`Runnable`] 70 | //! vanishes and only reappears when its [`Waker`][`core::task::Waker`] wakes the task, thus 71 | //! scheduling it to be run again. 72 | 73 | #![no_std] 74 | #![warn(missing_docs, missing_debug_implementations, rust_2018_idioms)] 75 | #![doc(test(attr(deny(rust_2018_idioms, warnings))))] 76 | #![doc(test(attr(allow(unused_variables))))] 77 | #![doc( 78 | html_favicon_url = "https://raw.githubusercontent.com/smol-rs/smol/master/assets/images/logo_fullsize_transparent.png" 79 | )] 80 | #![doc( 81 | html_logo_url = "https://raw.githubusercontent.com/smol-rs/smol/master/assets/images/logo_fullsize_transparent.png" 82 | )] 83 | 84 | extern crate alloc; 85 | #[cfg(feature = "std")] 86 | extern crate std; 87 | 88 | /// We can't use `?` in const contexts yet, so this macro acts 89 | /// as a workaround. 90 | macro_rules! leap { 91 | ($x: expr) => {{ 92 | match ($x) { 93 | Some(val) => val, 94 | None => return None, 95 | } 96 | }}; 97 | } 98 | 99 | macro_rules! leap_unwrap { 100 | ($x: expr) => {{ 101 | match ($x) { 102 | Some(val) => val, 103 | None => panic!("called `Option::unwrap()` on a `None` value"), 104 | } 105 | }}; 106 | } 107 | 108 | mod header; 109 | mod raw; 110 | mod runnable; 111 | mod state; 112 | mod task; 113 | mod utils; 114 | 115 | pub use crate::runnable::{ 116 | spawn, spawn_unchecked, Builder, Runnable, Schedule, ScheduleInfo, WithInfo, 117 | }; 118 | pub use crate::task::{FallibleTask, Task}; 119 | 120 | #[cfg(feature = "std")] 121 | pub use crate::runnable::spawn_local; 122 | -------------------------------------------------------------------------------- /examples/with-metadata.rs: -------------------------------------------------------------------------------- 1 | //! A single threaded executor that uses shortest-job-first scheduling. 2 | 3 | use std::cell::RefCell; 4 | use std::collections::BinaryHeap; 5 | use std::pin::Pin; 6 | use std::task::{Context, Poll}; 7 | use std::thread; 8 | use std::time::{Duration, Instant}; 9 | use std::{cell::Cell, future::Future}; 10 | 11 | use async_task::{Builder, Runnable, Task}; 12 | use pin_project_lite::pin_project; 13 | use smol::{channel, future}; 14 | 15 | struct ByDuration(Runnable); 16 | 17 | impl ByDuration { 18 | fn duration(&self) -> Duration { 19 | self.0.metadata().inner.get() 20 | } 21 | } 22 | 23 | impl PartialEq for ByDuration { 24 | fn eq(&self, other: &Self) -> bool { 25 | self.duration() == other.duration() 26 | } 27 | } 28 | 29 | impl Eq for ByDuration {} 30 | 31 | impl PartialOrd for ByDuration { 32 | fn partial_cmp(&self, other: &Self) -> Option { 33 | Some(self.cmp(other)) 34 | } 35 | } 36 | 37 | impl Ord for ByDuration { 38 | fn cmp(&self, other: &Self) -> std::cmp::Ordering { 39 | self.duration().cmp(&other.duration()).reverse() 40 | } 41 | } 42 | 43 | pin_project! { 44 | #[must_use = "futures do nothing unless you `.await` or poll them"] 45 | struct MeasureRuntime<'a, F> { 46 | #[pin] 47 | f: F, 48 | duration: &'a Cell 49 | } 50 | } 51 | 52 | impl Future for MeasureRuntime<'_, F> { 53 | type Output = F::Output; 54 | 55 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 56 | let this = self.project(); 57 | let duration_cell: &Cell = this.duration; 58 | let start = Instant::now(); 59 | let res = F::poll(this.f, cx); 60 | let new_duration = Instant::now() - start; 61 | duration_cell.set(duration_cell.get() / 2 + new_duration / 2); 62 | res 63 | } 64 | } 65 | 66 | pub struct DurationMetadata { 67 | inner: Cell, 68 | } 69 | 70 | thread_local! { 71 | // A queue that holds scheduled tasks. 72 | static QUEUE: RefCell> = RefCell::new(BinaryHeap::new()); 73 | } 74 | 75 | fn make_future_fn<'a, F>(future: F) -> impl FnOnce(&'a DurationMetadata) -> MeasureRuntime<'a, F> { 76 | move |duration_meta| MeasureRuntime { 77 | f: future, 78 | duration: &duration_meta.inner, 79 | } 80 | } 81 | 82 | fn ensure_safe_schedule(f: F) -> F { 83 | f 84 | } 85 | 86 | /// Spawns a future on the executor. 87 | pub fn spawn(future: F) -> Task 88 | where 89 | F: Future + 'static, 90 | T: 'static, 91 | { 92 | let spawn_thread_id = thread::current().id(); 93 | // Create a task that is scheduled by pushing it into the queue. 94 | let schedule = ensure_safe_schedule(move |runnable| { 95 | if thread::current().id() != spawn_thread_id { 96 | panic!("Task would be run on a different thread than spawned on."); 97 | } 98 | QUEUE.with(move |queue| queue.borrow_mut().push(ByDuration(runnable))); 99 | }); 100 | let future_fn = make_future_fn(future); 101 | let (runnable, task) = unsafe { 102 | Builder::new() 103 | .metadata(DurationMetadata { 104 | inner: Cell::new(Duration::default()), 105 | }) 106 | .spawn_unchecked(future_fn, schedule) 107 | }; 108 | 109 | // Schedule the task by pushing it into the queue. 110 | runnable.schedule(); 111 | 112 | task 113 | } 114 | 115 | pub fn block_on(future: F) 116 | where 117 | F: Future + 'static, 118 | { 119 | let task = spawn(future); 120 | while !task.is_finished() { 121 | let Some(runnable) = QUEUE.with(|queue| queue.borrow_mut().pop()) else { 122 | thread::yield_now(); 123 | continue; 124 | }; 125 | runnable.0.run(); 126 | } 127 | } 128 | 129 | fn main() { 130 | // Spawn a future and await its result. 131 | block_on(async { 132 | let (sender, receiver) = channel::bounded(1); 133 | let world = spawn(async move { 134 | receiver.recv().await.unwrap(); 135 | println!("world.") 136 | }); 137 | let hello = spawn(async move { 138 | sender.send(()).await.unwrap(); 139 | print!("Hello, ") 140 | }); 141 | future::zip(hello, world).await; 142 | }); 143 | } 144 | -------------------------------------------------------------------------------- /tests/cancel.rs: -------------------------------------------------------------------------------- 1 | use std::future::Future; 2 | use std::pin::Pin; 3 | use std::sync::atomic::{AtomicUsize, Ordering}; 4 | use std::task::{Context, Poll}; 5 | use std::thread; 6 | use std::time::Duration; 7 | 8 | use async_task::Runnable; 9 | use easy_parallel::Parallel; 10 | use smol::future; 11 | 12 | // Creates a future with event counters. 13 | // 14 | // Usage: `future!(f, POLL, DROP_F, DROP_T)` 15 | // 16 | // The future `f` outputs `Poll::Ready`. 17 | // When it gets polled, `POLL` is incremented. 18 | // When it gets dropped, `DROP_F` is incremented. 19 | // When the output gets dropped, `DROP_T` is incremented. 20 | macro_rules! future { 21 | ($name:pat, $poll:ident, $drop_f:ident, $drop_t:ident) => { 22 | static $poll: AtomicUsize = AtomicUsize::new(0); 23 | static $drop_f: AtomicUsize = AtomicUsize::new(0); 24 | static $drop_t: AtomicUsize = AtomicUsize::new(0); 25 | 26 | let $name = { 27 | struct Fut(#[allow(dead_code)] Box); 28 | 29 | impl Future for Fut { 30 | type Output = Out; 31 | 32 | fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { 33 | $poll.fetch_add(1, Ordering::SeqCst); 34 | thread::sleep(ms(400)); 35 | Poll::Ready(Out(Box::new(0), true)) 36 | } 37 | } 38 | 39 | impl Drop for Fut { 40 | fn drop(&mut self) { 41 | $drop_f.fetch_add(1, Ordering::SeqCst); 42 | } 43 | } 44 | 45 | #[derive(Default)] 46 | struct Out(#[allow(dead_code)] Box, bool); 47 | 48 | impl Drop for Out { 49 | fn drop(&mut self) { 50 | if self.1 { 51 | $drop_t.fetch_add(1, Ordering::SeqCst); 52 | } 53 | } 54 | } 55 | 56 | Fut(Box::new(0)) 57 | }; 58 | }; 59 | } 60 | 61 | // Creates a schedule function with event counters. 62 | // 63 | // Usage: `schedule!(s, SCHED, DROP)` 64 | // 65 | // The schedule function `s` does nothing. 66 | // When it gets invoked, `SCHED` is incremented. 67 | // When it gets dropped, `DROP` is incremented. 68 | macro_rules! schedule { 69 | ($name:pat, $sched:ident, $drop:ident) => { 70 | static $drop: AtomicUsize = AtomicUsize::new(0); 71 | static $sched: AtomicUsize = AtomicUsize::new(0); 72 | 73 | let $name = { 74 | struct Guard(#[allow(dead_code)] Box); 75 | 76 | impl Drop for Guard { 77 | fn drop(&mut self) { 78 | $drop.fetch_add(1, Ordering::SeqCst); 79 | } 80 | } 81 | 82 | let guard = Guard(Box::new(0)); 83 | move |runnable: Runnable| { 84 | let _ = &guard; 85 | runnable.schedule(); 86 | $sched.fetch_add(1, Ordering::SeqCst); 87 | } 88 | }; 89 | }; 90 | } 91 | 92 | fn ms(ms: u64) -> Duration { 93 | Duration::from_millis(ms) 94 | } 95 | 96 | #[test] 97 | fn run_and_cancel() { 98 | future!(f, POLL, DROP_F, DROP_T); 99 | schedule!(s, SCHEDULE, DROP_S); 100 | let (runnable, task) = async_task::spawn(f, s); 101 | 102 | runnable.run(); 103 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 104 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 105 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 106 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 107 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 108 | 109 | assert!(future::block_on(task.cancel()).is_some()); 110 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 111 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 112 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 113 | assert_eq!(DROP_T.load(Ordering::SeqCst), 1); 114 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 115 | } 116 | 117 | #[test] 118 | fn cancel_and_run() { 119 | future!(f, POLL, DROP_F, DROP_T); 120 | schedule!(s, SCHEDULE, DROP_S); 121 | let (runnable, task) = async_task::spawn(f, s); 122 | 123 | Parallel::new() 124 | .add(|| { 125 | thread::sleep(ms(200)); 126 | runnable.run(); 127 | 128 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 129 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 130 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 131 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 132 | 133 | thread::sleep(ms(200)); 134 | 135 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 136 | }) 137 | .add(|| { 138 | assert!(future::block_on(task.cancel()).is_none()); 139 | 140 | thread::sleep(ms(200)); 141 | 142 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 143 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 144 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 145 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 146 | 147 | thread::sleep(ms(200)); 148 | 149 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 150 | }) 151 | .run(); 152 | } 153 | 154 | #[test] 155 | fn cancel_during_run() { 156 | future!(f, POLL, DROP_F, DROP_T); 157 | schedule!(s, SCHEDULE, DROP_S); 158 | let (runnable, task) = async_task::spawn(f, s); 159 | 160 | Parallel::new() 161 | .add(|| { 162 | runnable.run(); 163 | 164 | thread::sleep(ms(200)); 165 | 166 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 167 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 168 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 169 | assert_eq!(DROP_T.load(Ordering::SeqCst), 1); 170 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 171 | }) 172 | .add(|| { 173 | thread::sleep(ms(200)); 174 | 175 | assert!(future::block_on(task.cancel()).is_none()); 176 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 177 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 178 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 179 | assert_eq!(DROP_T.load(Ordering::SeqCst), 1); 180 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 181 | }) 182 | .run(); 183 | } 184 | -------------------------------------------------------------------------------- /tests/ready.rs: -------------------------------------------------------------------------------- 1 | use std::future::Future; 2 | use std::pin::Pin; 3 | use std::sync::atomic::{AtomicUsize, Ordering}; 4 | use std::task::{Context, Poll}; 5 | use std::thread; 6 | use std::time::Duration; 7 | 8 | use async_task::Runnable; 9 | use easy_parallel::Parallel; 10 | use smol::future; 11 | 12 | // Creates a future with event counters. 13 | // 14 | // Usage: `future!(f, POLL, DROP_F, DROP_T)` 15 | // 16 | // The future `f` sleeps for 200 ms and outputs `Poll::Ready`. 17 | // When it gets polled, `POLL` is incremented. 18 | // When it gets dropped, `DROP_F` is incremented. 19 | // When the output gets dropped, `DROP_T` is incremented. 20 | macro_rules! future { 21 | ($name:pat, $poll:ident, $drop_f:ident, $drop_t:ident) => { 22 | static $poll: AtomicUsize = AtomicUsize::new(0); 23 | static $drop_f: AtomicUsize = AtomicUsize::new(0); 24 | static $drop_t: AtomicUsize = AtomicUsize::new(0); 25 | 26 | let $name = { 27 | struct Fut(#[allow(dead_code)] Box); 28 | 29 | impl Future for Fut { 30 | type Output = Out; 31 | 32 | fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { 33 | $poll.fetch_add(1, Ordering::SeqCst); 34 | thread::sleep(ms(400)); 35 | Poll::Ready(Out(Box::new(0), true)) 36 | } 37 | } 38 | 39 | impl Drop for Fut { 40 | fn drop(&mut self) { 41 | $drop_f.fetch_add(1, Ordering::SeqCst); 42 | } 43 | } 44 | 45 | #[derive(Default)] 46 | struct Out(#[allow(dead_code)] Box, bool); 47 | 48 | impl Drop for Out { 49 | fn drop(&mut self) { 50 | if self.1 { 51 | $drop_t.fetch_add(1, Ordering::SeqCst); 52 | } 53 | } 54 | } 55 | 56 | Fut(Box::new(0)) 57 | }; 58 | }; 59 | } 60 | 61 | // Creates a schedule function with event counters. 62 | // 63 | // Usage: `schedule!(s, SCHED, DROP)` 64 | // 65 | // The schedule function `s` does nothing. 66 | // When it gets invoked, `SCHED` is incremented. 67 | // When it gets dropped, `DROP` is incremented. 68 | macro_rules! schedule { 69 | ($name:pat, $sched:ident, $drop:ident) => { 70 | static $drop: AtomicUsize = AtomicUsize::new(0); 71 | static $sched: AtomicUsize = AtomicUsize::new(0); 72 | 73 | let $name = { 74 | struct Guard(#[allow(dead_code)] Box); 75 | 76 | impl Drop for Guard { 77 | fn drop(&mut self) { 78 | $drop.fetch_add(1, Ordering::SeqCst); 79 | } 80 | } 81 | 82 | let guard = Guard(Box::new(0)); 83 | move |_runnable: Runnable| { 84 | let _ = &guard; 85 | $sched.fetch_add(1, Ordering::SeqCst); 86 | } 87 | }; 88 | }; 89 | } 90 | 91 | fn ms(ms: u64) -> Duration { 92 | Duration::from_millis(ms) 93 | } 94 | 95 | #[test] 96 | fn cancel_during_run() { 97 | future!(f, POLL, DROP_F, DROP_T); 98 | schedule!(s, SCHEDULE, DROP_S); 99 | let (runnable, task) = async_task::spawn(f, s); 100 | 101 | Parallel::new() 102 | .add(|| { 103 | runnable.run(); 104 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 105 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 106 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 107 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 108 | assert_eq!(DROP_T.load(Ordering::SeqCst), 1); 109 | }) 110 | .add(|| { 111 | thread::sleep(ms(200)); 112 | 113 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 114 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 115 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 116 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 117 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 118 | 119 | drop(task); 120 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 121 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 122 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 123 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 124 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 125 | 126 | thread::sleep(ms(400)); 127 | 128 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 129 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 130 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 131 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 132 | assert_eq!(DROP_T.load(Ordering::SeqCst), 1); 133 | }) 134 | .run(); 135 | } 136 | 137 | #[test] 138 | fn join_during_run() { 139 | future!(f, POLL, DROP_F, DROP_T); 140 | schedule!(s, SCHEDULE, DROP_S); 141 | let (runnable, task) = async_task::spawn(f, s); 142 | 143 | Parallel::new() 144 | .add(|| { 145 | runnable.run(); 146 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 147 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 148 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 149 | 150 | thread::sleep(ms(200)); 151 | 152 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 153 | }) 154 | .add(|| { 155 | thread::sleep(ms(200)); 156 | 157 | future::block_on(task); 158 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 159 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 160 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 161 | assert_eq!(DROP_T.load(Ordering::SeqCst), 1); 162 | 163 | thread::sleep(ms(200)); 164 | 165 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 166 | }) 167 | .run(); 168 | } 169 | 170 | #[test] 171 | fn try_join_during_run() { 172 | future!(f, POLL, DROP_F, DROP_T); 173 | schedule!(s, SCHEDULE, DROP_S); 174 | let (runnable, mut task) = async_task::spawn(f, s); 175 | 176 | Parallel::new() 177 | .add(|| { 178 | runnable.run(); 179 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 180 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 181 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 182 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 183 | assert_eq!(DROP_T.load(Ordering::SeqCst), 1); 184 | }) 185 | .add(|| { 186 | thread::sleep(ms(200)); 187 | 188 | future::block_on(future::or(&mut task, future::ready(Default::default()))); 189 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 190 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 191 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 192 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 193 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 194 | drop(task); 195 | }) 196 | .run(); 197 | } 198 | 199 | #[test] 200 | fn detach_during_run() { 201 | future!(f, POLL, DROP_F, DROP_T); 202 | schedule!(s, SCHEDULE, DROP_S); 203 | let (runnable, task) = async_task::spawn(f, s); 204 | 205 | Parallel::new() 206 | .add(|| { 207 | runnable.run(); 208 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 209 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 210 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 211 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 212 | assert_eq!(DROP_T.load(Ordering::SeqCst), 1); 213 | }) 214 | .add(|| { 215 | thread::sleep(ms(200)); 216 | 217 | task.detach(); 218 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 219 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 220 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 221 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 222 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 223 | }) 224 | .run(); 225 | } 226 | -------------------------------------------------------------------------------- /tests/panic.rs: -------------------------------------------------------------------------------- 1 | use std::future::Future; 2 | use std::panic::catch_unwind; 3 | use std::pin::Pin; 4 | use std::sync::atomic::{AtomicUsize, Ordering}; 5 | use std::task::{Context, Poll}; 6 | use std::thread; 7 | use std::time::Duration; 8 | 9 | use async_task::Runnable; 10 | use easy_parallel::Parallel; 11 | use smol::future; 12 | 13 | // Creates a future with event counters. 14 | // 15 | // Usage: `future!(f, POLL, DROP)` 16 | // 17 | // The future `f` sleeps for 200 ms and then panics. 18 | // When it gets polled, `POLL` is incremented. 19 | // When it gets dropped, `DROP` is incremented. 20 | macro_rules! future { 21 | ($name:pat, $poll:ident, $drop:ident) => { 22 | static $poll: AtomicUsize = AtomicUsize::new(0); 23 | static $drop: AtomicUsize = AtomicUsize::new(0); 24 | 25 | let $name = { 26 | struct Fut(#[allow(dead_code)] Box); 27 | 28 | impl Future for Fut { 29 | type Output = (); 30 | 31 | fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { 32 | $poll.fetch_add(1, Ordering::SeqCst); 33 | thread::sleep(ms(400)); 34 | panic!() 35 | } 36 | } 37 | 38 | impl Drop for Fut { 39 | fn drop(&mut self) { 40 | $drop.fetch_add(1, Ordering::SeqCst); 41 | } 42 | } 43 | 44 | Fut(Box::new(0)) 45 | }; 46 | }; 47 | } 48 | 49 | // Creates a schedule function with event counters. 50 | // 51 | // Usage: `schedule!(s, SCHED, DROP)` 52 | // 53 | // The schedule function `s` does nothing. 54 | // When it gets invoked, `SCHED` is incremented. 55 | // When it gets dropped, `DROP` is incremented. 56 | macro_rules! schedule { 57 | ($name:pat, $sched:ident, $drop:ident) => { 58 | static $drop: AtomicUsize = AtomicUsize::new(0); 59 | static $sched: AtomicUsize = AtomicUsize::new(0); 60 | 61 | let $name = { 62 | struct Guard(#[allow(dead_code)] Box); 63 | 64 | impl Drop for Guard { 65 | fn drop(&mut self) { 66 | $drop.fetch_add(1, Ordering::SeqCst); 67 | } 68 | } 69 | 70 | let guard = Guard(Box::new(0)); 71 | move |_runnable: Runnable| { 72 | let _ = &guard; 73 | $sched.fetch_add(1, Ordering::SeqCst); 74 | } 75 | }; 76 | }; 77 | } 78 | 79 | fn ms(ms: u64) -> Duration { 80 | Duration::from_millis(ms) 81 | } 82 | 83 | #[test] 84 | fn cancel_during_run() { 85 | future!(f, POLL, DROP_F); 86 | schedule!(s, SCHEDULE, DROP_S); 87 | let (runnable, task) = async_task::spawn(f, s); 88 | 89 | Parallel::new() 90 | .add(|| { 91 | assert!(catch_unwind(|| runnable.run()).is_err()); 92 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 93 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 94 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 95 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 96 | }) 97 | .add(|| { 98 | thread::sleep(ms(200)); 99 | 100 | drop(task); 101 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 102 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 103 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 104 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 105 | }) 106 | .run(); 107 | } 108 | 109 | #[test] 110 | fn run_and_join() { 111 | future!(f, POLL, DROP_F); 112 | schedule!(s, SCHEDULE, DROP_S); 113 | let (runnable, task) = async_task::spawn(f, s); 114 | 115 | assert!(catch_unwind(|| runnable.run()).is_err()); 116 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 117 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 118 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 119 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 120 | 121 | assert!(catch_unwind(|| future::block_on(task)).is_err()); 122 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 123 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 124 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 125 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 126 | } 127 | 128 | #[test] 129 | fn try_join_and_run_and_join() { 130 | future!(f, POLL, DROP_F); 131 | schedule!(s, SCHEDULE, DROP_S); 132 | let (runnable, mut task) = async_task::spawn(f, s); 133 | 134 | future::block_on(future::or(&mut task, future::ready(()))); 135 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 136 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 137 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 138 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 139 | 140 | assert!(catch_unwind(|| runnable.run()).is_err()); 141 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 142 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 143 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 144 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 145 | 146 | assert!(catch_unwind(|| future::block_on(task)).is_err()); 147 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 148 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 149 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 150 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 151 | } 152 | 153 | #[test] 154 | fn join_during_run() { 155 | future!(f, POLL, DROP_F); 156 | schedule!(s, SCHEDULE, DROP_S); 157 | let (runnable, task) = async_task::spawn(f, s); 158 | 159 | Parallel::new() 160 | .add(|| { 161 | assert!(catch_unwind(|| runnable.run()).is_err()); 162 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 163 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 164 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 165 | 166 | thread::sleep(ms(200)); 167 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 168 | }) 169 | .add(|| { 170 | thread::sleep(ms(200)); 171 | 172 | assert!(catch_unwind(|| future::block_on(task)).is_err()); 173 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 174 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 175 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 176 | 177 | thread::sleep(ms(200)); 178 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 179 | }) 180 | .run(); 181 | } 182 | 183 | #[test] 184 | fn try_join_during_run() { 185 | future!(f, POLL, DROP_F); 186 | schedule!(s, SCHEDULE, DROP_S); 187 | let (runnable, mut task) = async_task::spawn(f, s); 188 | 189 | Parallel::new() 190 | .add(|| { 191 | assert!(catch_unwind(|| runnable.run()).is_err()); 192 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 193 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 194 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 195 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 196 | }) 197 | .add(|| { 198 | thread::sleep(ms(200)); 199 | 200 | future::block_on(future::or(&mut task, future::ready(()))); 201 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 202 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 203 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 204 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 205 | drop(task); 206 | }) 207 | .run(); 208 | } 209 | 210 | #[test] 211 | fn detach_during_run() { 212 | future!(f, POLL, DROP_F); 213 | schedule!(s, SCHEDULE, DROP_S); 214 | let (runnable, task) = async_task::spawn(f, s); 215 | 216 | Parallel::new() 217 | .add(|| { 218 | assert!(catch_unwind(|| runnable.run()).is_err()); 219 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 220 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 221 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 222 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 223 | }) 224 | .add(|| { 225 | thread::sleep(ms(200)); 226 | 227 | task.detach(); 228 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 229 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 230 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 231 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 232 | }) 233 | .run(); 234 | } 235 | -------------------------------------------------------------------------------- /src/header.rs: -------------------------------------------------------------------------------- 1 | use core::cell::UnsafeCell; 2 | use core::fmt; 3 | use core::task::{RawWaker, Waker}; 4 | 5 | #[cfg(not(feature = "portable-atomic"))] 6 | use core::sync::atomic::AtomicUsize; 7 | use core::sync::atomic::Ordering; 8 | #[cfg(feature = "portable-atomic")] 9 | use portable_atomic::AtomicUsize; 10 | 11 | use crate::raw::TaskVTable; 12 | use crate::state::*; 13 | use crate::utils::abort; 14 | use crate::utils::abort_on_panic; 15 | 16 | /// Actions to take upon calling [`Header::drop_waker`]. 17 | pub(crate) enum DropWakerAction { 18 | /// Re-schedule the task. 19 | Schedule, 20 | /// Destroy the task. 21 | Destroy, 22 | /// Do nothing. 23 | None, 24 | } 25 | 26 | pub(crate) struct Header { 27 | /// Current state of the task. 28 | /// 29 | /// Contains flags representing the current state and the reference count. 30 | pub(crate) state: AtomicUsize, 31 | 32 | /// The task that is blocked on the `Task` handle. 33 | /// 34 | /// This waker needs to be woken up once the task completes or is closed. 35 | pub(crate) awaiter: UnsafeCell>, 36 | 37 | /// The virtual table. 38 | /// 39 | /// In addition to the actual waker virtual table, it also contains pointers to several other 40 | /// methods necessary for bookkeeping the heap-allocated task. 41 | pub(crate) vtable: &'static TaskVTable, 42 | 43 | /// Whether or not a panic that occurs in the task should be propagated. 44 | #[cfg(feature = "std")] 45 | pub(crate) propagate_panic: bool, 46 | } 47 | 48 | impl Header { 49 | /// Notifies the awaiter blocked on this task. 50 | /// 51 | /// If the awaiter is the same as the current waker, it will not be notified. 52 | #[inline] 53 | pub(crate) fn notify(&self, current: Option<&Waker>) { 54 | if let Some(w) = self.take(current) { 55 | abort_on_panic(|| w.wake()); 56 | } 57 | } 58 | 59 | /// Takes the awaiter blocked on this task. 60 | /// 61 | /// If there is no awaiter or if it is the same as the current waker, returns `None`. 62 | #[inline] 63 | pub(crate) fn take(&self, current: Option<&Waker>) -> Option { 64 | // Set the bit indicating that the task is notifying its awaiter. 65 | let state = self.state.fetch_or(NOTIFYING, Ordering::AcqRel); 66 | 67 | // If the task was not notifying or registering an awaiter... 68 | if state & (NOTIFYING | REGISTERING) == 0 { 69 | // Take the waker out. 70 | let waker = unsafe { (*self.awaiter.get()).take() }; 71 | 72 | // Unset the bit indicating that the task is notifying its awaiter. 73 | self.state 74 | .fetch_and(!NOTIFYING & !AWAITER, Ordering::Release); 75 | 76 | // Finally, notify the waker if it's different from the current waker. 77 | if let Some(w) = waker { 78 | match current { 79 | None => return Some(w), 80 | Some(c) if !w.will_wake(c) => return Some(w), 81 | Some(_) => abort_on_panic(|| drop(w)), 82 | } 83 | } 84 | } 85 | 86 | None 87 | } 88 | 89 | /// Registers a new awaiter blocked on this task. 90 | /// 91 | /// This method is called when `Task` is polled and it has not yet completed. 92 | #[inline] 93 | pub(crate) fn register(&self, waker: &Waker) { 94 | // Load the state and synchronize with it. 95 | let mut state = self.state.fetch_or(0, Ordering::Acquire); 96 | 97 | loop { 98 | // There can't be two concurrent registrations because `Task` can only be polled 99 | // by a unique pinned reference. 100 | debug_assert!(state & REGISTERING == 0); 101 | 102 | // If we're in the notifying state at this moment, just wake and return without 103 | // registering. 104 | if state & NOTIFYING != 0 { 105 | abort_on_panic(|| waker.wake_by_ref()); 106 | return; 107 | } 108 | 109 | // Mark the state to let other threads know we're registering a new awaiter. 110 | match self.state.compare_exchange_weak( 111 | state, 112 | state | REGISTERING, 113 | Ordering::AcqRel, 114 | Ordering::Acquire, 115 | ) { 116 | Ok(_) => { 117 | state |= REGISTERING; 118 | break; 119 | } 120 | Err(s) => state = s, 121 | } 122 | } 123 | 124 | // Put the waker into the awaiter field. 125 | unsafe { 126 | abort_on_panic(|| (*self.awaiter.get()) = Some(waker.clone())); 127 | } 128 | 129 | // This variable will contain the newly registered waker if a notification comes in before 130 | // we complete registration. 131 | let mut waker = None; 132 | 133 | loop { 134 | // If there was a notification, take the waker out of the awaiter field. 135 | if state & NOTIFYING != 0 { 136 | if let Some(w) = unsafe { (*self.awaiter.get()).take() } { 137 | abort_on_panic(|| waker = Some(w)); 138 | } 139 | } 140 | 141 | // The new state is not being notified nor registered, but there might or might not be 142 | // an awaiter depending on whether there was a concurrent notification. 143 | let new = if waker.is_none() { 144 | (state & !NOTIFYING & !REGISTERING) | AWAITER 145 | } else { 146 | state & !NOTIFYING & !REGISTERING & !AWAITER 147 | }; 148 | 149 | match self 150 | .state 151 | .compare_exchange_weak(state, new, Ordering::AcqRel, Ordering::Acquire) 152 | { 153 | Ok(_) => break, 154 | Err(s) => state = s, 155 | } 156 | } 157 | 158 | // If there was a notification during registration, wake the awaiter now. 159 | if let Some(w) = waker { 160 | abort_on_panic(|| w.wake()); 161 | } 162 | } 163 | 164 | /// Clones a waker. 165 | pub(crate) unsafe fn clone_waker(ptr: *const ()) -> RawWaker { 166 | let header = ptr as *const Header; 167 | 168 | // Increment the reference count. With any kind of reference-counted data structure, 169 | // relaxed ordering is appropriate when incrementing the counter. 170 | let state = (*header).state.fetch_add(REFERENCE, Ordering::Relaxed); 171 | 172 | // If the reference count overflowed, abort. 173 | if state > isize::MAX as usize { 174 | abort(); 175 | } 176 | 177 | RawWaker::new(ptr, (*header).vtable.raw_waker_vtable) 178 | } 179 | 180 | #[inline(never)] 181 | pub(crate) unsafe fn drop_waker(ptr: *const ()) -> DropWakerAction { 182 | let header = ptr as *const Header; 183 | 184 | // Decrement the reference count. 185 | let new = (*header).state.fetch_sub(REFERENCE, Ordering::AcqRel) - REFERENCE; 186 | 187 | // If this was the last reference to the task and the `Task` has been dropped too, 188 | // then we need to decide how to destroy the task. 189 | if new & !(REFERENCE - 1) == 0 && new & TASK == 0 { 190 | if new & (COMPLETED | CLOSED) == 0 { 191 | // If the task was not completed nor closed, close it and schedule one more time so 192 | // that its future gets dropped by the executor. 193 | (*header) 194 | .state 195 | .store(SCHEDULED | CLOSED | REFERENCE, Ordering::Release); 196 | DropWakerAction::Schedule 197 | } else { 198 | // Otherwise, destroy the task right away. 199 | DropWakerAction::Destroy 200 | } 201 | } else { 202 | DropWakerAction::None 203 | } 204 | } 205 | } 206 | 207 | // SAFETY: repr(C) is explicitly used here so that casts between `Header` and `HeaderWithMetadata` 208 | // can be done safely without additional offsets. 209 | // 210 | /// The header of a task. 211 | /// 212 | /// This header is stored in memory at the beginning of the heap-allocated task. 213 | #[repr(C)] 214 | pub(crate) struct HeaderWithMetadata { 215 | pub(crate) header: Header, 216 | 217 | /// Metadata associated with the task. 218 | /// 219 | /// This metadata may be provided to the user. 220 | pub(crate) metadata: M, 221 | } 222 | 223 | impl fmt::Debug for HeaderWithMetadata { 224 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 225 | let state = self.header.state.load(Ordering::SeqCst); 226 | 227 | f.debug_struct("Header") 228 | .field("scheduled", &(state & SCHEDULED != 0)) 229 | .field("running", &(state & RUNNING != 0)) 230 | .field("completed", &(state & COMPLETED != 0)) 231 | .field("closed", &(state & CLOSED != 0)) 232 | .field("awaiter", &(state & AWAITER != 0)) 233 | .field("task", &(state & TASK != 0)) 234 | .field("ref_count", &(state / REFERENCE)) 235 | .field("metadata", &self.metadata) 236 | .finish() 237 | } 238 | } 239 | -------------------------------------------------------------------------------- /tests/waker_ready.rs: -------------------------------------------------------------------------------- 1 | use std::cell::Cell; 2 | use std::future::Future; 3 | use std::pin::Pin; 4 | use std::sync::atomic::{AtomicUsize, Ordering}; 5 | use std::task::{Context, Poll}; 6 | use std::thread; 7 | use std::time::Duration; 8 | 9 | use async_task::Runnable; 10 | use atomic_waker::AtomicWaker; 11 | 12 | // Creates a future with event counters. 13 | // 14 | // Usage: `future!(f, get_waker, POLL, DROP)` 15 | // 16 | // The future `f` always sleeps for 200 ms, and returns `Poll::Ready` the second time it is polled. 17 | // When it gets polled, `POLL` is incremented. 18 | // When it gets dropped, `DROP` is incremented. 19 | // 20 | // Every time the future is run, it stores the waker into a global variable. 21 | // This waker can be extracted using the `get_waker()` function. 22 | macro_rules! future { 23 | ($name:pat, $get_waker:pat, $poll:ident, $drop:ident) => { 24 | static $poll: AtomicUsize = AtomicUsize::new(0); 25 | static $drop: AtomicUsize = AtomicUsize::new(0); 26 | static WAKER: AtomicWaker = AtomicWaker::new(); 27 | 28 | let ($name, $get_waker) = { 29 | struct Fut(Cell, #[allow(dead_code)] Box); 30 | 31 | impl Future for Fut { 32 | type Output = Box; 33 | 34 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 35 | WAKER.register(cx.waker()); 36 | $poll.fetch_add(1, Ordering::SeqCst); 37 | thread::sleep(ms(200)); 38 | 39 | if self.0.get() { 40 | Poll::Ready(Box::new(0)) 41 | } else { 42 | self.0.set(true); 43 | Poll::Pending 44 | } 45 | } 46 | } 47 | 48 | impl Drop for Fut { 49 | fn drop(&mut self) { 50 | $drop.fetch_add(1, Ordering::SeqCst); 51 | } 52 | } 53 | 54 | (Fut(Cell::new(false), Box::new(0)), || WAKER.take().unwrap()) 55 | }; 56 | }; 57 | } 58 | 59 | // Creates a schedule function with event counters. 60 | // 61 | // Usage: `schedule!(s, chan, SCHED, DROP)` 62 | // 63 | // The schedule function `s` pushes the task into `chan`. 64 | // When it gets invoked, `SCHED` is incremented. 65 | // When it gets dropped, `DROP` is incremented. 66 | // 67 | // Receiver `chan` extracts the task when it is scheduled. 68 | macro_rules! schedule { 69 | ($name:pat, $chan:pat, $sched:ident, $drop:ident) => { 70 | static $drop: AtomicUsize = AtomicUsize::new(0); 71 | static $sched: AtomicUsize = AtomicUsize::new(0); 72 | 73 | let ($name, $chan) = { 74 | let (s, r) = flume::unbounded(); 75 | 76 | struct Guard(#[allow(dead_code)] Box); 77 | 78 | impl Drop for Guard { 79 | fn drop(&mut self) { 80 | $drop.fetch_add(1, Ordering::SeqCst); 81 | } 82 | } 83 | 84 | let guard = Guard(Box::new(0)); 85 | let sched = move |runnable: Runnable| { 86 | let _ = &guard; 87 | $sched.fetch_add(1, Ordering::SeqCst); 88 | s.send(runnable).unwrap(); 89 | }; 90 | 91 | (sched, r) 92 | }; 93 | }; 94 | } 95 | 96 | fn ms(ms: u64) -> Duration { 97 | Duration::from_millis(ms) 98 | } 99 | 100 | #[test] 101 | fn wake() { 102 | future!(f, get_waker, POLL, DROP_F); 103 | schedule!(s, chan, SCHEDULE, DROP_S); 104 | let (mut runnable, task) = async_task::spawn(f, s); 105 | task.detach(); 106 | 107 | assert!(chan.is_empty()); 108 | 109 | runnable.run(); 110 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 111 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 112 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 113 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 114 | assert_eq!(chan.len(), 0); 115 | 116 | get_waker().wake(); 117 | runnable = chan.recv().unwrap(); 118 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 119 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 120 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 121 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 122 | assert_eq!(chan.len(), 0); 123 | 124 | runnable.run(); 125 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 126 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 127 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 128 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 129 | assert_eq!(chan.len(), 0); 130 | 131 | get_waker().wake(); 132 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 133 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 134 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 135 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 136 | assert_eq!(chan.len(), 0); 137 | } 138 | 139 | #[test] 140 | fn wake_by_ref() { 141 | future!(f, get_waker, POLL, DROP_F); 142 | schedule!(s, chan, SCHEDULE, DROP_S); 143 | let (mut runnable, task) = async_task::spawn(f, s); 144 | task.detach(); 145 | 146 | assert!(chan.is_empty()); 147 | 148 | runnable.run(); 149 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 150 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 151 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 152 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 153 | assert_eq!(chan.len(), 0); 154 | 155 | get_waker().wake_by_ref(); 156 | runnable = chan.recv().unwrap(); 157 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 158 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 159 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 160 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 161 | assert_eq!(chan.len(), 0); 162 | 163 | runnable.run(); 164 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 165 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 166 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 167 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 168 | assert_eq!(chan.len(), 0); 169 | 170 | get_waker().wake_by_ref(); 171 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 172 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 173 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 174 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 175 | assert_eq!(chan.len(), 0); 176 | } 177 | 178 | #[allow(clippy::redundant_clone)] // This is intentional 179 | #[test] 180 | fn clone() { 181 | future!(f, get_waker, POLL, DROP_F); 182 | schedule!(s, chan, SCHEDULE, DROP_S); 183 | let (mut runnable, task) = async_task::spawn(f, s); 184 | task.detach(); 185 | 186 | runnable.run(); 187 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 188 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 189 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 190 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 191 | assert_eq!(chan.len(), 0); 192 | 193 | let w2 = get_waker().clone(); 194 | let w3 = w2.clone(); 195 | let w4 = w3.clone(); 196 | w4.wake(); 197 | 198 | runnable = chan.recv().unwrap(); 199 | runnable.run(); 200 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 201 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 202 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 203 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 204 | assert_eq!(chan.len(), 0); 205 | 206 | w3.wake(); 207 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 208 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 209 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 210 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 211 | assert_eq!(chan.len(), 0); 212 | 213 | drop(w2); 214 | drop(get_waker()); 215 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 216 | } 217 | 218 | #[test] 219 | fn wake_dropped() { 220 | future!(f, get_waker, POLL, DROP_F); 221 | schedule!(s, chan, SCHEDULE, DROP_S); 222 | let (runnable, task) = async_task::spawn(f, s); 223 | task.detach(); 224 | 225 | runnable.run(); 226 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 227 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 228 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 229 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 230 | assert_eq!(chan.len(), 0); 231 | 232 | let waker = get_waker(); 233 | 234 | waker.wake_by_ref(); 235 | drop(chan.recv().unwrap()); 236 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 237 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 238 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 239 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 240 | assert_eq!(chan.len(), 0); 241 | 242 | waker.wake(); 243 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 244 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 245 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 246 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 247 | assert_eq!(chan.len(), 0); 248 | } 249 | 250 | #[test] 251 | fn wake_completed() { 252 | future!(f, get_waker, POLL, DROP_F); 253 | schedule!(s, chan, SCHEDULE, DROP_S); 254 | let (runnable, task) = async_task::spawn(f, s); 255 | task.detach(); 256 | 257 | runnable.run(); 258 | let waker = get_waker(); 259 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 260 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 261 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 262 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 263 | assert_eq!(chan.len(), 0); 264 | 265 | waker.wake(); 266 | chan.recv().unwrap().run(); 267 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 268 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 269 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 270 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 271 | assert_eq!(chan.len(), 0); 272 | 273 | get_waker().wake(); 274 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 275 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 276 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 277 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 278 | assert_eq!(chan.len(), 0); 279 | } 280 | -------------------------------------------------------------------------------- /tests/basic.rs: -------------------------------------------------------------------------------- 1 | use std::future::Future; 2 | use std::pin::Pin; 3 | use std::ptr::NonNull; 4 | use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; 5 | use std::sync::Arc; 6 | use std::task::{Context, Poll}; 7 | 8 | use async_task::Runnable; 9 | use smol::future; 10 | 11 | // Creates a future with event counters. 12 | // 13 | // Usage: `future!(f, POLL, DROP)` 14 | // 15 | // The future `f` always returns `Poll::Ready`. 16 | // When it gets polled, `POLL` is incremented. 17 | // When it gets dropped, `DROP` is incremented. 18 | macro_rules! future { 19 | ($name:pat, $poll:ident, $drop:ident) => { 20 | static $poll: AtomicUsize = AtomicUsize::new(0); 21 | static $drop: AtomicUsize = AtomicUsize::new(0); 22 | 23 | let $name = { 24 | struct Fut(#[allow(dead_code)] Box); 25 | 26 | impl Future for Fut { 27 | type Output = Box; 28 | 29 | fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { 30 | $poll.fetch_add(1, Ordering::SeqCst); 31 | Poll::Ready(Box::new(0)) 32 | } 33 | } 34 | 35 | impl Drop for Fut { 36 | fn drop(&mut self) { 37 | $drop.fetch_add(1, Ordering::SeqCst); 38 | } 39 | } 40 | 41 | Fut(Box::new(0)) 42 | }; 43 | }; 44 | } 45 | 46 | // Creates a schedule function with event counters. 47 | // 48 | // Usage: `schedule!(s, SCHED, DROP)` 49 | // 50 | // The schedule function `s` does nothing. 51 | // When it gets invoked, `SCHED` is incremented. 52 | // When it gets dropped, `DROP` is incremented. 53 | macro_rules! schedule { 54 | ($name:pat, $sched:ident, $drop:ident) => { 55 | static $drop: AtomicUsize = AtomicUsize::new(0); 56 | static $sched: AtomicUsize = AtomicUsize::new(0); 57 | 58 | let $name = { 59 | struct Guard(#[allow(dead_code)] Box); 60 | 61 | impl Drop for Guard { 62 | fn drop(&mut self) { 63 | $drop.fetch_add(1, Ordering::SeqCst); 64 | } 65 | } 66 | 67 | let guard = Guard(Box::new(0)); 68 | move |_runnable| { 69 | let _ = &guard; 70 | $sched.fetch_add(1, Ordering::SeqCst); 71 | } 72 | }; 73 | }; 74 | } 75 | 76 | fn try_await(f: impl Future) -> Option { 77 | future::block_on(future::poll_once(f)) 78 | } 79 | 80 | #[test] 81 | fn drop_and_detach() { 82 | future!(f, POLL, DROP_F); 83 | schedule!(s, SCHEDULE, DROP_S); 84 | let (runnable, task) = async_task::spawn(f, s); 85 | 86 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 87 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 88 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 89 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 90 | 91 | drop(runnable); 92 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 93 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 94 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 95 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 96 | 97 | task.detach(); 98 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 99 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 100 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 101 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 102 | } 103 | 104 | #[test] 105 | fn detach_and_drop() { 106 | future!(f, POLL, DROP_F); 107 | schedule!(s, SCHEDULE, DROP_S); 108 | let (runnable, task) = async_task::spawn(f, s); 109 | 110 | task.detach(); 111 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 112 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 113 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 114 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 115 | 116 | drop(runnable); 117 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 118 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 119 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 120 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 121 | } 122 | 123 | #[test] 124 | fn detach_and_run() { 125 | future!(f, POLL, DROP_F); 126 | schedule!(s, SCHEDULE, DROP_S); 127 | let (runnable, task) = async_task::spawn(f, s); 128 | 129 | task.detach(); 130 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 131 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 132 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 133 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 134 | 135 | runnable.run(); 136 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 137 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 138 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 139 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 140 | } 141 | 142 | #[test] 143 | fn run_and_detach() { 144 | future!(f, POLL, DROP_F); 145 | schedule!(s, SCHEDULE, DROP_S); 146 | let (runnable, task) = async_task::spawn(f, s); 147 | 148 | runnable.run(); 149 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 150 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 151 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 152 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 153 | 154 | task.detach(); 155 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 156 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 157 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 158 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 159 | } 160 | 161 | #[test] 162 | fn cancel_and_run() { 163 | future!(f, POLL, DROP_F); 164 | schedule!(s, SCHEDULE, DROP_S); 165 | let (runnable, task) = async_task::spawn(f, s); 166 | 167 | drop(task); 168 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 169 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 170 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 171 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 172 | 173 | runnable.run(); 174 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 175 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 176 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 177 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 178 | } 179 | 180 | #[test] 181 | fn run_and_cancel() { 182 | future!(f, POLL, DROP_F); 183 | schedule!(s, SCHEDULE, DROP_S); 184 | let (runnable, task) = async_task::spawn(f, s); 185 | 186 | runnable.run(); 187 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 188 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 189 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 190 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 191 | 192 | drop(task); 193 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 194 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 195 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 196 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 197 | } 198 | 199 | #[test] 200 | fn cancel_join() { 201 | future!(f, POLL, DROP_F); 202 | schedule!(s, SCHEDULE, DROP_S); 203 | let (runnable, mut task) = async_task::spawn(f, s); 204 | 205 | assert!(try_await(&mut task).is_none()); 206 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 207 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 208 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 209 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 210 | 211 | runnable.run(); 212 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 213 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 214 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 215 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 216 | 217 | assert!(try_await(&mut task).is_some()); 218 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 219 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 220 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 221 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 222 | 223 | drop(task); 224 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 225 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 226 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 227 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 228 | } 229 | 230 | #[test] 231 | fn schedule() { 232 | let (s, r) = flume::unbounded(); 233 | let schedule = move |runnable| s.send(runnable).unwrap(); 234 | let (runnable, _task) = async_task::spawn(future::poll_fn(|_| Poll::<()>::Pending), schedule); 235 | 236 | assert!(r.is_empty()); 237 | runnable.schedule(); 238 | 239 | let runnable = r.recv().unwrap(); 240 | assert!(r.is_empty()); 241 | runnable.schedule(); 242 | 243 | let runnable = r.recv().unwrap(); 244 | assert!(r.is_empty()); 245 | runnable.schedule(); 246 | 247 | r.recv().unwrap(); 248 | } 249 | 250 | #[test] 251 | fn schedule_counter() { 252 | static COUNT: AtomicUsize = AtomicUsize::new(0); 253 | 254 | let (s, r) = flume::unbounded(); 255 | let schedule = move |runnable: Runnable| { 256 | COUNT.fetch_add(1, Ordering::SeqCst); 257 | s.send(runnable).unwrap(); 258 | }; 259 | let (runnable, _task) = async_task::spawn(future::poll_fn(|_| Poll::<()>::Pending), schedule); 260 | runnable.schedule(); 261 | 262 | r.recv().unwrap().schedule(); 263 | r.recv().unwrap().schedule(); 264 | assert_eq!(COUNT.load(Ordering::SeqCst), 3); 265 | r.recv().unwrap(); 266 | } 267 | 268 | #[test] 269 | fn drop_inside_schedule() { 270 | struct DropGuard(AtomicUsize); 271 | impl Drop for DropGuard { 272 | fn drop(&mut self) { 273 | self.0.fetch_add(1, Ordering::SeqCst); 274 | } 275 | } 276 | let guard = DropGuard(AtomicUsize::new(0)); 277 | 278 | let (runnable, _) = async_task::spawn(async {}, move |runnable| { 279 | assert_eq!(guard.0.load(Ordering::SeqCst), 0); 280 | drop(runnable); 281 | assert_eq!(guard.0.load(Ordering::SeqCst), 0); 282 | }); 283 | runnable.schedule(); 284 | } 285 | 286 | #[test] 287 | fn waker() { 288 | let (s, r) = flume::unbounded(); 289 | let schedule = move |runnable| s.send(runnable).unwrap(); 290 | let (runnable, _task) = async_task::spawn(future::poll_fn(|_| Poll::<()>::Pending), schedule); 291 | 292 | assert!(r.is_empty()); 293 | let waker = runnable.waker(); 294 | runnable.run(); 295 | waker.wake_by_ref(); 296 | 297 | let runnable = r.recv().unwrap(); 298 | runnable.run(); 299 | waker.wake(); 300 | r.recv().unwrap(); 301 | } 302 | 303 | #[test] 304 | fn raw() { 305 | // Dispatch schedules a function for execution at a later point. For tests, we execute it straight away. 306 | fn dispatch(trampoline: extern "C" fn(NonNull<()>), context: NonNull<()>) { 307 | trampoline(context) 308 | } 309 | extern "C" fn trampoline(runnable: NonNull<()>) { 310 | let task = unsafe { Runnable::<()>::from_raw(runnable) }; 311 | task.run(); 312 | } 313 | 314 | let task_got_executed = Arc::new(AtomicBool::new(false)); 315 | let (runnable, _handle) = async_task::spawn( 316 | { 317 | let task_got_executed = task_got_executed.clone(); 318 | async move { task_got_executed.store(true, Ordering::SeqCst) } 319 | }, 320 | |runnable: Runnable<()>| dispatch(trampoline, runnable.into_raw()), 321 | ); 322 | runnable.schedule(); 323 | 324 | assert!(task_got_executed.load(Ordering::SeqCst)); 325 | } 326 | -------------------------------------------------------------------------------- /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 [yyyy] [name of copyright owner] 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 | -------------------------------------------------------------------------------- /tests/waker_panic.rs: -------------------------------------------------------------------------------- 1 | use std::cell::Cell; 2 | use std::future::Future; 3 | use std::panic::{catch_unwind, AssertUnwindSafe}; 4 | use std::pin::Pin; 5 | use std::sync::atomic::{AtomicUsize, Ordering}; 6 | use std::task::{Context, Poll}; 7 | use std::thread; 8 | use std::time::Duration; 9 | 10 | use async_task::Runnable; 11 | use atomic_waker::AtomicWaker; 12 | use easy_parallel::Parallel; 13 | use smol::future; 14 | 15 | // Creates a future with event counters. 16 | // 17 | // Usage: `future!(f, get_waker, POLL, DROP)` 18 | // 19 | // The future `f` always sleeps for 200 ms, and panics the second time it is polled. 20 | // When it gets polled, `POLL` is incremented. 21 | // When it gets dropped, `DROP` is incremented. 22 | // 23 | // Every time the future is run, it stores the waker into a global variable. 24 | // This waker can be extracted using the `get_waker()` function. 25 | macro_rules! future { 26 | ($name:pat, $get_waker:pat, $poll:ident, $drop:ident) => { 27 | static $poll: AtomicUsize = AtomicUsize::new(0); 28 | static $drop: AtomicUsize = AtomicUsize::new(0); 29 | static WAKER: AtomicWaker = AtomicWaker::new(); 30 | 31 | let ($name, $get_waker) = { 32 | struct Fut(Cell, #[allow(dead_code)] Box); 33 | 34 | impl Future for Fut { 35 | type Output = (); 36 | 37 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 38 | WAKER.register(cx.waker()); 39 | $poll.fetch_add(1, Ordering::SeqCst); 40 | thread::sleep(ms(400)); 41 | 42 | if self.0.get() { 43 | panic!() 44 | } else { 45 | self.0.set(true); 46 | Poll::Pending 47 | } 48 | } 49 | } 50 | 51 | impl Drop for Fut { 52 | fn drop(&mut self) { 53 | $drop.fetch_add(1, Ordering::SeqCst); 54 | } 55 | } 56 | 57 | (Fut(Cell::new(false), Box::new(0)), || WAKER.take().unwrap()) 58 | }; 59 | }; 60 | } 61 | 62 | // Creates a schedule function with event counters. 63 | // 64 | // Usage: `schedule!(s, chan, SCHED, DROP)` 65 | // 66 | // The schedule function `s` pushes the task into `chan`. 67 | // When it gets invoked, `SCHED` is incremented. 68 | // When it gets dropped, `DROP` is incremented. 69 | // 70 | // Receiver `chan` extracts the task when it is scheduled. 71 | macro_rules! schedule { 72 | ($name:pat, $chan:pat, $sched:ident, $drop:ident) => { 73 | static $drop: AtomicUsize = AtomicUsize::new(0); 74 | static $sched: AtomicUsize = AtomicUsize::new(0); 75 | 76 | let ($name, $chan) = { 77 | let (s, r) = flume::unbounded(); 78 | 79 | struct Guard(#[allow(dead_code)] Box); 80 | 81 | impl Drop for Guard { 82 | fn drop(&mut self) { 83 | $drop.fetch_add(1, Ordering::SeqCst); 84 | } 85 | } 86 | 87 | let guard = Guard(Box::new(0)); 88 | let sched = move |runnable: Runnable| { 89 | let _ = &guard; 90 | $sched.fetch_add(1, Ordering::SeqCst); 91 | s.send(runnable).unwrap(); 92 | }; 93 | 94 | (sched, r) 95 | }; 96 | }; 97 | } 98 | 99 | fn ms(ms: u64) -> Duration { 100 | Duration::from_millis(ms) 101 | } 102 | 103 | fn try_await(f: impl Future) -> Option { 104 | future::block_on(future::poll_once(f)) 105 | } 106 | 107 | #[test] 108 | fn wake_during_run() { 109 | future!(f, get_waker, POLL, DROP_F); 110 | schedule!(s, chan, SCHEDULE, DROP_S); 111 | let (runnable, task) = async_task::spawn(f, s); 112 | 113 | runnable.run(); 114 | let waker = get_waker(); 115 | waker.wake_by_ref(); 116 | let runnable = chan.recv().unwrap(); 117 | 118 | Parallel::new() 119 | .add(|| { 120 | assert!(catch_unwind(|| runnable.run()).is_err()); 121 | drop(get_waker()); 122 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 123 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 124 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 125 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 126 | assert_eq!(chan.len(), 0); 127 | }) 128 | .add(|| { 129 | thread::sleep(ms(200)); 130 | 131 | waker.wake(); 132 | task.detach(); 133 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 134 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 135 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 136 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 137 | assert_eq!(chan.len(), 0); 138 | 139 | thread::sleep(ms(400)); 140 | 141 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 142 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 143 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 144 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 145 | assert_eq!(chan.len(), 0); 146 | }) 147 | .run(); 148 | } 149 | 150 | #[test] 151 | fn cancel_during_run() { 152 | future!(f, get_waker, POLL, DROP_F); 153 | schedule!(s, chan, SCHEDULE, DROP_S); 154 | let (runnable, task) = async_task::spawn(f, s); 155 | 156 | runnable.run(); 157 | let waker = get_waker(); 158 | waker.wake(); 159 | let runnable = chan.recv().unwrap(); 160 | 161 | Parallel::new() 162 | .add(|| { 163 | assert!(catch_unwind(|| runnable.run()).is_err()); 164 | drop(get_waker()); 165 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 166 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 167 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 168 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 169 | assert_eq!(chan.len(), 0); 170 | }) 171 | .add(|| { 172 | thread::sleep(ms(200)); 173 | 174 | drop(task); 175 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 176 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 177 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 178 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 179 | assert_eq!(chan.len(), 0); 180 | 181 | thread::sleep(ms(400)); 182 | 183 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 184 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 185 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 186 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 187 | assert_eq!(chan.len(), 0); 188 | }) 189 | .run(); 190 | } 191 | 192 | #[test] 193 | fn wake_and_cancel_during_run() { 194 | future!(f, get_waker, POLL, DROP_F); 195 | schedule!(s, chan, SCHEDULE, DROP_S); 196 | let (runnable, task) = async_task::spawn(f, s); 197 | 198 | runnable.run(); 199 | let waker = get_waker(); 200 | waker.wake_by_ref(); 201 | let runnable = chan.recv().unwrap(); 202 | 203 | Parallel::new() 204 | .add(|| { 205 | assert!(catch_unwind(|| runnable.run()).is_err()); 206 | drop(get_waker()); 207 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 208 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 209 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 210 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 211 | assert_eq!(chan.len(), 0); 212 | }) 213 | .add(|| { 214 | thread::sleep(ms(200)); 215 | 216 | waker.wake(); 217 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 218 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 219 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 220 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 221 | assert_eq!(chan.len(), 0); 222 | 223 | drop(task); 224 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 225 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 226 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 227 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 228 | assert_eq!(chan.len(), 0); 229 | 230 | thread::sleep(ms(400)); 231 | 232 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 233 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 234 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 235 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 236 | assert_eq!(chan.len(), 0); 237 | }) 238 | .run(); 239 | } 240 | 241 | #[flaky_test::flaky_test] 242 | fn cancel_and_wake_during_run() { 243 | future!(f, get_waker, POLL, DROP_F); 244 | schedule!(s, chan, SCHEDULE, DROP_S); 245 | POLL.store(0, Ordering::SeqCst); 246 | DROP_F.store(0, Ordering::SeqCst); 247 | SCHEDULE.store(0, Ordering::SeqCst); 248 | DROP_S.store(0, Ordering::SeqCst); 249 | 250 | let (runnable, task) = async_task::spawn(f, s); 251 | 252 | runnable.run(); 253 | let waker = get_waker(); 254 | waker.wake_by_ref(); 255 | let runnable = chan.recv().unwrap(); 256 | 257 | Parallel::new() 258 | .add(|| { 259 | assert!(catch_unwind(|| runnable.run()).is_err()); 260 | drop(get_waker()); 261 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 262 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 263 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 264 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 265 | assert_eq!(chan.len(), 0); 266 | }) 267 | .add(|| { 268 | thread::sleep(ms(200)); 269 | 270 | drop(task); 271 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 272 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 273 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 274 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 275 | assert_eq!(chan.len(), 0); 276 | 277 | waker.wake(); 278 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 279 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 280 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 281 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 282 | assert_eq!(chan.len(), 0); 283 | 284 | thread::sleep(ms(400)); 285 | 286 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 287 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 288 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 289 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 290 | assert_eq!(chan.len(), 0); 291 | }) 292 | .run(); 293 | } 294 | 295 | #[test] 296 | fn panic_and_poll() { 297 | future!(f, get_waker, POLL, DROP_F); 298 | schedule!(s, chan, SCHEDULE, DROP_S); 299 | let (runnable, task) = async_task::spawn(f, s); 300 | 301 | runnable.run(); 302 | get_waker().wake(); 303 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 304 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 305 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 306 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 307 | 308 | let mut task = task; 309 | assert!(try_await(&mut task).is_none()); 310 | 311 | let runnable = chan.recv().unwrap(); 312 | assert!(catch_unwind(|| runnable.run()).is_err()); 313 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 314 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 315 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 316 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 317 | 318 | assert!(catch_unwind(AssertUnwindSafe(|| try_await(&mut task))).is_err()); 319 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 320 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 321 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 322 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 323 | 324 | drop(get_waker()); 325 | drop(task); 326 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 327 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 328 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 329 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 330 | } 331 | -------------------------------------------------------------------------------- /tests/waker_pending.rs: -------------------------------------------------------------------------------- 1 | use std::future::Future; 2 | use std::pin::Pin; 3 | use std::sync::atomic::{AtomicUsize, Ordering}; 4 | use std::task::{Context, Poll}; 5 | use std::thread; 6 | use std::time::Duration; 7 | 8 | use async_task::Runnable; 9 | use atomic_waker::AtomicWaker; 10 | use easy_parallel::Parallel; 11 | 12 | // Creates a future with event counters. 13 | // 14 | // Usage: `future!(f, get_waker, POLL, DROP)` 15 | // 16 | // The future `f` always sleeps for 200 ms and returns `Poll::Pending`. 17 | // When it gets polled, `POLL` is incremented. 18 | // When it gets dropped, `DROP` is incremented. 19 | // 20 | // Every time the future is run, it stores the waker into a global variable. 21 | // This waker can be extracted using the `get_waker()` function. 22 | macro_rules! future { 23 | ($name:pat, $get_waker:pat, $poll:ident, $drop:ident) => { 24 | static $poll: AtomicUsize = AtomicUsize::new(0); 25 | static $drop: AtomicUsize = AtomicUsize::new(0); 26 | static WAKER: AtomicWaker = AtomicWaker::new(); 27 | 28 | let ($name, $get_waker) = { 29 | struct Fut(#[allow(dead_code)] Box); 30 | 31 | impl Future for Fut { 32 | type Output = (); 33 | 34 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 35 | WAKER.register(cx.waker()); 36 | $poll.fetch_add(1, Ordering::SeqCst); 37 | thread::sleep(ms(400)); 38 | Poll::Pending 39 | } 40 | } 41 | 42 | impl Drop for Fut { 43 | fn drop(&mut self) { 44 | $drop.fetch_add(1, Ordering::SeqCst); 45 | } 46 | } 47 | 48 | (Fut(Box::new(0)), || WAKER.take().unwrap()) 49 | }; 50 | }; 51 | } 52 | 53 | // Creates a schedule function with event counters. 54 | // 55 | // Usage: `schedule!(s, chan, SCHED, DROP)` 56 | // 57 | // The schedule function `s` pushes the task into `chan`. 58 | // When it gets invoked, `SCHED` is incremented. 59 | // When it gets dropped, `DROP` is incremented. 60 | // 61 | // Receiver `chan` extracts the task when it is scheduled. 62 | macro_rules! schedule { 63 | ($name:pat, $chan:pat, $sched:ident, $drop:ident) => { 64 | static $drop: AtomicUsize = AtomicUsize::new(0); 65 | static $sched: AtomicUsize = AtomicUsize::new(0); 66 | 67 | let ($name, $chan) = { 68 | let (s, r) = flume::unbounded(); 69 | 70 | struct Guard(#[allow(dead_code)] Box); 71 | 72 | impl Drop for Guard { 73 | fn drop(&mut self) { 74 | $drop.fetch_add(1, Ordering::SeqCst); 75 | } 76 | } 77 | 78 | let guard = Guard(Box::new(0)); 79 | let sched = move |runnable: Runnable| { 80 | let _ = &guard; 81 | $sched.fetch_add(1, Ordering::SeqCst); 82 | s.send(runnable).unwrap(); 83 | }; 84 | 85 | (sched, r) 86 | }; 87 | }; 88 | } 89 | 90 | fn ms(ms: u64) -> Duration { 91 | Duration::from_millis(ms) 92 | } 93 | 94 | #[test] 95 | fn wake_during_run() { 96 | future!(f, get_waker, POLL, DROP_F); 97 | schedule!(s, chan, SCHEDULE, DROP_S); 98 | let (runnable, _task) = async_task::spawn(f, s); 99 | 100 | runnable.run(); 101 | let waker = get_waker(); 102 | waker.wake_by_ref(); 103 | let runnable = chan.recv().unwrap(); 104 | 105 | Parallel::new() 106 | .add(|| { 107 | runnable.run(); 108 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 109 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 2); 110 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 111 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 112 | assert_eq!(chan.len(), 1); 113 | }) 114 | .add(|| { 115 | thread::sleep(ms(200)); 116 | 117 | waker.wake_by_ref(); 118 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 119 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 120 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 121 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 122 | assert_eq!(chan.len(), 0); 123 | 124 | thread::sleep(ms(400)); 125 | 126 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 127 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 2); 128 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 129 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 130 | assert_eq!(chan.len(), 1); 131 | }) 132 | .run(); 133 | 134 | chan.recv().unwrap(); 135 | drop(get_waker()); 136 | } 137 | 138 | #[test] 139 | fn cancel_during_run() { 140 | future!(f, get_waker, POLL, DROP_F); 141 | schedule!(s, chan, SCHEDULE, DROP_S); 142 | let (runnable, task) = async_task::spawn(f, s); 143 | 144 | runnable.run(); 145 | let waker = get_waker(); 146 | waker.wake(); 147 | let runnable = chan.recv().unwrap(); 148 | 149 | Parallel::new() 150 | .add(|| { 151 | runnable.run(); 152 | drop(get_waker()); 153 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 154 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 155 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 156 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 157 | assert_eq!(chan.len(), 0); 158 | }) 159 | .add(|| { 160 | thread::sleep(ms(200)); 161 | 162 | drop(task); 163 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 164 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 165 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 166 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 167 | assert_eq!(chan.len(), 0); 168 | 169 | thread::sleep(ms(400)); 170 | 171 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 172 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 173 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 174 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 175 | assert_eq!(chan.len(), 0); 176 | }) 177 | .run(); 178 | } 179 | 180 | #[test] 181 | fn wake_and_cancel_during_run() { 182 | future!(f, get_waker, POLL, DROP_F); 183 | schedule!(s, chan, SCHEDULE, DROP_S); 184 | let (runnable, task) = async_task::spawn(f, s); 185 | 186 | runnable.run(); 187 | let waker = get_waker(); 188 | waker.wake_by_ref(); 189 | let runnable = chan.recv().unwrap(); 190 | 191 | Parallel::new() 192 | .add(|| { 193 | runnable.run(); 194 | drop(get_waker()); 195 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 196 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 197 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 198 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 199 | assert_eq!(chan.len(), 0); 200 | }) 201 | .add(|| { 202 | thread::sleep(ms(200)); 203 | 204 | waker.wake(); 205 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 206 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 207 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 208 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 209 | assert_eq!(chan.len(), 0); 210 | 211 | drop(task); 212 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 213 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 214 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 215 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 216 | assert_eq!(chan.len(), 0); 217 | 218 | thread::sleep(ms(400)); 219 | 220 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 221 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 222 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 223 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 224 | assert_eq!(chan.len(), 0); 225 | }) 226 | .run(); 227 | } 228 | 229 | #[test] 230 | fn cancel_and_wake_during_run() { 231 | future!(f, get_waker, POLL, DROP_F); 232 | schedule!(s, chan, SCHEDULE, DROP_S); 233 | let (runnable, task) = async_task::spawn(f, s); 234 | 235 | runnable.run(); 236 | let waker = get_waker(); 237 | waker.wake_by_ref(); 238 | let runnable = chan.recv().unwrap(); 239 | 240 | Parallel::new() 241 | .add(|| { 242 | runnable.run(); 243 | drop(get_waker()); 244 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 245 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 246 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 247 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 248 | assert_eq!(chan.len(), 0); 249 | }) 250 | .add(|| { 251 | thread::sleep(ms(200)); 252 | 253 | drop(task); 254 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 255 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 256 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 257 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 258 | assert_eq!(chan.len(), 0); 259 | 260 | waker.wake(); 261 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 262 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 263 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 264 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 265 | assert_eq!(chan.len(), 0); 266 | 267 | thread::sleep(ms(400)); 268 | 269 | assert_eq!(POLL.load(Ordering::SeqCst), 2); 270 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 271 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 272 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 273 | assert_eq!(chan.len(), 0); 274 | }) 275 | .run(); 276 | } 277 | 278 | #[test] 279 | fn drop_last_waker() { 280 | future!(f, get_waker, POLL, DROP_F); 281 | schedule!(s, chan, SCHEDULE, DROP_S); 282 | let (runnable, task) = async_task::spawn(f, s); 283 | 284 | runnable.run(); 285 | let waker = get_waker(); 286 | 287 | task.detach(); 288 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 289 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 290 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 291 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 292 | assert_eq!(chan.len(), 0); 293 | 294 | drop(waker); 295 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 296 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 297 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 298 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 299 | assert_eq!(chan.len(), 1); 300 | 301 | chan.recv().unwrap().run(); 302 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 303 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 304 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 305 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 306 | assert_eq!(chan.len(), 0); 307 | } 308 | 309 | #[test] 310 | fn cancel_last_task() { 311 | future!(f, get_waker, POLL, DROP_F); 312 | schedule!(s, chan, SCHEDULE, DROP_S); 313 | let (runnable, task) = async_task::spawn(f, s); 314 | 315 | runnable.run(); 316 | drop(get_waker()); 317 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 318 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 319 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 320 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 321 | assert_eq!(chan.len(), 0); 322 | 323 | drop(task); 324 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 325 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 326 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 327 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 328 | assert_eq!(chan.len(), 1); 329 | 330 | chan.recv().unwrap().run(); 331 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 332 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 333 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 334 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 335 | assert_eq!(chan.len(), 0); 336 | } 337 | 338 | #[test] 339 | fn drop_last_task() { 340 | future!(f, get_waker, POLL, DROP_F); 341 | schedule!(s, chan, SCHEDULE, DROP_S); 342 | let (runnable, task) = async_task::spawn(f, s); 343 | 344 | runnable.run(); 345 | drop(get_waker()); 346 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 347 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 348 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 349 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 350 | assert_eq!(chan.len(), 0); 351 | 352 | task.detach(); 353 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 354 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 355 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 356 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 357 | assert_eq!(chan.len(), 1); 358 | 359 | chan.recv().unwrap().run(); 360 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 361 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); 362 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 363 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 364 | assert_eq!(chan.len(), 0); 365 | } 366 | -------------------------------------------------------------------------------- /tests/join.rs: -------------------------------------------------------------------------------- 1 | use std::cell::Cell; 2 | use std::future::Future; 3 | use std::panic::{catch_unwind, AssertUnwindSafe}; 4 | use std::pin::Pin; 5 | use std::sync::atomic::{AtomicUsize, Ordering}; 6 | use std::task::{Context, Poll}; 7 | use std::thread; 8 | use std::time::Duration; 9 | 10 | use async_task::Runnable; 11 | use easy_parallel::Parallel; 12 | use smol::future; 13 | 14 | // Creates a future with event counters. 15 | // 16 | // Usage: `future!(f, POLL, DROP_F, DROP_T)` 17 | // 18 | // The future `f` outputs `Poll::Ready`. 19 | // When it gets polled, `POLL` is incremented. 20 | // When it gets dropped, `DROP_F` is incremented. 21 | // When the output gets dropped, `DROP_T` is incremented. 22 | macro_rules! future { 23 | ($name:pat, $poll:ident, $drop_f:ident, $drop_t:ident) => { 24 | static $poll: AtomicUsize = AtomicUsize::new(0); 25 | static $drop_f: AtomicUsize = AtomicUsize::new(0); 26 | static $drop_t: AtomicUsize = AtomicUsize::new(0); 27 | 28 | let $name = { 29 | struct Fut(#[allow(dead_code)] Box); 30 | 31 | impl Future for Fut { 32 | type Output = Out; 33 | 34 | fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { 35 | $poll.fetch_add(1, Ordering::SeqCst); 36 | Poll::Ready(Out(Box::new(0), true)) 37 | } 38 | } 39 | 40 | impl Drop for Fut { 41 | fn drop(&mut self) { 42 | $drop_f.fetch_add(1, Ordering::SeqCst); 43 | } 44 | } 45 | 46 | #[derive(Default)] 47 | struct Out(#[allow(dead_code)] Box, bool); 48 | 49 | impl Drop for Out { 50 | fn drop(&mut self) { 51 | if self.1 { 52 | $drop_t.fetch_add(1, Ordering::SeqCst); 53 | } 54 | } 55 | } 56 | 57 | Fut(Box::new(0)) 58 | }; 59 | }; 60 | } 61 | 62 | // Creates a schedule function with event counters. 63 | // 64 | // Usage: `schedule!(s, SCHED, DROP)` 65 | // 66 | // The schedule function `s` does nothing. 67 | // When it gets invoked, `SCHED` is incremented. 68 | // When it gets dropped, `DROP` is incremented. 69 | macro_rules! schedule { 70 | ($name:pat, $sched:ident, $drop:ident) => { 71 | static $drop: AtomicUsize = AtomicUsize::new(0); 72 | static $sched: AtomicUsize = AtomicUsize::new(0); 73 | 74 | let $name = { 75 | struct Guard(#[allow(dead_code)] Box); 76 | 77 | impl Drop for Guard { 78 | fn drop(&mut self) { 79 | $drop.fetch_add(1, Ordering::SeqCst); 80 | } 81 | } 82 | 83 | let guard = Guard(Box::new(0)); 84 | move |runnable: Runnable| { 85 | let _ = &guard; 86 | runnable.schedule(); 87 | $sched.fetch_add(1, Ordering::SeqCst); 88 | } 89 | }; 90 | }; 91 | } 92 | 93 | fn ms(ms: u64) -> Duration { 94 | Duration::from_millis(ms) 95 | } 96 | 97 | #[test] 98 | fn drop_and_join() { 99 | future!(f, POLL, DROP_F, DROP_T); 100 | schedule!(s, SCHEDULE, DROP_S); 101 | let (runnable, task) = async_task::spawn(f, s); 102 | 103 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 104 | 105 | drop(runnable); 106 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 107 | 108 | assert!(catch_unwind(|| future::block_on(task)).is_err()); 109 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 110 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 111 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 112 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 113 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 114 | } 115 | 116 | #[test] 117 | fn run_and_join() { 118 | future!(f, POLL, DROP_F, DROP_T); 119 | schedule!(s, SCHEDULE, DROP_S); 120 | let (runnable, task) = async_task::spawn(f, s); 121 | 122 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 123 | 124 | runnable.run(); 125 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 126 | 127 | assert!(catch_unwind(|| future::block_on(task)).is_ok()); 128 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 129 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 130 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 131 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 132 | assert_eq!(DROP_T.load(Ordering::SeqCst), 1); 133 | } 134 | 135 | #[test] 136 | fn detach_and_run() { 137 | future!(f, POLL, DROP_F, DROP_T); 138 | schedule!(s, SCHEDULE, DROP_S); 139 | let (runnable, task) = async_task::spawn(f, s); 140 | 141 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 142 | 143 | task.detach(); 144 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 145 | 146 | runnable.run(); 147 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 148 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 149 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 150 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 151 | assert_eq!(DROP_T.load(Ordering::SeqCst), 1); 152 | } 153 | 154 | #[test] 155 | fn join_twice() { 156 | future!(f, POLL, DROP_F, DROP_T); 157 | schedule!(s, SCHEDULE, DROP_S); 158 | let (runnable, mut task) = async_task::spawn(f, s); 159 | 160 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 161 | 162 | runnable.run(); 163 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 164 | 165 | future::block_on(&mut task); 166 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 167 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 168 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 169 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 170 | assert_eq!(DROP_T.load(Ordering::SeqCst), 1); 171 | 172 | assert!(catch_unwind(AssertUnwindSafe(|| future::block_on(&mut task))).is_err()); 173 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 174 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 175 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 176 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 177 | assert_eq!(DROP_T.load(Ordering::SeqCst), 1); 178 | 179 | task.detach(); 180 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 181 | } 182 | 183 | #[test] 184 | fn join_and_cancel() { 185 | future!(f, POLL, DROP_F, DROP_T); 186 | schedule!(s, SCHEDULE, DROP_S); 187 | let (runnable, task) = async_task::spawn(f, s); 188 | 189 | Parallel::new() 190 | .add(|| { 191 | thread::sleep(ms(200)); 192 | drop(runnable); 193 | 194 | thread::sleep(ms(400)); 195 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 196 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 197 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 198 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 199 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 200 | }) 201 | .add(|| { 202 | assert!(catch_unwind(|| future::block_on(task)).is_err()); 203 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 204 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 205 | 206 | thread::sleep(ms(200)); 207 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 208 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 209 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 210 | }) 211 | .run(); 212 | } 213 | 214 | #[test] 215 | fn join_and_run() { 216 | future!(f, POLL, DROP_F, DROP_T); 217 | schedule!(s, SCHEDULE, DROP_S); 218 | let (runnable, task) = async_task::spawn(f, s); 219 | 220 | Parallel::new() 221 | .add(|| { 222 | thread::sleep(ms(400)); 223 | 224 | runnable.run(); 225 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 226 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 227 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 228 | 229 | thread::sleep(ms(200)); 230 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 231 | }) 232 | .add(|| { 233 | future::block_on(task); 234 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 235 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 236 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 237 | assert_eq!(DROP_T.load(Ordering::SeqCst), 1); 238 | 239 | thread::sleep(ms(200)); 240 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 241 | }) 242 | .run(); 243 | } 244 | 245 | #[test] 246 | fn try_join_and_run_and_join() { 247 | future!(f, POLL, DROP_F, DROP_T); 248 | schedule!(s, SCHEDULE, DROP_S); 249 | let (runnable, mut task) = async_task::spawn(f, s); 250 | 251 | Parallel::new() 252 | .add(|| { 253 | thread::sleep(ms(400)); 254 | 255 | runnable.run(); 256 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 257 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 258 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 259 | 260 | thread::sleep(ms(200)); 261 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 262 | }) 263 | .add(|| { 264 | future::block_on(future::or(&mut task, future::ready(Default::default()))); 265 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 266 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 267 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 268 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 269 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 270 | 271 | future::block_on(task); 272 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 273 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 274 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 275 | assert_eq!(DROP_T.load(Ordering::SeqCst), 1); 276 | 277 | thread::sleep(ms(200)); 278 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 279 | }) 280 | .run(); 281 | } 282 | 283 | #[test] 284 | fn try_join_and_cancel_and_run() { 285 | future!(f, POLL, DROP_F, DROP_T); 286 | schedule!(s, SCHEDULE, DROP_S); 287 | let (runnable, mut task) = async_task::spawn(f, s); 288 | 289 | Parallel::new() 290 | .add(|| { 291 | thread::sleep(ms(200)); 292 | 293 | runnable.run(); 294 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 295 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 296 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 297 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 298 | }) 299 | .add(|| { 300 | future::block_on(future::or(&mut task, future::ready(Default::default()))); 301 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 302 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 303 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 304 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 305 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 306 | 307 | drop(task); 308 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 309 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 310 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 311 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 312 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 313 | }) 314 | .run(); 315 | } 316 | 317 | #[test] 318 | fn try_join_and_run_and_cancel() { 319 | future!(f, POLL, DROP_F, DROP_T); 320 | schedule!(s, SCHEDULE, DROP_S); 321 | let (runnable, mut task) = async_task::spawn(f, s); 322 | 323 | Parallel::new() 324 | .add(|| { 325 | thread::sleep(ms(200)); 326 | 327 | runnable.run(); 328 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 329 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 330 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 331 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 332 | }) 333 | .add(|| { 334 | future::block_on(future::or(&mut task, future::ready(Default::default()))); 335 | assert_eq!(POLL.load(Ordering::SeqCst), 0); 336 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 337 | assert_eq!(DROP_F.load(Ordering::SeqCst), 0); 338 | assert_eq!(DROP_S.load(Ordering::SeqCst), 0); 339 | assert_eq!(DROP_T.load(Ordering::SeqCst), 0); 340 | 341 | thread::sleep(ms(400)); 342 | 343 | drop(task); 344 | assert_eq!(POLL.load(Ordering::SeqCst), 1); 345 | assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); 346 | assert_eq!(DROP_F.load(Ordering::SeqCst), 1); 347 | assert_eq!(DROP_S.load(Ordering::SeqCst), 1); 348 | assert_eq!(DROP_T.load(Ordering::SeqCst), 1); 349 | }) 350 | .run(); 351 | } 352 | 353 | #[test] 354 | fn await_output() { 355 | struct Fut(Cell>); 356 | 357 | impl Fut { 358 | fn new(t: T) -> Fut { 359 | Fut(Cell::new(Some(t))) 360 | } 361 | } 362 | 363 | impl Future for Fut { 364 | type Output = T; 365 | 366 | fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { 367 | Poll::Ready(self.0.take().unwrap()) 368 | } 369 | } 370 | 371 | for i in 0..10 { 372 | let (runnable, task) = async_task::spawn(Fut::new(i), drop); 373 | runnable.run(); 374 | assert_eq!(future::block_on(task), i); 375 | } 376 | 377 | for i in 0..10 { 378 | let (runnable, task) = async_task::spawn(Fut::new(vec![7; i]), drop); 379 | runnable.run(); 380 | assert_eq!(future::block_on(task), vec![7; i]); 381 | } 382 | 383 | let (runnable, task) = async_task::spawn(Fut::new("foo".to_string()), drop); 384 | runnable.run(); 385 | assert_eq!(future::block_on(task), "foo"); 386 | } 387 | -------------------------------------------------------------------------------- /src/task.rs: -------------------------------------------------------------------------------- 1 | use core::fmt; 2 | use core::future::Future; 3 | use core::marker::PhantomData; 4 | use core::mem; 5 | use core::pin::Pin; 6 | use core::ptr::NonNull; 7 | use core::sync::atomic::Ordering; 8 | use core::task::{Context, Poll}; 9 | 10 | use crate::header::{Header, HeaderWithMetadata}; 11 | use crate::raw::Panic; 12 | use crate::state::*; 13 | use crate::ScheduleInfo; 14 | 15 | /// A spawned task. 16 | /// 17 | /// A [`Task`] can be awaited to retrieve the output of its future. 18 | /// 19 | /// Dropping a [`Task`] cancels it, which means its future won't be polled again. To drop the 20 | /// [`Task`] handle without canceling it, use [`detach()`][`Task::detach()`] instead. To cancel a 21 | /// task gracefully and wait until it is fully destroyed, use the [`cancel()`][Task::cancel()] 22 | /// method. 23 | /// 24 | /// Note that canceling a task actually wakes it and reschedules one last time. Then, the executor 25 | /// can destroy the task by simply dropping its [`Runnable`][`super::Runnable`] or by invoking 26 | /// [`run()`][`super::Runnable::run()`]. 27 | /// 28 | /// # Examples 29 | /// 30 | /// ``` 31 | /// use smol::{future, Executor}; 32 | /// use std::thread; 33 | /// 34 | /// let ex = Executor::new(); 35 | /// 36 | /// // Spawn a future onto the executor. 37 | /// let task = ex.spawn(async { 38 | /// println!("Hello from a task!"); 39 | /// 1 + 2 40 | /// }); 41 | /// 42 | /// // Run an executor thread. 43 | /// thread::spawn(move || future::block_on(ex.run(future::pending::<()>()))); 44 | /// 45 | /// // Wait for the task's output. 46 | /// assert_eq!(future::block_on(task), 3); 47 | /// ``` 48 | #[must_use = "tasks get canceled when dropped, use `.detach()` to run them in the background"] 49 | pub struct Task { 50 | /// A raw task pointer. 51 | pub(crate) ptr: NonNull<()>, 52 | 53 | /// A marker capturing generic types `T` and `M`. 54 | pub(crate) _marker: PhantomData<(T, M)>, 55 | } 56 | 57 | unsafe impl Send for Task {} 58 | unsafe impl Sync for Task {} 59 | 60 | impl Unpin for Task {} 61 | 62 | #[cfg(feature = "std")] 63 | impl std::panic::UnwindSafe for Task {} 64 | #[cfg(feature = "std")] 65 | impl std::panic::RefUnwindSafe for Task {} 66 | 67 | impl Task { 68 | /// Detaches the task to let it keep running in the background. 69 | /// 70 | /// # Examples 71 | /// 72 | /// ``` 73 | /// use smol::{Executor, Timer}; 74 | /// use std::time::Duration; 75 | /// 76 | /// let ex = Executor::new(); 77 | /// 78 | /// // Spawn a daemon future. 79 | /// ex.spawn(async { 80 | /// loop { 81 | /// println!("I'm a daemon task looping forever."); 82 | /// Timer::after(Duration::from_secs(1)).await; 83 | /// } 84 | /// }) 85 | /// .detach(); 86 | /// ``` 87 | pub fn detach(self) { 88 | let this = self; 89 | let _out = set_detached::(this.ptr.as_ptr()); 90 | mem::forget(this); 91 | } 92 | 93 | /// Cancels the task and waits for it to stop running. 94 | /// 95 | /// Returns the task's output if it was completed just before it got canceled, or [`None`] if 96 | /// it didn't complete. 97 | /// 98 | /// While it's possible to simply drop the [`Task`] to cancel it, this is a cleaner way of 99 | /// canceling because it also waits for the task to stop running. 100 | /// 101 | /// # Examples 102 | /// 103 | /// ``` 104 | /// # if cfg!(miri) { return; } // Miri does not support epoll 105 | /// use smol::{future, Executor, Timer}; 106 | /// use std::thread; 107 | /// use std::time::Duration; 108 | /// 109 | /// let ex = Executor::new(); 110 | /// 111 | /// // Spawn a daemon future. 112 | /// let task = ex.spawn(async { 113 | /// loop { 114 | /// println!("Even though I'm in an infinite loop, you can still cancel me!"); 115 | /// Timer::after(Duration::from_secs(1)).await; 116 | /// } 117 | /// }); 118 | /// 119 | /// // Run an executor thread. 120 | /// thread::spawn(move || future::block_on(ex.run(future::pending::<()>()))); 121 | /// 122 | /// future::block_on(async { 123 | /// Timer::after(Duration::from_secs(3)).await; 124 | /// task.cancel().await; 125 | /// }); 126 | /// ``` 127 | pub async fn cancel(self) -> Option { 128 | let this = self; 129 | set_canceled(this.ptr.as_ptr()); 130 | this.fallible().await 131 | } 132 | 133 | /// Converts this task into a [`FallibleTask`]. 134 | /// 135 | /// Like [`Task`], a fallible task will poll the task's output until it is 136 | /// completed or cancelled due to its [`Runnable`][`super::Runnable`] being 137 | /// dropped without being run. Resolves to the task's output when completed, 138 | /// or [`None`] if it didn't complete. 139 | /// 140 | /// # Examples 141 | /// 142 | /// ``` 143 | /// use smol::{future, Executor}; 144 | /// use std::thread; 145 | /// 146 | /// let ex = Executor::new(); 147 | /// 148 | /// // Spawn a future onto the executor. 149 | /// let task = ex.spawn(async { 150 | /// println!("Hello from a task!"); 151 | /// 1 + 2 152 | /// }) 153 | /// .fallible(); 154 | /// 155 | /// // Run an executor thread. 156 | /// thread::spawn(move || future::block_on(ex.run(future::pending::<()>()))); 157 | /// 158 | /// // Wait for the task's output. 159 | /// assert_eq!(future::block_on(task), Some(3)); 160 | /// ``` 161 | /// 162 | /// ``` 163 | /// use smol::future; 164 | /// 165 | /// // Schedule function which drops the runnable without running it. 166 | /// let schedule = move |runnable| drop(runnable); 167 | /// 168 | /// // Create a task with the future and the schedule function. 169 | /// let (runnable, task) = async_task::spawn(async { 170 | /// println!("Hello from a task!"); 171 | /// 1 + 2 172 | /// }, schedule); 173 | /// runnable.schedule(); 174 | /// 175 | /// // Wait for the task's output. 176 | /// assert_eq!(future::block_on(task.fallible()), None); 177 | /// ``` 178 | pub fn fallible(self) -> FallibleTask { 179 | FallibleTask { task: self } 180 | } 181 | 182 | fn header_with_metadata(&self) -> &HeaderWithMetadata { 183 | let ptr = self.ptr.as_ptr(); 184 | let header = ptr as *const HeaderWithMetadata; 185 | unsafe { &*header } 186 | } 187 | 188 | /// Returns `true` if the current task is finished. 189 | /// 190 | /// Note that in a multithreaded environment, this task can change finish immediately after calling this function. 191 | pub fn is_finished(&self) -> bool { 192 | let ptr = self.ptr.as_ptr(); 193 | let header = ptr as *const Header; 194 | 195 | unsafe { 196 | let state = (*header).state.load(Ordering::Acquire); 197 | state & (CLOSED | COMPLETED) != 0 198 | } 199 | } 200 | 201 | /// Get the metadata associated with this task. 202 | /// 203 | /// Tasks can be created with a metadata object associated with them; by default, this 204 | /// is a `()` value. See the [`Builder::metadata()`] method for more information. 205 | pub fn metadata(&self) -> &M { 206 | let ptr = self.ptr.as_ptr(); 207 | let header = ptr as *const HeaderWithMetadata; 208 | &unsafe { &*header }.metadata 209 | } 210 | } 211 | 212 | /// Puts the task in detached state. 213 | #[inline(never)] 214 | fn set_detached(ptr: *const ()) -> Option> { 215 | let header = ptr as *const Header; 216 | 217 | unsafe { 218 | // A place where the output will be stored in case it needs to be dropped. 219 | let mut output = None; 220 | 221 | // Optimistically assume the `Task` is being detached just after creating the task. 222 | // This is a common case so if the `Task` is datached, the overhead of it is only one 223 | // compare-exchange operation. 224 | if let Err(mut state) = (*header).state.compare_exchange_weak( 225 | SCHEDULED | TASK | REFERENCE, 226 | SCHEDULED | REFERENCE, 227 | Ordering::AcqRel, 228 | Ordering::Acquire, 229 | ) { 230 | loop { 231 | // If the task has been completed but not yet closed, that means its output 232 | // must be dropped. 233 | if state & COMPLETED != 0 && state & CLOSED == 0 { 234 | // Mark the task as closed in order to grab its output. 235 | match (*header).state.compare_exchange_weak( 236 | state, 237 | state | CLOSED, 238 | Ordering::AcqRel, 239 | Ordering::Acquire, 240 | ) { 241 | Ok(_) => { 242 | // Read the output. 243 | output = Some( 244 | ((*header).vtable.get_output(ptr) as *mut Result).read(), 245 | ); 246 | 247 | // Update the state variable because we're continuing the loop. 248 | state |= CLOSED; 249 | } 250 | Err(s) => state = s, 251 | } 252 | } else { 253 | // If this is the last reference to the task and it's not closed, then 254 | // close it and schedule one more time so that its future gets dropped by 255 | // the executor. 256 | let new = if state & (!(REFERENCE - 1) | CLOSED) == 0 { 257 | SCHEDULED | CLOSED | REFERENCE 258 | } else { 259 | state & !TASK 260 | }; 261 | 262 | // Unset the `TASK` flag. 263 | match (*header).state.compare_exchange_weak( 264 | state, 265 | new, 266 | Ordering::AcqRel, 267 | Ordering::Acquire, 268 | ) { 269 | Ok(_) => { 270 | // If this is the last reference to the task, we need to either 271 | // schedule dropping its future or destroy it. 272 | if state & !(REFERENCE - 1) == 0 { 273 | if state & CLOSED == 0 { 274 | ((*header).vtable.schedule)(ptr, ScheduleInfo::new(false)); 275 | } else { 276 | ((*header).vtable.destroy)(ptr); 277 | } 278 | } 279 | 280 | break; 281 | } 282 | Err(s) => state = s, 283 | } 284 | } 285 | } 286 | } 287 | 288 | output 289 | } 290 | } 291 | 292 | /// Puts the task in canceled state. 293 | #[inline(never)] 294 | fn set_canceled(ptr: *const ()) { 295 | let header = ptr as *const Header; 296 | 297 | unsafe { 298 | let mut state = (*header).state.load(Ordering::Acquire); 299 | 300 | loop { 301 | // If the task has been completed or closed, it can't be canceled. 302 | if state & (COMPLETED | CLOSED) != 0 { 303 | break; 304 | } 305 | 306 | // If the task is not scheduled nor running, we'll need to schedule it. 307 | let new = if state & (SCHEDULED | RUNNING) == 0 { 308 | (state | SCHEDULED | CLOSED) + REFERENCE 309 | } else { 310 | state | CLOSED 311 | }; 312 | 313 | // Mark the task as closed. 314 | match (*header).state.compare_exchange_weak( 315 | state, 316 | new, 317 | Ordering::AcqRel, 318 | Ordering::Acquire, 319 | ) { 320 | Ok(_) => { 321 | // If the task is not scheduled nor running, schedule it one more time so 322 | // that its future gets dropped by the executor. 323 | if state & (SCHEDULED | RUNNING) == 0 { 324 | ((*header).vtable.schedule)(ptr, ScheduleInfo::new(false)); 325 | } 326 | 327 | // Notify the awaiter that the task has been closed. 328 | if state & AWAITER != 0 { 329 | (*header).notify(None); 330 | } 331 | 332 | break; 333 | } 334 | Err(s) => state = s, 335 | } 336 | } 337 | } 338 | } 339 | 340 | /// Polls the task to retrieve its output. 341 | /// 342 | /// Returns `Some` if the task has completed or `None` if it was closed. 343 | /// 344 | /// A task becomes closed in the following cases: 345 | /// 346 | /// 1. It gets canceled by `Runnable::drop()`, `Task::drop()`, or `Task::cancel()`. 347 | /// 2. Its output gets awaited by the `Task`. 348 | /// 3. It panics while polling the future. 349 | /// 4. It is completed and the `Task` gets dropped. 350 | fn poll_task(ptr: *const (), cx: &mut Context<'_>) -> Poll> { 351 | let header = ptr as *const Header; 352 | 353 | unsafe { 354 | let mut state = (*header).state.load(Ordering::Acquire); 355 | 356 | loop { 357 | // If the task has been closed, notify the awaiter and return `None`. 358 | if state & CLOSED != 0 { 359 | // If the task is scheduled or running, we need to wait until its future is 360 | // dropped. 361 | if state & (SCHEDULED | RUNNING) != 0 { 362 | // Replace the waker with one associated with the current task. 363 | (*header).register(cx.waker()); 364 | 365 | // Reload the state after registering. It is possible changes occurred just 366 | // before registration so we need to check for that. 367 | state = (*header).state.load(Ordering::Acquire); 368 | 369 | // If the task is still scheduled or running, we need to wait because its 370 | // future is not dropped yet. 371 | if state & (SCHEDULED | RUNNING) != 0 { 372 | return Poll::Pending; 373 | } 374 | } 375 | 376 | // Even though the awaiter is most likely the current task, it could also be 377 | // another task. 378 | (*header).notify(Some(cx.waker())); 379 | return Poll::Ready(None); 380 | } 381 | 382 | // If the task is not completed, register the current task. 383 | if state & COMPLETED == 0 { 384 | // Replace the waker with one associated with the current task. 385 | (*header).register(cx.waker()); 386 | 387 | // Reload the state after registering. It is possible that the task became 388 | // completed or closed just before registration so we need to check for that. 389 | state = (*header).state.load(Ordering::Acquire); 390 | 391 | // If the task has been closed, restart. 392 | if state & CLOSED != 0 { 393 | continue; 394 | } 395 | 396 | // If the task is still not completed, we're blocked on it. 397 | if state & COMPLETED == 0 { 398 | return Poll::Pending; 399 | } 400 | } 401 | 402 | // Since the task is now completed, mark it as closed in order to grab its output. 403 | match (*header).state.compare_exchange( 404 | state, 405 | state | CLOSED, 406 | Ordering::AcqRel, 407 | Ordering::Acquire, 408 | ) { 409 | Ok(_) => { 410 | // Notify the awaiter. Even though the awaiter is most likely the current 411 | // task, it could also be another task. 412 | if state & AWAITER != 0 { 413 | (*header).notify(Some(cx.waker())); 414 | } 415 | 416 | // Take the output from the task. 417 | let output = (*header).vtable.get_output(ptr) as *mut Result; 418 | let output = output.read(); 419 | 420 | // Propagate the panic if the task panicked. 421 | let output = match output { 422 | Ok(output) => output, 423 | Err(panic) => { 424 | #[cfg(feature = "std")] 425 | std::panic::resume_unwind(panic); 426 | 427 | #[cfg(not(feature = "std"))] 428 | match panic {} 429 | } 430 | }; 431 | 432 | return Poll::Ready(Some(output)); 433 | } 434 | Err(s) => state = s, 435 | } 436 | } 437 | } 438 | } 439 | 440 | impl Drop for Task { 441 | fn drop(&mut self) { 442 | let ptr = self.ptr.as_ptr(); 443 | set_canceled(ptr); 444 | set_detached::(ptr); 445 | } 446 | } 447 | 448 | impl Future for Task { 449 | type Output = T; 450 | 451 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 452 | match poll_task::(self.ptr.as_ptr(), cx) { 453 | Poll::Ready(t) => Poll::Ready(t.expect("Task polled after completion")), 454 | Poll::Pending => Poll::Pending, 455 | } 456 | } 457 | } 458 | 459 | impl fmt::Debug for Task { 460 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 461 | f.debug_struct("Task") 462 | .field("header", self.header_with_metadata()) 463 | .finish() 464 | } 465 | } 466 | 467 | /// A spawned task with a fallible response. 468 | /// 469 | /// This type behaves like [`Task`], however it produces an `Option` when 470 | /// polled and will return `None` if the executor dropped its 471 | /// [`Runnable`][`super::Runnable`] without being run. 472 | /// 473 | /// This can be useful to avoid the panic produced when polling the `Task` 474 | /// future if the executor dropped its `Runnable`. 475 | #[must_use = "tasks get canceled when dropped, use `.detach()` to run them in the background"] 476 | pub struct FallibleTask { 477 | task: Task, 478 | } 479 | 480 | impl FallibleTask { 481 | /// Detaches the task to let it keep running in the background. 482 | /// 483 | /// # Examples 484 | /// 485 | /// ``` 486 | /// use smol::{Executor, Timer}; 487 | /// use std::time::Duration; 488 | /// 489 | /// let ex = Executor::new(); 490 | /// 491 | /// // Spawn a daemon future. 492 | /// ex.spawn(async { 493 | /// loop { 494 | /// println!("I'm a daemon task looping forever."); 495 | /// Timer::after(Duration::from_secs(1)).await; 496 | /// } 497 | /// }) 498 | /// .fallible() 499 | /// .detach(); 500 | /// ``` 501 | pub fn detach(self) { 502 | self.task.detach() 503 | } 504 | 505 | /// Cancels the task and waits for it to stop running. 506 | /// 507 | /// Returns the task's output if it was completed just before it got canceled, or [`None`] if 508 | /// it didn't complete. 509 | /// 510 | /// While it's possible to simply drop the [`Task`] to cancel it, this is a cleaner way of 511 | /// canceling because it also waits for the task to stop running. 512 | /// 513 | /// # Examples 514 | /// 515 | /// ``` 516 | /// # if cfg!(miri) { return; } // Miri does not support epoll 517 | /// use smol::{future, Executor, Timer}; 518 | /// use std::thread; 519 | /// use std::time::Duration; 520 | /// 521 | /// let ex = Executor::new(); 522 | /// 523 | /// // Spawn a daemon future. 524 | /// let task = ex.spawn(async { 525 | /// loop { 526 | /// println!("Even though I'm in an infinite loop, you can still cancel me!"); 527 | /// Timer::after(Duration::from_secs(1)).await; 528 | /// } 529 | /// }) 530 | /// .fallible(); 531 | /// 532 | /// // Run an executor thread. 533 | /// thread::spawn(move || future::block_on(ex.run(future::pending::<()>()))); 534 | /// 535 | /// future::block_on(async { 536 | /// Timer::after(Duration::from_secs(3)).await; 537 | /// task.cancel().await; 538 | /// }); 539 | /// ``` 540 | pub async fn cancel(self) -> Option { 541 | self.task.cancel().await 542 | } 543 | 544 | /// Returns `true` if the current task is finished. 545 | /// 546 | /// Note that in a multithreaded environment, this task can change finish immediately after calling this function. 547 | pub fn is_finished(&self) -> bool { 548 | self.task.is_finished() 549 | } 550 | } 551 | 552 | impl Future for FallibleTask { 553 | type Output = Option; 554 | 555 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 556 | poll_task::(self.task.ptr.as_ptr(), cx) 557 | } 558 | } 559 | 560 | impl fmt::Debug for FallibleTask { 561 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 562 | f.debug_struct("FallibleTask") 563 | .field("header", self.task.header_with_metadata()) 564 | .finish() 565 | } 566 | } 567 | -------------------------------------------------------------------------------- /src/raw.rs: -------------------------------------------------------------------------------- 1 | use alloc::alloc::Layout as StdLayout; 2 | use core::future::Future; 3 | use core::marker::PhantomData; 4 | use core::mem::{self, ManuallyDrop}; 5 | use core::pin::Pin; 6 | use core::ptr::NonNull; 7 | use core::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; 8 | 9 | use core::sync::atomic::Ordering; 10 | 11 | use crate::header::{DropWakerAction, Header, HeaderWithMetadata}; 12 | use crate::runnable::{Schedule, ScheduleInfo}; 13 | use crate::state::*; 14 | use crate::utils::{abort, abort_on_panic, max, Layout}; 15 | use crate::Runnable; 16 | 17 | #[cfg(feature = "std")] 18 | pub(crate) type Panic = alloc::boxed::Box; 19 | 20 | #[cfg(not(feature = "std"))] 21 | pub(crate) type Panic = core::convert::Infallible; 22 | 23 | /// The vtable for a task. 24 | pub(crate) struct TaskVTable { 25 | pub(crate) raw_waker_vtable: &'static RawWakerVTable, 26 | 27 | /// Schedules the task. 28 | pub(crate) schedule: unsafe fn(*const (), ScheduleInfo), 29 | 30 | /// Drops the future inside the task. 31 | pub(crate) drop_future: unsafe fn(*const (), &TaskLayout), 32 | 33 | /// Destroys the task. 34 | pub(crate) destroy: unsafe fn(*const ()), 35 | 36 | /// Runs the task. 37 | pub(crate) run: unsafe fn(*const ()) -> bool, 38 | 39 | /// The memory layout of the task. This information enables 40 | /// debuggers to decode raw task memory blobs. Do not remove 41 | /// the field, even if it appears to be unused. 42 | pub(crate) layout_info: &'static TaskLayout, 43 | } 44 | 45 | impl TaskVTable { 46 | /// Returns a pointer to the output inside a task. 47 | pub(crate) unsafe fn get_output(&self, ptr: *const ()) -> *const () { 48 | ptr.add_byte(self.layout_info.offset_r) 49 | } 50 | } 51 | 52 | /// Memory layout of a task. 53 | /// 54 | /// This struct contains the following information: 55 | /// 56 | /// 1. How to allocate and deallocate the task. 57 | /// 2. How to access the fields inside the task. 58 | #[derive(Clone, Copy)] 59 | pub(crate) struct TaskLayout { 60 | /// Memory layout of the whole task. 61 | pub(crate) layout: StdLayout, 62 | 63 | /// Offset into the task at which the schedule function is stored. 64 | pub(crate) offset_s: usize, 65 | 66 | /// Offset into the task at which the future is stored. 67 | pub(crate) offset_f: usize, 68 | 69 | /// Offset into the task at which the output is stored. 70 | pub(crate) offset_r: usize, 71 | } 72 | 73 | /// Raw pointers to the fields inside a task. 74 | pub(crate) struct RawTask { 75 | /// The task header. 76 | pub(crate) header: *const HeaderWithMetadata, 77 | 78 | /// The schedule function. 79 | pub(crate) schedule: *const S, 80 | 81 | /// The future. 82 | pub(crate) future: *mut F, 83 | 84 | /// The output of the future. 85 | pub(crate) output: *mut Result, 86 | } 87 | 88 | impl Copy for RawTask {} 89 | 90 | impl Clone for RawTask { 91 | fn clone(&self) -> Self { 92 | *self 93 | } 94 | } 95 | 96 | impl RawTask { 97 | pub(crate) const TASK_LAYOUT: TaskLayout = Self::eval_task_layout(); 98 | 99 | /// Computes the memory layout for a task. 100 | #[inline] 101 | const fn eval_task_layout() -> TaskLayout { 102 | // Compute the layouts for `Header`, `S`, `F`, and `T`. 103 | let layout_header = Layout::new::>(); 104 | let layout_s = Layout::new::(); 105 | let layout_f = Layout::new::(); 106 | let layout_r = Layout::new::>(); 107 | 108 | // Compute the layout for `union { F, T }`. 109 | let size_union = max(layout_f.size(), layout_r.size()); 110 | let align_union = max(layout_f.align(), layout_r.align()); 111 | let layout_union = Layout::from_size_align(size_union, align_union); 112 | 113 | // Compute the layout for `Header` followed by `S` and `union { F, T }`. 114 | let layout = layout_header; 115 | let (layout, offset_s) = leap_unwrap!(layout.extend(layout_s)); 116 | let (layout, offset_union) = leap_unwrap!(layout.extend(layout_union)); 117 | let offset_f = offset_union; 118 | let offset_r = offset_union; 119 | 120 | TaskLayout { 121 | layout: unsafe { layout.into_std() }, 122 | offset_s, 123 | offset_f, 124 | offset_r, 125 | } 126 | } 127 | } 128 | 129 | /// Allocates a task with the given `future` and `schedule` function. 130 | /// 131 | /// It is assumed that initially only the `Runnable` and the `Task` exist. 132 | /// 133 | /// Use a macro to brute force inlining to minimize stack copies of potentially 134 | /// large futures. 135 | macro_rules! allocate_task { 136 | ($f:tt, $s:tt, $m:tt, $builder:ident, $schedule:ident, $raw:ident => $future:block) => {{ 137 | let allocation = 138 | alloc::alloc::alloc(RawTask::<$f, <$f as Future>::Output, $s, $m>::TASK_LAYOUT.layout); 139 | // Allocate enough space for the entire task. 140 | let ptr = NonNull::new(allocation as *mut ()).unwrap_or_else(|| crate::utils::abort()); 141 | 142 | let $raw = RawTask::<$f, <$f as Future>::Output, $s, $m>::from_ptr(ptr.as_ptr()); 143 | 144 | let crate::Builder { 145 | metadata, 146 | #[cfg(feature = "std")] 147 | propagate_panic, 148 | } = $builder; 149 | 150 | // Write the header as the first field of the task. 151 | ($raw.header as *mut HeaderWithMetadata<$m>).write(HeaderWithMetadata { 152 | header: Header { 153 | #[cfg(not(feature = "portable-atomic"))] 154 | state: core::sync::atomic::AtomicUsize::new(SCHEDULED | TASK | REFERENCE), 155 | #[cfg(feature = "portable-atomic")] 156 | state: portable_atomic::AtomicUsize::new(SCHEDULED | TASK | REFERENCE), 157 | awaiter: core::cell::UnsafeCell::new(None), 158 | vtable: &RawTask::<$f, <$f as Future>::Output, $s, $m>::TASK_VTABLE, 159 | #[cfg(feature = "std")] 160 | propagate_panic, 161 | }, 162 | metadata, 163 | }); 164 | 165 | // Write the schedule function as the third field of the task. 166 | ($raw.schedule as *mut S).write($schedule); 167 | 168 | // Explicitly avoid using abort_on_panic here to avoid extra stack 169 | // copies of the future on lower optimization levels. 170 | let bomb = crate::utils::Bomb; 171 | 172 | // Generate the future, now that the metadata has been pinned in place. 173 | // Write the future as the fourth field of the task. 174 | $raw.future.write($future); 175 | // (&(*raw.header).metadata) 176 | 177 | mem::forget(bomb); 178 | ptr 179 | }}; 180 | } 181 | 182 | pub(crate) use allocate_task; 183 | 184 | impl RawTask 185 | where 186 | F: Future, 187 | S: Schedule, 188 | { 189 | pub(crate) const RAW_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new( 190 | Header::clone_waker, 191 | wake::, 192 | wake_by_ref::, 193 | drop_waker, 194 | ); 195 | 196 | pub(crate) const TASK_VTABLE: TaskVTable = TaskVTable { 197 | raw_waker_vtable: &Self::RAW_WAKER_VTABLE, 198 | schedule: schedule::, 199 | drop_future: drop_future::, 200 | destroy: destroy::, 201 | run: Self::run, 202 | layout_info: &Self::TASK_LAYOUT, 203 | }; 204 | 205 | /// Creates a `RawTask` from a raw task pointer. 206 | #[inline] 207 | pub(crate) fn from_ptr(ptr: *const ()) -> Self { 208 | unsafe { 209 | Self { 210 | header: ptr as *const HeaderWithMetadata, 211 | schedule: ptr.add_byte(Self::TASK_LAYOUT.offset_s) as *const S, 212 | future: ptr.add_byte(Self::TASK_LAYOUT.offset_f) as *mut F, 213 | output: ptr.add_byte(Self::TASK_LAYOUT.offset_r) as *mut Result, 214 | } 215 | } 216 | } 217 | 218 | /// Runs a task. 219 | /// 220 | /// If polling its future panics, the task will be closed and the panic will be propagated into 221 | /// the caller. 222 | unsafe fn run(ptr: *const ()) -> bool { 223 | let raw = Self::from_ptr(ptr); 224 | let header = ptr as *const Header; 225 | 226 | // Create a context from the raw task pointer and the vtable inside the its header. 227 | let waker = ManuallyDrop::new(Waker::from_raw(RawWaker::new(ptr, &Self::RAW_WAKER_VTABLE))); 228 | let cx = &mut Context::from_waker(&waker); 229 | 230 | let mut state = (*header).state.load(Ordering::Acquire); 231 | 232 | // Update the task's state before polling its future. 233 | loop { 234 | // If the task has already been closed, drop the task reference and return. 235 | if state & CLOSED != 0 { 236 | // Drop the future. 237 | drop_future::(ptr, &Self::TASK_LAYOUT); 238 | 239 | // Mark the task as unscheduled. 240 | let state = (*header).state.fetch_and(!SCHEDULED, Ordering::AcqRel); 241 | 242 | // Take the awaiter out. 243 | let mut awaiter = None; 244 | if state & AWAITER != 0 { 245 | awaiter = (*header).take(None); 246 | } 247 | 248 | // Drop the task reference. 249 | drop_ref(ptr); 250 | 251 | // Notify the awaiter that the future has been dropped. 252 | if let Some(w) = awaiter { 253 | abort_on_panic(|| w.wake()); 254 | } 255 | return false; 256 | } 257 | 258 | // Mark the task as unscheduled and running. 259 | match (*header).state.compare_exchange_weak( 260 | state, 261 | (state & !SCHEDULED) | RUNNING, 262 | Ordering::AcqRel, 263 | Ordering::Acquire, 264 | ) { 265 | Ok(_) => { 266 | // Update the state because we're continuing with polling the future. 267 | state = (state & !SCHEDULED) | RUNNING; 268 | break; 269 | } 270 | Err(s) => state = s, 271 | } 272 | } 273 | 274 | // Poll the inner future, but surround it with a guard that closes the task in case polling 275 | // panics. 276 | // If available, we should also try to catch the panic so that it is propagated correctly. 277 | let guard = Guard::(ptr, &Self::TASK_LAYOUT, PhantomData); 278 | 279 | // Panic propagation is not available for no_std. 280 | #[cfg(not(feature = "std"))] 281 | let poll = ::poll(Pin::new_unchecked(&mut *raw.future), cx).map(Ok); 282 | 283 | #[cfg(feature = "std")] 284 | let poll = { 285 | // Check if we should propagate panics. 286 | if (*header).propagate_panic { 287 | // Use catch_unwind to catch the panic. 288 | match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { 289 | ::poll(Pin::new_unchecked(&mut *raw.future), cx) 290 | })) { 291 | Ok(Poll::Ready(v)) => Poll::Ready(Ok(v)), 292 | Ok(Poll::Pending) => Poll::Pending, 293 | Err(e) => Poll::Ready(Err(e)), 294 | } 295 | } else { 296 | ::poll(Pin::new_unchecked(&mut *raw.future), cx).map(Ok) 297 | } 298 | }; 299 | 300 | mem::forget(guard); 301 | 302 | match poll { 303 | Poll::Ready(out) => { 304 | // Replace the future with its output. 305 | drop_future::(ptr, &Self::TASK_LAYOUT); 306 | raw.output.write(out); 307 | 308 | // The task is now completed. 309 | loop { 310 | // If the `Task` is dropped, we'll need to close it and drop the output. 311 | let new = if state & TASK == 0 { 312 | (state & !RUNNING & !SCHEDULED) | COMPLETED | CLOSED 313 | } else { 314 | (state & !RUNNING & !SCHEDULED) | COMPLETED 315 | }; 316 | 317 | // Mark the task as not running and completed. 318 | match (*header).state.compare_exchange_weak( 319 | state, 320 | new, 321 | Ordering::AcqRel, 322 | Ordering::Acquire, 323 | ) { 324 | Ok(_) => { 325 | // If the `Task` is dropped or if the task was closed while running, 326 | // now it's time to drop the output. 327 | if state & TASK == 0 || state & CLOSED != 0 { 328 | // Drop the output. 329 | abort_on_panic(|| raw.output.drop_in_place()); 330 | } 331 | 332 | // Take the awaiter out. 333 | let mut awaiter = None; 334 | if state & AWAITER != 0 { 335 | awaiter = (*header).take(None); 336 | } 337 | 338 | // Drop the task reference. 339 | drop_ref(ptr); 340 | 341 | // Notify the awaiter that the future has been dropped. 342 | if let Some(w) = awaiter { 343 | abort_on_panic(|| w.wake()); 344 | } 345 | break; 346 | } 347 | Err(s) => state = s, 348 | } 349 | } 350 | } 351 | Poll::Pending => { 352 | let mut future_dropped = false; 353 | 354 | // The task is still not completed. 355 | loop { 356 | // If the task was closed while running, we'll need to unschedule in case it 357 | // was woken up and then destroy it. 358 | let new = if state & CLOSED != 0 { 359 | state & !RUNNING & !SCHEDULED 360 | } else { 361 | state & !RUNNING 362 | }; 363 | 364 | if state & CLOSED != 0 && !future_dropped { 365 | // The thread that closed the task didn't drop the future because it was 366 | // running so now it's our responsibility to do so. 367 | drop_future::(ptr, &Self::TASK_LAYOUT); 368 | future_dropped = true; 369 | } 370 | 371 | // Mark the task as not running. 372 | match (*header).state.compare_exchange_weak( 373 | state, 374 | new, 375 | Ordering::AcqRel, 376 | Ordering::Acquire, 377 | ) { 378 | Ok(state) => { 379 | // If the task was closed while running, we need to notify the awaiter. 380 | // If the task was woken up while running, we need to schedule it. 381 | // Otherwise, we just drop the task reference. 382 | if state & CLOSED != 0 { 383 | // Take the awaiter out. 384 | let mut awaiter = None; 385 | if state & AWAITER != 0 { 386 | awaiter = (*header).take(None); 387 | } 388 | 389 | // Drop the task reference. 390 | drop_ref(ptr); 391 | 392 | // Notify the awaiter that the future has been dropped. 393 | if let Some(w) = awaiter { 394 | abort_on_panic(|| w.wake()); 395 | } 396 | } else if state & SCHEDULED != 0 { 397 | // The thread that woke the task up didn't reschedule it because 398 | // it was running so now it's our responsibility to do so. 399 | schedule::(ptr, ScheduleInfo::new(true)); 400 | return true; 401 | } else { 402 | // Drop the task reference. 403 | drop_ref(ptr); 404 | } 405 | break; 406 | } 407 | Err(s) => state = s, 408 | } 409 | } 410 | } 411 | } 412 | 413 | false 414 | } 415 | } 416 | 417 | /// A guard that closes the task if polling its future panics. 418 | struct Guard(*const (), &'static TaskLayout, PhantomData F>) 419 | where 420 | F: Future; 421 | 422 | impl Drop for Guard 423 | where 424 | F: Future, 425 | { 426 | fn drop(&mut self) { 427 | let ptr = self.0; 428 | let task_layout = self.1; 429 | let header = ptr as *const Header; 430 | 431 | unsafe { 432 | let header = &*header; 433 | let mut state = header.state.load(Ordering::Acquire); 434 | 435 | loop { 436 | // If the task was closed while running, then unschedule it, drop its 437 | // future, and drop the task reference. 438 | if state & CLOSED != 0 { 439 | // The thread that closed the task didn't drop the future because it 440 | // was running so now it's our responsibility to do so. 441 | drop_future::(ptr, task_layout); 442 | 443 | // Mark the task as not running and not scheduled. 444 | header 445 | .state 446 | .fetch_and(!RUNNING & !SCHEDULED, Ordering::AcqRel); 447 | 448 | // Take the awaiter out. 449 | let mut awaiter = None; 450 | if state & AWAITER != 0 { 451 | awaiter = header.take(None); 452 | } 453 | 454 | // Drop the task reference. 455 | drop_ref(ptr); 456 | 457 | // Notify the awaiter that the future has been dropped. 458 | if let Some(w) = awaiter { 459 | abort_on_panic(|| w.wake()); 460 | } 461 | break; 462 | } 463 | 464 | // Mark the task as not running, not scheduled, and closed. 465 | match header.state.compare_exchange_weak( 466 | state, 467 | (state & !RUNNING & !SCHEDULED) | CLOSED, 468 | Ordering::AcqRel, 469 | Ordering::Acquire, 470 | ) { 471 | Ok(state) => { 472 | // Drop the future because the task is now closed. 473 | drop_future::(ptr, task_layout); 474 | 475 | // Take the awaiter out. 476 | let mut awaiter = None; 477 | if state & AWAITER != 0 { 478 | awaiter = header.take(None); 479 | } 480 | 481 | // Drop the task reference. 482 | drop_ref(ptr); 483 | 484 | // Notify the awaiter that the future has been dropped. 485 | if let Some(w) = awaiter { 486 | abort_on_panic(|| w.wake()); 487 | } 488 | break; 489 | } 490 | Err(s) => state = s, 491 | } 492 | } 493 | } 494 | } 495 | } 496 | 497 | /// Schedules a task for running. 498 | /// 499 | /// This function doesn't modify the state of the task. It only passes the task reference to 500 | /// its schedule function. 501 | unsafe fn schedule, M>(ptr: *const (), info: ScheduleInfo) { 502 | let header = ptr as *const Header; 503 | let task_layout = (*header).vtable.layout_info; 504 | let schedule = ptr.add_byte(task_layout.offset_s) as *mut S; 505 | 506 | // If the schedule function has captured variables, create a temporary waker that prevents 507 | // the task from getting deallocated while the function is being invoked. 508 | let _waker; 509 | if mem::size_of::() > 0 { 510 | _waker = Waker::from_raw(Header::clone_waker(ptr)); 511 | } 512 | 513 | let task = Runnable::from_raw(NonNull::new_unchecked(ptr as *mut ())); 514 | (*schedule).schedule(task, info); 515 | } 516 | 517 | /// Drops a waker. 518 | /// 519 | /// This function will decrement the reference count. If it drops down to zero, the associated 520 | /// `Task` has been dropped too, and the task has not been completed, then it will get 521 | /// scheduled one more time so that its future gets dropped by the executor. 522 | #[inline] 523 | unsafe fn drop_waker(ptr: *const ()) { 524 | let header = ptr as *const Header; 525 | match Header::drop_waker(ptr) { 526 | DropWakerAction::Schedule => ((*header).vtable.schedule)(ptr, ScheduleInfo::new(false)), 527 | DropWakerAction::Destroy => ((*header).vtable.destroy)(ptr), 528 | DropWakerAction::None => {} 529 | } 530 | } 531 | 532 | /// Drops the future inside a task. 533 | #[inline] 534 | unsafe fn drop_future(ptr: *const (), task_layout: &TaskLayout) { 535 | let future_ptr = ptr.add_byte(task_layout.offset_f) as *mut F; 536 | 537 | // We need a safeguard against panics because the destructor can panic. 538 | abort_on_panic(|| { 539 | future_ptr.drop_in_place(); 540 | }) 541 | } 542 | 543 | /// Wakes a waker. 544 | unsafe fn wake, M>(ptr: *const ()) { 545 | // This is just an optimization. If the schedule function has captured variables, then 546 | // we'll do less reference counting if we wake the waker by reference and then drop it. 547 | if mem::size_of::() > 0 { 548 | wake_by_ref::(ptr); 549 | drop_waker(ptr); 550 | return; 551 | } 552 | 553 | let header = ptr as *const Header; 554 | 555 | let mut state = (*header).state.load(Ordering::Acquire); 556 | 557 | loop { 558 | // If the task is completed or closed, it can't be woken up. 559 | if state & (COMPLETED | CLOSED) != 0 { 560 | // Drop the waker. 561 | drop_waker(ptr); 562 | break; 563 | } 564 | 565 | // If the task is already scheduled, we just need to synchronize with the thread that 566 | // will run the task by "publishing" our current view of the memory. 567 | if state & SCHEDULED != 0 { 568 | // Update the state without actually modifying it. 569 | match (*header).state.compare_exchange_weak( 570 | state, 571 | state, 572 | Ordering::AcqRel, 573 | Ordering::Acquire, 574 | ) { 575 | Ok(_) => { 576 | // Drop the waker. 577 | drop_waker(ptr); 578 | break; 579 | } 580 | Err(s) => state = s, 581 | } 582 | } else { 583 | // Mark the task as scheduled. 584 | match (*header).state.compare_exchange_weak( 585 | state, 586 | state | SCHEDULED, 587 | Ordering::AcqRel, 588 | Ordering::Acquire, 589 | ) { 590 | Ok(_) => { 591 | // If the task is not yet scheduled and isn't currently running, now is the 592 | // time to schedule it. 593 | if state & RUNNING == 0 { 594 | // Schedule the task. 595 | schedule::(ptr, ScheduleInfo::new(false)); 596 | } else { 597 | // Drop the waker. 598 | drop_waker(ptr); 599 | } 600 | 601 | break; 602 | } 603 | Err(s) => state = s, 604 | } 605 | } 606 | } 607 | } 608 | 609 | /// Wakes a waker by reference. 610 | unsafe fn wake_by_ref, M>(ptr: *const ()) { 611 | let header = ptr as *const Header; 612 | let header = &*header; 613 | let task_layout = header.vtable.layout_info; 614 | 615 | let mut state = header.state.load(Ordering::Acquire); 616 | 617 | loop { 618 | // If the task is completed or closed, it can't be woken up. 619 | if state & (COMPLETED | CLOSED) != 0 { 620 | break; 621 | } 622 | 623 | // If the task is already scheduled, we just need to synchronize with the thread that 624 | // will run the task by "publishing" our current view of the memory. 625 | if state & SCHEDULED != 0 { 626 | // Update the state without actually modifying it. 627 | match header.state.compare_exchange_weak( 628 | state, 629 | state, 630 | Ordering::AcqRel, 631 | Ordering::Acquire, 632 | ) { 633 | Ok(_) => break, 634 | Err(s) => state = s, 635 | } 636 | } else { 637 | // If the task is not running, we can schedule right away. 638 | let new = if state & RUNNING == 0 { 639 | (state | SCHEDULED) + REFERENCE 640 | } else { 641 | state | SCHEDULED 642 | }; 643 | 644 | // Mark the task as scheduled. 645 | match header.state.compare_exchange_weak( 646 | state, 647 | new, 648 | Ordering::AcqRel, 649 | Ordering::Acquire, 650 | ) { 651 | Ok(_) => { 652 | // If the task is not running, now is the time to schedule. 653 | if state & RUNNING == 0 { 654 | // If the reference count overflowed, abort. 655 | if state > isize::MAX as usize { 656 | abort(); 657 | } 658 | 659 | let schedule = ptr.add_byte(task_layout.offset_s) as *mut S; 660 | 661 | // Schedule the task. There is no need to call `Self::schedule(ptr)` 662 | // because the schedule function cannot be destroyed while the waker is 663 | // still alive. 664 | let task = Runnable::from_raw(NonNull::new_unchecked(ptr as *mut ())); 665 | (*schedule).schedule(task, ScheduleInfo::new(false)); 666 | } 667 | 668 | break; 669 | } 670 | Err(s) => state = s, 671 | } 672 | } 673 | } 674 | } 675 | 676 | /// Cleans up task's resources and deallocates it. 677 | /// 678 | /// The schedule function will be dropped, and the task will then get deallocated. 679 | /// The task must be closed before this function is called. 680 | #[inline] 681 | unsafe fn destroy(ptr: *const ()) { 682 | let header = ptr as *const Header; 683 | let task_layout = (*header).vtable.layout_info; 684 | let schedule = ptr.add_byte(task_layout.offset_s); 685 | 686 | // We need a safeguard against panics because destructors can panic. 687 | abort_on_panic(|| { 688 | // Drop the header along with the metadata. 689 | (ptr as *mut HeaderWithMetadata).drop_in_place(); 690 | 691 | // Drop the schedule function. 692 | (schedule as *mut S).drop_in_place(); 693 | }); 694 | 695 | // Finally, deallocate the memory reserved by the task. 696 | alloc::alloc::dealloc(ptr as *mut u8, task_layout.layout); 697 | } 698 | 699 | /// Drops a task reference (`Runnable` or `Waker`). 700 | /// 701 | /// This function will decrement the reference count. If it drops down to zero and the 702 | /// associated `Task` handle has been dropped too, then the task gets destroyed. 703 | #[inline] 704 | pub(crate) unsafe fn drop_ref(ptr: *const ()) { 705 | let header = ptr as *const Header; 706 | let header = &*header; 707 | 708 | // Decrement the reference count. 709 | let new = header.state.fetch_sub(REFERENCE, Ordering::AcqRel) - REFERENCE; 710 | 711 | // If this was the last reference to the task and the `Task` has been dropped too, 712 | // then destroy the task. 713 | if new & !(REFERENCE - 1) == 0 && new & TASK == 0 { 714 | (header.vtable.destroy)(ptr); 715 | } 716 | } 717 | 718 | trait PointerPolyfill { 719 | // Polyfill for `byte_add`. 720 | // TODO: Replace this with `byte_add` once the MSRV should be bumped past 1.75 721 | /// Adds an unsigned offset in bytes to a pointer. 722 | /// 723 | /// `count` is in units of bytes. 724 | /// 725 | /// This is purely a convenience for casting to a `u8` pointer and 726 | /// using [add][pointer::add] on it. See that method for documentation 727 | /// and safety requirements. 728 | /// 729 | /// # Safety 730 | /// If any of the following conditions are violated, the result is Undefined Behavior: 731 | /// 732 | /// - The offset in bytes, count * size_of::(), computed on mathematical integers 733 | /// (without “wrapping around”), must fit in an isize. 734 | /// - If the computed offset is non-zero, then self must be derived from a pointer to 735 | /// some allocation, and the entire memory range between self and the result must be 736 | /// in bounds of that allocation. In particular, this range must not “wrap around” 737 | /// the edge of the address space. 738 | unsafe fn add_byte(self, size: usize) -> Self; 739 | } 740 | 741 | impl PointerPolyfill for *const T { 742 | #[inline] 743 | unsafe fn add_byte(self, size: usize) -> Self { 744 | (self.cast::().add(size)).cast::() 745 | } 746 | } 747 | -------------------------------------------------------------------------------- /src/runnable.rs: -------------------------------------------------------------------------------- 1 | use core::fmt; 2 | use core::future::Future; 3 | use core::marker::PhantomData; 4 | use core::mem; 5 | use core::ptr::NonNull; 6 | use core::sync::atomic::Ordering; 7 | use core::task::Waker; 8 | 9 | use crate::header::Header; 10 | use crate::header::HeaderWithMetadata; 11 | use crate::raw::drop_ref; 12 | use crate::raw::{allocate_task, RawTask}; 13 | use crate::state::*; 14 | use crate::Task; 15 | 16 | mod sealed { 17 | use super::*; 18 | pub trait Sealed {} 19 | 20 | impl Sealed for F where F: Fn(Runnable) {} 21 | 22 | impl Sealed for WithInfo where F: Fn(Runnable, ScheduleInfo) {} 23 | } 24 | 25 | /// A builder that creates a new task. 26 | #[derive(Debug)] 27 | pub struct Builder { 28 | /// The metadata associated with the task. 29 | pub(crate) metadata: M, 30 | 31 | /// Whether or not a panic that occurs in the task should be propagated. 32 | #[cfg(feature = "std")] 33 | pub(crate) propagate_panic: bool, 34 | } 35 | 36 | impl Default for Builder { 37 | fn default() -> Self { 38 | Builder::new().metadata(M::default()) 39 | } 40 | } 41 | 42 | /// Extra scheduling information that can be passed to the scheduling function. 43 | /// 44 | /// The data source of this struct is directly from the actual implementation 45 | /// of the crate itself, different from [`Runnable`]'s metadata, which is 46 | /// managed by the caller. 47 | /// 48 | /// # Examples 49 | /// 50 | /// ``` 51 | /// use async_task::{Runnable, ScheduleInfo, WithInfo}; 52 | /// use std::sync::{Arc, Mutex}; 53 | /// 54 | /// // The future inside the task. 55 | /// let future = async { 56 | /// println!("Hello, world!"); 57 | /// }; 58 | /// 59 | /// // If the task gets woken up while running, it will be sent into this channel. 60 | /// let (s, r) = flume::unbounded(); 61 | /// // Otherwise, it will be placed into this slot. 62 | /// let lifo_slot = Arc::new(Mutex::new(None)); 63 | /// let schedule = move |runnable: Runnable, info: ScheduleInfo| { 64 | /// if info.woken_while_running { 65 | /// s.send(runnable).unwrap() 66 | /// } else { 67 | /// let last = lifo_slot.lock().unwrap().replace(runnable); 68 | /// if let Some(last) = last { 69 | /// s.send(last).unwrap() 70 | /// } 71 | /// } 72 | /// }; 73 | /// 74 | /// // Create the actual scheduler to be spawned with some future. 75 | /// let scheduler = WithInfo(schedule); 76 | /// // Create a task with the future and the scheduler. 77 | /// let (runnable, task) = async_task::spawn(future, scheduler); 78 | /// ``` 79 | #[derive(Debug, Copy, Clone)] 80 | #[non_exhaustive] 81 | pub struct ScheduleInfo { 82 | /// Indicates whether the task gets woken up while running. 83 | /// 84 | /// It is set to true usually because the task has yielded itself to the 85 | /// scheduler. 86 | pub woken_while_running: bool, 87 | } 88 | 89 | impl ScheduleInfo { 90 | pub(crate) fn new(woken_while_running: bool) -> Self { 91 | ScheduleInfo { 92 | woken_while_running, 93 | } 94 | } 95 | } 96 | 97 | /// The trait for scheduling functions. 98 | pub trait Schedule: sealed::Sealed { 99 | /// The actual scheduling procedure. 100 | fn schedule(&self, runnable: Runnable, info: ScheduleInfo); 101 | } 102 | 103 | impl Schedule for F 104 | where 105 | F: Fn(Runnable), 106 | { 107 | fn schedule(&self, runnable: Runnable, _: ScheduleInfo) { 108 | self(runnable) 109 | } 110 | } 111 | 112 | /// Pass a scheduling function with more scheduling information - a.k.a. 113 | /// [`ScheduleInfo`]. 114 | /// 115 | /// Sometimes, it's useful to pass the runnable's state directly to the 116 | /// scheduling function, such as whether it's woken up while running. The 117 | /// scheduler can thus use the information to determine its scheduling 118 | /// strategy. 119 | /// 120 | /// The data source of [`ScheduleInfo`] is directly from the actual 121 | /// implementation of the crate itself, different from [`Runnable`]'s metadata, 122 | /// which is managed by the caller. 123 | /// 124 | /// # Examples 125 | /// 126 | /// ``` 127 | /// use async_task::{ScheduleInfo, WithInfo}; 128 | /// use std::sync::{Arc, Mutex}; 129 | /// 130 | /// // The future inside the task. 131 | /// let future = async { 132 | /// println!("Hello, world!"); 133 | /// }; 134 | /// 135 | /// // If the task gets woken up while running, it will be sent into this channel. 136 | /// let (s, r) = flume::unbounded(); 137 | /// // Otherwise, it will be placed into this slot. 138 | /// let lifo_slot = Arc::new(Mutex::new(None)); 139 | /// let schedule = move |runnable, info: ScheduleInfo| { 140 | /// if info.woken_while_running { 141 | /// s.send(runnable).unwrap() 142 | /// } else { 143 | /// let last = lifo_slot.lock().unwrap().replace(runnable); 144 | /// if let Some(last) = last { 145 | /// s.send(last).unwrap() 146 | /// } 147 | /// } 148 | /// }; 149 | /// 150 | /// // Create a task with the future and the schedule function. 151 | /// let (runnable, task) = async_task::spawn(future, WithInfo(schedule)); 152 | /// ``` 153 | #[derive(Debug)] 154 | pub struct WithInfo(pub F); 155 | 156 | impl From for WithInfo { 157 | fn from(value: F) -> Self { 158 | WithInfo(value) 159 | } 160 | } 161 | 162 | impl Schedule for WithInfo 163 | where 164 | F: Fn(Runnable, ScheduleInfo), 165 | { 166 | fn schedule(&self, runnable: Runnable, info: ScheduleInfo) { 167 | (self.0)(runnable, info) 168 | } 169 | } 170 | 171 | impl Builder<()> { 172 | /// Creates a new task builder. 173 | /// 174 | /// By default, this task builder has no metadata. Use the [`metadata`] method to 175 | /// set the metadata. 176 | /// 177 | /// # Examples 178 | /// 179 | /// ``` 180 | /// use async_task::Builder; 181 | /// 182 | /// let (runnable, task) = Builder::new().spawn(|()| async {}, |_| {}); 183 | /// ``` 184 | pub fn new() -> Builder<()> { 185 | Builder { 186 | metadata: (), 187 | #[cfg(feature = "std")] 188 | propagate_panic: false, 189 | } 190 | } 191 | 192 | /// Adds metadata to the task. 193 | /// 194 | /// In certain cases, it may be useful to associate some metadata with a task. For instance, 195 | /// you may want to associate a name with a task, or a priority for a priority queue. This 196 | /// method allows the user to attach arbitrary metadata to a task that is available through 197 | /// the [`Runnable`] or the [`Task`]. 198 | /// 199 | /// # Examples 200 | /// 201 | /// This example creates an executor that associates a "priority" number with each task, and 202 | /// then runs the tasks in order of priority. 203 | /// 204 | /// ``` 205 | /// use async_task::{Builder, Runnable}; 206 | /// use once_cell::sync::Lazy; 207 | /// use std::cmp; 208 | /// use std::collections::BinaryHeap; 209 | /// use std::sync::Mutex; 210 | /// 211 | /// # smol::future::block_on(async { 212 | /// /// A wrapper around a `Runnable` that implements `Ord` so that it can be used in a 213 | /// /// priority queue. 214 | /// struct TaskWrapper(Runnable); 215 | /// 216 | /// impl PartialEq for TaskWrapper { 217 | /// fn eq(&self, other: &Self) -> bool { 218 | /// self.0.metadata() == other.0.metadata() 219 | /// } 220 | /// } 221 | /// 222 | /// impl Eq for TaskWrapper {} 223 | /// 224 | /// impl PartialOrd for TaskWrapper { 225 | /// fn partial_cmp(&self, other: &Self) -> Option { 226 | /// Some(self.cmp(other)) 227 | /// } 228 | /// } 229 | /// 230 | /// impl Ord for TaskWrapper { 231 | /// fn cmp(&self, other: &Self) -> cmp::Ordering { 232 | /// self.0.metadata().cmp(other.0.metadata()) 233 | /// } 234 | /// } 235 | /// 236 | /// static EXECUTOR: Lazy>> = Lazy::new(|| { 237 | /// Mutex::new(BinaryHeap::new()) 238 | /// }); 239 | /// 240 | /// let schedule = |runnable| { 241 | /// EXECUTOR.lock().unwrap().push(TaskWrapper(runnable)); 242 | /// }; 243 | /// 244 | /// // Spawn a few tasks with different priorities. 245 | /// let spawn_task = move |priority| { 246 | /// let (runnable, task) = Builder::new().metadata(priority).spawn( 247 | /// move |_| async move { priority }, 248 | /// schedule, 249 | /// ); 250 | /// runnable.schedule(); 251 | /// task 252 | /// }; 253 | /// 254 | /// let t1 = spawn_task(1); 255 | /// let t2 = spawn_task(2); 256 | /// let t3 = spawn_task(3); 257 | /// 258 | /// // Run the tasks in order of priority. 259 | /// let mut metadata_seen = vec![]; 260 | /// while let Some(TaskWrapper(runnable)) = EXECUTOR.lock().unwrap().pop() { 261 | /// metadata_seen.push(*runnable.metadata()); 262 | /// runnable.run(); 263 | /// } 264 | /// 265 | /// assert_eq!(metadata_seen, vec![3, 2, 1]); 266 | /// assert_eq!(t1.await, 1); 267 | /// assert_eq!(t2.await, 2); 268 | /// assert_eq!(t3.await, 3); 269 | /// # }); 270 | /// ``` 271 | pub fn metadata(self, metadata: M) -> Builder { 272 | Builder { 273 | metadata, 274 | #[cfg(feature = "std")] 275 | propagate_panic: self.propagate_panic, 276 | } 277 | } 278 | } 279 | 280 | // Use a macro to brute force inlining to minimize stack copies of potentially 281 | // large futures. 282 | macro_rules! spawn_unchecked { 283 | ($f:tt, $s:tt, $m:tt, $builder:ident, $schedule:ident, $raw:ident => $future:block) => {{ 284 | let ptr = allocate_task!($f, $s, $m, $builder, $schedule, $raw => $future); 285 | 286 | #[allow(unused_unsafe)] 287 | // SAFTETY: The task was just allocated above. 288 | let runnable = unsafe { Runnable::from_raw(ptr) }; 289 | let task = Task { 290 | ptr, 291 | _marker: PhantomData, 292 | }; 293 | (runnable, task) 294 | }}; 295 | } 296 | 297 | impl Builder { 298 | /// Propagates panics that occur in the task. 299 | /// 300 | /// When this is `true`, panics that occur in the task will be propagated to the caller of 301 | /// the [`Task`]. When this is false, no special action is taken when a panic occurs in the 302 | /// task, meaning that the caller of [`Runnable::run`] will observe a panic. 303 | /// 304 | /// This is only available when the `std` feature is enabled. By default, this is `false`. 305 | /// 306 | /// # Examples 307 | /// 308 | /// ``` 309 | /// use async_task::Builder; 310 | /// use futures_lite::future::poll_fn; 311 | /// use std::future::Future; 312 | /// use std::panic; 313 | /// use std::pin::Pin; 314 | /// use std::task::{Context, Poll}; 315 | /// 316 | /// fn did_panic(f: F) -> bool { 317 | /// panic::catch_unwind(panic::AssertUnwindSafe(f)).is_err() 318 | /// } 319 | /// 320 | /// # smol::future::block_on(async { 321 | /// let (runnable1, mut task1) = Builder::new() 322 | /// .propagate_panic(true) 323 | /// .spawn(|()| async move { panic!() }, |_| {}); 324 | /// 325 | /// let (runnable2, mut task2) = Builder::new() 326 | /// .propagate_panic(false) 327 | /// .spawn(|()| async move { panic!() }, |_| {}); 328 | /// 329 | /// assert!(!did_panic(|| { runnable1.run(); })); 330 | /// assert!(did_panic(|| { runnable2.run(); })); 331 | /// 332 | /// let waker = poll_fn(|cx| Poll::Ready(cx.waker().clone())).await; 333 | /// let mut cx = Context::from_waker(&waker); 334 | /// assert!(did_panic(|| { let _ = Pin::new(&mut task1).poll(&mut cx); })); 335 | /// assert!(did_panic(|| { let _ = Pin::new(&mut task2).poll(&mut cx); })); 336 | /// # }); 337 | /// ``` 338 | #[cfg(feature = "std")] 339 | pub fn propagate_panic(self, propagate_panic: bool) -> Builder { 340 | Builder { 341 | metadata: self.metadata, 342 | propagate_panic, 343 | } 344 | } 345 | 346 | /// Creates a new task. 347 | /// 348 | /// The returned [`Runnable`] is used to poll the `future`, and the [`Task`] is used to await its 349 | /// output. 350 | /// 351 | /// Method [`run()`][`Runnable::run()`] polls the task's future once. Then, the [`Runnable`] 352 | /// vanishes and only reappears when its [`Waker`] wakes the task, thus scheduling it to be run 353 | /// again. 354 | /// 355 | /// When the task is woken, its [`Runnable`] is passed to the `schedule` function. 356 | /// The `schedule` function should not attempt to run the [`Runnable`] nor to drop it. Instead, it 357 | /// should push it into a task queue so that it can be processed later. 358 | /// 359 | /// If you need to spawn a future that does not implement [`Send`] or isn't `'static`, consider 360 | /// using [`spawn_local()`] or [`spawn_unchecked()`] instead. 361 | /// 362 | /// # Examples 363 | /// 364 | /// ``` 365 | /// use async_task::Builder; 366 | /// 367 | /// // The future inside the task. 368 | /// let future = async { 369 | /// println!("Hello, world!"); 370 | /// }; 371 | /// 372 | /// // A function that schedules the task when it gets woken up. 373 | /// let (s, r) = flume::unbounded(); 374 | /// let schedule = move |runnable| s.send(runnable).unwrap(); 375 | /// 376 | /// // Create a task with the future and the schedule function. 377 | /// let (runnable, task) = Builder::new().spawn(|()| future, schedule); 378 | /// ``` 379 | pub fn spawn(self, future: F, schedule: S) -> (Runnable, Task) 380 | where 381 | F: FnOnce(&M) -> Fut, 382 | Fut: Future + Send + 'static, 383 | Fut::Output: Send + 'static, 384 | S: Schedule + Send + Sync + 'static, 385 | { 386 | unsafe { 387 | spawn_unchecked!(Fut, S, M, self, schedule, raw => { future(&(*raw.header).metadata) }) 388 | } 389 | } 390 | 391 | /// Creates a new thread-local task. 392 | /// 393 | /// This function is same as [`spawn()`], except it does not require [`Send`] on `future`. If the 394 | /// [`Runnable`] is used or dropped on another thread, a panic will occur. 395 | /// 396 | /// This function is only available when the `std` feature for this crate is enabled. 397 | /// 398 | /// # Examples 399 | /// 400 | /// ``` 401 | /// use async_task::{Builder, Runnable}; 402 | /// use flume::{Receiver, Sender}; 403 | /// use std::rc::Rc; 404 | /// 405 | /// thread_local! { 406 | /// // A queue that holds scheduled tasks. 407 | /// static QUEUE: (Sender, Receiver) = flume::unbounded(); 408 | /// } 409 | /// 410 | /// // Make a non-Send future. 411 | /// let msg: Rc = "Hello, world!".into(); 412 | /// let future = async move { 413 | /// println!("{}", msg); 414 | /// }; 415 | /// 416 | /// // A function that schedules the task when it gets woken up. 417 | /// let s = QUEUE.with(|(s, _)| s.clone()); 418 | /// let schedule = move |runnable| s.send(runnable).unwrap(); 419 | /// 420 | /// // Create a task with the future and the schedule function. 421 | /// let (runnable, task) = Builder::new().spawn_local(move |()| future, schedule); 422 | /// ``` 423 | #[cfg(feature = "std")] 424 | pub fn spawn_local( 425 | self, 426 | future: F, 427 | schedule: S, 428 | ) -> (Runnable, Task) 429 | where 430 | F: FnOnce(&M) -> Fut, 431 | Fut: Future + 'static, 432 | Fut::Output: 'static, 433 | S: Schedule + Send + Sync + 'static, 434 | { 435 | use std::mem::ManuallyDrop; 436 | use std::pin::Pin; 437 | use std::task::{Context, Poll}; 438 | use std::thread::{self, ThreadId}; 439 | 440 | #[inline] 441 | fn thread_id() -> ThreadId { 442 | std::thread_local! { 443 | static ID: ThreadId = thread::current().id(); 444 | } 445 | ID.try_with(|id| *id) 446 | .unwrap_or_else(|_| thread::current().id()) 447 | } 448 | 449 | struct Checked { 450 | id: ThreadId, 451 | inner: ManuallyDrop, 452 | } 453 | 454 | impl Drop for Checked { 455 | fn drop(&mut self) { 456 | assert!( 457 | self.id == thread_id(), 458 | "local task dropped by a thread that didn't spawn it" 459 | ); 460 | unsafe { 461 | ManuallyDrop::drop(&mut self.inner); 462 | } 463 | } 464 | } 465 | 466 | impl Future for Checked { 467 | type Output = F::Output; 468 | 469 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 470 | assert!( 471 | self.id == thread_id(), 472 | "local task polled by a thread that didn't spawn it" 473 | ); 474 | unsafe { self.map_unchecked_mut(|c| &mut *c.inner).poll(cx) } 475 | } 476 | } 477 | 478 | // Wrap the future into one that checks which thread it's on. 479 | let future = move |meta| { 480 | let future = future(meta); 481 | 482 | Checked { 483 | id: thread_id(), 484 | inner: ManuallyDrop::new(future), 485 | } 486 | }; 487 | 488 | unsafe { self.spawn_unchecked(future, schedule) } 489 | } 490 | 491 | /// Creates a new task without [`Send`], [`Sync`], and `'static` bounds. 492 | /// 493 | /// This function is same as [`spawn()`], except it does not require [`Send`], [`Sync`], and 494 | /// `'static` on `future` and `schedule`. 495 | /// 496 | /// # Safety 497 | /// 498 | /// - If `Fut` is not [`Send`], its [`Runnable`] must be used and dropped on the original 499 | /// thread. 500 | /// - If `Fut` is not `'static`, borrowed non-metadata variables must outlive its [`Runnable`]. 501 | /// - If `schedule` is not [`Send`] and [`Sync`], all instances of the [`Runnable`]'s [`Waker`] 502 | /// must be used and dropped on the original thread. 503 | /// - If `schedule` is not `'static`, borrowed variables must outlive all instances of the 504 | /// [`Runnable`]'s [`Waker`]. 505 | /// 506 | /// # Examples 507 | /// 508 | /// ``` 509 | /// use async_task::Builder; 510 | /// 511 | /// // The future inside the task. 512 | /// let future = async { 513 | /// println!("Hello, world!"); 514 | /// }; 515 | /// 516 | /// // If the task gets woken up, it will be sent into this channel. 517 | /// let (s, r) = flume::unbounded(); 518 | /// let schedule = move |runnable| s.send(runnable).unwrap(); 519 | /// 520 | /// // Create a task with the future and the schedule function. 521 | /// let (runnable, task) = unsafe { Builder::new().spawn_unchecked(move |()| future, schedule) }; 522 | /// ``` 523 | pub unsafe fn spawn_unchecked<'a, F, Fut, S>( 524 | self, 525 | future: F, 526 | schedule: S, 527 | ) -> (Runnable, Task) 528 | where 529 | F: FnOnce(&'a M) -> Fut, 530 | Fut: Future + 'a, 531 | S: Schedule, 532 | M: 'a, 533 | { 534 | spawn_unchecked!(Fut, S, M, self, schedule, raw => { future(&(*raw.header).metadata) }) 535 | } 536 | } 537 | 538 | /// Creates a new task. 539 | /// 540 | /// The returned [`Runnable`] is used to poll the `future`, and the [`Task`] is used to await its 541 | /// output. 542 | /// 543 | /// Method [`run()`][`Runnable::run()`] polls the task's future once. Then, the [`Runnable`] 544 | /// vanishes and only reappears when its [`Waker`] wakes the task, thus scheduling it to be run 545 | /// again. 546 | /// 547 | /// When the task is woken, its [`Runnable`] is passed to the `schedule` function. 548 | /// The `schedule` function should not attempt to run the [`Runnable`] nor to drop it. Instead, it 549 | /// should push it into a task queue so that it can be processed later. 550 | /// 551 | /// If you need to spawn a future that does not implement [`Send`] or isn't `'static`, consider 552 | /// using [`spawn_local()`] or [`spawn_unchecked()`] instead. 553 | /// 554 | /// # Examples 555 | /// 556 | /// ``` 557 | /// // The future inside the task. 558 | /// let future = async { 559 | /// println!("Hello, world!"); 560 | /// }; 561 | /// 562 | /// // A function that schedules the task when it gets woken up. 563 | /// let (s, r) = flume::unbounded(); 564 | /// let schedule = move |runnable| s.send(runnable).unwrap(); 565 | /// 566 | /// // Create a task with the future and the schedule function. 567 | /// let (runnable, task) = async_task::spawn(future, schedule); 568 | /// ``` 569 | pub fn spawn(future: F, schedule: S) -> (Runnable, Task) 570 | where 571 | F: Future + Send + 'static, 572 | F::Output: Send + 'static, 573 | S: Schedule + Send + Sync + 'static, 574 | { 575 | let builder = Builder::new(); 576 | unsafe { spawn_unchecked!(F, S, (), builder, schedule, raw => { future }) } 577 | } 578 | 579 | /// Creates a new thread-local task. 580 | /// 581 | /// This function is same as [`spawn()`], except it does not require [`Send`] on `future`. If the 582 | /// [`Runnable`] is used or dropped on another thread, a panic will occur. 583 | /// 584 | /// This function is only available when the `std` feature for this crate is enabled. 585 | /// 586 | /// # Examples 587 | /// 588 | /// ``` 589 | /// use async_task::Runnable; 590 | /// use flume::{Receiver, Sender}; 591 | /// use std::rc::Rc; 592 | /// 593 | /// thread_local! { 594 | /// // A queue that holds scheduled tasks. 595 | /// static QUEUE: (Sender, Receiver) = flume::unbounded(); 596 | /// } 597 | /// 598 | /// // Make a non-Send future. 599 | /// let msg: Rc = "Hello, world!".into(); 600 | /// let future = async move { 601 | /// println!("{}", msg); 602 | /// }; 603 | /// 604 | /// // A function that schedules the task when it gets woken up. 605 | /// let s = QUEUE.with(|(s, _)| s.clone()); 606 | /// let schedule = move |runnable| s.send(runnable).unwrap(); 607 | /// 608 | /// // Create a task with the future and the schedule function. 609 | /// let (runnable, task) = async_task::spawn_local(future, schedule); 610 | /// ``` 611 | #[cfg(feature = "std")] 612 | pub fn spawn_local(future: F, schedule: S) -> (Runnable, Task) 613 | where 614 | F: Future + 'static, 615 | F::Output: 'static, 616 | S: Schedule + Send + Sync + 'static, 617 | { 618 | Builder::new().spawn_local(move |()| future, schedule) 619 | } 620 | 621 | /// Creates a new task without [`Send`], [`Sync`], and `'static` bounds. 622 | /// 623 | /// This function is same as [`spawn()`], except it does not require [`Send`], [`Sync`], and 624 | /// `'static` on `future` and `schedule`. 625 | /// 626 | /// # Safety 627 | /// 628 | /// - If `future` is not [`Send`], its [`Runnable`] must be used and dropped on the original 629 | /// thread. 630 | /// - If `future` is not `'static`, borrowed variables must outlive its [`Runnable`]. 631 | /// - If `schedule` is not [`Send`] and [`Sync`], all instances of the [`Runnable`]'s [`Waker`] 632 | /// must be used and dropped on the original thread. 633 | /// - If `schedule` is not `'static`, borrowed variables must outlive all instances of the 634 | /// [`Runnable`]'s [`Waker`]. 635 | /// 636 | /// # Examples 637 | /// 638 | /// ``` 639 | /// // The future inside the task. 640 | /// let future = async { 641 | /// println!("Hello, world!"); 642 | /// }; 643 | /// 644 | /// // If the task gets woken up, it will be sent into this channel. 645 | /// let (s, r) = flume::unbounded(); 646 | /// let schedule = move |runnable| s.send(runnable).unwrap(); 647 | /// 648 | /// // Create a task with the future and the schedule function. 649 | /// let (runnable, task) = unsafe { async_task::spawn_unchecked(future, schedule) }; 650 | /// ``` 651 | pub unsafe fn spawn_unchecked(future: F, schedule: S) -> (Runnable, Task) 652 | where 653 | F: Future, 654 | S: Schedule, 655 | { 656 | let builder = Builder::new(); 657 | spawn_unchecked!(F, S, (), builder, schedule, raw => { future }) 658 | } 659 | 660 | /// A handle to a runnable task. 661 | /// 662 | /// Every spawned task has a single [`Runnable`] handle, which only exists when the task is 663 | /// scheduled for running. 664 | /// 665 | /// Method [`run()`][`Runnable::run()`] polls the task's future once. Then, the [`Runnable`] 666 | /// vanishes and only reappears when its [`Waker`] wakes the task, thus scheduling it to be run 667 | /// again. 668 | /// 669 | /// Dropping a [`Runnable`] cancels the task, which means its future won't be polled again, and 670 | /// awaiting the [`Task`] after that will result in a panic. 671 | /// 672 | /// # Examples 673 | /// 674 | /// ``` 675 | /// use async_task::Runnable; 676 | /// use once_cell::sync::Lazy; 677 | /// use std::{panic, thread}; 678 | /// 679 | /// // A simple executor. 680 | /// static QUEUE: Lazy> = Lazy::new(|| { 681 | /// let (sender, receiver) = flume::unbounded::(); 682 | /// thread::spawn(|| { 683 | /// for runnable in receiver { 684 | /// let _ignore_panic = panic::catch_unwind(|| runnable.run()); 685 | /// } 686 | /// }); 687 | /// sender 688 | /// }); 689 | /// 690 | /// // Create a task with a simple future. 691 | /// let schedule = |runnable| QUEUE.send(runnable).unwrap(); 692 | /// let (runnable, task) = async_task::spawn(async { 1 + 2 }, schedule); 693 | /// 694 | /// // Schedule the task and await its output. 695 | /// runnable.schedule(); 696 | /// assert_eq!(smol::future::block_on(task), 3); 697 | /// ``` 698 | pub struct Runnable { 699 | /// A pointer to the heap-allocated task. 700 | pub(crate) ptr: NonNull<()>, 701 | 702 | /// A marker capturing generic type `M`. 703 | pub(crate) _marker: PhantomData, 704 | } 705 | 706 | unsafe impl Send for Runnable {} 707 | unsafe impl Sync for Runnable {} 708 | 709 | #[cfg(feature = "std")] 710 | impl std::panic::UnwindSafe for Runnable {} 711 | #[cfg(feature = "std")] 712 | impl std::panic::RefUnwindSafe for Runnable {} 713 | 714 | impl Runnable { 715 | /// Get the metadata associated with this task. 716 | /// 717 | /// Tasks can be created with a metadata object associated with them; by default, this 718 | /// is a `()` value. See the [`Builder::metadata()`] method for more information. 719 | pub fn metadata(&self) -> &M { 720 | &self.header_with_metadata().metadata 721 | } 722 | 723 | /// Schedules the task. 724 | /// 725 | /// This is a convenience method that passes the [`Runnable`] to the schedule function. 726 | /// 727 | /// # Examples 728 | /// 729 | /// ``` 730 | /// // A function that schedules the task when it gets woken up. 731 | /// let (s, r) = flume::unbounded(); 732 | /// let schedule = move |runnable| s.send(runnable).unwrap(); 733 | /// 734 | /// // Create a task with a simple future and the schedule function. 735 | /// let (runnable, task) = async_task::spawn(async {}, schedule); 736 | /// 737 | /// // Schedule the task. 738 | /// assert_eq!(r.len(), 0); 739 | /// runnable.schedule(); 740 | /// assert_eq!(r.len(), 1); 741 | /// # let handle = std::thread::spawn(move || { for runnable in r { runnable.run(); }}); 742 | /// # smol::future::block_on(task); 743 | /// # handle.join().unwrap(); 744 | /// ``` 745 | pub fn schedule(self) { 746 | let ptr = self.ptr.as_ptr(); 747 | let header = ptr as *const Header; 748 | mem::forget(self); 749 | 750 | unsafe { 751 | ((*header).vtable.schedule)(ptr, ScheduleInfo::new(false)); 752 | } 753 | } 754 | 755 | /// Runs the task by polling its future. 756 | /// 757 | /// Returns `true` if the task was woken while running, in which case the [`Runnable`] gets 758 | /// rescheduled at the end of this method invocation. Otherwise, returns `false` and the 759 | /// [`Runnable`] vanishes until the task is woken. 760 | /// The return value is just a hint: `true` usually indicates that the task has yielded, i.e. 761 | /// it woke itself and then gave the control back to the executor. 762 | /// 763 | /// If the [`Task`] handle was dropped or if [`cancel()`][`Task::cancel()`] was called, then 764 | /// this method simply destroys the task. 765 | /// 766 | /// If the polled future panics, this method propagates the panic, and awaiting the [`Task`] 767 | /// after that will also result in a panic. 768 | /// 769 | /// # Examples 770 | /// 771 | /// ``` 772 | /// // A function that schedules the task when it gets woken up. 773 | /// let (s, r) = flume::unbounded(); 774 | /// let schedule = move |runnable| s.send(runnable).unwrap(); 775 | /// 776 | /// // Create a task with a simple future and the schedule function. 777 | /// let (runnable, task) = async_task::spawn(async { 1 + 2 }, schedule); 778 | /// 779 | /// // Run the task and check its output. 780 | /// runnable.run(); 781 | /// assert_eq!(smol::future::block_on(task), 3); 782 | /// ``` 783 | pub fn run(self) -> bool { 784 | let ptr = self.ptr.as_ptr(); 785 | let header = ptr as *const Header; 786 | mem::forget(self); 787 | 788 | unsafe { ((*header).vtable.run)(ptr) } 789 | } 790 | 791 | /// Returns a waker associated with this task. 792 | /// 793 | /// # Examples 794 | /// 795 | /// ``` 796 | /// use smol::future; 797 | /// 798 | /// // A function that schedules the task when it gets woken up. 799 | /// let (s, r) = flume::unbounded(); 800 | /// let schedule = move |runnable| s.send(runnable).unwrap(); 801 | /// 802 | /// // Create a task with a simple future and the schedule function. 803 | /// let (runnable, task) = async_task::spawn(future::pending::<()>(), schedule); 804 | /// 805 | /// // Take a waker and run the task. 806 | /// let waker = runnable.waker(); 807 | /// runnable.run(); 808 | /// 809 | /// // Reschedule the task by waking it. 810 | /// assert_eq!(r.len(), 0); 811 | /// waker.wake(); 812 | /// assert_eq!(r.len(), 1); 813 | /// # let handle = std::thread::spawn(move || { for runnable in r { runnable.run(); }}); 814 | /// # smol::future::block_on(task.cancel()); // cancel because the future is future::pending 815 | /// # handle.join().unwrap(); 816 | /// ``` 817 | pub fn waker(&self) -> Waker { 818 | unsafe { 819 | let raw_waker = Header::clone_waker(self.ptr.as_ptr()); 820 | Waker::from_raw(raw_waker) 821 | } 822 | } 823 | 824 | fn header(&self) -> &Header { 825 | unsafe { &*(self.ptr.as_ptr() as *const Header) } 826 | } 827 | 828 | fn header_with_metadata(&self) -> &HeaderWithMetadata { 829 | unsafe { &*(self.ptr.as_ptr() as *const HeaderWithMetadata) } 830 | } 831 | 832 | /// Converts this task into a raw pointer. 833 | /// 834 | /// To avoid a memory leak the pointer must be converted back to a Runnable using [`Runnable::from_raw`][from_raw]. 835 | /// 836 | /// `into_raw` does not change the state of the [`Task`], but there is no guarantee that it will be in the same state after calling [`Runnable::from_raw`][from_raw], 837 | /// as the corresponding [`Task`] might have been dropped or cancelled. 838 | /// 839 | /// # Examples 840 | /// 841 | /// ```rust 842 | /// use async_task::{Runnable, spawn}; 843 | /// 844 | /// let (runnable, task) = spawn(async {}, |_| {}); 845 | /// let runnable_pointer = runnable.into_raw(); 846 | /// 847 | /// unsafe { 848 | /// // Convert back to an `Runnable` to prevent leak. 849 | /// let runnable = Runnable::<()>::from_raw(runnable_pointer); 850 | /// runnable.run(); 851 | /// // Further calls to `Runnable::from_raw(runnable_pointer)` would be memory-unsafe. 852 | /// } 853 | /// // The memory was freed when `x` went out of scope above, so `runnable_pointer` is now dangling! 854 | /// ``` 855 | /// [from_raw]: #method.from_raw 856 | pub fn into_raw(self) -> NonNull<()> { 857 | let ptr = self.ptr; 858 | mem::forget(self); 859 | ptr 860 | } 861 | 862 | /// Converts a raw pointer into a Runnable. 863 | /// 864 | /// # Safety 865 | /// 866 | /// This method should only be used with raw pointers returned from [`Runnable::into_raw`][into_raw]. 867 | /// It is not safe to use the provided pointer once it is passed to `from_raw`. 868 | /// Crucially, it is unsafe to call `from_raw` multiple times with the same pointer - even if the resulting [`Runnable`] is not used - 869 | /// as internally `async-task` uses reference counting. 870 | /// 871 | /// It is however safe to call [`Runnable::into_raw`][into_raw] on a [`Runnable`] created with `from_raw` or 872 | /// after the [`Task`] associated with a given Runnable has been dropped or cancelled. 873 | /// 874 | /// The state of the [`Runnable`] created with `from_raw` is not specified. 875 | /// 876 | /// # Examples 877 | /// 878 | /// ```rust 879 | /// use async_task::{Runnable, spawn}; 880 | /// 881 | /// let (runnable, task) = spawn(async {}, |_| {}); 882 | /// let runnable_pointer = runnable.into_raw(); 883 | /// 884 | /// drop(task); 885 | /// unsafe { 886 | /// // Convert back to an `Runnable` to prevent leak. 887 | /// let runnable = Runnable::<()>::from_raw(runnable_pointer); 888 | /// let did_poll = runnable.run(); 889 | /// assert!(!did_poll); 890 | /// // Further calls to `Runnable::from_raw(runnable_pointer)` would be memory-unsafe. 891 | /// } 892 | /// // The memory was freed when `x` went out of scope above, so `runnable_pointer` is now dangling! 893 | /// ``` 894 | /// 895 | /// [into_raw]: #method.into_raw 896 | pub unsafe fn from_raw(ptr: NonNull<()>) -> Self { 897 | Self { 898 | ptr, 899 | _marker: Default::default(), 900 | } 901 | } 902 | } 903 | 904 | impl Drop for Runnable { 905 | fn drop(&mut self) { 906 | let ptr = self.ptr.as_ptr(); 907 | let header = self.header(); 908 | 909 | unsafe { 910 | let mut state = header.state.load(Ordering::Acquire); 911 | 912 | loop { 913 | // If the task has been completed or closed, it can't be canceled. 914 | if state & (COMPLETED | CLOSED) != 0 { 915 | break; 916 | } 917 | 918 | // Mark the task as closed. 919 | match header.state.compare_exchange_weak( 920 | state, 921 | state | CLOSED, 922 | Ordering::AcqRel, 923 | Ordering::Acquire, 924 | ) { 925 | Ok(_) => break, 926 | Err(s) => state = s, 927 | } 928 | } 929 | 930 | // Drop the future. 931 | (header.vtable.drop_future)(ptr, header.vtable.layout_info); 932 | 933 | // Mark the task as unscheduled. 934 | let state = header.state.fetch_and(!SCHEDULED, Ordering::AcqRel); 935 | 936 | // Notify the awaiter that the future has been dropped. 937 | if state & AWAITER != 0 { 938 | (*header).notify(None); 939 | } 940 | 941 | // Drop the task reference. 942 | drop_ref(ptr); 943 | } 944 | } 945 | } 946 | 947 | impl fmt::Debug for Runnable { 948 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 949 | let ptr = self.ptr.as_ptr(); 950 | let header = ptr as *const HeaderWithMetadata; 951 | 952 | f.debug_struct("Runnable") 953 | .field("header", unsafe { &(*header) }) 954 | .finish() 955 | } 956 | } 957 | --------------------------------------------------------------------------------