Add a facility for delaying quit until critical tasks finish

Co-Authored-By: Antonio Scandurra <me@as-cii.com>

Cargo.lock

@@ -2309,6 +2309,7 @@ dependencies = [
  "etagere",
  "font-kit",
  "foreign-types",
+ "futures",
  "gpui_macros",
  "image 0.23.14",
  "lazy_static",

crates/gpui/Cargo.toml

@@ -15,6 +15,7 @@ backtrace = "0.3"
 ctor = "0.1"
 env_logger = { version = "0.8", optional = true }
 etagere = "0.2"
+futures = "0.3"
 image = "0.23"
 lazy_static = "1.4.0"
 log = "0.4"

crates/gpui/src/app.rs

@@ -265,6 +265,18 @@ impl App {
         self
     }
 
+    pub fn on_quit<F>(self, mut callback: F) -> Self
+    where
+        F: 'static + FnMut(&mut MutableAppContext),
+    {
+        let cx = self.0.clone();
+        self.0
+            .borrow_mut()
+            .foreground_platform
+            .on_quit(Box::new(move || callback(&mut *cx.borrow_mut())));
+        self
+    }
+
     pub fn on_event<F>(self, mut callback: F) -> Self
     where
         F: 'static + FnMut(Event, &mut MutableAppContext) -> bool,

crates/gpui/src/executor.rs

@@ -38,9 +38,13 @@ pub enum Foreground {
 }
 
 pub enum Background {
-    Deterministic(Arc<Deterministic>),
+    Deterministic {
+        executor: Arc<Deterministic>,
+        critical_tasks: Mutex<Vec<Task<()>>>,
+    },
     Production {
         executor: Arc<smol::Executor<'static>>,
+        critical_tasks: Mutex<Vec<Task<()>>>,
         _stop: channel::Sender<()>,
     },
 }
@@ -500,6 +504,7 @@ impl Background {
 
         Self::Production {
             executor,
+            critical_tasks: Default::default(),
             _stop: stop.0,
         }
     }
@@ -516,11 +521,36 @@ impl Background {
         let future = any_future(future);
         let any_task = match self {
             Self::Production { executor, .. } => executor.spawn(future),
-            Self::Deterministic(executor) => executor.spawn(future),
+            Self::Deterministic { executor, .. } => executor.spawn(future),
         };
         Task::send(any_task)
     }
 
+    pub fn spawn_critical<T, F>(&self, future: F)
+    where
+        T: 'static + Send,
+        F: Send + Future<Output = T> + 'static,
+    {
+        let task = self.spawn(async move {
+            future.await;
+        });
+        match self {
+            Self::Production { critical_tasks, .. }
+            | Self::Deterministic { critical_tasks, .. } => critical_tasks.lock().push(task),
+        }
+    }
+
+    pub fn block_on_critical_tasks(&self, timeout: Duration) -> bool {
+        match self {
+            Background::Production { critical_tasks, .. }
+            | Self::Deterministic { critical_tasks, .. } => {
+                let tasks = mem::take(&mut *critical_tasks.lock());
+                self.block_with_timeout(timeout, futures::future::join_all(tasks))
+                    .is_ok()
+            }
+        }
+    }
+
     pub fn block_with_timeout<F, T>(
         &self,
         timeout: Duration,
@@ -534,7 +564,7 @@ impl Background {
         if !timeout.is_zero() {
             let output = match self {
                 Self::Production { .. } => smol::block_on(util::timeout(timeout, &mut future)).ok(),
-                Self::Deterministic(executor) => executor.block_on(&mut future),
+                Self::Deterministic { executor, .. } => executor.block_on(&mut future),
             };
             if let Some(output) = output {
                 return Ok(*output.downcast().unwrap());
@@ -587,7 +617,10 @@ pub fn deterministic(seed: u64) -> (Rc<Foreground>, Arc<Background>) {
     let executor = Arc::new(Deterministic::new(seed));
     (
         Rc::new(Foreground::Deterministic(executor.clone())),
-        Arc::new(Background::Deterministic(executor)),
+        Arc::new(Background::Deterministic {
+            executor,
+            critical_tasks: Default::default(),
+        }),
     )
 }
 

crates/gpui/src/platform.rs

@@ -60,6 +60,7 @@ pub trait Platform: Send + Sync {
 pub(crate) trait ForegroundPlatform {
     fn on_become_active(&self, callback: Box<dyn FnMut()>);
     fn on_resign_active(&self, callback: Box<dyn FnMut()>);
+    fn on_quit(&self, callback: Box<dyn FnMut()>);
     fn on_event(&self, callback: Box<dyn FnMut(Event) -> bool>);
     fn on_open_files(&self, callback: Box<dyn FnMut(Vec<PathBuf>)>);
     fn run(&self, on_finish_launching: Box<dyn FnOnce() -> ()>);

crates/gpui/src/platform/mac/platform.rs

@@ -81,6 +81,10 @@ unsafe fn build_classes() {
             sel!(applicationDidResignActive:),
             did_resign_active as extern "C" fn(&mut Object, Sel, id),
         );
+        decl.add_method(
+            sel!(applicationWillTerminate:),
+            will_terminate as extern "C" fn(&mut Object, Sel, id),
+        );
         decl.add_method(
             sel!(handleGPUIMenuItem:),
             handle_menu_item as extern "C" fn(&mut Object, Sel, id),
@@ -100,6 +104,7 @@ pub struct MacForegroundPlatform(RefCell<MacForegroundPlatformState>);
 pub struct MacForegroundPlatformState {
     become_active: Option<Box<dyn FnMut()>>,
     resign_active: Option<Box<dyn FnMut()>>,
+    quit: Option<Box<dyn FnMut()>>,
     event: Option<Box<dyn FnMut(crate::Event) -> bool>>,
     menu_command: Option<Box<dyn FnMut(&dyn AnyAction)>>,
     open_files: Option<Box<dyn FnMut(Vec<PathBuf>)>>,
@@ -196,6 +201,10 @@ impl platform::ForegroundPlatform for MacForegroundPlatform {
         self.0.borrow_mut().resign_active = Some(callback);
     }
 
+    fn on_quit(&self, callback: Box<dyn FnMut()>) {
+        self.0.borrow_mut().quit = Some(callback);
+    }
+
     fn on_event(&self, callback: Box<dyn FnMut(crate::Event) -> bool>) {
         self.0.borrow_mut().event = Some(callback);
     }
@@ -664,6 +673,13 @@ extern "C" fn did_resign_active(this: &mut Object, _: Sel, _: id) {
     }
 }
 
+extern "C" fn will_terminate(this: &mut Object, _: Sel, _: id) {
+    let platform = unsafe { get_foreground_platform(this) };
+    if let Some(callback) = platform.0.borrow_mut().quit.as_mut() {
+        callback();
+    }
+}
+
 extern "C" fn open_files(this: &mut Object, _: Sel, _: id, paths: id) {
     let paths = unsafe {
         (0..paths.count())

crates/gpui/src/platform/test.rs

@@ -58,6 +58,8 @@ impl super::ForegroundPlatform for ForegroundPlatform {
 
     fn on_resign_active(&self, _: Box<dyn FnMut()>) {}
 
+    fn on_quit(&self, _: Box<dyn FnMut()>) {}
+
     fn on_event(&self, _: Box<dyn FnMut(crate::Event) -> bool>) {}
 
     fn on_open_files(&self, _: Box<dyn FnMut(Vec<std::path::PathBuf>)>) {}