I want to make a flag object that controls whether some tasks can be executed. It works like this:
- The flag has two states: enabled and disabled.
- There are some worker tasks that will only be executed only if:
- the current state of the flag is enabled,
- or wait for the next enabled state.
- There may be some other tasks that enable or disable the flag as they wish.
So the API of the flag could be:
pub struct Flag { ... }
impl Flag {
/// Creates a new flag object.
pub fn new(enabled: bool) -> Self { ... }
/// Enables the flag.
pub fn enable(&self) { ... }
/// Disables the flag.
pub fn disable(&self) { ... }
/// Waits the flag to become enabled.
pub async fn wait_enabled(&self) { ... }
}
Here is my implementation:
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::task::{Context, Poll};
use tokio::sync::futures::Notified;
use tokio::sync::Notify;
pin_project_lite::pin_project! {
#[project = WaitEnabledProjection]
pub struct WaitEnabled<'a> {
#[pin]
notified: Option<Notified<'a>>,
}
}
impl Future for WaitEnabled<'_> {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
match self.project().notified.as_pin_mut() {
None => Poll::Ready(()),
Some(notified) => notified.poll(cx),
}
}
}
pub struct Flag {
enabled: AtomicBool,
notify: Notify,
}
impl Flag {
/// Creates a new flag object.
pub fn new(enabled: bool) -> Self {
Self {
enabled: AtomicBool::new(enabled),
notify: Notify::new(),
}
}
/// Enables the flag.
pub fn enable(&self) {
self.enabled.store(true, Ordering::Relaxed);
self.notify.notify_waiters();
}
/// Disables the flag.
pub fn disable(&self) {
self.enabled.store(false, Ordering::Relaxed);
}
/// Waits the flag to become enabled.
pub fn wait_enabled(&self) -> WaitEnabled {
if !self.enabled.load(Ordering::Relaxed) {
let notified = self.notify.notified();
// This check prevents losing wake up calls.
if !self.enabled.load(Ordering::Relaxed) {
return WaitEnabled {
notified: Some(notified),
};
}
}
WaitEnabled { notified: None }
}
/// An alternative simpler implementation of `wait_enabled` but outputs bigger `Future` and the output `Future`
/// cannot be named.
pub async fn wait_enabled_2(&self) {
if !self.enabled.load(Ordering::Relaxed) {
let notified = self.notify.notified();
// This check prevents losing wake up calls.
if !self.enabled.load(Ordering::Relaxed) {
notified.await;
}
}
}
}
Here is a playground link for my implementation: https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=501b09d254598b128b80d810ac2f0655.
Here are my questions:
- Is my implementation correct?
- I use
Relaxed
ordering for checking the current state, is it strong enough to ensure correctness? - After the
Notified
future becomes ready, I do not check the state again, which may be a problem if there are spurious wake ups. DoesNotified
future guarantee the wake up comes fromnotify_waiters
calls?
- I use
- Is there an existing solution to this problem so I don’t have to implement it myself?
- What is the best way to implement such a flag type?