feature_flags: Add FeatureFlagAppExt::wait_for_flag_or_timeout (#24055)

This PR adds a new `wait_for_flag_or_timeout` method to the
`FeatureFlagAppExt` trait.

This encapsulates the somewhat gnarly code for using `wait_for_flag`
with a timeout.

A side benefit of this is that the tasks waiting on the feature flags
run in parallel, so in the case where the feature flags do not resolve
we don't end up having to wait on consecutive timeouts. This should help
a bit with https://github.com/zed-industries/zed/issues/23922.

Release Notes:

- N/A

crates/feature_flags/Cargo.toml

@@ -12,5 +12,6 @@ workspace = true
 path = "src/feature_flags.rs"
 
 [dependencies]
-gpui.workspace = true
 futures.workspace = true
+gpui.workspace = true
+smol.workspace = true

crates/feature_flags/src/feature_flags.rs

@@ -1,5 +1,7 @@
-use futures::{channel::oneshot, FutureExt as _};
-use gpui::{App, Context, Global, Subscription, Window};
+use futures::channel::oneshot;
+use futures::{select_biased, FutureExt};
+use gpui::{App, Context, Global, Subscription, Task, Window};
+use std::time::Duration;
 use std::{future::Future, pin::Pin, task::Poll};
 
 #[derive(Default)]
@@ -114,6 +116,10 @@ where
 
 pub trait FeatureFlagAppExt {
     fn wait_for_flag<T: FeatureFlag>(&mut self) -> WaitForFlag;
+
+    /// Waits for the specified feature flag to resolve, up to the given timeout.
+    fn wait_for_flag_or_timeout<T: FeatureFlag>(&mut self, timeout: Duration) -> Task<bool>;
+
     fn update_flags(&mut self, staff: bool, flags: Vec<String>);
     fn set_staff(&mut self, staff: bool);
     fn has_flag<T: FeatureFlag>(&self) -> bool;
@@ -180,6 +186,20 @@ impl FeatureFlagAppExt for App {
 
         WaitForFlag(rx, subscription)
     }
+
+    fn wait_for_flag_or_timeout<T: FeatureFlag>(&mut self, timeout: Duration) -> Task<bool> {
+        let wait_for_flag = self.wait_for_flag::<T>();
+
+        self.spawn(|_cx| async move {
+            let mut wait_for_flag = wait_for_flag.fuse();
+            let mut timeout = FutureExt::fuse(smol::Timer::after(timeout));
+
+            select_biased! {
+                is_enabled = wait_for_flag => is_enabled,
+                _ = timeout => false,
+            }
+        })
+    }
 }
 
 pub struct WaitForFlag(oneshot::Receiver<bool>, Option<Subscription>);