From c4dbdadb64638f672c83f831138addb75c603575 Mon Sep 17 00:00:00 2001 From: Stjepan Glavina Date: Sun, 19 Jul 2020 17:43:22 +0200 Subject: [PATCH] Initial commit --- .github/FUNDING.yml | 1 + .github/workflows/build-and-test.yaml | 51 ++++ .github/workflows/lint.yaml | 26 ++ .github/workflows/security.yaml | 20 ++ .gitignore | 2 + CHANGELOG.md | 3 + Cargo.toml | 16 + LICENSE-APACHE | 201 +++++++++++++ LICENSE-MIT | 23 ++ README.md | 27 ++ src/lib.rs | 414 ++++++++++++++++++++++++++ tests/basic.rs | 50 ++++ 12 files changed, 834 insertions(+) create mode 100644 .github/FUNDING.yml create mode 100644 .github/workflows/build-and-test.yaml create mode 100644 .github/workflows/lint.yaml create mode 100644 .github/workflows/security.yaml create mode 100644 .gitignore create mode 100644 CHANGELOG.md create mode 100644 Cargo.toml create mode 100644 LICENSE-APACHE create mode 100644 LICENSE-MIT create mode 100644 README.md create mode 100644 src/lib.rs create mode 100644 tests/basic.rs diff --git a/.github/FUNDING.yml b/.github/FUNDING.yml new file mode 100644 index 0000000..e1f75e7 --- /dev/null +++ b/.github/FUNDING.yml @@ -0,0 +1 @@ +github: stjepang diff --git a/.github/workflows/build-and-test.yaml b/.github/workflows/build-and-test.yaml new file mode 100644 index 0000000..19d1ebb --- /dev/null +++ b/.github/workflows/build-and-test.yaml @@ -0,0 +1,51 @@ +name: Build and test + +on: + push: + branches: + - master + pull_request: + +jobs: + build_and_test: + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest] + rust: [nightly, beta, stable] + steps: + - uses: actions/checkout@v2 + + - name: Set current week of the year in environnement + if: startsWith(matrix.os, 'ubuntu') || startsWith(matrix.os, 'macOS') + run: echo "::set-env name=CURRENT_WEEK::$(date +%V)" + + - name: Set current week of the year in environnement + if: startsWith(matrix.os, 'windows') + run: echo "::set-env name=CURRENT_WEEK::$(Get-Date -UFormat %V)" + + - name: Install latest ${{ matrix.rust }} + uses: actions-rs/toolchain@v1 + with: + toolchain: ${{ matrix.rust }} + profile: minimal + override: true + + - name: Run cargo check + uses: actions-rs/cargo@v1 + with: + command: check + args: --all --bins --examples --tests --all-features + + - name: Run cargo check (without dev-dependencies to catch missing feature flags) + if: startsWith(matrix.rust, 'nightly') + uses: actions-rs/cargo@v1 + with: + command: check + args: -Z features=dev_dep + + - name: Run cargo test + uses: actions-rs/cargo@v1 + with: + command: test diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml new file mode 100644 index 0000000..7e9bd98 --- /dev/null +++ b/.github/workflows/lint.yaml @@ -0,0 +1,26 @@ +name: Lint + +on: + push: + branches: + - master + pull_request: + +jobs: + clippy: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v2 + + - name: Set current week of the year in environnement + run: echo "::set-env name=CURRENT_WEEK::$(date +%V)" + + - uses: actions-rs/toolchain@v1 + with: + toolchain: stable + profile: minimal + components: clippy + - uses: actions-rs/clippy-check@v1 + with: + token: ${{ secrets.GITHUB_TOKEN }} + args: --all-features -- -W clippy::all diff --git a/.github/workflows/security.yaml b/.github/workflows/security.yaml new file mode 100644 index 0000000..8f722e7 --- /dev/null +++ b/.github/workflows/security.yaml @@ -0,0 +1,20 @@ +name: Security audit + +on: + push: + branches: + - master + pull_request: + +jobs: + security_audit: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v2 + + - name: Set current week of the year in environnement + run: echo "::set-env name=CURRENT_WEEK::$(date +%V)" + + - uses: actions-rs/audit-check@v1 + with: + token: ${{ secrets.GITHUB_TOKEN }} diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..96ef6c0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +/target +Cargo.lock diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..af42cea --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,3 @@ +# Version 1.0.0 + +- Initial version diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..3dbf8d7 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "atomic-waker" +version = "1.0.0" +authors = ["Stjepan Glavina "] +edition = "2018" +description = "A synchronization primitive for task wakeup" +license = "Apache-2.0 OR MIT" +repository = "https://github.com/stjepang/futures-lite" +homepage = "https://github.com/stjepang/futures-lite" +documentation = "https://docs.rs/futures-lite" +keywords = ["waker", "notify", "wake", "futures", "async"] +categories = ["asynchronous", "concurrency"] +readme = "README.md" + +[dev-dependencies] +futures = "0.3.5" diff --git a/LICENSE-APACHE b/LICENSE-APACHE new file mode 100644 index 0000000..16fe87b --- /dev/null +++ b/LICENSE-APACHE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + +TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + +1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + +2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + +3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + +4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + +5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + +6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + +7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + +8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + +9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + +END OF TERMS AND CONDITIONS + +APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + +Copyright [yyyy] [name of copyright owner] + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. diff --git a/LICENSE-MIT b/LICENSE-MIT new file mode 100644 index 0000000..31aa793 --- /dev/null +++ b/LICENSE-MIT @@ -0,0 +1,23 @@ +Permission is hereby granted, free of charge, to any +person obtaining a copy of this software and associated +documentation files (the "Software"), to deal in the +Software without restriction, including without +limitation the rights to use, copy, modify, merge, +publish, distribute, sublicense, and/or sell copies of +the Software, and to permit persons to whom the Software +is furnished to do so, subject to the following +conditions: + +The above copyright notice and this permission notice +shall be included in all copies or substantial portions +of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF +ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED +TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A +PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT +SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR +IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +DEALINGS IN THE SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..2c0ec04 --- /dev/null +++ b/README.md @@ -0,0 +1,27 @@ +# atomic-waker + +[![Build](https://github.com/stjepang/atomic-waker/workflows/Build%20and%20test/badge.svg)]( +https://github.com/stjepang/atomic-waker/actions) +[![License](https://img.shields.io/badge/license-MIT%2FApache--2.0-blue.svg)]( +https://github.com/stjepang/atomic-waker) +[![Cargo](https://img.shields.io/crates/v/atomic-waker.svg)]( +https://crates.io/crates/atomic-waker) +[![Documentation](https://docs.rs/atomic-waker/badge.svg)]( +https://docs.rs/atomic-waker) + +`futures::task::AtomicWaker` extracted into its own crate. + +## License + +Licensed under either of + + * Apache License, Version 2.0 ([LICENSE-APACHE](LICENSE-APACHE) or http://www.apache.org/licenses/LICENSE-2.0) + * MIT license ([LICENSE-MIT](LICENSE-MIT) or http://opensource.org/licenses/MIT) + +at your option. + +#### Contribution + +Unless you explicitly state otherwise, any contribution intentionally submitted +for inclusion in the work by you, as defined in the Apache-2.0 license, shall be +dual licensed as above, without any additional terms or conditions. diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..17f87f4 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,414 @@ +//! `futures::task::AtomicWaker` extracted into its own crate. + +use std::cell::UnsafeCell; +use std::fmt; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering::{Acquire, Release, AcqRel}; +use std::task::Waker; + +/// A synchronization primitive for task wakeup. +/// +/// Sometimes the task interested in a given event will change over time. +/// An `AtomicWaker` can coordinate concurrent notifications with the consumer +/// potentially "updating" the underlying task to wake up. This is useful in +/// scenarios where a computation completes in another thread and wants to +/// notify the consumer, but the consumer is in the process of being migrated to +/// a new logical task. +/// +/// Consumers should call `register` before checking the result of a computation +/// and producers should call `wake` after producing the computation (this +/// differs from the usual `thread::park` pattern). It is also permitted for +/// `wake` to be called **before** `register`. This results in a no-op. +/// +/// A single `AtomicWaker` may be reused for any number of calls to `register` or +/// `wake`. +/// +/// # Memory ordering +/// +/// Calling `register` "acquires" all memory "released" by calls to `wake` +/// before the call to `register`. Later calls to `wake` will wake the +/// registered waker (on contention this wake might be triggered in `register`). +/// +/// For concurrent calls to `register` (should be avoided) the ordering is only +/// guaranteed for the winning call. +/// +/// # Examples +/// +/// Here is a simple example providing a `Flag` that can be signalled manually +/// when it is ready. +/// +/// ``` +/// use futures::future::Future; +/// use futures::task::{Context, Poll, AtomicWaker}; +/// use std::sync::Arc; +/// use std::sync::atomic::AtomicBool; +/// use std::sync::atomic::Ordering::Relaxed; +/// use std::pin::Pin; +/// +/// struct Inner { +/// waker: AtomicWaker, +/// set: AtomicBool, +/// } +/// +/// #[derive(Clone)] +/// struct Flag(Arc); +/// +/// impl Flag { +/// pub fn new() -> Self { +/// Flag(Arc::new(Inner { +/// waker: AtomicWaker::new(), +/// set: AtomicBool::new(false), +/// })) +/// } +/// +/// pub fn signal(&self) { +/// self.0.set.store(true, Relaxed); +/// self.0.waker.wake(); +/// } +/// } +/// +/// impl Future for Flag { +/// type Output = (); +/// +/// fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { +/// // quick check to avoid registration if already done. +/// if self.0.set.load(Relaxed) { +/// return Poll::Ready(()); +/// } +/// +/// self.0.waker.register(cx.waker()); +/// +/// // Need to check condition **after** `register` to avoid a race +/// // condition that would result in lost notifications. +/// if self.0.set.load(Relaxed) { +/// Poll::Ready(()) +/// } else { +/// Poll::Pending +/// } +/// } +/// } +/// ``` +pub struct AtomicWaker { + state: AtomicUsize, + waker: UnsafeCell>, +} + +// `AtomicWaker` is a multi-consumer, single-producer transfer cell. The cell +// stores a `Waker` value produced by calls to `register` and many threads can +// race to take the waker (to wake it) by calling `wake`. +// +// If a new `Waker` instance is produced by calling `register` before an +// existing one is consumed, then the existing one is overwritten. +// +// While `AtomicWaker` is single-producer, the implementation ensures memory +// safety. In the event of concurrent calls to `register`, there will be a +// single winner whose waker will get stored in the cell. The losers will not +// have their tasks woken. As such, callers should ensure to add synchronization +// to calls to `register`. +// +// The implementation uses a single `AtomicUsize` value to coordinate access to +// the `Waker` cell. There are two bits that are operated on independently. +// These are represented by `REGISTERING` and `WAKING`. +// +// The `REGISTERING` bit is set when a producer enters the critical section. The +// `WAKING` bit is set when a consumer enters the critical section. Neither bit +// being set is represented by `WAITING`. +// +// A thread obtains an exclusive lock on the waker cell by transitioning the +// state from `WAITING` to `REGISTERING` or `WAKING`, depending on the operation +// the thread wishes to perform. When this transition is made, it is guaranteed +// that no other thread will access the waker cell. +// +// # Registering +// +// On a call to `register`, an attempt to transition the state from WAITING to +// REGISTERING is made. On success, the caller obtains a lock on the waker cell. +// +// If the lock is obtained, then the thread sets the waker cell to the waker +// provided as an argument. Then it attempts to transition the state back from +// `REGISTERING` -> `WAITING`. +// +// If this transition is successful, then the registering process is complete +// and the next call to `wake` will observe the waker. +// +// If the transition fails, then there was a concurrent call to `wake` that was +// unable to access the waker cell (due to the registering thread holding the +// lock). To handle this, the registering thread removes the waker it just set +// from the cell and calls `wake` on it. This call to wake represents the +// attempt to wake by the other thread (that set the `WAKING` bit). The state is +// then transitioned from `REGISTERING | WAKING` back to `WAITING`. This +// transition must succeed because, at this point, the state cannot be +// transitioned by another thread. +// +// # Waking +// +// On a call to `wake`, an attempt to transition the state from `WAITING` to +// `WAKING` is made. On success, the caller obtains a lock on the waker cell. +// +// If the lock is obtained, then the thread takes ownership of the current value +// in the waker cell, and calls `wake` on it. The state is then transitioned +// back to `WAITING`. This transition must succeed as, at this point, the state +// cannot be transitioned by another thread. +// +// If the thread is unable to obtain the lock, the `WAKING` bit is still. This +// is because it has either been set by the current thread but the previous +// value included the `REGISTERING` bit **or** a concurrent thread is in the +// `WAKING` critical section. Either way, no action must be taken. +// +// If the current thread is the only concurrent call to `wake` and another +// thread is in the `register` critical section, when the other thread **exits** +// the `register` critical section, it will observe the `WAKING` bit and handle +// the wake itself. +// +// If another thread is in the `wake` critical section, then it will handle +// waking the task. +// +// # A potential race (is safely handled). +// +// Imagine the following situation: +// +// * Thread A obtains the `wake` lock and wakes a task. +// +// * Before thread A releases the `wake` lock, the woken task is scheduled. +// +// * Thread B attempts to wake the task. In theory this should result in the +// task being woken, but it cannot because thread A still holds the wake lock. +// +// This case is handled by requiring users of `AtomicWaker` to call `register` +// **before** attempting to observe the application state change that resulted +// in the task being awoken. The wakers also change the application state before +// calling wake. +// +// Because of this, the waker will do one of two things. +// +// 1) Observe the application state change that Thread B is woken for. In this +// case, it is OK for Thread B's wake to be lost. +// +// 2) Call register before attempting to observe the application state. Since +// Thread A still holds the `wake` lock, the call to `register` will result +// in the task waking itself and get scheduled again. + +/// Idle state +const WAITING: usize = 0; + +/// A new waker value is being registered with the `AtomicWaker` cell. +const REGISTERING: usize = 0b01; + +/// The waker currently registered with the `AtomicWaker` cell is being woken. +const WAKING: usize = 0b10; + +impl AtomicWaker { + /// Create an `AtomicWaker`. + pub const fn new() -> Self { + // Make sure that task is Sync + trait AssertSync: Sync {} + impl AssertSync for Waker {} + + AtomicWaker { + state: AtomicUsize::new(WAITING), + waker: UnsafeCell::new(None), + } + } + + /// Registers the waker to be notified on calls to `wake`. + /// + /// The new task will take place of any previous tasks that were registered + /// by previous calls to `register`. Any calls to `wake` that happen after + /// a call to `register` (as defined by the memory ordering rules), will + /// notify the `register` caller's task and deregister the waker from future + /// notifications. Because of this, callers should ensure `register` gets + /// invoked with a new `Waker` **each** time they require a wakeup. + /// + /// It is safe to call `register` with multiple other threads concurrently + /// calling `wake`. This will result in the `register` caller's current + /// task being notified once. + /// + /// This function is safe to call concurrently, but this is generally a bad + /// idea. Concurrent calls to `register` will attempt to register different + /// tasks to be notified. One of the callers will win and have its task set, + /// but there is no guarantee as to which caller will succeed. + /// + /// # Examples + /// + /// Here is how `register` is used when implementing a flag. + /// + /// ``` + /// use futures::future::Future; + /// use futures::task::{Context, Poll, AtomicWaker}; + /// use std::sync::atomic::AtomicBool; + /// use std::sync::atomic::Ordering::Relaxed; + /// use std::pin::Pin; + /// + /// struct Flag { + /// waker: AtomicWaker, + /// set: AtomicBool, + /// } + /// + /// impl Future for Flag { + /// type Output = (); + /// + /// fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + /// // Register **before** checking `set` to avoid a race condition + /// // that would result in lost notifications. + /// self.waker.register(cx.waker()); + /// + /// if self.set.load(Relaxed) { + /// Poll::Ready(()) + /// } else { + /// Poll::Pending + /// } + /// } + /// } + /// ``` + pub fn register(&self, waker: &Waker) { + match self.state.compare_and_swap(WAITING, REGISTERING, Acquire) { + WAITING => { + unsafe { + // Locked acquired, update the waker cell + *self.waker.get() = Some(waker.clone()); + + // Release the lock. If the state transitioned to include + // the `WAKING` bit, this means that at least one wake has + // been called concurrently. + // + // Start by assuming that the state is `REGISTERING` as this + // is what we just set it to. If this holds, we know that no + // other writes were performed in the meantime, so there is + // nothing to acquire, only release. In case of concurrent + // wakers, we need to acquire their releases, so success needs + // to do both. + let res = self.state.compare_exchange( + REGISTERING, WAITING, AcqRel, Acquire); + + match res { + Ok(_) => { + // memory ordering: acquired self.state during CAS + // - if previous wakes went through it syncs with + // their final release (`fetch_and`) + // - if there was no previous wake the next wake + // will wake us, no sync needed. + } + Err(actual) => { + // This branch can only be reached if at least one + // concurrent thread called `wake`. In this + // case, `actual` **must** be `REGISTERING | + // `WAKING`. + debug_assert_eq!(actual, REGISTERING | WAKING); + + // Take the waker to wake once the atomic operation has + // completed. + let waker = (*self.waker.get()).take().unwrap(); + + // We need to return to WAITING state (clear our lock and + // concurrent WAKING flag). This needs to acquire all + // WAKING fetch_or releases and it needs to release our + // update to self.waker, so we need a `swap` operation. + self.state.swap(WAITING, AcqRel); + + // memory ordering: we acquired the state for all + // concurrent wakes, but future wakes might still + // need to wake us in case we can't make progress + // from the pending wakes. + // + // So we simply schedule to come back later (we could + // also simply leave the registration in place above). + waker.wake(); + } + } + } + } + WAKING => { + // Currently in the process of waking the task, i.e., + // `wake` is currently being called on the old task handle. + // + // memory ordering: we acquired the state for all + // concurrent wakes, but future wakes might still + // need to wake us in case we can't make progress + // from the pending wakes. + // + // So we simply schedule to come back later (we + // could also spin here trying to acquire the lock + // to register). + waker.wake_by_ref(); + } + state => { + // In this case, a concurrent thread is holding the + // "registering" lock. This probably indicates a bug in the + // caller's code as racing to call `register` doesn't make much + // sense. + // + // memory ordering: don't care. a concurrent register() is going + // to succeed and provide proper memory ordering. + // + // We just want to maintain memory safety. It is ok to drop the + // call to `register`. + debug_assert!( + state == REGISTERING || + state == REGISTERING | WAKING); + } + } + } + + /// Calls `wake` on the last `Waker` passed to `register`. + /// + /// If `register` has not been called yet, then this does nothing. + pub fn wake(&self) { + if let Some(waker) = self.take() { + waker.wake(); + } + } + + /// Returns the last `Waker` passed to `register`, so that the user can wake it. + /// + /// + /// Sometimes, just waking the AtomicWaker is not fine grained enough. This allows the user + /// to take the waker and then wake it separately, rather than performing both steps in one + /// atomic action. + /// + /// If a waker has not been registered, this returns `None`. + pub fn take(&self) -> Option { + // AcqRel ordering is used in order to acquire the value of the `task` + // cell as well as to establish a `release` ordering with whatever + // memory the `AtomicWaker` is associated with. + match self.state.fetch_or(WAKING, AcqRel) { + WAITING => { + // The waking lock has been acquired. + let waker = unsafe { (*self.waker.get()).take() }; + + // Release the lock + self.state.fetch_and(!WAKING, Release); + + waker + } + state => { + // There is a concurrent thread currently updating the + // associated task. + // + // Nothing more to do as the `WAKING` bit has been set. It + // doesn't matter if there are concurrent registering threads or + // not. + // + debug_assert!( + state == REGISTERING || + state == REGISTERING | WAKING || + state == WAKING); + None + } + } + } +} + +impl Default for AtomicWaker { + fn default() -> Self { + AtomicWaker::new() + } +} + +impl fmt::Debug for AtomicWaker { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "AtomicWaker") + } +} + +unsafe impl Send for AtomicWaker {} +unsafe impl Sync for AtomicWaker {} diff --git a/tests/basic.rs b/tests/basic.rs new file mode 100644 index 0000000..09c3f92 --- /dev/null +++ b/tests/basic.rs @@ -0,0 +1,50 @@ +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::sync::Arc; +use std::task::Poll; +use std::thread; + +use atomic_waker::AtomicWaker; +use futures::executor::block_on; +use futures::future::poll_fn; + +#[test] +fn basic() { + let atomic_waker = Arc::new(AtomicWaker::new()); + let atomic_waker_copy = atomic_waker.clone(); + + let returned_pending = Arc::new(AtomicUsize::new(0)); + let returned_pending_copy = returned_pending.clone(); + + let woken = Arc::new(AtomicUsize::new(0)); + let woken_copy = woken.clone(); + + let t = thread::spawn(move || { + let mut pending_count = 0; + + block_on(poll_fn(move |cx| { + if woken_copy.load(Ordering::Relaxed) == 1 { + Poll::Ready(()) + } else { + // Assert we return pending exactly once + assert_eq!(0, pending_count); + pending_count += 1; + atomic_waker_copy.register(cx.waker()); + + returned_pending_copy.store(1, Ordering::Relaxed); + + Poll::Pending + } + })) + }); + + while returned_pending.load(Ordering::Relaxed) == 0 {} + + // give spawned thread some time to sleep in `block_on` + thread::yield_now(); + + woken.store(1, Ordering::Relaxed); + atomic_waker.wake(); + + t.join().unwrap(); +}