Merge pull request #207 from Dstack-TEE/gw-rt

gw: Seperate proxy runtime from Rocket

Author: kvinwang

Cargo.lock

@@ -87,6 +87,7 @@ tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
 safe-write = "0.1.2"
 nix = "0.29.0"
 sd-notify = "0.4.5"
+jemallocator = "0.5.4"
 
 # Serialization/Parsing
 bon = "3.4.0"

Cargo.toml

@@ -45,6 +45,7 @@ serde-duration.workspace = true
 reqwest = { workspace = true, features = ["json"] }
 hyper = { workspace = true, features = ["server", "http1"] }
 hyper-util = { version = "0.1", features = ["tokio"] }
+jemallocator.workspace = true
 
 [target.'cfg(unix)'.dependencies]
 nix = { workspace = true, features = ["resource"] }

gateway/Cargo.toml

@@ -60,6 +60,7 @@ buffer_size = 8192
 connect_top_n = 3
 localhost_enabled = false
 app_address_ns_prefix = "_dstack-app-address"
+workers = 32
 
 [core.proxy.timeouts]
 # Timeout for establishing a connection to the target app.

gateway/gateway.toml

@@ -77,6 +77,7 @@ pub struct ProxyConfig {
     pub buffer_size: usize,
     pub connect_top_n: usize,
     pub localhost_enabled: bool,
+    pub workers: usize,
     pub app_address_ns_prefix: String,
 }
 

gateway/src/config.rs

@@ -20,6 +20,9 @@ mod models;
 mod proxy;
 mod web_routes;
 
+#[global_allocator]
+static ALLOCATOR: jemallocator::Jemalloc = jemallocator::Jemalloc;
+
 fn app_version() -> String {
     const CARGO_PKG_VERSION: &str = env!("CARGO_PKG_VERSION");
     const VERSION: &str = git_version::git_version!(

gateway/src/main.rs

@@ -145,7 +145,15 @@ async fn handle_connection(
     }
 }
 
-pub async fn run(config: &ProxyConfig, proxy: Proxy) -> Result<()> {
+#[inline(never)]
+pub async fn proxy_main(config: &ProxyConfig, proxy: Proxy) -> Result<()> {
+    let workers_rt = tokio::runtime::Builder::new_multi_thread()
+        .thread_name("proxy-worker")
+        .enable_all()
+        .worker_threads(config.workers)
+        .build()
+        .expect("Failed to build Tokio runtime");
+
     let dotted_base_domain = {
         let base_domain = config.base_domain.as_str();
         let base_domain = base_domain.strip_prefix(".").unwrap_or(base_domain);
@@ -174,7 +182,7 @@ pub async fn run(config: &ProxyConfig, proxy: Proxy) -> Result<()> {
                 info!(%from, "new connection");
                 let proxy = proxy.clone();
                 let dotted_base_domain = dotted_base_domain.clone();
-                tokio::spawn(
+                workers_rt.spawn(
                     async move {
                         let _conn_entered = conn_entered;
                         let timeouts = &proxy.config.proxy.timeouts;
@@ -211,14 +219,24 @@ fn next_connection_id() -> usize {
 }
 
 pub fn start(config: ProxyConfig, app_state: Proxy) {
-    tokio::spawn(async move {
-        if let Err(err) = run(&config, app_state).await {
-            error!(
-                "error on {}:{}: {err:?}",
-                config.listen_addr, config.listen_port
-            );
-        }
-    });
+    std::thread::Builder::new()
+        .name("proxy-main".to_string())
+        .spawn(move || {
+            // Create a new single-threaded runtime
+            let rt = tokio::runtime::Builder::new_current_thread()
+                .enable_all()
+                .build()
+                .expect("Failed to build Tokio runtime");
+
+            // Run the proxy_main function in this runtime
+            if let Err(err) = rt.block_on(proxy_main(&config, app_state)) {
+                error!(
+                    "error on {}:{}: {err:?}",
+                    config.listen_addr, config.listen_port
+                );
+            }
+        })
+        .expect("Failed to spawn proxy-main thread");
 }
 
 #[cfg(test)]