Allow external handles to be provided to gpui_tokio (#42795)

William Whittaker created 1 month ago

This PR allows for a handle to an existing Tokio runtime to be passed to
gpui_tokio's initialization function, which means that Tokio runtimes
created externally can be used.

Mikayla suggested that the function simply take the runtime from
whatever context the initialization function is called from but I think
there could reasonably be situations where that isn't the case and this
shouldn't have a meaningful impact to code complexity. If you want to
use the current context's runtime you can just do
`gpui_tokio::init_from_handle(cx, Handle::current());`.

This doesn't have an impact on the current users of the crate - the
existing `init()` function is functionally unchanged.

Release Notes:

- N/A

Change summary

crates/gpui_tokio/src/gpui_tokio.rs | 47 +++++++++++++++++++++++-------
1 file changed, 35 insertions(+), 12 deletions(-)

Detailed changes

crates/gpui_tokio/src/gpui_tokio.rs

@@ -5,25 +5,48 @@ use util::defer;
 
 pub use tokio::task::JoinError;
 
+/// Initializes the Tokio wrapper using a new Tokio runtime with 2 worker threads.
+///
+/// If you need more threads (or access to the runtime outside of GPUI), you can create the runtime
+/// yourself and pass a Handle to `init_from_handle`.
 pub fn init(cx: &mut App) {
-    cx.set_global(GlobalTokio::new());
+    let runtime = tokio::runtime::Builder::new_multi_thread()
+        // Since we now have two executors, let's try to keep our footprint small
+        .worker_threads(2)
+        .enable_all()
+        .build()
+        .expect("Failed to initialize Tokio");
+
+    cx.set_global(GlobalTokio::new(RuntimeHolder::Owned(runtime)));
+}
+
+/// Initializes the Tokio wrapper using a Tokio runtime handle.
+pub fn init_from_handle(cx: &mut App, handle: tokio::runtime::Handle) {
+    cx.set_global(GlobalTokio::new(RuntimeHolder::Shared(handle)));
+}
+
+enum RuntimeHolder {
+    Owned(tokio::runtime::Runtime),
+    Shared(tokio::runtime::Handle),
+}
+
+impl RuntimeHolder {
+    pub fn handle(&self) -> &tokio::runtime::Handle {
+        match self {
+            RuntimeHolder::Owned(runtime) => runtime.handle(),
+            RuntimeHolder::Shared(handle) => handle,
+        }
+    }
 }
 
 struct GlobalTokio {
-    runtime: tokio::runtime::Runtime,
+    runtime: RuntimeHolder,
 }
 
 impl Global for GlobalTokio {}
 
 impl GlobalTokio {
-    fn new() -> Self {
-        let runtime = tokio::runtime::Builder::new_multi_thread()
-            // Since we now have two executors, let's try to keep our footprint small
-            .worker_threads(2)
-            .enable_all()
-            .build()
-            .expect("Failed to initialize Tokio");
-
+    fn new(runtime: RuntimeHolder) -> Self {
         Self { runtime }
     }
 }
@@ -40,7 +63,7 @@ impl Tokio {
         R: Send + 'static,
     {
         cx.read_global(|tokio: &GlobalTokio, cx| {
-            let join_handle = tokio.runtime.spawn(f);
+            let join_handle = tokio.runtime.handle().spawn(f);
             let abort_handle = join_handle.abort_handle();
             let cancel = defer(move || {
                 abort_handle.abort();
@@ -62,7 +85,7 @@ impl Tokio {
         R: Send + 'static,
     {
         cx.read_global(|tokio: &GlobalTokio, cx| {
-            let join_handle = tokio.runtime.spawn(f);
+            let join_handle = tokio.runtime.handle().spawn(f);
             let abort_handle = join_handle.abort_handle();
             let cancel = defer(move || {
                 abort_handle.abort();