From 3b9fc883a392689b233ab5348fcee38742548c94 Mon Sep 17 00:00:00 2001
From: Mike Aizatsky <maizatskyi@cloudflare.com>
Date: Tue, 5 Nov 2024 10:43:26 -0800
Subject: [PATCH] [rust] expose init_tokio (#3061)

EW has more complicated setup process and needs to initialize tokio
within sbox.
---
 src/rust/cxx-integration/lib.rs   |  8 +++++++-
 src/rust/cxx-integration/tokio.rs | 12 +++++++++---
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/src/rust/cxx-integration/lib.rs b/src/rust/cxx-integration/lib.rs
index 0ec149af0ee..cc58882253d 100644
--- a/src/rust/cxx-integration/lib.rs
+++ b/src/rust/cxx-integration/lib.rs
@@ -14,7 +14,13 @@ mod ffi {
 }
 
 pub fn init() {
-    tokio::init();
+    init_tokio(None);
+}
+
+/// Initialize tokio runtime.
+/// Should not be called directly but as a part of a downstream cxx-integration init.
+pub fn init_tokio(worker_threads: Option<usize>) {
+    tokio::init(worker_threads);
 }
 
 fn trigger_panic(msg: &str) {
diff --git a/src/rust/cxx-integration/tokio.rs b/src/rust/cxx-integration/tokio.rs
index 1e684116f1b..09355cd63b3 100644
--- a/src/rust/cxx-integration/tokio.rs
+++ b/src/rust/cxx-integration/tokio.rs
@@ -8,10 +8,16 @@ static TOKIO_RUNTIME: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
 
 /// Initialize tokio runtime.
 /// Must be called after all forking and sandbox setup is finished.
-pub(crate) fn init() {
+pub(crate) fn init(worker_threads: Option<usize>) {
     assert!(TOKIO_RUNTIME.get().is_none());
 
-    let runtime = tokio::runtime::Builder::new_multi_thread()
+    let mut builder = tokio::runtime::Builder::new_multi_thread();
+
+    if let Some(worker_threads) = worker_threads {
+        builder.worker_threads(worker_threads);
+    }
+
+    let runtime = builder
         .enable_time()
         .enable_io()
         .build()
@@ -58,7 +64,7 @@ mod test {
 
     #[test]
     fn test_tokio_init() {
-        init();
+        init(None);
         let join = spawn(async move {
             tokio::time::sleep(Duration::from_millis(1)).await;
             42