Start proxy_server on a dedicated Tokio runtime

proxy_agent/config/GuestProxyAgent.linux.json

@@ -4,6 +4,7 @@
     "latchKeyFolder": "/var/lib/azure-proxy-agent/keys",
     "monitorIntervalInSeconds": 60,
     "pollKeyStatusIntervalInSeconds": 15,
+    "proxyServerRuntimeWorkerThreads": 2,
     "hostGAPluginSupport": 1,
     "ebpfProgramName": "ebpf_cgroup.o",
     "cgroupRoot": "/sys/fs/cgroup",

proxy_agent/config/GuestProxyAgent.windows.json

@@ -4,6 +4,7 @@
     "latchKeyFolder": "%SYSTEMDRIVE%\\WindowsAzure\\ProxyAgent\\Keys",
     "monitorIntervalInSeconds": 60,
     "pollKeyStatusIntervalInSeconds": 15,
+    "proxyServerRuntimeWorkerThreads": 2,
     "hostGAPluginSupport": 1,
     "ebpfProgramName": "redirect.bpf.sys",
     "fileLogLevel": "Trace",

proxy_agent/src/common/config.rs

@@ -19,6 +19,7 @@
 use crate::common::constants;
 use crate::common::logger;
 use once_cell::sync::Lazy;
+use proxy_agent_shared::current_info;
 use proxy_agent_shared::{logger::LoggerLevel, misc_helpers};
 use serde_derive::{Deserialize, Serialize};
 use std::str::FromStr;
@@ -55,6 +56,10 @@ pub fn get_max_event_file_count() -> usize {
     SYSTEM_CONFIG.get_max_event_file_count()
 }
 
+pub fn get_proxy_server_runtime_worker_threads() -> usize {
+    SYSTEM_CONFIG.get_proxy_server_runtime_worker_threads()
+}
+
 pub fn get_ebpf_file_full_path() -> Option<PathBuf> {
     SYSTEM_CONFIG.get_ebpf_file_full_path()
 }
