tenant.openagents/omega

1use std::future::Future;
2
3use gpui::{App, AppContext, Global, ReadGlobal, Task};
4use gpui_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    let handle = runtime.handle().clone();
21    cx.set_global(GlobalTokio {
22        owned_runtime: Some(runtime),
23        handle,
24    });
25}
26
27/// Initializes the Tokio wrapper using a Tokio runtime handle.
28pub fn init_from_handle(cx: &mut App, handle: tokio::runtime::Handle) {
29    cx.set_global(GlobalTokio {
30        owned_runtime: None,
31        handle,
32    });
33}
34
35struct GlobalTokio {
36    owned_runtime: Option<tokio::runtime::Runtime>,
37    handle: tokio::runtime::Handle,
38}
39
40impl Global for GlobalTokio {}
41
42impl Drop for GlobalTokio {
43    fn drop(&mut self) {
44        if let Some(runtime) = self.owned_runtime.take() {
45            runtime.shutdown_background();
46        }
47    }
48}
49
50pub struct Tokio {}
51
52impl Tokio {
53    /// Spawns the given future on Tokio's thread pool, and returns it via a GPUI task
54    /// Note that the Tokio task will be cancelled if the GPUI task is dropped
55    pub fn spawn<C, Fut, R>(cx: &C, f: Fut) -> Task<Result<R, JoinError>>
56    where
57        C: AppContext,
58        Fut: Future<Output = R> + Send + 'static,
59        R: Send + 'static,
60    {
61        cx.read_global(|tokio: &GlobalTokio, cx| {
62            let join_handle = tokio.handle.spawn(f);
63            let abort_handle = join_handle.abort_handle();
64            let cancel = defer(move || {
65                abort_handle.abort();
66            });
67            cx.background_spawn(async move {
68                let result = join_handle.await;
69                drop(cancel);
70                result
71            })
72        })
73    }
74
75    /// Spawns the given future on Tokio's thread pool, and returns it via a GPUI task
76    /// Note that the Tokio task will be cancelled if the GPUI task is dropped
77    pub fn spawn_result<C, Fut, R>(cx: &C, f: Fut) -> Task<anyhow::Result<R>>
78    where
79        C: AppContext,
80        Fut: Future<Output = anyhow::Result<R>> + Send + 'static,
81        R: Send + 'static,
82    {
83        cx.read_global(|tokio: &GlobalTokio, cx| {
84            let join_handle = tokio.handle.spawn(f);
85            let abort_handle = join_handle.abort_handle();
86            let cancel = defer(move || {
87                abort_handle.abort();
88            });
89            cx.background_spawn(async move {
90                let result = join_handle.await?;
91                drop(cancel);
92                result
93            })
94        })
95    }
96
97    pub fn handle(cx: &App) -> tokio::runtime::Handle {
98        GlobalTokio::global(cx).handle.clone()
99    }
100}
101