Skip to content

Commit

Permalink
Make it possible to time out after a deadline
Browse files Browse the repository at this point in the history
  • Loading branch information
sbarral committed May 15, 2024
1 parent b1596a3 commit 75d214e
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 1 deletion.
3 changes: 3 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@ An efficient async condition variable for lock-free algorithms.
categories = ["asynchronous", "concurrency"]
keywords = ["async", "event", "atomic", "futures"]

[dependencies]
pin-project-lite = "0.2"

[dev-dependencies]
tokio = { version = "1", features = ["full"] }
futures-executor = "0.3"
Expand Down
76 changes: 75 additions & 1 deletion src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,12 @@
//! });
//! ```
// Temporary workaround until the `async_event_loom` flag can be whitelisted
// without a `build.rs` [1].
//
// [1]: (https://github.com/rust-lang/rust/issues/124800).
#![allow(unexpected_cfgs)]

mod loom_exports;

use std::future::Future;
Expand All @@ -81,6 +87,7 @@ use std::task::{Context, Poll, Waker};
use loom_exports::cell::UnsafeCell;
use loom_exports::sync::atomic::{self, AtomicBool};
use loom_exports::sync::Mutex;
use pin_project_lite::pin_project;

/// An object that can receive or send notifications.
pub struct Event {
Expand Down Expand Up @@ -130,9 +137,29 @@ impl Event {

/// Returns a future that can be `await`ed until the provided predicate is
/// satisfied.
pub fn wait_until<F: FnMut() -> Option<T>, T>(&self, predicate: F) -> WaitUntil<F, T> {
pub fn wait_until<F, T>(&self, predicate: F) -> WaitUntil<F, T>
where
F: FnMut() -> Option<T>,
{
WaitUntil::new(&self.wait_set, predicate)
}

/// Returns a future that can be `await`ed until the provided predicate is
/// satisfied or until the provided future completes.
///
/// The deadline is specified as a `Future` that is expected to resolves to
/// `()` after some duration, such as a `tokio::time::Sleep` future.
pub fn wait_until_or_timeout<F, T, D>(
&self,
predicate: F,
deadline: D,
) -> WaitUntilOrTimeout<F, T, D>
where
F: FnMut() -> Option<T>,
D: Future<Output = ()>,
{
WaitUntilOrTimeout::new(&self.wait_set, predicate, deadline)
}
}

impl Default for Event {
Expand Down Expand Up @@ -378,6 +405,53 @@ enum WaitUntilState {
Completed,
}

pin_project! {
/// A future that can be `await`ed until a predicate is satisfied or until a
/// deadline elapses.
pub struct WaitUntilOrTimeout<'a, F: FnMut() -> Option<T>, T, D: Future<Output = ()>> {
wait_until: WaitUntil<'a, F, T>,
#[pin]
deadline: D,
}
}

impl<'a, F, T, D> WaitUntilOrTimeout<'a, F, T, D>
where
F: FnMut() -> Option<T>,
D: Future<Output = ()>,
{
/// Creates a future associated with the specified event sink that can be
/// `await`ed until the specified predicate is satisfied, or until the
/// specified timeout future completes.
fn new(wait_set: &'a WaitSet, predicate: F, deadline: D) -> Self {
Self {
wait_until: WaitUntil::new(wait_set, predicate),
deadline,
}
}
}

impl<'a, F, T, D> Future for WaitUntilOrTimeout<'a, F, T, D>
where
F: FnMut() -> Option<T>,
D: Future<Output = ()>,
{
type Output = Option<T>;

#[inline]
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();

if let Poll::Ready(value) = Pin::new(this.wait_until).poll(cx) {
Poll::Ready(Some(value))
} else if this.deadline.poll(cx).is_ready() {
Poll::Ready(None)
} else {
Poll::Pending
}
}
}

/// A set of notifiers.
///
/// The set wraps a Mutex-protected list of notifiers and manages a flag for
Expand Down

0 comments on commit 75d214e

Please sign in to comment.