From ba1b1fc5501db5d8bcdc8cacdc4d19ea5a2b4b83 Mon Sep 17 00:00:00 2001 From: Zhidong Peng Date: Wed, 2 Sep 2026 14:33:27 -0700 Subject: [PATCH 1/5] Set tokio runtime worker threads between 2 and 10 instead default to number of cpus --- proxy_agent/src/common/helpers.rs | 13 ++++++++++++- proxy_agent/src/main.rs | 17 +++++++++++++++-- 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/proxy_agent/src/common/helpers.rs b/proxy_agent/src/common/helpers.rs index 0469773c..866f49a8 100644 --- a/proxy_agent/src/common/helpers.rs +++ b/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 = 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 +/// Limit the number of worker threads to a maximum of 10 and minimum of 2 +static TOKIO_RUNTIME_WORKER_THREADS: Lazy = Lazy::new(|| { + let cpu_count = current_info::get_cpu_count(); + cpu_count.clamp(2, 10) +}); + +pub fn get_worker_threads() -> usize { + *TOKIO_RUNTIME_WORKER_THREADS +} diff --git a/proxy_agent/src/main.rs b/proxy_agent/src/main.rs index 5f090fd4..00d98ed5 100644 --- a/proxy_agent/src/main.rs +++ b/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::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_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 From 2c8e6aa7f538728847734d41cf450e5d70532771 Mon Sep 17 00:00:00 2001 From: "Zhidong Peng (HE/HIM)" Date: Thu, 3 Sep 2026 22:53:37 -0700 Subject: [PATCH 2/5] Update start_message --- proxy_agent/src/service.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/proxy_agent/src/service.rs b/proxy_agent/src/service.rs index b1a0c1b6..c3603f08 100644 --- a/proxy_agent/src/service.rs +++ b/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_worker_threads(), helpers::get_elapsed_time_in_millisec() ); logger::write_information(start_message.clone()); From 4be149ca87023932a09bde386479083b19cfd243 Mon Sep 17 00:00:00 2001 From: Zhidong Peng Date: Wed, 9 Sep 2026 16:24:26 +0000 Subject: [PATCH 3/5] update tokio runtime worker between 1 and 4. --- proxy_agent/src/common/helpers.rs | 12 ++++++------ proxy_agent/src/main.rs | 2 +- proxy_agent/src/service.rs | 2 +- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/proxy_agent/src/common/helpers.rs b/proxy_agent/src/common/helpers.rs index 866f49a8..a87fcca8 100644 --- a/proxy_agent/src/common/helpers.rs +++ b/proxy_agent/src/common/helpers.rs @@ -22,13 +22,13 @@ pub fn write_startup_event( message } -/// Determine the number of worker threads for the tokio runtime -/// Limit the number of worker threads to a maximum of 10 and minimum of 2 -static TOKIO_RUNTIME_WORKER_THREADS: Lazy = Lazy::new(|| { +/// 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 = Lazy::new(|| { let cpu_count = current_info::get_cpu_count(); - cpu_count.clamp(2, 10) + cpu_count.clamp(1, 4) }); -pub fn get_worker_threads() -> usize { - *TOKIO_RUNTIME_WORKER_THREADS +pub fn get_tokio_main_worker_threads() -> usize { + *TOKIO_MAIN_RUNTIME_WORKER_THREADS } diff --git a/proxy_agent/src/main.rs b/proxy_agent/src/main.rs index 00d98ed5..f586880a 100644 --- a/proxy_agent/src/main.rs +++ b/proxy_agent/src/main.rs @@ -41,7 +41,7 @@ static ASYNC_RUNTIME_HANDLE: tokio::sync::OnceCell = /// It also determines the number of worker threads for the tokio runtime fn main() { tokio::runtime::Builder::new_multi_thread() - .worker_threads(helpers::get_worker_threads()) + .worker_threads(helpers::get_tokio_main_worker_threads()) .enable_all() .build() .unwrap() diff --git a/proxy_agent/src/service.rs b/proxy_agent/src/service.rs index c3603f08..2689b17b 100644 --- a/proxy_agent/src/service.rs +++ b/proxy_agent/src/service.rs @@ -61,7 +61,7 @@ pub async fn start_service(shared_state: SharedState) { current_info::get_current_exe_version(), current_info::get_long_os_version(), current_info::get_cpu_arch(), - helpers::get_worker_threads(), + helpers::get_tokio_main_worker_threads(), helpers::get_elapsed_time_in_millisec() ); logger::write_information(start_message.clone()); From f6f1e24df19f6a172186dbf38c1c7c235d10645e Mon Sep 17 00:00:00 2001 From: "Zhidong Peng (HE/HIM)" Date: Tue, 8 Sep 2026 13:06:43 -0700 Subject: [PATCH 4/5] start proxy_server on delicated tokio runtime. --- proxy_agent/config/GuestProxyAgent.linux.json | 1 + .../config/GuestProxyAgent.windows.json | 1 + proxy_agent/src/common/config.rs | 58 ++++++++++++++++++- proxy_agent/src/common/constants.rs | 1 + proxy_agent/src/proxy/proxy_server.rs | 44 +++++++++++++- proxy_agent/src/service.rs | 12 ++-- 6 files changed, 109 insertions(+), 8 deletions(-) diff --git a/proxy_agent/config/GuestProxyAgent.linux.json b/proxy_agent/config/GuestProxyAgent.linux.json index 90b6d0b5..c580e3c4 100644 --- a/proxy_agent/config/GuestProxyAgent.linux.json +++ b/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", diff --git a/proxy_agent/config/GuestProxyAgent.windows.json b/proxy_agent/config/GuestProxyAgent.windows.json index b43e5768..17e80205 100644 --- a/proxy_agent/config/GuestProxyAgent.windows.json +++ b/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", diff --git a/proxy_agent/src/common/config.rs b/proxy_agent/src/common/config.rs index 4365f148..854447cd 100644 --- a/proxy_agent/src/common/config.rs +++ b/proxy_agent/src/common/config.rs @@ -17,8 +17,10 @@ //! ``` use crate::common::constants; +use crate::common::helpers; 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 +57,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 { SYSTEM_CONFIG.get_ebpf_file_full_path() } @@ -100,6 +106,8 @@ pub struct Config { #[serde(skip_serializing_if = "Option::is_none")] maxEventFileCount: Option, #[serde(skip_serializing_if = "Option::is_none")] + proxyServerRuntimeWorkerThreads: Option, + #[serde(skip_serializing_if = "Option::is_none")] ebpfFileFullPath: Option, ebpfProgramName: String, #[serde(skip_serializing_if = "Option::is_none")] @@ -193,6 +201,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,11 +278,12 @@ 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; use std::{env, fs}; + use sysinfo::Cpu; #[test] fn config_struct_test() { @@ -316,6 +333,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 +375,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#"{ diff --git a/proxy_agent/src/common/constants.rs b/proxy_agent/src/common/constants.rs index 6091b60d..dee9466a 100644 --- a/proxy_agent/src/common/constants.rs +++ b/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"; diff --git a/proxy_agent/src/proxy/proxy_server.rs b/proxy_agent/src/proxy/proxy_server.rs index a40ba44f..9f977c66 100644 --- a/proxy_agent/src/proxy/proxy_server.rs +++ b/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; @@ -44,11 +45,11 @@ use hyper::{Request, Response}; use hyper_util::rt::TokioIo; use proxy_agent_shared::common_state::CommonState; use proxy_agent_shared::error::HyperErrorType; -use proxy_agent_shared::hyper_client; use proxy_agent_shared::logger::LoggerLevel; use proxy_agent_shared::misc_helpers; use proxy_agent_shared::proxy_agent_aggregate_status::ModuleState; use proxy_agent_shared::telemetry::event_logger; +use proxy_agent_shared::{current_info, hyper_client}; use std::time::Duration; use tokio::net::TcpListener; use tokio::net::TcpStream; @@ -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> { + 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, @@ -1205,6 +1235,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 diff --git a/proxy_agent/src/service.rs b/proxy_agent/src/service.rs index 2689b17b..b980b590 100644 --- a/proxy_agent/src/service.rs +++ b/proxy_agent/src/service.rs @@ -91,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)] From b561ab98c93706bfec54a8dccab6afd7fbb6c8cd Mon Sep 17 00:00:00 2001 From: Zhidong Peng Date: Wed, 9 Sep 2026 16:48:03 +0000 Subject: [PATCH 5/5] fix build --- proxy_agent/src/common/config.rs | 2 -- proxy_agent/src/proxy/proxy_server.rs | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/proxy_agent/src/common/config.rs b/proxy_agent/src/common/config.rs index 854447cd..2b1a035a 100644 --- a/proxy_agent/src/common/config.rs +++ b/proxy_agent/src/common/config.rs @@ -17,7 +17,6 @@ //! ``` use crate::common::constants; -use crate::common::helpers; use crate::common::logger; use once_cell::sync::Lazy; use proxy_agent_shared::current_info; @@ -283,7 +282,6 @@ mod tests { use std::io::Write; use std::path::PathBuf; use std::{env, fs}; - use sysinfo::Cpu; #[test] fn config_struct_test() { diff --git a/proxy_agent/src/proxy/proxy_server.rs b/proxy_agent/src/proxy/proxy_server.rs index 9f977c66..5303eb2b 100644 --- a/proxy_agent/src/proxy/proxy_server.rs +++ b/proxy_agent/src/proxy/proxy_server.rs @@ -45,11 +45,11 @@ use hyper::{Request, Response}; use hyper_util::rt::TokioIo; use proxy_agent_shared::common_state::CommonState; use proxy_agent_shared::error::HyperErrorType; +use proxy_agent_shared::hyper_client; use proxy_agent_shared::logger::LoggerLevel; use proxy_agent_shared::misc_helpers; use proxy_agent_shared::proxy_agent_aggregate_status::ModuleState; use proxy_agent_shared::telemetry::event_logger; -use proxy_agent_shared::{current_info, hyper_client}; use std::time::Duration; use tokio::net::TcpListener; use tokio::net::TcpStream;