Box future outputs before passing them to async_task

Co-Authored-By: Nathan Sobo <nathan@zed.dev>

gpui/src/app.rs

@@ -1,6 +1,6 @@
 use crate::{
     elements::ElementBox,
-    executor,
+    executor::{self, Task},
     keymap::{self, Keystroke},
     platform::{self, CursorStyle, Platform, PromptLevel, WindowOptions},
     presenter::Presenter,
@@ -8,7 +8,6 @@ use crate::{
     AssetCache, AssetSource, ClipboardItem, FontCache, PathPromptOptions, TextLayoutCache,
 };
 use anyhow::{anyhow, Result};
-use async_task::Task;
 use keymap::MatchResult;
 use parking_lot::Mutex;
 use platform::Event;

gpui/src/executor.rs

@@ -1,12 +1,12 @@
 use anyhow::{anyhow, Result};
 use async_task::Runnable;
-pub use async_task::Task;
 use backtrace::{Backtrace, BacktraceFmt, BytesOrWideString};
 use parking_lot::Mutex;
 use postage::{barrier, prelude::Stream as _};
 use rand::prelude::*;
 use smol::{channel, prelude::*, Executor, Timer};
 use std::{
+    any::Any,
     fmt::{self, Debug},
     marker::PhantomData,
     mem,
@@ -42,6 +42,24 @@ pub enum Background {
     },
 }
 
+type AnyLocalFuture = Pin<Box<dyn 'static + Future<Output = Box<dyn Any + 'static>>>>;
+type AnyFuture = Pin<Box<dyn 'static + Send + Future<Output = Box<dyn Any + Send + 'static>>>>;
+type AnyTask = async_task::Task<Box<dyn Any + Send + 'static>>;
+type AnyLocalTask = async_task::Task<Box<dyn Any + 'static>>;
+
+pub enum Task<T> {
+    Local {
+        any_task: AnyLocalTask,
+        result_type: PhantomData<T>,
+    },
+    Send {
+        any_task: AnyTask,
+        result_type: PhantomData<T>,
+    },
+}
+
+unsafe impl<T: Send> Send for Task<T> {}
+
 struct DeterministicState {
     rng: StdRng,
     seed: u64,
@@ -77,10 +95,7 @@ impl Deterministic {
         }
     }
 
-    fn spawn_from_foreground<T>(&self, future: Pin<Box<dyn Future<Output = T>>>) -> Task<T>
-    where
-        T: 'static,
-    {
+    fn spawn_from_foreground(&self, future: AnyLocalFuture) -> AnyLocalTask {
         let backtrace = Backtrace::new_unresolved();
         let scheduled_once = AtomicBool::new(false);
         let state = self.state.clone();
@@ -99,10 +114,7 @@ impl Deterministic {
         task
     }
 
-    fn spawn<T>(&self, future: Pin<Box<dyn Send + Future<Output = T>>>) -> Task<T>
-    where
-        T: 'static + Send,
-    {
+    fn spawn(&self, future: AnyFuture) -> AnyTask {
         let backtrace = Backtrace::new_unresolved();
         let state = self.state.clone();
         let unparker = self.parker.lock().unparker();
@@ -117,10 +129,7 @@ impl Deterministic {
         task
     }
 
-    fn run<T>(&self, mut future: Pin<Box<dyn Future<Output = T>>>) -> T
-    where
-        T: 'static,
-    {
+    fn run(&self, mut future: AnyLocalFuture) -> Box<dyn Any> {
         let woken = Arc::new(AtomicBool::new(false));
         loop {
             if let Some(result) = self.run_internal(woken.clone(), &mut future) {
@@ -138,18 +147,15 @@ impl Deterministic {
 
     fn run_until_parked(&self) {
         let woken = Arc::new(AtomicBool::new(false));
-        let mut future = std::future::pending::<()>().boxed_local();
+        let mut future = any_local_future(std::future::pending::<()>());
         self.run_internal(woken, &mut future);
     }
 
-    fn run_internal<T>(
+    fn run_internal(
         &self,
         woken: Arc<AtomicBool>,
-        future: &mut Pin<Box<dyn Future<Output = T>>>,
-    ) -> Option<T>
-    where
-        T: 'static,
-    {
+        future: &mut AnyLocalFuture,
+    ) -> Option<Box<dyn Any>> {
         let unparker = self.parker.lock().unparker();
         let waker = waker_fn(move || {
             woken.store(true, SeqCst);
@@ -203,13 +209,7 @@ impl Deterministic {
         }
     }
 
-    pub fn block_on<F, T>(&self, future: F) -> Option<T>
-    where
-        T: 'static,
-        F: Future<Output = T>,
-    {
-        smol::pin!(future);
-
+    fn block_on(&self, future: &mut AnyLocalFuture) -> Option<Box<dyn Any>> {
         let unparker = self.parker.lock().unparker();
         let waker = waker_fn(move || {
             unparker.unpark();
@@ -394,8 +394,9 @@ impl Foreground {
     }
 
     pub fn spawn<T: 'static>(&self, future: impl Future<Output = T> + 'static) -> Task<T> {
-        let future = future.boxed_local();
-        match self {
+        let future = any_local_future(future);
+        let any_task = match self {
+            Self::Deterministic(executor) => executor.spawn_from_foreground(future),
             Self::Platform { dispatcher, .. } => {
                 let dispatcher = dispatcher.clone();
                 let schedule = move |runnable: Runnable| dispatcher.run_on_main_thread(runnable);
@@ -404,17 +405,18 @@ impl Foreground {
                 task
             }
             Self::Test(executor) => executor.spawn(future),
-            Self::Deterministic(executor) => executor.spawn_from_foreground(future),
-        }
+        };
+        Task::local(any_task)
     }
 
     pub fn run<T: 'static>(&self, future: impl 'static + Future<Output = T>) -> T {
-        let future = future.boxed_local();
-        match self {
+        let future = any_local_future(future);
+        let any_value = match self {
+            Self::Deterministic(executor) => executor.run(future),
             Self::Platform { .. } => panic!("you can't call run on a platform foreground executor"),
             Self::Test(executor) => smol::block_on(executor.run(future)),
-            Self::Deterministic(executor) => executor.run(future),
-        }
+        };
+        *any_value.downcast().unwrap()
     }
 
     pub fn forbid_parking(&self) {
@@ -500,33 +502,34 @@ impl Background {
         T: 'static + Send,
         F: Send + Future<Output = T> + 'static,
     {
-        let future = future.boxed();
-        match self {
+        let future = any_future(future);
+        let any_task = match self {
             Self::Production { executor, .. } => executor.spawn(future),
             Self::Deterministic(executor) => executor.spawn(future),
-        }
+        };
+        Task::send(any_task)
     }
 
     pub fn block_with_timeout<F, T>(
         &self,
         timeout: Duration,
         future: F,
-    ) -> Result<T, Pin<Box<dyn Future<Output = T>>>>
+    ) -> Result<T, impl Future<Output = T>>
     where
         T: 'static,
         F: 'static + Unpin + Future<Output = T>,
     {
-        let mut future = future.boxed_local();
+        let mut future = any_local_future(future);
         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),
             };
             if let Some(output) = output {
-                return Ok(output);
+                return Ok(*output.downcast().unwrap());
             }
         }
-        Err(future)
+        Err(async { *future.await.downcast().unwrap() })
     }
 
     pub async fn scoped<'scope, F>(&self, scheduler: F)
@@ -576,3 +579,68 @@ pub fn deterministic(seed: u64) -> (Rc<Foreground>, Arc<Background>) {
         Arc::new(Background::Deterministic(executor)),
     )
 }
+
+impl<T> Task<T> {
+    fn local(any_task: AnyLocalTask) -> Self {
+        Self::Local {
+            any_task,
+            result_type: PhantomData,
+        }
+    }
+
+    pub fn detach(self) {
+        match self {
+            Task::Local { any_task, .. } => any_task.detach(),
+            Task::Send { any_task, .. } => any_task.detach(),
+        }
+    }
+}
+
+impl<T: Send> Task<T> {
+    fn send(any_task: AnyTask) -> Self {
+        Self::Send {
+            any_task,
+            result_type: PhantomData,
+        }
+    }
+}
+
+impl<T: fmt::Debug> fmt::Debug for Task<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Task::Local { any_task, .. } => any_task.fmt(f),
+            Task::Send { any_task, .. } => any_task.fmt(f),
+        }
+    }
+}
+
+impl<T: 'static> Future for Task<T> {
+    type Output = T;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match unsafe { self.get_unchecked_mut() } {
+            Task::Local { any_task, .. } => {
+                any_task.poll(cx).map(|value| *value.downcast().unwrap())
+            }
+            Task::Send { any_task, .. } => {
+                any_task.poll(cx).map(|value| *value.downcast().unwrap())
+            }
+        }
+    }
+}
+
+fn any_future<T, F>(future: F) -> AnyFuture
+where
+    T: 'static + Send,
+    F: Future<Output = T> + Send + 'static,
+{
+    async { Box::new(future.await) as Box<dyn Any + Send> }.boxed()
+}
+
+fn any_local_future<T, F>(future: F) -> AnyLocalFuture
+where
+    T: 'static,
+    F: Future<Output = T> + 'static,
+{
+    async { Box::new(future.await) as Box<dyn Any> }.boxed_local()
+}