├── .github └── workflows │ ├── ci.yml │ └── loom.yml ├── .gitignore ├── CHANGELOG.md ├── Cargo.toml ├── LICENSE-APACHE ├── LICENSE-MIT ├── README.md └── src ├── arc_waker.rs ├── borrowed_waker.rs ├── lib.rs ├── loom_exports.rs ├── primitives.rs └── waker.rs /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | 3 | on: 4 | pull_request: 5 | push: 6 | branches: [ main ] 7 | 8 | env: 9 | RUSTFLAGS: -Dwarnings 10 | 11 | jobs: 12 | check: 13 | name: Check 14 | runs-on: ubuntu-latest 15 | strategy: 16 | fail-fast: false 17 | matrix: 18 | rust: 19 | - stable 20 | - 1.56.1 21 | steps: 22 | - name: Checkout sources 23 | uses: actions/checkout@v4 24 | 25 | - name: Install toolchain 26 | uses: dtolnay/rust-toolchain@master 27 | with: 28 | toolchain: ${{ matrix.rust }} 29 | 30 | - name: Run cargo check 31 | run: cargo check 32 | 33 | test: 34 | name: Test suite 35 | runs-on: ubuntu-latest 36 | steps: 37 | - name: Checkout sources 38 | uses: actions/checkout@v4 39 | 40 | - name: Install toolchain 41 | uses: dtolnay/rust-toolchain@stable 42 | 43 | - name: Run cargo test 44 | run: cargo test --all-features 45 | 46 | loom-light: 47 | name: Test suite (Loom, low preemption bound) 48 | runs-on: ubuntu-latest 49 | steps: 50 | - name: Checkout sources 51 | uses: actions/checkout@v4 52 | 53 | - name: Install toolchain 54 | uses: dtolnay/rust-toolchain@stable 55 | 56 | - name: Run cargo test (Loom) 57 | run: cargo test --lib --release --all-features 58 | env: 59 | RUSTFLAGS: --cfg diatomic_waker_loom -Dwarnings 60 | LOOM_MAX_PREEMPTIONS: 2 61 | 62 | lints: 63 | name: Lints 64 | runs-on: ubuntu-latest 65 | steps: 66 | - name: Checkout sources 67 | uses: actions/checkout@v4 68 | 69 | - name: Install toolchain 70 | uses: dtolnay/rust-toolchain@stable 71 | with: 72 | components: rustfmt, clippy 73 | 74 | - name: Run cargo fmt 75 | run: cargo fmt --all -- --check 76 | 77 | - name: Run cargo clippy 78 | run: cargo clippy 79 | 80 | docs: 81 | name: Docs 82 | runs-on: ubuntu-latest 83 | steps: 84 | - name: Checkout sources 85 | uses: actions/checkout@v4 86 | 87 | - name: Install toolchain 88 | uses: dtolnay/rust-toolchain@nightly 89 | 90 | - name: Run cargo doc 91 | run: cargo doc --no-deps --document-private-items --all-features 92 | env: 93 | RUSTDOCFLAGS: --cfg docsrs -Dwarnings -------------------------------------------------------------------------------- /.github/workflows/loom.yml: -------------------------------------------------------------------------------- 1 | name: Loom 2 | 3 | on: 4 | workflow_dispatch: 5 | 6 | jobs: 7 | loom: 8 | name: Test suite (Loom, default preemption bounds) 9 | runs-on: ubuntu-latest 10 | steps: 11 | - name: Checkout sources 12 | uses: actions/checkout@v4 13 | 14 | - name: Install toolchain 15 | uses: dtolnay/rust-toolchain@stable 16 | 17 | - name: Run cargo test (Loom) 18 | run: cargo test --lib --release --all-features 19 | env: 20 | RUSTFLAGS: --cfg diatomic_waker_loom 21 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | Cargo.lock 3 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # 0.2.3 (2024-09-08) 2 | 3 | - Remove `Unpin` bound on `wait_until` methods ([#15]). 4 | 5 | [#15]: https://github.com/asynchronics/diatomic-waker/pull/15 6 | 7 | 8 | # 0.2.2 (2024-08-24) 9 | 10 | - Make `loom` a `dev-dependency` to prevent transitive dependencies issues ([#9], [#11]). 11 | - Upgrade CI checkout action ([#10]). 12 | 13 | [#9]: https://github.com/asynchronics/diatomic-waker/pull/9 14 | [#10]: https://github.com/asynchronics/diatomic-waker/pull/10 15 | [#11]: https://github.com/asynchronics/diatomic-waker/pull/11 16 | 17 | 18 | # 0.2.1 (2024-08-21) 19 | 20 | - Upgrade `loom` ([#7]). 21 | 22 | [#7]: https://github.com/asynchronics/diatomic-waker/pull/7 23 | 24 | 25 | # 0.2.0 (2024-07-28) 26 | 27 | - Remove unnecessary `Unpin` bound on `WaitUntil`'s closures. 28 | - Make the crate embedded-friendly with `alloc` as a default, optional feature 29 | ([#1]). 30 | - Add non-owned counterparts to `WakeSink` and `WakeSource` that can be used 31 | with `no-alloc` ([#2], [#3]). 32 | - Update and make CI more strict ([#5], [#6]). 33 | - Move `DiatomicWaker` and `WaitUntil` to the root module ([#6]). 34 | 35 | 36 | [#1]: https://github.com/asynchronics/diatomic-waker/pull/1 37 | [#2]: https://github.com/asynchronics/diatomic-waker/pull/2 38 | [#3]: https://github.com/asynchronics/diatomic-waker/pull/3 39 | [#5]: https://github.com/asynchronics/diatomic-waker/pull/5 40 | [#6]: https://github.com/asynchronics/diatomic-waker/pull/6 41 | 42 | 43 | # 0.1.0 (2022-10-12) 44 | 45 | Initial release 46 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "diatomic-waker" 3 | # When incrementing version and releasing to crates.io: 4 | # - Update crate version in README.md 5 | # - Update CHANGELOG.md 6 | # - Update if necessary copyright notice in LICENSE-MIT 7 | # - Create a "vX.Y.Z" git tag 8 | version = "0.2.3" 9 | authors = ["Asynchronics and contributors"] 10 | edition = "2021" 11 | rust-version = "1.56" 12 | license = "MIT OR Apache-2.0" 13 | repository = "https://github.com/asynchronics/diatomic-waker" 14 | readme = "README.md" 15 | description = """ 16 | An async, lock-free synchronization primitive for task wakeup. 17 | """ 18 | categories = ["asynchronous", "concurrency"] 19 | keywords = ["async", "waker", "atomic", "no_std", "no_alloc"] 20 | 21 | [features] 22 | default = ["alloc"] 23 | alloc = [] 24 | 25 | [dev-dependencies] 26 | pollster = "0.3" 27 | 28 | [target.'cfg(diatomic_waker_loom)'.dev-dependencies] 29 | waker-fn = "1.1" 30 | loom = "0.7" 31 | 32 | [lints.rust] 33 | unexpected_cfgs = { level = "warn", check-cfg = ['cfg(diatomic_waker_loom)'] } 34 | 35 | [package.metadata.docs.rs] 36 | all-features = true 37 | rustdoc-args = ["--cfg", "docsrs"] 38 | -------------------------------------------------------------------------------- /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 | 203 | -------------------------------------------------------------------------------- /LICENSE-MIT: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2024 Asynchronics and contributors 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in 13 | all copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # diatomic-waker 2 | 3 | An async, fast synchronization primitives for task wakeup. 4 | 5 | [![Cargo](https://img.shields.io/crates/v/diatomic-waker.svg)](https://crates.io/crates/diatomic-waker) 6 | [![Documentation](https://docs.rs/diatomic-waker/badge.svg)](https://docs.rs/diatomic-waker) 7 | [![License](https://img.shields.io/badge/license-MIT%2FApache--2.0-blue.svg)](https://github.com/asynchronics/diatomic-waker#license) 8 | 9 | ## Overview 10 | 11 | `diatomic-waker` is similar to [`atomic-waker`][atomic-waker] in that it enables 12 | concurrent updates and notifications to a wrapped `Waker`. Unlike the latter, 13 | however, it does not use spinlocks[^spinlocks] and is significantly faster, in 14 | particular when the consumer needs to be notified periodically rather than just 15 | once. It can in particular be used as a very fast, single-consumer [eventcount] 16 | to turn a non-blocking data structure into an asynchronous one (see MPSC channel 17 | receiver example). 18 | 19 | This library is an offshoot of [Asynchronix][asynchronix], an ongoing effort at 20 | a high performance asynchronous computation framework for system simulation. 21 | 22 | [atomic-waker]: https://docs.rs/atomic-waker/latest/atomic_waker/ 23 | [eventcount]: https://www.1024cores.net/home/lock-free-algorithms/eventcounts 24 | [asynchronix]: https://github.com/asynchronics/asynchronix 25 | [^spinlocks]: The implementation of [AtomicWaker][atomic-waker] yields to the 26 | runtime on contention, which is in effect an executor-mediated spinlock. 27 | 28 | ## Usage 29 | 30 | Add this to your `Cargo.toml`: 31 | 32 | ```toml 33 | [dependencies] 34 | diatomic-waker = "0.2.3" 35 | ``` 36 | 37 | ## Features flags 38 | 39 | By default, this crate enables the `alloc` feature to provide the owned 40 | `WakeSink` and `WakeSource`. It can be made `no-std`-compatible by specifying 41 | `default-features = false`. 42 | 43 | ## Example 44 | 45 | A multi-producer, single-consumer channel of capacity 1 for sending 46 | `NonZeroUsize` values, with an asynchronous receiver: 47 | 48 | ```rust 49 | use std::num::NonZeroUsize; 50 | use std::sync::atomic::{AtomicUsize, Ordering}; 51 | use std::sync::Arc; 52 | 53 | use diatomic_waker::{WakeSink, WakeSource}; 54 | 55 | // The sending side of the channel. 56 | #[derive(Clone)] 57 | struct Sender { 58 | wake_src: WakeSource, 59 | value: Arc, 60 | } 61 | 62 | // The receiving side of the channel. 63 | struct Receiver { 64 | wake_sink: WakeSink, 65 | value: Arc, 66 | } 67 | 68 | // Creates an empty channel. 69 | fn channel() -> (Sender, Receiver) { 70 | let value = Arc::new(AtomicUsize::new(0)); 71 | let wake_sink = WakeSink::new(); 72 | let wake_src = wake_sink.source(); 73 | ( 74 | Sender { 75 | wake_src, 76 | value: value.clone(), 77 | }, 78 | Receiver { wake_sink, value }, 79 | ) 80 | } 81 | 82 | impl Sender { 83 | // Sends a value if the channel is empty. 84 | fn try_send(&self, value: NonZeroUsize) -> bool { 85 | let success = self 86 | .value 87 | .compare_exchange(0, value.get(), Ordering::Relaxed, Ordering::Relaxed) 88 | .is_ok(); 89 | if success { 90 | self.wake_src.notify() 91 | }; 92 | success 93 | } 94 | } 95 | 96 | impl Receiver { 97 | // Receives a value asynchronously. 98 | async fn recv(&mut self) -> NonZeroUsize { 99 | // Wait until the predicate returns `Some(value)`, i.e. when the atomic 100 | // value becomes non-zero. 101 | self.wake_sink 102 | .wait_until(|| NonZeroUsize::new(self.value.swap(0, Ordering::Relaxed))) 103 | .await 104 | } 105 | } 106 | ``` 107 | 108 | ## Safety 109 | 110 | This is a low-level primitive and as such its implementation relies on `unsafe`. 111 | The test suite makes extensive use of [Loom] to assess its correctness. As 112 | amazing as it is, however, Loom is only a tool: it cannot formally prove the 113 | absence of data races. 114 | 115 | [Loom]: https://github.com/tokio-rs/loom 116 | 117 | 118 | ## Implementation details 119 | 120 | A distinguishing feature of `diatomic-waker` is its use of two waker storage 121 | slots (hence its name) rather than one. This makes it possible to achieve 122 | lock-freedom in situations where waker registration and notification are 123 | performed concurrently. In the case of concurrent notifications, even though one 124 | notifier does hold a notification lock, other notifiers never block: they merely 125 | request the holder of the lock to send another notification, which is a 126 | wait-free operation. 127 | 128 | Compared to `atomic-waker`, dummy notifications (with no waker registered) are 129 | much cheaper. The overall cost of a successful notification (registration + 130 | notification itself) is also much cheaper in the common case where the 131 | registered/unregistered waker is always the same, because the last waker is 132 | always cached to avoid undue cloning. Quantitatively, the costs in terms of 133 | atomic Read-Modify-Write (RMW) operations are: 134 | 135 | * dummy notification due to no waker being registered: 1 RMW vs 2 RMWs for 136 | `atomic-waker`, 137 | * registration of the same waker as the last registered waker + notification: 138 | 1+3=4 RMWs vs 3+4=7 RMWs for `atomic-waker` (this assumes 1 RMW for 139 | `Waker::wake_by_ref`, 2 RMWs for `Waker::wake`, 1 RMW for `Waker::clone`). 140 | * registration of a new waker + notification: 3+3=6 RMWs vs 3+4=7 RMWs for 141 | `atomic-waker` (same assumptions as above + 1 RMW for `Waker::drop`); this is 142 | typically only necessary for the very first registration, 143 | * very few RMWs and predictable cost on contention due to the absence of 144 | spinlocks. 145 | 146 | 147 | ## License 148 | 149 | This software is licensed under the [Apache License, Version 2.0](LICENSE-APACHE) or the 150 | [MIT license](LICENSE-MIT), at your option. 151 | 152 | 153 | ### Contribution 154 | 155 | Unless you explicitly state otherwise, any contribution intentionally submitted 156 | for inclusion in the work by you, as defined in the Apache-2.0 license, shall be 157 | dual licensed as above, without any additional terms or conditions. 158 | -------------------------------------------------------------------------------- /src/arc_waker.rs: -------------------------------------------------------------------------------- 1 | use alloc::sync::Arc; 2 | use core::task::Waker; 3 | 4 | use crate::DiatomicWaker; 5 | use crate::WaitUntil; 6 | 7 | /// An owned object that can await notifications from one or several 8 | /// [`WakeSource`]s. 9 | /// 10 | /// See the [crate-level documentation](crate) for usage. 11 | #[derive(Debug, Default)] 12 | pub struct WakeSink { 13 | /// The shared data. 14 | inner: Arc, 15 | } 16 | 17 | impl WakeSink { 18 | /// Creates a new sink. 19 | pub fn new() -> Self { 20 | Self { 21 | inner: Arc::new(DiatomicWaker::new()), 22 | } 23 | } 24 | 25 | /// Creates an owned source. 26 | #[inline] 27 | pub fn source(&self) -> WakeSource { 28 | WakeSource { 29 | inner: self.inner.clone(), 30 | } 31 | } 32 | 33 | /// Registers a new waker. 34 | /// 35 | /// Registration is lazy: the waker is cloned only if it differs from the 36 | /// last registered waker (note that the last registered waker is cached 37 | /// even if it was unregistered). 38 | #[inline] 39 | pub fn register(&mut self, waker: &Waker) { 40 | // Safety: `DiatomicWaker::register`, `DiatomicWaker::unregister` and 41 | // `DiatomicWaker::wait_until` cannot be used concurrently from multiple 42 | // thread since `WakeSink` does not implement `Clone` and the wrappers 43 | // of the above methods require exclusive ownership to `WakeSink`. 44 | unsafe { self.inner.register(waker) }; 45 | } 46 | 47 | /// Unregisters the waker. 48 | /// 49 | /// After the waker is unregistered, subsequent calls to 50 | /// `WakeSource::notify` will be ignored. 51 | /// 52 | /// Note that the previously-registered waker (if any) remains cached. 53 | #[inline] 54 | pub fn unregister(&mut self) { 55 | // Safety: `DiatomicWaker::register`, `DiatomicWaker::unregister` and 56 | // `DiatomicWaker::wait_until` cannot be used concurrently from multiple 57 | // thread since `WakeSink` does not implement `Clone` and the wrappers 58 | // of the above methods require exclusive ownership to `WakeSink`. 59 | unsafe { self.inner.unregister() }; 60 | } 61 | 62 | /// Returns a future that can be `await`ed until the provided predicate 63 | /// returns a value. 64 | /// 65 | /// The predicate is checked each time a notification is received. 66 | #[inline] 67 | pub fn wait_until(&mut self, predicate: P) -> WaitUntil<'_, P, T> 68 | where 69 | P: FnMut() -> Option, 70 | { 71 | // Safety: `DiatomicWaker::register`, `DiatomicWaker::unregister` and 72 | // `DiatomicWaker::wait_until` cannot be used concurrently from multiple 73 | // thread since `WakeSink` does not implement `Clone` and the wrappers 74 | // of the above methods require exclusive ownership to `WakeSink`. 75 | unsafe { self.inner.wait_until(predicate) } 76 | } 77 | } 78 | 79 | /// An owned object that can send notifications to a [`WakeSink`]. 80 | /// 81 | /// See the [crate-level documentation](crate) for usage. 82 | #[derive(Clone, Debug)] 83 | pub struct WakeSource { 84 | /// The shared data. 85 | inner: Arc, 86 | } 87 | 88 | impl WakeSource { 89 | /// Notifies the sink if a waker is registered. 90 | /// 91 | /// This automatically unregisters any waker that may have been previously 92 | /// registered. 93 | #[inline] 94 | pub fn notify(&self) { 95 | self.inner.notify(); 96 | } 97 | } 98 | -------------------------------------------------------------------------------- /src/borrowed_waker.rs: -------------------------------------------------------------------------------- 1 | use core::task::Waker; 2 | 3 | use crate::{DiatomicWaker, WaitUntil}; 4 | 5 | /// A non-owned object that can await notifications from one or several 6 | /// [`WakeSourceRef`]s. 7 | /// 8 | /// See the [crate-level documentation](crate) for usage. 9 | #[derive(Debug)] 10 | pub struct WakeSinkRef<'a> { 11 | /// The shared data. 12 | pub(crate) inner: &'a DiatomicWaker, 13 | } 14 | 15 | impl<'a> WakeSinkRef<'a> { 16 | /// Creates a new `WakeSourceRef` associated to this sink with the same 17 | /// lifetime. 18 | #[inline] 19 | pub fn source_ref(&self) -> WakeSourceRef<'a> { 20 | WakeSourceRef { inner: self.inner } 21 | } 22 | 23 | /// Registers a new waker. 24 | /// 25 | /// Registration is lazy: the waker is cloned only if it differs from the 26 | /// last registered waker (note that the last registered waker is cached 27 | /// even if it was unregistered). 28 | #[inline] 29 | pub fn register(&mut self, waker: &Waker) { 30 | // Safety: `DiatomicWaker::register`, `DiatomicWaker::unregister` and 31 | // `DiatomicWaker::wait_until` cannot be used concurrently from multiple 32 | // thread since `WakeSinkRef` does not implement `Clone` and the 33 | // wrappers of the above methods require exclusive ownership to 34 | // `WakeSinkRef`. 35 | unsafe { self.inner.register(waker) }; 36 | } 37 | 38 | /// Unregisters the waker. 39 | /// 40 | /// After the waker is unregistered, subsequent calls to 41 | /// `WakeSourceRef::notify` will be ignored. 42 | /// 43 | /// Note that the previously-registered waker (if any) remains cached. 44 | #[inline] 45 | pub fn unregister(&mut self) { 46 | // Safety: `DiatomicWaker::register`, `DiatomicWaker::unregister` and 47 | // `DiatomicWaker::wait_until` cannot be used concurrently from multiple 48 | // thread since `WakeSinkRef` does not implement `Clone` and the 49 | // wrappers of the above methods require exclusive ownership to 50 | // `WakeSinkRef`. 51 | unsafe { self.inner.unregister() }; 52 | } 53 | 54 | /// Returns a future that can be `await`ed until the provided predicate 55 | /// returns a value. 56 | /// 57 | /// The predicate is checked each time a notification is received. 58 | #[inline] 59 | pub fn wait_until(&mut self, predicate: P) -> WaitUntil<'_, P, T> 60 | where 61 | P: FnMut() -> Option, 62 | { 63 | // Safety: `DiatomicWaker::register`, `DiatomicWaker::unregister` and 64 | // `DiatomicWaker::wait_until` cannot be used concurrently from multiple 65 | // thread since `WakeSinkRef` does not implement `Clone` and the 66 | // wrappers of the above methods require exclusive ownership to 67 | // `WakeSinkRef`. 68 | unsafe { self.inner.wait_until(predicate) } 69 | } 70 | } 71 | 72 | /// A non-owned object that can send notifications to a [`WakeSinkRef`]. 73 | /// 74 | /// See the [crate-level documentation](crate) for usage. 75 | #[derive(Clone, Debug)] 76 | pub struct WakeSourceRef<'a> { 77 | /// The shared data. 78 | pub(crate) inner: &'a DiatomicWaker, 79 | } 80 | 81 | impl WakeSourceRef<'_> { 82 | /// Notifies the sink if a waker is registered. 83 | /// 84 | /// This automatically unregisters any waker that may have been previously 85 | /// registered. 86 | #[inline] 87 | pub fn notify(&self) { 88 | self.inner.notify(); 89 | } 90 | } 91 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | //! Async, fast synchronization primitives for task wakeup. 2 | //! 3 | //! `diatomic-waker` is similar to [`atomic-waker`][atomic-waker] in that it 4 | //! enables concurrent updates and notifications to a wrapped `Waker`. Unlike 5 | //! the latter, however, it does not use spinlocks[^spinlocks] and is faster, in 6 | //! particular when the consumer is notified periodically rather than just once. 7 | //! It can in particular be used as a very fast, single-consumer [eventcount] to 8 | //! turn a non-blocking data structure into an asynchronous one (see MPSC 9 | //! channel receiver example). 10 | //! 11 | //! The API distinguishes between the entity that registers wakers ([`WakeSink`] 12 | //! or [`WakeSinkRef`]) and the possibly many entities that notify the waker 13 | //! ([`WakeSource`]s or [`WakeSourceRef`]s). 14 | //! 15 | //! Most users will prefer to use [`WakeSink`] and [`WakeSource`], which readily 16 | //! store a shared [`DiatomicWaker`] within an `Arc`. You may otherwise elect to 17 | //! allocate a [`DiatomicWaker`] yourself, but will then need to use the 18 | //! lifetime-bounded [`WakeSinkRef`] and [`WakeSourceRef`], or ensure by other 19 | //! means that waker registration is not performed concurrently. 20 | //! 21 | //! [atomic-waker]: https://docs.rs/atomic-waker/latest/atomic_waker/ 22 | //! [eventcount]: 23 | //! https://www.1024cores.net/home/lock-free-algorithms/eventcounts 24 | //! [^spinlocks]: The implementation of [AtomicWaker][atomic-waker] yields to the 25 | //! runtime on contention, which is in effect an executor-mediated spinlock. 26 | //! 27 | //! # Features flags 28 | //! 29 | //! By default, this crate enables the `alloc` feature to provide the owned 30 | //! [`WakeSink`] and [`WakeSource`]. It can be made `no-std`-compatible by 31 | //! specifying `default-features = false`. 32 | //! 33 | //! 34 | //! # Examples 35 | //! 36 | //! A multi-producer, single-consumer channel of capacity 1 for sending 37 | //! `NonZeroUsize` values, with an asynchronous receiver: 38 | //! 39 | //! ``` 40 | //! use std::num::NonZeroUsize; 41 | //! use std::sync::atomic::{AtomicUsize, Ordering}; 42 | //! use std::sync::Arc; 43 | //! 44 | //! use diatomic_waker::{WakeSink, WakeSource}; 45 | //! 46 | //! // The sending side of the channel. 47 | //! #[derive(Clone)] 48 | //! struct Sender { 49 | //! wake_src: WakeSource, 50 | //! value: Arc, 51 | //! } 52 | //! 53 | //! // The receiving side of the channel. 54 | //! struct Receiver { 55 | //! wake_sink: WakeSink, 56 | //! value: Arc, 57 | //! } 58 | //! 59 | //! // Creates an empty channel. 60 | //! fn channel() -> (Sender, Receiver) { 61 | //! let value = Arc::new(AtomicUsize::new(0)); 62 | //! let wake_sink = WakeSink::new(); 63 | //! let wake_src = wake_sink.source(); 64 | //! 65 | //! ( 66 | //! Sender { 67 | //! wake_src, 68 | //! value: value.clone(), 69 | //! }, 70 | //! Receiver { wake_sink, value }, 71 | //! ) 72 | //! } 73 | //! 74 | //! impl Sender { 75 | //! // Sends a value if the channel is empty. 76 | //! fn try_send(&self, value: NonZeroUsize) -> bool { 77 | //! let success = self 78 | //! .value 79 | //! .compare_exchange(0, value.get(), Ordering::Relaxed, Ordering::Relaxed) 80 | //! .is_ok(); 81 | //! if success { 82 | //! self.wake_src.notify() 83 | //! }; 84 | //! 85 | //! success 86 | //! } 87 | //! } 88 | //! 89 | //! impl Receiver { 90 | //! // Receives a value asynchronously. 91 | //! async fn recv(&mut self) -> NonZeroUsize { 92 | //! // Wait until the predicate returns `Some(value)`, i.e. when the atomic 93 | //! // value becomes non-zero. 94 | //! self.wake_sink 95 | //! .wait_until(|| NonZeroUsize::new(self.value.swap(0, Ordering::Relaxed))) 96 | //! .await 97 | //! } 98 | //! } 99 | //! ``` 100 | //! 101 | //! 102 | //! In some case, it may be necessary to use the lower-level 103 | //! [`register`](WakeSink::register) and [`unregister`](WakeSink::unregister) 104 | //! methods rather than the [`wait_until`](WakeSink::wait_until) convenience 105 | //! method. 106 | //! 107 | //! This is how the behavior of the above `recv` method could be 108 | //! reproduced with a hand-coded future: 109 | //! 110 | //! ``` 111 | //! use std::future::Future; 112 | //! # use std::num::NonZeroUsize; 113 | //! use std::pin::Pin; 114 | //! # use std::sync::atomic::{AtomicUsize, Ordering}; 115 | //! # use std::sync::Arc; 116 | //! use std::task::{Context, Poll}; 117 | //! # use diatomic_waker::WakeSink; 118 | //! 119 | //! # struct Receiver { 120 | //! # wake_sink: WakeSink, 121 | //! # value: Arc, 122 | //! # } 123 | //! struct Recv<'a> { 124 | //! receiver: &'a mut Receiver, 125 | //! } 126 | //! 127 | //! impl Future for Recv<'_> { 128 | //! type Output = NonZeroUsize; 129 | //! 130 | //! fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 131 | //! // Avoid waker registration if a value is readily available. 132 | //! let value = NonZeroUsize::new(self.receiver.value.swap(0, Ordering::Relaxed)); 133 | //! if let Some(value) = value { 134 | //! return Poll::Ready(value); 135 | //! } 136 | //! 137 | //! // Register the waker to be polled again once a value is available. 138 | //! self.receiver.wake_sink.register(cx.waker()); 139 | //! 140 | //! // Check again after registering the waker to prevent a race condition. 141 | //! let value = NonZeroUsize::new(self.receiver.value.swap(0, Ordering::Relaxed)); 142 | //! if let Some(value) = value { 143 | //! // Avoid a spurious wake-up. 144 | //! self.receiver.wake_sink.unregister(); 145 | //! 146 | //! return Poll::Ready(value); 147 | //! } 148 | //! 149 | //! Poll::Pending 150 | //! } 151 | //! } 152 | //! ``` 153 | #![warn(missing_docs, missing_debug_implementations, unreachable_pub)] 154 | #![cfg_attr(not(test), no_std)] 155 | #![cfg_attr(docsrs, feature(doc_auto_cfg, doc_cfg_hide))] 156 | 157 | #[cfg(feature = "alloc")] 158 | extern crate alloc; 159 | 160 | #[cfg(feature = "alloc")] 161 | mod arc_waker; 162 | mod borrowed_waker; 163 | mod loom_exports; 164 | #[deprecated( 165 | since = "0.2.0", 166 | note = "items from this module are now available in the root module" 167 | )] 168 | pub mod primitives; 169 | mod waker; 170 | 171 | #[cfg(feature = "alloc")] 172 | pub use arc_waker::{WakeSink, WakeSource}; 173 | pub use borrowed_waker::{WakeSinkRef, WakeSourceRef}; 174 | pub use waker::{DiatomicWaker, WaitUntil}; 175 | 176 | /// Tests. 177 | #[cfg(all(test, not(diatomic_waker_loom)))] 178 | mod tests { 179 | use std::sync::atomic::{AtomicBool, Ordering}; 180 | use std::thread; 181 | use std::time::Duration; 182 | 183 | use pollster::block_on; 184 | 185 | use super::*; 186 | 187 | #[test] 188 | fn waker_wait_until() { 189 | let mut sink = WakeSink::new(); 190 | let source = sink.source(); 191 | static FLAG: AtomicBool = AtomicBool::new(false); 192 | 193 | let t1 = thread::spawn(move || { 194 | std::thread::sleep(Duration::from_millis(10)); 195 | source.notify(); // force a spurious notification 196 | std::thread::sleep(Duration::from_millis(10)); 197 | FLAG.store(true, Ordering::Relaxed); 198 | source.notify(); 199 | }); 200 | 201 | let t2 = thread::spawn(move || { 202 | block_on(sink.wait_until(|| { 203 | if FLAG.load(Ordering::Relaxed) { 204 | Some(()) 205 | } else { 206 | None 207 | } 208 | })); 209 | 210 | assert!(FLAG.load(Ordering::Relaxed)); 211 | }); 212 | 213 | t1.join().unwrap(); 214 | t2.join().unwrap(); 215 | } 216 | 217 | #[test] 218 | fn waker_ref_wait_until() { 219 | let mut w = DiatomicWaker::new(); 220 | let mut sink = w.sink_ref(); 221 | let source = sink.source_ref(); 222 | static FLAG: AtomicBool = AtomicBool::new(false); 223 | 224 | thread::scope(|s| { 225 | s.spawn(move || { 226 | std::thread::sleep(Duration::from_millis(10)); 227 | source.notify(); // force a spurious notification 228 | std::thread::sleep(Duration::from_millis(10)); 229 | FLAG.store(true, Ordering::Relaxed); 230 | source.notify(); 231 | }); 232 | 233 | s.spawn(move || { 234 | block_on(sink.wait_until(|| { 235 | if FLAG.load(Ordering::Relaxed) { 236 | Some(()) 237 | } else { 238 | None 239 | } 240 | })); 241 | 242 | assert!(FLAG.load(Ordering::Relaxed)); 243 | }); 244 | }); 245 | } 246 | } 247 | 248 | /// Loom tests. 249 | #[cfg(all(test, diatomic_waker_loom))] 250 | mod tests { 251 | use super::*; 252 | 253 | use core::task::Waker; 254 | use std::future::Future; 255 | use std::pin::Pin; 256 | use std::sync::atomic::Ordering; 257 | use std::sync::Arc; 258 | use std::task::{Context, Poll}; 259 | 260 | use loom::model::Builder; 261 | use loom::sync::atomic::{AtomicU32, AtomicUsize}; 262 | use loom::thread; 263 | 264 | use waker_fn::waker_fn; 265 | 266 | /// A waker factory that registers notifications from the newest waker only. 267 | #[derive(Clone, Default)] 268 | struct MultiWaker { 269 | state: Arc, 270 | } 271 | impl MultiWaker { 272 | /// Clears the notification flag and returns the former notification 273 | /// status. 274 | /// 275 | /// This operation has Acquire semantic when a notification is indeed 276 | /// present, and Relaxed otherwise. It is therefore appropriate to 277 | /// simulate a scheduler receiving a notification as it ensures that all 278 | /// memory operations preceding the notification of a task are visible. 279 | fn take_notification(&self) -> bool { 280 | // Clear the notification flag. 281 | let mut state = self.state.load(Ordering::Relaxed); 282 | loop { 283 | // This is basically a `fetch_or` but with an atomic memory 284 | // ordering that depends on the LSB. 285 | let notified_stated = state | 1; 286 | let unnotified_stated = state & !1; 287 | match self.state.compare_exchange_weak( 288 | notified_stated, 289 | unnotified_stated, 290 | Ordering::Acquire, 291 | Ordering::Relaxed, 292 | ) { 293 | Ok(_) => return true, 294 | Err(s) => { 295 | state = s; 296 | if state == unnotified_stated { 297 | return false; 298 | } 299 | } 300 | } 301 | } 302 | } 303 | 304 | /// Clears the notification flag and creates a new waker. 305 | fn new_waker(&self) -> Waker { 306 | // Increase the epoch and clear the notification flag. 307 | let mut state = self.state.load(Ordering::Relaxed); 308 | let mut epoch; 309 | loop { 310 | // Increase the epoch by 2. 311 | epoch = (state & !1) + 2; 312 | match self.state.compare_exchange_weak( 313 | state, 314 | epoch, 315 | Ordering::Relaxed, 316 | Ordering::Relaxed, 317 | ) { 318 | Ok(_) => break, 319 | Err(s) => state = s, 320 | } 321 | } 322 | 323 | // Create a waker that only notifies if it is the newest waker. 324 | let waker_state = self.state.clone(); 325 | waker_fn(move || { 326 | let mut state = waker_state.load(Ordering::Relaxed); 327 | loop { 328 | let new_state = if state & !1 == epoch { 329 | epoch | 1 330 | } else { 331 | break; 332 | }; 333 | match waker_state.compare_exchange( 334 | state, 335 | new_state, 336 | Ordering::Release, 337 | Ordering::Relaxed, 338 | ) { 339 | Ok(_) => break, 340 | Err(s) => state = s, 341 | } 342 | } 343 | }) 344 | } 345 | } 346 | 347 | // A simple counter that can be used to simulate the availability of a 348 | // certain number of tokens. In order to model the weakest possible 349 | // predicate from the viewpoint of atomic memory ordering, only Relaxed 350 | // atomic operations are used. 351 | #[derive(Clone, Default)] 352 | struct Counter { 353 | count: Arc, 354 | } 355 | impl Counter { 356 | fn increment(&self) { 357 | self.count.fetch_add(1, Ordering::Relaxed); 358 | } 359 | fn try_decrement(&self) -> bool { 360 | let mut count = self.count.load(Ordering::Relaxed); 361 | loop { 362 | if count == 0 { 363 | return false; 364 | } 365 | match self.count.compare_exchange( 366 | count, 367 | count - 1, 368 | Ordering::Relaxed, 369 | Ordering::Relaxed, 370 | ) { 371 | Ok(_) => return true, 372 | Err(c) => count = c, 373 | } 374 | } 375 | } 376 | } 377 | 378 | /// Test whether notifications may be lost. 379 | /// 380 | /// Make a certain amount of tokens available and notify the sink each time 381 | /// a token is made available. Optionally, it is possible to: 382 | /// - request that `max_spurious_wake` threads will simulate a spurious 383 | /// wake-up, 384 | /// - change the waker each time it is polled. 385 | /// 386 | /// A default preemption bound will be applied if none was specified through 387 | /// an environment variable. 388 | fn loom_notify( 389 | token_count: usize, 390 | max_spurious_wake: usize, 391 | change_waker: bool, 392 | preemption_bound: usize, 393 | ) { 394 | // Only set the preemption bound if it wasn't already specified via a environment variable. 395 | let mut builder = Builder::new(); 396 | if builder.preemption_bound.is_none() { 397 | builder.preemption_bound = Some(preemption_bound); 398 | } 399 | 400 | builder.check(move || { 401 | let token_counter = Counter::default(); 402 | let mut wake_sink = WakeSink::new(); 403 | 404 | for src_id in 0..token_count { 405 | thread::spawn({ 406 | let token_counter = token_counter.clone(); 407 | let wake_src = wake_sink.source(); 408 | 409 | move || { 410 | if src_id < max_spurious_wake { 411 | wake_src.notify(); 412 | } 413 | token_counter.increment(); 414 | wake_src.notify(); 415 | } 416 | }); 417 | } 418 | 419 | let multi_waker = MultiWaker::default(); 420 | let mut waker = multi_waker.new_waker(); 421 | let mut satisfied_predicates_count = 0; 422 | 423 | // Iterate until all tokens are "received". 424 | // 425 | // Note: the loop does not have any assertion. This is by design: 426 | // missed notifications will be caught by Loom with a `Model 427 | // exceeded maximum number of branches` error because the spin loop 428 | // will then spin forever. 429 | while satisfied_predicates_count < token_count { 430 | let mut wait_until = wake_sink.wait_until(|| { 431 | if token_counter.try_decrement() { 432 | Some(()) 433 | } else { 434 | None 435 | } 436 | }); 437 | 438 | // Poll the predicate until it is satisfied. 439 | loop { 440 | let mut cx = Context::from_waker(&waker); 441 | let poll_state = Pin::new(&mut wait_until).poll(&mut cx); 442 | 443 | if poll_state == Poll::Ready(()) { 444 | satisfied_predicates_count += 1; 445 | break; 446 | } 447 | 448 | // Simulate the scheduler by spinning until the next 449 | // notification. 450 | while !multi_waker.take_notification() { 451 | thread::yield_now(); 452 | } 453 | 454 | if change_waker { 455 | waker = multi_waker.new_waker(); 456 | } 457 | } 458 | } 459 | }); 460 | } 461 | 462 | #[test] 463 | fn loom_notify_two_tokens() { 464 | const DEFAULT_PREEMPTION_BOUND: usize = 4; 465 | 466 | loom_notify(2, 0, false, DEFAULT_PREEMPTION_BOUND); 467 | } 468 | 469 | #[test] 470 | fn loom_notify_two_tokens_one_spurious() { 471 | const DEFAULT_PREEMPTION_BOUND: usize = 4; 472 | 473 | loom_notify(2, 1, false, DEFAULT_PREEMPTION_BOUND); 474 | } 475 | 476 | #[test] 477 | fn loom_notify_two_tokens_change_waker() { 478 | const DEFAULT_PREEMPTION_BOUND: usize = 3; 479 | 480 | loom_notify(2, 0, true, DEFAULT_PREEMPTION_BOUND); 481 | } 482 | 483 | #[test] 484 | fn loom_notify_two_tokens_one_spurious_change_waker() { 485 | const DEFAULT_PREEMPTION_BOUND: usize = 3; 486 | 487 | loom_notify(2, 1, true, DEFAULT_PREEMPTION_BOUND); 488 | } 489 | 490 | #[test] 491 | fn loom_notify_three_tokens() { 492 | const DEFAULT_PREEMPTION_BOUND: usize = 2; 493 | 494 | loom_notify(3, 0, false, DEFAULT_PREEMPTION_BOUND); 495 | } 496 | 497 | #[test] 498 | /// Test whether concurrent read and write access to the waker is possible. 499 | /// 500 | /// 3 different wakers are registered to force a waker slot to be re-used. 501 | fn loom_waker_slot_reuse() { 502 | // This tests require a high preemption bound to catch typical atomic 503 | // memory ordering mistakes. 504 | const DEFAULT_PREEMPTION_BOUND: usize = 5; 505 | 506 | // Only set the preemption bound if it wasn't already specified via a 507 | // environment variable. 508 | let mut builder = Builder::new(); 509 | if builder.preemption_bound.is_none() { 510 | builder.preemption_bound = Some(DEFAULT_PREEMPTION_BOUND); 511 | } 512 | 513 | builder.check(move || { 514 | let mut wake_sink = WakeSink::new(); 515 | 516 | thread::spawn({ 517 | let wake_src = wake_sink.source(); 518 | 519 | move || { 520 | wake_src.notify(); 521 | } 522 | }); 523 | thread::spawn({ 524 | let wake_src = wake_sink.source(); 525 | 526 | move || { 527 | wake_src.notify(); 528 | wake_src.notify(); 529 | } 530 | }); 531 | 532 | let multi_waker = MultiWaker::default(); 533 | for _ in 0..3 { 534 | let waker = multi_waker.new_waker(); 535 | wake_sink.register(&waker); 536 | } 537 | }); 538 | } 539 | } 540 | -------------------------------------------------------------------------------- /src/loom_exports.rs: -------------------------------------------------------------------------------- 1 | pub(crate) mod sync { 2 | pub(crate) mod atomic { 3 | #[cfg(not(all(test, diatomic_waker_loom)))] 4 | pub(crate) use core::sync::atomic::AtomicUsize; 5 | #[cfg(all(test, diatomic_waker_loom))] 6 | pub(crate) use loom::sync::atomic::AtomicUsize; 7 | } 8 | } 9 | 10 | #[cfg(all(test, diatomic_waker_loom))] 11 | pub(crate) mod cell { 12 | pub(crate) use loom::cell::UnsafeCell; 13 | } 14 | #[cfg(not(all(test, diatomic_waker_loom)))] 15 | pub(crate) mod cell { 16 | #[derive(Debug)] 17 | pub(crate) struct UnsafeCell(core::cell::UnsafeCell); 18 | 19 | impl UnsafeCell { 20 | pub(crate) const fn new(data: T) -> UnsafeCell { 21 | UnsafeCell(core::cell::UnsafeCell::new(data)) 22 | } 23 | pub(crate) fn with(&self, f: impl FnOnce(*const T) -> R) -> R { 24 | f(self.0.get()) 25 | } 26 | pub(crate) fn with_mut(&self, f: impl FnOnce(*mut T) -> R) -> R { 27 | f(self.0.get()) 28 | } 29 | } 30 | } 31 | -------------------------------------------------------------------------------- /src/primitives.rs: -------------------------------------------------------------------------------- 1 | //! Primitives for task wakeup. 2 | 3 | pub use crate::{DiatomicWaker, WaitUntil}; 4 | -------------------------------------------------------------------------------- /src/waker.rs: -------------------------------------------------------------------------------- 1 | use core::future::Future; 2 | use core::pin::Pin; 3 | use core::sync::atomic::Ordering; 4 | use core::task::{Context, Poll, Waker}; 5 | 6 | use crate::loom_exports::cell::UnsafeCell; 7 | use crate::loom_exports::sync::atomic::AtomicUsize; 8 | use crate::WakeSinkRef; 9 | 10 | // The state of the waker is tracked by the following bit flags: 11 | // 12 | // * INDEX [I]: slot index of the current waker, if any (0 or 1), 13 | // * UPDATE [U]: an updated waker has been registered in the redundant slot at 14 | // index 1 - INDEX, 15 | // * REGISTERED [R]: a waker is registered and awaits a notification 16 | // * LOCKED [L]: a notifier has taken the notifier lock and is in the process of 17 | // sending a notification, 18 | // * NOTIFICATION [N]: a notifier has failed to take the lock when a waker was 19 | // registered and has requested the notifier holding the lock to send a 20 | // notification on its behalf (implies REGISTERED and LOCKED). 21 | // 22 | // The waker stored in the slot at INDEX ("current" waker) is shared between the 23 | // sink (entity which registers wakers) and the source that holds the notifier 24 | // lock (if any). For this reason, this waker may only be accessed by shared 25 | // reference. The waker at 1 - INDEX is exclusively owned by the sink, which is 26 | // free to mutate it. 27 | 28 | // Summary of valid states: 29 | // 30 | // | N L R U I | 31 | // |---------------------| 32 | // | 0 any any any any | 33 | // | 1 1 1 any any | 34 | 35 | // [I] Index of the current waker (0 or 1). 36 | const INDEX: usize = 0b00001; 37 | // [U] Indicates that an updated waker is available at 1 - INDEX. 38 | const UPDATE: usize = 0b00010; 39 | // [R] Indicates that a waker has been registered. 40 | const REGISTERED: usize = 0b00100; 41 | // [L] Indicates that a notifier holds the notifier lock to the waker at INDEX. 42 | const LOCKED: usize = 0b01000; 43 | // [N] Indicates that a notifier has failed to acquire the lock and has 44 | // requested the notifier holding the lock to notify on its behalf. 45 | const NOTIFICATION: usize = 0b10000; 46 | 47 | /// A primitive that can send or await notifications. 48 | /// 49 | /// It is almost always preferable to use the [`WakeSink`](crate::WakeSink) and 50 | /// [`WakeSource`](crate::WakeSource) which offer more convenience at the cost 51 | /// of an allocation in an `Arc`. 52 | /// 53 | /// If allocation is not possible or desirable, the 54 | /// [`sink_ref`](DiatomicWaker::sink_ref) method can be used to create a 55 | /// [`WakeSinkRef`] handle and one or more 56 | /// [`WakeSourceRef`](crate::borrowed_waker::WakeSourceRef)s, the non-owned 57 | /// counterparts to `WakeSink` and `WakeSource`. 58 | /// 59 | /// Finally, `DiatomicWaker` exposes `unsafe` methods that can be used to create 60 | /// custom synchronization primitives. 61 | #[derive(Debug)] 62 | pub struct DiatomicWaker { 63 | /// A bit field for `INDEX`, `UPDATE`, `REGISTERED`, `LOCKED` and `NOTIFICATION`. 64 | state: AtomicUsize, 65 | /// Redundant slots for the waker. 66 | waker: [UnsafeCell>; 2], 67 | } 68 | 69 | impl DiatomicWaker { 70 | /// Creates a new `DiatomicWaker`. 71 | #[cfg(not(all(test, diatomic_waker_loom)))] 72 | pub const fn new() -> Self { 73 | Self { 74 | state: AtomicUsize::new(0), 75 | waker: [UnsafeCell::new(None), UnsafeCell::new(None)], 76 | } 77 | } 78 | 79 | #[cfg(all(test, diatomic_waker_loom))] 80 | pub fn new() -> Self { 81 | Self { 82 | state: AtomicUsize::new(0), 83 | waker: [UnsafeCell::new(None), UnsafeCell::new(None)], 84 | } 85 | } 86 | 87 | /// Returns a sink with a lifetime bound to this `DiatomicWaker`. 88 | /// 89 | /// This mutably borrows the waker, thus ensuring that at most one 90 | /// associated sink can be active at a time. 91 | pub fn sink_ref(&mut self) -> WakeSinkRef<'_> { 92 | WakeSinkRef { inner: self } 93 | } 94 | 95 | /// Sends a notification if a waker is registered. 96 | /// 97 | /// This automatically unregisters any waker that may have been previously 98 | /// registered. 99 | pub fn notify(&self) { 100 | // Transitions: see `try_lock` and `try_unlock`. 101 | 102 | let mut state = if let Ok(s) = try_lock(&self.state) { 103 | s 104 | } else { 105 | return; 106 | }; 107 | 108 | loop { 109 | let idx = state & INDEX; 110 | 111 | // Safety: the notifier lock has been acquired, which guarantees 112 | // exclusive access to the waker at `INDEX`. 113 | unsafe { 114 | self.wake_by_ref(idx); 115 | } 116 | 117 | if let Err(s) = try_unlock(&self.state, state) { 118 | state = s; 119 | } else { 120 | return; 121 | } 122 | 123 | // One more loop iteration is necessary because the waker was 124 | // registered again and another notifier has failed to send a 125 | // notification while the notifier lock was taken. 126 | } 127 | } 128 | 129 | /// Registers a new waker. 130 | /// 131 | /// Registration is lazy: the waker is cloned only if it differs from the 132 | /// last registered waker (note that the last registered waker is cached 133 | /// even if it was unregistered). 134 | /// 135 | /// # Safety 136 | /// 137 | /// The `register`, `unregister` and `wait_until` methods cannot be used 138 | /// concurrently from multiple threads. 139 | pub unsafe fn register(&self, waker: &Waker) { 140 | // Transitions if the new waker is the same as the one currently stored. 141 | // 142 | // | N L R U I | N L R U I | 143 | // |-----------------|-----------------| 144 | // | n l r u i | n l 1 u i | 145 | // 146 | // 147 | // Transitions if the new waker needs to be stored: 148 | // 149 | // Step 1 (only necessary if the state initially indicates R=U=1): 150 | // 151 | // | N L R U I | N L R U I | 152 | // |-----------------|-----------------| 153 | // | n l r u i | 0 l 0 u i | 154 | // 155 | // Step 2: 156 | // 157 | // | N L R U I | N L R U I | 158 | // |-----------------|-----------------| 159 | // | n l r u i | n l 1 1 i | 160 | 161 | // Ordering: Acquire ordering is necessary to synchronize with the 162 | // Release unlocking operation in `notify`, which ensures that all calls 163 | // to the waker in the redundant slot have completed. 164 | let state = self.state.load(Ordering::Acquire); 165 | 166 | // Compute the index of the waker that was most recently updated. Note 167 | // that the value of `recent_idx` as computed below remains correct even 168 | // if the state is stale since only this thread can store new wakers. 169 | let mut idx = state & INDEX; 170 | let recent_idx = if state & UPDATE == 0 { 171 | idx 172 | } else { 173 | INDEX - idx 174 | }; 175 | 176 | // Safety: it is safe to call `will_wake` since the registering thread 177 | // is the only one allowed to mutate the wakers so there can be no 178 | // concurrent mutable access to the waker. 179 | let is_up_to_date = self.will_wake(recent_idx, waker); 180 | 181 | // Fast path in case the waker is up to date. 182 | if is_up_to_date { 183 | // Set the `REGISTERED` flag. Ideally, the `NOTIFICATION` flag would 184 | // be cleared at the same time to avoid a spurious wake-up, but it 185 | // probably isn't worth the overhead of a CAS loop because having 186 | // this flag set when calling `register` is very unlikely: it would 187 | // mean that since the last call to `register`: 188 | // 1) a notifier has been holding the lock continuously, 189 | // 2) another notifier has tried and failed to take the lock, and 190 | // 3) `unregister` was never called. 191 | // 192 | // Ordering: Acquire ordering synchronizes with the Release and 193 | // AcqRel RMWs in `try_lock` (called by `notify`) and ensures that 194 | // either the predicate set before the call to `notify` will be 195 | // visible after the call to `register`, or the registered waker 196 | // will be visible during the call to `notify` (or both). Note that 197 | // Release ordering is not necessary since the waker has not changed 198 | // and this RMW takes part in a release sequence headed by the 199 | // initial registration of the waker. 200 | self.state.fetch_or(REGISTERED, Ordering::Acquire); 201 | 202 | return; 203 | } 204 | 205 | // The waker needs to be stored in the redundant slot. 206 | // 207 | // It is necessary to make sure that either the `UPDATE` or the 208 | // `REGISTERED` flag is cleared to prevent concurrent access by a notifier 209 | // to the redundant waker slot while the waker is updated. 210 | // 211 | // Note that only the thread registering the waker can set `REGISTERED` 212 | // and `UPDATE` so even if the state is stale, observing `REGISTERED` or 213 | // `UPDATE` as cleared guarantees that such flag is and will remain 214 | // cleared until this thread sets them. 215 | if state & (UPDATE | REGISTERED) == (UPDATE | REGISTERED) { 216 | // Clear the `REGISTERED` and `NOTIFICATION` flags. 217 | // 218 | // Ordering: Acquire ordering is necessary to synchronize with the 219 | // Release unlocking operation in `notify`, which ensures that all 220 | // calls to the waker in the redundant slot have completed. 221 | let state = self 222 | .state 223 | .fetch_and(!(REGISTERED | NOTIFICATION), Ordering::Acquire); 224 | 225 | // It is possible that `UPDATE` was cleared and `INDEX` was switched 226 | // by a notifier after the initial load of the state, so the waker 227 | // index needs to be updated. 228 | idx = state & INDEX; 229 | } 230 | 231 | // Always store the new waker in the redundant slot to avoid racing with 232 | // a notifier. 233 | let redundant_idx = 1 - idx; 234 | 235 | // Store the new waker. 236 | // 237 | // Safety: it is safe to store the new waker in the redundant slot 238 | // because the `REGISTERED` flag and/or the `UPDATE` flag are/is cleared 239 | // so the notifier will not attempt to switch the waker. 240 | self.set_waker(redundant_idx, waker.clone()); 241 | 242 | // Make the waker visible. 243 | // 244 | // Ordering: Acquire ordering synchronizes with the Release and AcqRel 245 | // RMWs in `try_lock` (called by `notify`) and ensures that either the 246 | // predicate set before the call to `notify` will be visible after the 247 | // call to `register`, or the registered waker will be visible during 248 | // the call to `notify` (or both). Since the waker has been modified 249 | // above, Release ordering is also necessary to synchronize with the 250 | // AcqRel RMW in `try_lock` (success case) and ensure that the 251 | // modification to the waker is fully visible when notifying. 252 | self.state.fetch_or(UPDATE | REGISTERED, Ordering::AcqRel); 253 | } 254 | 255 | /// Unregisters the waker. 256 | /// 257 | /// After the waker is unregistered, subsequent calls to `notify` will be 258 | /// ignored. 259 | /// 260 | /// Note that the previously-registered waker (if any) remains cached. 261 | /// 262 | /// # Safety 263 | /// 264 | /// The `register`, `unregister` and `wait_until` methods cannot be used 265 | /// concurrently from multiple threads. 266 | pub unsafe fn unregister(&self) { 267 | // Transitions: 268 | // 269 | // | N L R U I | N L R U I | 270 | // |-----------------|-----------------| 271 | // | n l r u i | 0 l 0 u i | 272 | 273 | // Modify the state. Note that the waker is not dropped: caching it can 274 | // avoid a waker drop/cloning cycle (typically, 2 RMWs) in the frequent 275 | // case when the next waker to be registered will be the same as the one 276 | // being unregistered. 277 | // 278 | // Ordering: no waker was modified so Relaxed ordering is sufficient. 279 | self.state 280 | .fetch_and(!(REGISTERED | NOTIFICATION), Ordering::Relaxed); 281 | } 282 | 283 | /// Returns a future that can be `await`ed until the provided predicate 284 | /// returns a value. 285 | /// 286 | /// The predicate is checked each time a notification is received. 287 | /// 288 | /// # Safety 289 | /// 290 | /// The `register`, `unregister` and `wait_until` methods cannot be used 291 | /// concurrently from multiple threads. 292 | pub unsafe fn wait_until(&self, predicate: P) -> WaitUntil<'_, P, T> 293 | where 294 | P: FnMut() -> Option, 295 | { 296 | WaitUntil::new(self, predicate) 297 | } 298 | 299 | /// Sets the waker at index `idx`. 300 | /// 301 | /// # Safety 302 | /// 303 | /// The caller must have exclusive access to the waker at index `idx`. 304 | unsafe fn set_waker(&self, idx: usize, new: Waker) { 305 | self.waker[idx].with_mut(|waker| (*waker) = Some(new)); 306 | } 307 | 308 | /// Notify the waker at index `idx`. 309 | /// 310 | /// # Safety 311 | /// 312 | /// The waker at index `idx` cannot be modified concurrently. 313 | unsafe fn wake_by_ref(&self, idx: usize) { 314 | self.waker[idx].with(|waker| { 315 | if let Some(waker) = &*waker { 316 | waker.wake_by_ref(); 317 | } 318 | }); 319 | } 320 | 321 | /// Check whether the waker at index `idx` will wake the same task as the 322 | /// provided waker. 323 | /// 324 | /// # Safety 325 | /// 326 | /// The waker at index `idx` cannot be modified concurrently. 327 | unsafe fn will_wake(&self, idx: usize, other: &Waker) -> bool { 328 | self.waker[idx].with(|waker| match &*waker { 329 | Some(waker) => waker.will_wake(other), 330 | None => false, 331 | }) 332 | } 333 | } 334 | 335 | impl Default for DiatomicWaker { 336 | fn default() -> Self { 337 | Self::new() 338 | } 339 | } 340 | 341 | unsafe impl Send for DiatomicWaker {} 342 | unsafe impl Sync for DiatomicWaker {} 343 | 344 | /// Attempts to acquire the notifier lock and returns the current state upon 345 | /// success. 346 | /// 347 | /// Acquisition of the lock will fail in the following cases: 348 | /// 349 | /// * the `REGISTERED` flag is cleared, meaning that there is no need to wake 350 | /// and therefore no need to lock, 351 | /// * the lock is already taken, in which case the `NOTIFICATION` flag will be 352 | /// set if the `REGISTERED` flag is set. 353 | /// 354 | /// If acquisition of the lock succeeds, the `REGISTERED` flag is cleared. If 355 | /// additionally the `UPDATE` flag was set, it is cleared and `INDEX` is 356 | /// flipped. 357 | /// 358 | /// Transition table: 359 | /// 360 | /// | N L R U I | N L R U I | 361 | /// |-----------------|-----------------| 362 | /// | 0 0 0 u i | 0 0 0 u i | (failure) 363 | /// | 0 0 1 0 i | 0 1 0 0 i | (success) 364 | /// | 0 0 1 1 i | 0 1 0 0 !i | (success) 365 | /// | 0 1 0 u i | 0 1 0 u i | (failure) 366 | /// | n 1 1 u i | 1 1 1 u i | (failure) 367 | /// 368 | fn try_lock(state: &AtomicUsize) -> Result { 369 | let mut old_state = state.load(Ordering::Relaxed); 370 | 371 | loop { 372 | if old_state & (LOCKED | REGISTERED) == REGISTERED { 373 | // Success path. 374 | 375 | // If `UPDATE` is set, clear `UPDATE` and flip `INDEX` with the xor 376 | // mask. 377 | let update_bit = old_state & UPDATE; 378 | let xor_mask = update_bit | (update_bit >> 1); 379 | 380 | // Set `LOCKED` and clear `REGISTERED` with the xor mask. 381 | let xor_mask = xor_mask | LOCKED | REGISTERED; 382 | 383 | let new_state = old_state ^ xor_mask; 384 | 385 | // Ordering: Acquire is necessary to synchronize with the Release 386 | // ordering in `register` so that the new waker, if any, is visible. 387 | // Release ordering synchronizes with the Acquire and AcqRel RMWs in 388 | // `register` and ensures that either the predicate set before the 389 | // call to `notify` will be visible after the call to `register`, or 390 | // the registered waker will be visible during the call to `notify` 391 | // (or both). 392 | match state.compare_exchange_weak( 393 | old_state, 394 | new_state, 395 | Ordering::AcqRel, 396 | Ordering::Relaxed, 397 | ) { 398 | Ok(_) => return Ok(new_state), 399 | Err(s) => old_state = s, 400 | } 401 | } else { 402 | // Failure path. 403 | 404 | // Set the `NOTIFICATION` bit if `REGISTERED` was set. 405 | let registered_bit = old_state & REGISTERED; 406 | let new_state = old_state | (registered_bit << 2); 407 | 408 | // Ordering: Release ordering synchronizes with the Acquire and 409 | // AcqRel RMWs in `register` and ensures that either the predicate 410 | // set before the call to `notify` will be visible after the call to 411 | // `register`, or the registered waker will be visible during the 412 | // call to `notify` (or both). 413 | match state.compare_exchange_weak( 414 | old_state, 415 | new_state, 416 | Ordering::Release, 417 | Ordering::Relaxed, 418 | ) { 419 | Ok(_) => return Err(()), 420 | Err(s) => old_state = s, 421 | } 422 | }; 423 | } 424 | } 425 | 426 | /// Attempts to release the notifier lock and returns the current state upon 427 | /// failure. 428 | /// 429 | /// Release of the lock will fail if the `NOTIFICATION` flag is set because it 430 | /// means that, after the lock was taken, the registering thread has requested 431 | /// to be notified again and another notifier has subsequently requested that 432 | /// such notification be sent on its behalf; if additionally the `UPDATE` flag 433 | /// was set (i.e. a new waker is available), it is cleared and `INDEX` is 434 | /// flipped. 435 | /// 436 | /// Transition table: 437 | /// 438 | /// | N L R U I | N L R U I | 439 | /// |-----------------|-----------------| 440 | /// | 0 1 r u i | 0 0 r u i | (success) 441 | /// | 1 1 1 0 i | 0 1 0 0 i | (failure) 442 | /// | 1 1 1 1 i | 0 1 0 0 !i | (failure) 443 | /// 444 | fn try_unlock(state: &AtomicUsize, mut old_state: usize) -> Result<(), usize> { 445 | loop { 446 | if old_state & NOTIFICATION == 0 { 447 | // Success path. 448 | 449 | let new_state = old_state & !LOCKED; 450 | 451 | // Ordering: Release is necessary to synchronize with the Acquire 452 | // ordering in `register` and ensure that the waker call has 453 | // completed before a new waker is stored. 454 | match state.compare_exchange_weak( 455 | old_state, 456 | new_state, 457 | Ordering::Release, 458 | Ordering::Relaxed, 459 | ) { 460 | Ok(_) => return Ok(()), 461 | Err(s) => old_state = s, 462 | } 463 | } else { 464 | // Failure path. 465 | 466 | // If `UPDATE` is set, clear `UPDATE` and flip `INDEX` with the xor mask. 467 | let update_bit = old_state & UPDATE; 468 | let xor_mask = update_bit | (update_bit >> 1); 469 | 470 | // Clear `NOTIFICATION` and `REGISTERED` with the xor mask. 471 | let xor_mask = xor_mask | NOTIFICATION | REGISTERED; 472 | 473 | let new_state = old_state ^ xor_mask; 474 | 475 | // Ordering: Release is necessary to synchronize with the Acquire 476 | // ordering in `register` and ensure that the call to 477 | // `Waker::wake_by_ref` has completed before a new waker is stored. 478 | // Acquire ordering is in turn necessary to ensure that any newly 479 | // registered waker is visible. 480 | match state.compare_exchange_weak( 481 | old_state, 482 | new_state, 483 | Ordering::AcqRel, 484 | Ordering::Relaxed, 485 | ) { 486 | Ok(_) => return Err(new_state), 487 | Err(s) => old_state = s, 488 | } 489 | }; 490 | } 491 | } 492 | 493 | /// A future that can be `await`ed until a predicate is satisfied. 494 | #[derive(Debug)] 495 | pub struct WaitUntil<'a, P, T> 496 | where 497 | P: FnMut() -> Option, 498 | { 499 | predicate: P, 500 | wake: &'a DiatomicWaker, 501 | } 502 | 503 | impl<'a, P, T> WaitUntil<'a, P, T> 504 | where 505 | P: FnMut() -> Option, 506 | { 507 | /// Creates a future associated to the specified wake that can be `await`ed 508 | /// until the specified predicate is satisfied. 509 | fn new(wake: &'a DiatomicWaker, predicate: P) -> Self { 510 | Self { predicate, wake } 511 | } 512 | } 513 | 514 | impl Option, T> Unpin for WaitUntil<'_, P, T> {} 515 | 516 | impl<'a, P, T> Future for WaitUntil<'a, P, T> 517 | where 518 | P: FnMut() -> Option, 519 | { 520 | type Output = T; 521 | 522 | fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 523 | // Safety: the safety of this method is contingent on the safety of the 524 | // `register` and `unregister` methods. Since a `WaitUntil` future can 525 | // only be created from the unsafe `wait_until` method, however, the 526 | // user must uphold the contract that `register`, `unregister` and 527 | // `wait_until` cannot be used concurrently from multiple threads. 528 | unsafe { 529 | if let Some(value) = (self.predicate)() { 530 | return Poll::Ready(value); 531 | } 532 | self.wake.register(cx.waker()); 533 | 534 | if let Some(value) = (self.predicate)() { 535 | self.wake.unregister(); 536 | return Poll::Ready(value); 537 | } 538 | } 539 | 540 | Poll::Pending 541 | } 542 | } 543 | --------------------------------------------------------------------------------