fix: gateway shutdown cancellation

ThisSeanZhang · ThisSeanZhang · commit 2588eedfa133 · 2026-03-10T03:57:37.000+08:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/landscape-gateway/Cargo.toml b/landscape-gateway/Cargo.toml
@@ -8,6 +8,7 @@ landscape-common = { path = "../landscape-common" }
 landscape-database = { path = "../landscape-database" }
 pingora = { workspace = true }
 tokio = { workspace = true }
+tokio-util = { workspace = true }
 tokio-rustls = { workspace = true }
 rustls = { workspace = true }
 tracing = { workspace = true }
diff --git a/landscape-gateway/src/lib.rs b/landscape-gateway/src/lib.rs
@@ -28,6 +28,7 @@ use tokio::runtime::Builder as RuntimeBuilder;
 use tokio::sync::watch;
 use tokio::task::JoinSet;
 use tokio_rustls::{server::TlsStream as TokioTlsStream, TlsAcceptor};
+use tokio_util::sync::CancellationToken;
 
 use crate::sni_proxy::{parse_sni_from_client_hello, proxy_tls_passthrough, SniProxyRouter};
 
@@ -41,11 +42,16 @@ pub struct GatewayTlsConfig {
 pub struct GatewayManager {
     rules: SharedRules,
     status: WatchService,
-    state: Mutex<Option<JoinHandle<()>>>,
+    state: Mutex<Option<GatewayRuntimeState>>,
     config: GatewayRuntimeConfig,
     tls_config: Option<GatewayTlsConfig>,
 }
 
+struct GatewayRuntimeState {
+    thread: JoinHandle<()>,
+    cancel: CancellationToken,
+}
+
 impl GatewayManager {
     pub fn new(
         initial_rules: Vec<HttpUpstreamRuleConfig>,
@@ -84,12 +90,22 @@ impl GatewayManager {
         let https_port = self.config.https_port;
         let tls_config = self.tls_config.clone();
         let status = self.status.clone();
-
-        let handle = std::thread::spawn(move || {
-            run_pingora_server(rules, http_port, https_port, tls_config, status);
+        let cancel = CancellationToken::new();
+        let thread_cancel = cancel.clone();
+
+        let thread = std::thread::spawn(move || {
+            run_pingora_server(
+                rules,
+                http_port,
+                https_port,
+                tls_config,
+                status.clone(),
+                thread_cancel,
+            );
+            status.just_change_status(ServiceStatus::Stop);
         });
 
-        *state = Some(handle);
+        *state = Some(GatewayRuntimeState { thread, cancel });
         self.status.just_change_status(ServiceStatus::Running);
         if self.tls_config.is_some() {
             tracing::info!(
@@ -113,17 +129,19 @@ impl GatewayManager {
         }
         tracing::info!("Signalling gateway to stop...");
         self.status.just_change_status(ServiceStatus::Stopping);
+        if let Some(state) = self.state.lock().unwrap().as_ref() {
+            state.cancel.cancel();
+        }
     }
 
     /// Block until the Pingora thread has exited. Call after shutdown().
     pub fn join(&self) {
         let mut state = self.state.lock().unwrap();
-        if let Some(handle) = state.take() {
+        if let Some(runtime_state) = state.take() {
             tracing::info!("Waiting for gateway thread to finish...");
-            if let Err(e) = handle.join() {
+            if let Err(e) = runtime_state.thread.join() {
                 tracing::error!("Gateway thread panicked: {:?}", e);
             }
-            self.status.just_change_status(ServiceStatus::Stop);
             tracing::info!("Gateway stopped");
         }
     }
@@ -152,7 +170,17 @@ impl GatewayManager {
 impl Drop for GatewayManager {
     fn drop(&mut self) {
         self.shutdown();
-        self.join();
+        if self.status.is_stop() {
+            self.join();
+            return;
+        }
+
+        if let Some(runtime_state) = self.state.lock().unwrap().take() {
+            runtime_state.cancel.cancel();
+            tracing::warn!(
+                "Dropping gateway manager before gateway thread fully stopped; detaching thread"
+            );
+        }
     }
 }
 
@@ -162,6 +190,7 @@ fn run_pingora_server(
     https_port: u16,
     tls_config: Option<GatewayTlsConfig>,
     status: WatchService,
+    cancel: CancellationToken,
 ) {
     use pingora::server::Server;
     use proxy_service::LandscapeReverseProxy;
@@ -178,8 +207,9 @@ fn run_pingora_server(
     let https_handle = tls_config.map(|tls_config| {
         let rules = rules.clone();
         let status = status.clone();
+        let cancel = cancel.child_token();
         std::thread::spawn(move || {
-            run_https_server(rules, https_port, tls_config, server_conf, status);
+            run_https_server(rules, https_port, tls_config, server_conf, status, cancel);
         })
     });
 
@@ -201,6 +231,7 @@ fn run_https_server(
     tls_config: GatewayTlsConfig,
     server_conf: Arc<pingora::server::configuration::ServerConf>,
     status: WatchService,
+    cancel: CancellationToken,
 ) {
     let runtime = RuntimeBuilder::new_multi_thread()
         .enable_all()
@@ -210,7 +241,7 @@ fn run_https_server(
 
     runtime.block_on(async move {
         if let Err(e) =
-            run_https_server_inner(rules, https_port, tls_config, server_conf, status).await
+            run_https_server_inner(rules, https_port, tls_config, server_conf, status, cancel).await
         {
             tracing::error!("Gateway HTTPS listener exited with error: {e}");
         }
@@ -223,6 +254,7 @@ async fn run_https_server_inner(
     tls_config: GatewayTlsConfig,
     server_conf: Arc<pingora::server::configuration::ServerConf>,
     status: WatchService,
+    cancel: CancellationToken,
 ) -> std::io::Result<()> {
     use proxy_service::LandscapeReverseProxy;
 
@@ -245,6 +277,10 @@ async fn run_https_server_inner(
                     break;
                 }
             }
+            _ = cancel.cancelled() => {
+                let _ = shutdown_tx.send(true);
+                break;
+            }
             accept_result = listener.accept() => {
                 let (stream, peer_addr) = match accept_result {
                     Ok(pair) => pair,
@@ -258,11 +294,15 @@ async fn run_https_server_inner(
                 let app = app.clone();
                 let sni_proxy_router = sni_proxy_router.clone();
                 let connection_shutdown = shutdown_rx.clone();
+                let connection_cancel = cancel.child_token();
 
                 tasks.spawn(async move {
                     if sni_proxy_router.has_sni_proxy_rules() {
                         let mut peek_buf = vec![0u8; 4096];
-                        match stream.peek(&mut peek_buf).await {
+                        match tokio::select! {
+                            _ = connection_cancel.cancelled() => return,
+                            result = stream.peek(&mut peek_buf) => result,
+                        } {
                             Ok(size) if size > 0 => {
                                 if let Some(sni) = parse_sni_from_client_hello(&peek_buf[..size]) {
                                     if let Some(target) = sni_proxy_router.match_target(&sni) {
@@ -273,7 +313,7 @@ async fn run_https_server_inner(
                                             target.target.address,
                                             target.target.port
                                         );
-                                        if let Err(e) = proxy_tls_passthrough(stream, &target).await {
+                                        if let Err(e) = proxy_tls_passthrough(stream, &target, connection_cancel.clone()).await {
                                             tracing::warn!(
                                                 "Gateway TLS passthrough failed for '{}' via rule '{}': {}",
                                                 target.sni,
@@ -293,7 +333,10 @@ async fn run_https_server_inner(
                         }
                     }
 
-                    let tls_result = tokio::time::timeout(Duration::from_secs(60), acceptor.accept(stream)).await;
+                    let tls_result = tokio::select! {
+                        _ = connection_cancel.cancelled() => return,
+                        result = tokio::time::timeout(Duration::from_secs(60), acceptor.accept(stream)) => result,
+                    };
                     let tls_stream = match tls_result {
                         Ok(Ok(stream)) => stream,
                         Ok(Err(e)) => {
@@ -307,7 +350,10 @@ async fn run_https_server_inner(
                     };
 
                     let stream: Stream = Box::new(GatewayTlsStream::new(tls_stream));
-                    let _ = app.process_new(stream, &connection_shutdown).await;
+                    tokio::select! {
+                        _ = connection_cancel.cancelled() => {}
+                        _ = app.process_new(stream, &connection_shutdown) => {}
+                    }
                 });
             }
         }
diff --git a/landscape-gateway/src/service.rs b/landscape-gateway/src/service.rs
@@ -61,13 +61,21 @@ impl GatewayService {
     /// to exit cleanly, without blocking the async runtime.
     pub async fn shutdown_and_wait(&self, timeout: std::time::Duration) {
         self.manager.shutdown();
-        let manager = self.manager.clone();
-        let join_task = tokio::task::spawn_blocking(move || {
-            manager.join();
-        });
-        match tokio::time::timeout(timeout, join_task).await {
-            Ok(Ok(())) => tracing::info!("Gateway thread exited cleanly."),
-            Ok(Err(e)) => tracing::error!("Gateway join task panicked: {:?}", e),
+        let mut status_rx = self.manager.watch_service().subscribe();
+        if matches!(*status_rx.borrow(), ServiceStatus::Stop) {
+            tracing::info!("Gateway thread exited cleanly.");
+            return;
+        }
+
+        let wait_result = tokio::time::timeout(
+            timeout,
+            status_rx.wait_for(|status| matches!(status, ServiceStatus::Stop)),
+        )
+        .await;
+
+        match wait_result {
+            Ok(Ok(_)) => tracing::info!("Gateway thread exited cleanly."),
+            Ok(Err(e)) => tracing::error!("Gateway status watch closed during shutdown: {:?}", e),
             Err(_) => tracing::warn!(
                 "Gateway did not stop within {}s timeout, proceeding.",
                 timeout.as_secs()
diff --git a/landscape-gateway/src/sni_proxy.rs b/landscape-gateway/src/sni_proxy.rs
@@ -3,8 +3,9 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use landscape_common::config::gateway::{
     HttpUpstreamMatchRule, HttpUpstreamRuleConfig, HttpUpstreamTarget, LoadBalanceMethod,
 };
-use tokio::io;
+use tokio::io::{self, AsyncWriteExt};
 use tokio::net::TcpStream;
+use tokio_util::sync::CancellationToken;
 
 use crate::SharedRules;
 
@@ -84,14 +85,53 @@ impl MatchedSniTarget {
 }
 
 pub async fn proxy_tls_passthrough(
-    mut downstream: TcpStream,
+    downstream: TcpStream,
     target: &MatchedSniTarget,
+    cancel: CancellationToken,
 ) -> io::Result<()> {
     let upstream_addr = format!("{}:{}", target.target.address, target.target.port);
-    let mut upstream = TcpStream::connect(&upstream_addr).await?;
+    let upstream = tokio::select! {
+        _ = cancel.cancelled() => return Ok(()),
+        result = TcpStream::connect(&upstream_addr) => result?,
+    };
+
+    let (mut downstream_read, mut downstream_write) = downstream.into_split();
+    let (mut upstream_read, mut upstream_write) = upstream.into_split();
+
+    let client_to_upstream = tokio::spawn(async move {
+        let result = io::copy(&mut downstream_read, &mut upstream_write).await;
+        let _ = upstream_write.shutdown().await;
+        result
+    });
+
+    let upstream_to_client = tokio::spawn(async move {
+        let result = io::copy(&mut upstream_read, &mut downstream_write).await;
+        let _ = downstream_write.shutdown().await;
+        result
+    });
+
+    tokio::pin!(client_to_upstream);
+    tokio::pin!(upstream_to_client);
+
+    let result = tokio::select! {
+        _ = cancel.cancelled() => Ok(()),
+        result = &mut client_to_upstream => join_copy_task(result),
+        result = &mut upstream_to_client => join_copy_task(result),
+    };
+
+    client_to_upstream.abort();
+    upstream_to_client.abort();
+
+    result
+}
 
-    let _ = io::copy_bidirectional(&mut downstream, &mut upstream).await?;
-    Ok(())
+fn join_copy_task(result: Result<io::Result<u64>, tokio::task::JoinError>) -> io::Result<()> {
+    match result {
+        Ok(Ok(_)) => Ok(()),
+        Ok(Err(e)) => Err(e),
+        Err(e) if e.is_cancelled() => Ok(()),
+        Err(e) => Err(io::Error::other(e)),
+    }
 }
 
 /// Parse the SNI (Server Name Indication) extension from a TLS ClientHello message.
