diff --git a/crates/gpui/src/executor.rs b/crates/gpui/src/executor.rs index 9c4b9e90e0..f16a88d2dd 100644 --- a/crates/gpui/src/executor.rs +++ b/crates/gpui/src/executor.rs @@ -306,6 +306,32 @@ impl Deterministic { None } + pub fn timer(&self, duration: Duration) -> impl Future { + let (tx, mut rx) = postage::barrier::channel(); + let timer_id; + { + let mut state = self.state.lock(); + let wakeup_at = state.now + duration; + timer_id = util::post_inc(&mut state.next_timer_id); + state.pending_timers.push((timer_id, wakeup_at, tx)); + } + + let remove_timer = util::defer({ + let state = self.state.clone(); + move || { + state + .lock() + .pending_timers + .retain(|(id, _, _)| *id != timer_id); + } + }); + + async move { + postage::prelude::Stream::recv(&mut rx).await; + drop(remove_timer); + } + } + pub fn advance_clock(&self, duration: Duration) { let mut state = self.state.lock(); state.now += duration; @@ -438,41 +464,18 @@ impl Foreground { } } - pub async fn timer(&self, duration: Duration) { - match self { - #[cfg(any(test, feature = "test-support"))] - Self::Deterministic { executor, .. } => { - use postage::prelude::Stream as _; + pub fn timer(&self, duration: Duration) -> impl Future { + let mut timer = None; - let (tx, mut rx) = postage::barrier::channel(); - let timer_id; - { - let mut state = executor.state.lock(); - let wakeup_at = state.now + duration; - timer_id = util::post_inc(&mut state.next_timer_id); - state.pending_timers.push((timer_id, wakeup_at, tx)); - } + #[cfg(any(test, feature = "test-support"))] + if let Self::Deterministic { executor, .. } = self { + timer = Some(executor.timer(duration)); + } - struct DropTimer<'a>(usize, &'a Foreground); - impl<'a> Drop for DropTimer<'a> { - fn drop(&mut self) { - match self.1 { - Foreground::Deterministic { executor, .. } => { - executor - .state - .lock() - .pending_timers - .retain(|(timer_id, _, _)| *timer_id != self.0); - } - _ => unreachable!(), - } - } - } - - let _guard = DropTimer(timer_id, self); - rx.recv().await; - } - _ => { + async move { + if let Some(timer) = timer { + timer.await; + } else { Timer::after(duration).await; } } @@ -600,6 +603,23 @@ impl Background { } } + pub fn timer(&self, duration: Duration) -> impl Future { + let mut timer = None; + + #[cfg(any(test, feature = "test-support"))] + if let Self::Deterministic { executor, .. } = self { + timer = Some(executor.timer(duration)); + } + + async move { + if let Some(timer) = timer { + timer.await; + } else { + Timer::after(duration).await; + } + } + } + #[cfg(any(test, feature = "test-support"))] pub async fn simulate_random_delay(&self) { use rand::prelude::*; diff --git a/crates/gpui/src/util.rs b/crates/gpui/src/util.rs index 9e59c387e8..dc857b4c66 100644 --- a/crates/gpui/src/util.rs +++ b/crates/gpui/src/util.rs @@ -1,5 +1,6 @@ use smol::future::FutureExt; use std::{future::Future, time::Duration}; +pub use util::*; pub fn post_inc(value: &mut usize) -> usize { let prev = *value; diff --git a/crates/util/src/lib.rs b/crates/util/src/lib.rs index 919fecf8f9..adefa732e7 100644 --- a/crates/util/src/lib.rs +++ b/crates/util/src/lib.rs @@ -123,6 +123,18 @@ where } } +struct Defer(Option); + +impl Drop for Defer { + fn drop(&mut self) { + self.0.take().map(|f| f()); + } +} + +pub fn defer(f: F) -> impl Drop { + Defer(Some(f)) +} + #[cfg(test)] mod tests { use super::*;