diff --git a/proxy/src/serverless/mod.rs b/proxy/src/serverless/mod.rs index 364490e70d..a080d8c4d6 100644 --- a/proxy/src/serverless/mod.rs +++ b/proxy/src/serverless/mod.rs @@ -41,7 +41,7 @@ use tokio::net::{TcpListener, TcpStream}; use tokio::time::timeout; use tokio_rustls::TlsAcceptor; use tokio_util::sync::CancellationToken; -use tokio_util::task::TaskTracker; +use tokio_util::task::task_tracker::TaskTrackerToken; use tracing::{Instrument, info, warn}; use crate::cancellation::CancellationHandler; @@ -149,10 +149,11 @@ pub async fn task_main( let conn_token = cancellation_token.child_token(); let tls_acceptor = tls_acceptor.clone(); let backend = backend.clone(); - let connections2 = connections.clone(); let cancellation_handler = cancellation_handler.clone(); let endpoint_rate_limiter = endpoint_rate_limiter.clone(); - connections.spawn( + + let tracker = connections.token(); + tokio::spawn( async move { let conn_token2 = conn_token.clone(); let _cancel_guard = config.http_config.cancel_set.insert(conn_id, conn_token2); @@ -179,7 +180,7 @@ pub async fn task_main( Box::pin(connection_handler( config, backend, - connections2, + tracker, cancellation_handler, endpoint_rate_limiter, conn_token, @@ -302,7 +303,7 @@ async fn connection_startup( async fn connection_handler( config: &'static ProxyConfig, backend: Arc, - connections: TaskTracker, + tracker: TaskTrackerToken, cancellation_handler: Arc, endpoint_rate_limiter: Arc, cancellation_token: CancellationToken, @@ -343,12 +344,12 @@ async fn connection_handler( // `request_handler` is not cancel safe. It expects to be cancelled only at specific times. // By spawning the future, we ensure it never gets cancelled until it decides to. - let handler = connections.spawn( + let handler = tokio::spawn( request_handler( req, config, backend.clone(), - connections.clone(), + tracker.clone(), cancellation_handler.clone(), session_id, conn_info2.clone(), @@ -394,7 +395,7 @@ async fn request_handler( mut request: hyper::Request, config: &'static ProxyConfig, backend: Arc, - ws_connections: TaskTracker, + tracker: TaskTrackerToken, cancellation_handler: Arc, session_id: uuid::Uuid, conn_info: ConnectionInfo, @@ -434,7 +435,6 @@ async fn request_handler( let (response, websocket) = framed_websockets::upgrade::upgrade(&mut request) .map_err(|e| ApiError::BadRequest(e.into()))?; - let tracker = ws_connections.token(); tokio::spawn( async move { if let Err(e) = websocket::serve_websocket(