From 4d9659adc4f4704c8a80411b22d2036d1b9e30fe Mon Sep 17 00:00:00 2001
From: Marshall Bowers <git@maxdeviant.com>
Date: Fri, 31 Jan 2025 17:35:23 -0500
Subject: [PATCH] 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
---
 Cargo.lock                                |  1 +
 crates/feature_flags/Cargo.toml           |  3 ++-
 crates/feature_flags/src/feature_flags.rs | 24 +++++++++++++++++--
 crates/zed/src/zed.rs                     | 28 ++++++-----------------
 4 files changed, 32 insertions(+), 24 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 35dac08333..211c27c9b2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4547,6 +4547,7 @@ version = "0.1.0"
 dependencies = [
  "futures 0.3.31",
  "gpui",
+ "smol",
 ]
 
 [[package]]
diff --git a/crates/feature_flags/Cargo.toml b/crates/feature_flags/Cargo.toml
index a2f991b6f5..65d6942d50 100644
--- a/crates/feature_flags/Cargo.toml
+++ b/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
diff --git a/crates/feature_flags/src/feature_flags.rs b/crates/feature_flags/src/feature_flags.rs
index 13e19674aa..eb014be881 100644
--- a/crates/feature_flags/src/feature_flags.rs
+++ b/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>);
diff --git a/crates/zed/src/zed.rs b/crates/zed/src/zed.rs
index 1cfc9460d4..7bedb1b994 100644
--- a/crates/zed/src/zed.rs
+++ b/crates/zed/src/zed.rs
@@ -20,7 +20,6 @@ use command_palette_hooks::CommandPaletteFilter;
 use editor::ProposedChangesEditorToolbar;
 use editor::{scroll::Autoscroll, Editor, MultiBuffer};
 use feature_flags::FeatureFlagAppExt;
-use futures::FutureExt;
 use futures::{channel::mpsc, select_biased, StreamExt};
 use gpui::{
     actions, point, px, Action, App, AppContext as _, AsyncApp, Context, DismissEvent, Element,
@@ -45,6 +44,7 @@ use settings::{
 };
 use std::any::TypeId;
 use std::path::PathBuf;
+use std::time::Duration;
 use std::{borrow::Cow, ops::Deref, path::Path, sync::Arc};
 use terminal_view::terminal_panel::{self, TerminalPanel};
 use theme::{ActiveTheme, ThemeSettings};
@@ -362,8 +362,10 @@ fn initialize_panels(
     window: &mut Window,
     cx: &mut Context<Workspace>,
 ) {
-    let assistant2_feature_flag = cx.wait_for_flag::<feature_flags::Assistant2FeatureFlag>();
-    let git_ui_feature_flag = cx.wait_for_flag::<feature_flags::GitUiFeatureFlag>();
+    let assistant2_feature_flag =
+        cx.wait_for_flag_or_timeout::<feature_flags::Assistant2FeatureFlag>(Duration::from_secs(5));
+    let git_ui_feature_flag =
+        cx.wait_for_flag_or_timeout::<feature_flags::GitUiFeatureFlag>(Duration::from_secs(5));
 
     let prompt_builder = prompt_builder.clone();
 
@@ -405,16 +407,7 @@ fn initialize_panels(
             workspace.add_panel(notification_panel, window, cx);
         })?;
 
-        let git_ui_enabled = {
-            let mut git_ui_feature_flag = git_ui_feature_flag.fuse();
-            let mut timeout =
-                FutureExt::fuse(smol::Timer::after(std::time::Duration::from_secs(5)));
-
-            select_biased! {
-                is_git_ui_enabled = git_ui_feature_flag => is_git_ui_enabled,
-                _ = timeout => false,
-            }
-        };
+        let git_ui_enabled = git_ui_feature_flag.await;
 
         let git_panel = if git_ui_enabled {
             Some(git_ui::git_panel::GitPanel::load(workspace_handle.clone(), cx.clone()).await?)
@@ -430,14 +423,7 @@ fn initialize_panels(
         let is_assistant2_enabled = if cfg!(test) {
             false
         } else {
-            let mut assistant2_feature_flag = assistant2_feature_flag.fuse();
-            let mut timeout =
-                FutureExt::fuse(smol::Timer::after(std::time::Duration::from_secs(5)));
-
-            select_biased! {
-                is_assistant2_enabled = assistant2_feature_flag => is_assistant2_enabled,
-                _ = timeout => false,
-            }
+            assistant2_feature_flag.await
         };
 
         let (assistant_panel, assistant2_panel) = if is_assistant2_enabled {