1use std::future::Future;
  2
  3use gpui::{App, AppContext, Global, ReadGlobal, Task};
  4use util::defer;
  5
  6pub use tokio::task::JoinError;
  7
  8/// Initializes the Tokio wrapper using a new Tokio runtime with 2 worker threads.
  9///
 10/// If you need more threads (or access to the runtime outside of GPUI), you can create the runtime
 11/// yourself and pass a Handle to `init_from_handle`.
 12pub fn init(cx: &mut App) {
 13    let runtime = tokio::runtime::Builder::new_multi_thread()
 14        // Since we now have two executors, let's try to keep our footprint small
 15        .worker_threads(2)
 16        .enable_all()
 17        .build()
 18        .expect("Failed to initialize Tokio");
 19
 20    cx.set_global(GlobalTokio::new(RuntimeHolder::Owned(runtime)));
 21}
 22
 23/// Initializes the Tokio wrapper using a Tokio runtime handle.
 24pub fn init_from_handle(cx: &mut App, handle: tokio::runtime::Handle) {
 25    cx.set_global(GlobalTokio::new(RuntimeHolder::Shared(handle)));
 26}
 27
 28enum RuntimeHolder {
 29    Owned(tokio::runtime::Runtime),
 30    Shared(tokio::runtime::Handle),
 31}
 32
 33impl RuntimeHolder {
 34    pub fn handle(&self) -> &tokio::runtime::Handle {
 35        match self {
 36            RuntimeHolder::Owned(runtime) => runtime.handle(),
 37            RuntimeHolder::Shared(handle) => handle,
 38        }
 39    }
 40}
 41
 42struct GlobalTokio {
 43    runtime: RuntimeHolder,
 44}
 45
 46impl Global for GlobalTokio {}
 47
 48impl GlobalTokio {
 49    fn new(runtime: RuntimeHolder) -> Self {
 50        Self { runtime }
 51    }
 52}
 53
 54pub struct Tokio {}
 55
 56impl Tokio {
 57    /// Spawns the given future on Tokio's thread pool, and returns it via a GPUI task
 58    /// Note that the Tokio task will be cancelled if the GPUI task is dropped
 59    pub fn spawn<C, Fut, R>(cx: &C, f: Fut) -> C::Result<Task<Result<R, JoinError>>>
 60    where
 61        C: AppContext,
 62        Fut: Future<Output = R> + Send + 'static,
 63        R: Send + 'static,
 64    {
 65        cx.read_global(|tokio: &GlobalTokio, cx| {
 66            let join_handle = tokio.runtime.handle().spawn(f);
 67            let abort_handle = join_handle.abort_handle();
 68            let cancel = defer(move || {
 69                abort_handle.abort();
 70            });
 71            cx.background_spawn(async move {
 72                let result = join_handle.await;
 73                drop(cancel);
 74                result
 75            })
 76        })
 77    }
 78
 79    /// Spawns the given future on Tokio's thread pool, and returns it via a GPUI task
 80    /// Note that the Tokio task will be cancelled if the GPUI task is dropped
 81    pub fn spawn_result<C, Fut, R>(cx: &C, f: Fut) -> C::Result<Task<anyhow::Result<R>>>
 82    where
 83        C: AppContext,
 84        Fut: Future<Output = anyhow::Result<R>> + Send + 'static,
 85        R: Send + 'static,
 86    {
 87        cx.read_global(|tokio: &GlobalTokio, cx| {
 88            let join_handle = tokio.runtime.handle().spawn(f);
 89            let abort_handle = join_handle.abort_handle();
 90            let cancel = defer(move || {
 91                abort_handle.abort();
 92            });
 93            cx.background_spawn(async move {
 94                let result = join_handle.await?;
 95                drop(cancel);
 96                result
 97            })
 98        })
 99    }
100
101    pub fn handle(cx: &App) -> tokio::runtime::Handle {
102        GlobalTokio::global(cx).runtime.handle().clone()
103    }
104}