Add FutureExt::with_timeout and use it for for Room::maintain_connection (#36175)

Release Notes:

- N/A

---------

Co-authored-by: Antonio Scandurra <me@as-cii.com>

crates/gpui/src/app/test_context.rs

@@ -585,7 +585,7 @@ impl<V: 'static> Entity<V> {
         cx.executor().advance_clock(advance_clock_by);
 
         async move {
-            let notification = crate::util::timeout(duration, rx.recv())
+            let notification = crate::util::smol_timeout(duration, rx.recv())
                 .await
                 .expect("next notification timed out");
             drop(subscription);
@@ -629,7 +629,7 @@ impl<V> Entity<V> {
         let handle = self.downgrade();
 
         async move {
-            crate::util::timeout(Duration::from_secs(1), async move {
+            crate::util::smol_timeout(Duration::from_secs(1), async move {
                 loop {
                     {
                         let cx = cx.borrow();

crates/gpui/src/gpui.rs

@@ -157,7 +157,7 @@ pub use taffy::{AvailableSpace, LayoutId};
 #[cfg(any(test, feature = "test-support"))]
 pub use test::*;
 pub use text_system::*;
-pub use util::arc_cow::ArcCow;
+pub use util::{FutureExt, Timeout, arc_cow::ArcCow};
 pub use view::*;
 pub use window::*;
 

crates/gpui/src/util.rs

@@ -1,13 +1,11 @@
-use std::sync::atomic::AtomicUsize;
-use std::sync::atomic::Ordering::SeqCst;
-#[cfg(any(test, feature = "test-support"))]
-use std::time::Duration;
-
-#[cfg(any(test, feature = "test-support"))]
-use futures::Future;
-
-#[cfg(any(test, feature = "test-support"))]
-use smol::future::FutureExt;
+use crate::{BackgroundExecutor, Task};
+use std::{
+    future::Future,
+    pin::Pin,
+    sync::atomic::{AtomicUsize, Ordering::SeqCst},
+    task,
+    time::Duration,
+};
 
 pub use util::*;
 
@@ -70,8 +68,59 @@ pub trait FluentBuilder {
     }
 }
 
+/// Extensions for Future types that provide additional combinators and utilities.
+pub trait FutureExt {
+    /// Requires a Future to complete before the specified duration has elapsed.
+    /// Similar to tokio::timeout.
+    fn with_timeout(self, timeout: Duration, executor: &BackgroundExecutor) -> WithTimeout<Self>
+    where
+        Self: Sized;
+}
+
+impl<T: Future> FutureExt for T {
+    fn with_timeout(self, timeout: Duration, executor: &BackgroundExecutor) -> WithTimeout<Self>
+    where
+        Self: Sized,
+    {
+        WithTimeout {
+            future: self,
+            timer: executor.timer(timeout),
+        }
+    }
+}
+
+pub struct WithTimeout<T> {
+    future: T,
+    timer: Task<()>,
+}
+
+#[derive(Debug, thiserror::Error)]
+#[error("Timed out before future resolved")]
+/// Error returned by with_timeout when the timeout duration elapsed before the future resolved
+pub struct Timeout;
+
+impl<T: Future> Future for WithTimeout<T> {
+    type Output = Result<T::Output, Timeout>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut task::Context) -> task::Poll<Self::Output> {
+        // SAFETY: the fields of Timeout are private and we never move the future ourselves
+        // And its already pinned since we are being polled (all futures need to be pinned to be polled)
+        let this = unsafe { self.get_unchecked_mut() };
+        let future = unsafe { Pin::new_unchecked(&mut this.future) };
+        let timer = unsafe { Pin::new_unchecked(&mut this.timer) };
+
+        if let task::Poll::Ready(output) = future.poll(cx) {
+            task::Poll::Ready(Ok(output))
+        } else if timer.poll(cx).is_ready() {
+            task::Poll::Ready(Err(Timeout))
+        } else {
+            task::Poll::Pending
+        }
+    }
+}
+
 #[cfg(any(test, feature = "test-support"))]
-pub async fn timeout<F, T>(timeout: Duration, f: F) -> Result<T, ()>
+pub async fn smol_timeout<F, T>(timeout: Duration, f: F) -> Result<T, ()>
 where
     F: Future<Output = T>,
 {
@@ -80,7 +129,7 @@ where
         Err(())
     };
     let future = async move { Ok(f.await) };
-    timer.race(future).await
+    smol::future::FutureExt::race(timer, future).await
 }
 
 /// Increment the given atomic counter if it is not zero.