From fa20870c4e44d1c65b3b78150d92953a975ba560 Mon Sep 17 00:00:00 2001 From: "Zhidong Peng (HE/HIM)" Date: Fri, 2 Oct 2026 15:46:57 -0700 Subject: [PATCH 1/2] Refactor the proxy_sever code could stop its listener cleanly. --- proxy_agent/src/proxy/proxy_server.rs | 25 +++++++++++++++++-------- proxy_agent/src/service.rs | 7 ++++++- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/proxy_agent/src/proxy/proxy_server.rs b/proxy_agent/src/proxy/proxy_server.rs index c3ec8426..faa07bf2 100644 --- a/proxy_agent/src/proxy/proxy_server.rs +++ b/proxy_agent/src/proxy/proxy_server.rs @@ -66,7 +66,8 @@ const START_LISTENER_RETRY_SLEEP_DURATION: Duration = Duration::from_secs(1); #[derive(Clone)] pub struct ProxyServer { port: u16, - cancellation_token: CancellationToken, + service_cancellation_token: CancellationToken, + listener_cancellation_token: CancellationToken, key_keeper_shared_state: KeyKeeperSharedState, common_state: CommonState, provision_shared_state: ProvisionSharedState, @@ -78,10 +79,15 @@ pub struct ProxyServer { } impl ProxyServer { - pub fn new(port: u16, shared_state: &SharedState) -> Self { + pub fn new( + port: u16, + shared_state: &SharedState, + listener_cancellation_token: CancellationToken, + ) -> Self { ProxyServer { port, - cancellation_token: shared_state.get_cancellation_token(), + service_cancellation_token: shared_state.get_cancellation_token(), + listener_cancellation_token, key_keeper_shared_state: shared_state.get_key_keeper_shared_state(), common_state: shared_state.get_common_state(), provision_shared_state: shared_state.get_provision_shared_state(), @@ -221,7 +227,7 @@ impl ProxyServer { logger::write_warning(format!("Failed to set module state: {e}")); } provision::listener_started(EventThreadsSharedState { - cancellation_token: self.cancellation_token.clone(), + cancellation_token: self.service_cancellation_token.clone(), common_state: self.common_state.clone(), access_control_shared_state: self.access_control_shared_state.clone(), redirector_shared_state: self.redirector_shared_state.clone(), @@ -235,7 +241,7 @@ impl ProxyServer { // We start a loop to continuously accept incoming connections loop { tokio::select! { - _ = self.cancellation_token.cancelled() => { + _ = self.listener_cancellation_token.cancelled() => { logger::write_warning("cancellation token signal received, stop the listener.".to_string()); let _= self.agent_status_shared_state .set_module_state(ModuleState::STOPPED, AgentStatusModule::ProxyServer) @@ -1207,7 +1213,8 @@ mod tests { 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 proxy_server = + proxy_server::ProxyServer::new(0, &shared_state, shared_state.get_cancellation_token()); let runtime_thread = proxy_server.start_on_dedicated_runtime().unwrap(); tokio::task::spawn_blocking(move || runtime_thread.join().unwrap()) @@ -1223,7 +1230,8 @@ mod tests { let shared_state = shared_state::SharedState::start_all(); let key_keeper_shared_state = shared_state.get_key_keeper_shared_state(); let cancellation_token = shared_state.get_cancellation_token(); - let proxy_server = proxy_server::ProxyServer::new(port, &shared_state); + let proxy_server = + proxy_server::ProxyServer::new(port, &shared_state, cancellation_token.clone()); tokio::spawn({ let proxy_server = proxy_server.clone(); @@ -1361,7 +1369,8 @@ mod tests { let port: u16 = 8092; // distinct from other tests let shared_state = shared_state::SharedState::start_all(); let cancellation_token = shared_state.get_cancellation_token(); - let proxy_server = proxy_server::ProxyServer::new(port, &shared_state); + let proxy_server = + proxy_server::ProxyServer::new(port, &shared_state, cancellation_token.clone()); tokio::spawn({ let proxy_server = proxy_server.clone(); diff --git a/proxy_agent/src/service.rs b/proxy_agent/src/service.rs index b980b590..75cfd31e 100644 --- a/proxy_agent/src/service.rs +++ b/proxy_agent/src/service.rs @@ -91,7 +91,12 @@ pub async fn start_service(shared_state: SharedState) { } }); - let proxy_server = ProxyServer::new(constants::PROXY_AGENT_PORT, &shared_state); + let proxy_cancellation_token = shared_state.get_cancellation_token().child_token(); + let proxy_server = ProxyServer::new( + constants::PROXY_AGENT_PORT, + &shared_state, + proxy_cancellation_token, + ); if let Err(e) = proxy_server.start_on_dedicated_runtime() { logger::write_error(format!( "Failed to start the proxy server runtime thread: {e}" From f7f94004456c0932f4e7318ae0901ddb5332fc4d Mon Sep 17 00:00:00 2001 From: "Zhidong Peng (HE/HIM)" Date: Mon, 5 Oct 2026 10:31:41 -0700 Subject: [PATCH 2/2] fix test build failure. --- proxy_agent/src/provision.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/proxy_agent/src/provision.rs b/proxy_agent/src/provision.rs index 016a4ac4..7121ea9e 100644 --- a/proxy_agent/src/provision.rs +++ b/proxy_agent/src/provision.rs @@ -746,7 +746,8 @@ mod tests { // initialize key keeper secure channel state to UNKNOWN let port: u16 = 8092; - let proxy_server = proxy_server::ProxyServer::new(port, &shared_state); + let proxy_server = + proxy_server::ProxyServer::new(port, &shared_state, cancellation_token.clone()); tokio::spawn({ let proxy_server = proxy_server.clone();