@@ -100,6 +105,8 @@ pub struct Config {
     #[serde(skip_serializing_if = "Option::is_none")]
     maxEventFileCount: Option<usize>,
     #[serde(skip_serializing_if = "Option::is_none")]
+    proxyServerRuntimeWorkerThreads: Option<usize>,
+    #[serde(skip_serializing_if = "Option::is_none")]
     ebpfFileFullPath: Option<String>,
     ebpfProgramName: String,
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -193,6 +200,14 @@ impl Config {
             .unwrap_or(constants::DEFAULT_MAX_EVENT_FILE_COUNT)
     }
 
+    /// Gets the number of worker threads for the proxy server runtime.
+    /// The value is clamped between 1 and the number of CPU cores available.
+    pub fn get_proxy_server_runtime_worker_threads(&self) -> usize {
+        self.proxyServerRuntimeWorkerThreads
+            .unwrap_or(constants::DEFAULT_PROXY_SERVER_RUNTIME_WORKER_THREADS)
+            .clamp(1, current_info::get_cpu_count())
+    }
+
     pub fn get_ebpf_program_name(&self) -> &str {
         &self.ebpfProgramName
     }
@@ -262,7 +277,7 @@ impl Config {
 mod tests {
     use crate::common::config::Config;
     use crate::common::constants;
-    use proxy_agent_shared::misc_helpers;
+    use proxy_agent_shared::{current_info, misc_helpers};
     use std::fs::File;
     use std::io::Write;
     use std::path::PathBuf;
@@ -316,6 +331,12 @@ mod tests {
             "get_max_event_file_count mismatch"
         );
 
+        assert_eq!(
+            constants::DEFAULT_PROXY_SERVER_RUNTIME_WORKER_THREADS,
+            config.get_proxy_server_runtime_worker_threads(),
+            "get_proxy_server_runtime_worker_threads mismatch"
+        );
+
         assert_eq!(
             "ebpfProgramName",
             config.get_ebpf_program_name(),
@@ -352,6 +373,39 @@ mod tests {
         _ = fs::remove_dir_all(&temp_test_path);
     }
 
+    #[test]
+    fn proxy_server_runtime_worker_threads_are_bounded() {
+        let config_file_path = env::temp_dir().join("proxy_runtime_worker_config.json");
+        let mut config = create_config_file(config_file_path.clone());
+
+        config.proxyServerRuntimeWorkerThreads = Some(0);
+        assert_eq!(
+            1,
+            config.get_proxy_server_runtime_worker_threads(),
+            "proxy server runtime worker threads lower bound mismatch"
+        );
+
+        let cpu_count = current_info::get_cpu_count();
+        if cpu_count > 1 {
+            let value = cpu_count - 1;
+            config.proxyServerRuntimeWorkerThreads = Some(value);
+            assert_eq!(
+                value,
+                config.get_proxy_server_runtime_worker_threads(),
+                "proxy server runtime worker threads mismatch"
+            );
+        }
+
+        config.proxyServerRuntimeWorkerThreads = Some(cpu_count + 1);
+        assert_eq!(
+            cpu_count,
+            config.get_proxy_server_runtime_worker_threads(),
+            "proxy server runtime worker threads upper bound mismatch"
+        );
+
+        _ = fs::remove_file(config_file_path);
+    }
+
     fn create_config_file(file_path: PathBuf) -> Config {
         let data = if cfg!(not(windows)) {
             r#"{

proxy_agent/src/common/constants.rs

@@ -26,6 +26,7 @@ pub const NOTIFY_HEADER: &str = "x-ms-azure-notify";
 
 // Default Config Settings
 pub const DEFAULT_MAX_EVENT_FILE_COUNT: usize = 30;
+pub const DEFAULT_PROXY_SERVER_RUNTIME_WORKER_THREADS: usize = 2;
 
 pub const CGROUP_ROOT: &str = "/sys/fs/cgroup";
 

proxy_agent/src/common/helpers.rs

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: MIT
 
 use once_cell::sync::Lazy;
-use proxy_agent_shared::telemetry::span::SimpleSpan;
+use proxy_agent_shared::{current_info, telemetry::span::SimpleSpan};
 
 static START: Lazy<SimpleSpan> = Lazy::new(SimpleSpan::new);
 
@@ -21,3 +21,14 @@ pub fn write_startup_event(
     crate::common::logger::write_serial_console_log(message.clone(), None);
     message
 }
+
+/// Determine the number of worker threads for the tokio runtime in main
+/// Limit the number of worker threads to a maximum of 4 and minimum of 1  
+static TOKIO_MAIN_RUNTIME_WORKER_THREADS: Lazy<usize> = Lazy::new(|| {
+    let cpu_count = current_info::get_cpu_count();
+    cpu_count.clamp(1, 4)
+});
+
+pub fn get_tokio_main_worker_threads() -> usize {
+    *TOKIO_MAIN_RUNTIME_WORKER_THREADS
+}

proxy_agent/src/main.rs

@@ -36,8 +36,21 @@ define_windows_service!(ffi_service_main, proxy_agent_windows_service_main);
 static ASYNC_RUNTIME_HANDLE: tokio::sync::OnceCell<tokio::runtime::Handle> =
     tokio::sync::OnceCell::const_new();
 
-#[tokio::main(flavor = "multi_thread")]
-async fn main() {
+/// The main entry point of the GPA process.
+/// It initializes the tokio runtime and calls the async main function.
+/// It also determines the number of worker threads for the tokio runtime
+fn main() {
+    tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(helpers::get_tokio_main_worker_threads())
+        .enable_all()
+        .build()
+        .unwrap()
+        .block_on(async {
+            async_main().await;
+        });
+}
+
+async fn async_main() {
     // set the tokio runtime handle
     #[cfg(windows)]
     ASYNC_RUNTIME_HANDLE

proxy_agent/src/proxy/proxy_server.rs

@@ -22,6 +22,7 @@
 
 use super::proxy_authorizer::AuthorizeResult;
 use super::proxy_connection::{ConnectionLogger, HttpConnectionContext, TcpConnectionContext};
+use crate::common::config;
 use crate::common::{constants, error::Error, helpers, logger, result::Result};
 use crate::proxy::{proxy_authorizer, proxy_summary::ProxySummary, Claims};
 use crate::shared_state::access_control_wrapper::AccessControlSharedState;
@@ -92,6 +93,35 @@ impl ProxyServer {
         }
     }
 
+    /// Starts the proxy server on an isolated Tokio runtime.
+    pub fn start_on_dedicated_runtime(self) -> std::io::Result<std::thread::JoinHandle<()>> {
+        let worker_threads = config::get_proxy_server_runtime_worker_threads();
+
+        std::thread::Builder::new()
+            .name("proxy-server-runtime".to_string())
+            .spawn(move || {
+                let runtime = match tokio::runtime::Builder::new_multi_thread()
+                    .worker_threads(worker_threads)
+                    .thread_name("proxy-server-worker")
+                    .enable_all()
+                    .build()
+                {
+                    Ok(runtime) => runtime,
+                    Err(e) => {
+                        logger::write_error(format!(
+                            "Failed to create the proxy server Tokio runtime: {e}"
+                        ));
+                        return;
+                    }
+                };
+
+                logger::write_information(format!(
+                    "Started dedicated proxy server Tokio runtime with {worker_threads} worker threads."
+                ));
+                runtime.block_on(self.start());
+            })
+    }
+
     /// start listener at the given address with retry logic if the address is in use
     async fn start_listener_with_retry(
         addr: &str,
@@ -1173,6 +1203,18 @@ mod tests {
     use std::collections::HashMap;
     use std::time::Duration;
 
+    #[tokio::test]
+    async fn dedicated_runtime_stops_on_cancellation() {
+        let shared_state = shared_state::SharedState::start_all();
+        shared_state.cancel_cancellation_token();
+        let proxy_server = proxy_server::ProxyServer::new(0, &shared_state);
+
+        let runtime_thread = proxy_server.start_on_dedicated_runtime().unwrap();
+        tokio::task::spawn_blocking(move || runtime_thread.join().unwrap())
+            .await
+            .unwrap();
+    }
+
     #[tokio::test]
     async fn direct_request_test() {
         // start listener, the port must different from the one used in production code

proxy_agent/src/service.rs

@@ -57,10 +57,11 @@ pub async fn start_service(shared_state: SharedState) {
     }
 
     let start_message = format!(
-        "============== GuestProxyAgent ({}) is starting on {}({}), elapsed: {}",
+        "============== GuestProxyAgent ({}) is starting on {}({}) with {} worker threads, elapsed: {}",
         current_info::get_current_exe_version(),
         current_info::get_long_os_version(),
         current_info::get_cpu_arch(),
+        helpers::get_tokio_main_worker_threads(),
         helpers::get_elapsed_time_in_millisec()
     );
     logger::write_information(start_message.clone());
@@ -90,12 +91,12 @@ pub async fn start_service(shared_state: SharedState) {
         }
     });
 
-    tokio::spawn({
-        let proxy_server = ProxyServer::new(constants::PROXY_AGENT_PORT, &shared_state);
-        async move {
-            proxy_server.start().await;
-        }
-    });
+    let proxy_server = ProxyServer::new(constants::PROXY_AGENT_PORT, &shared_state);
+    if let Err(e) = proxy_server.start_on_dedicated_runtime() {
+        logger::write_error(format!(
+            "Failed to start the proxy server runtime thread: {e}"
+        ));
+    }
 }
 
 #[cfg(windows)]