Refactor the proxy_sever code could stop its listener cleanly.

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();

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)
@@ -1213,7 +1219,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())
@@ -1229,7 +1236,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();
@@ -1367,7 +1375,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();

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}"