Wait for inflight requests to complete on worker shutdown (#520)

Author: gi0baro

Cargo.lock

@@ -53,7 +53,7 @@ tls-listener = { version = "=0.11", features = ["rustls-ring"] }
 tokio = { version = "1.40", features = ["full"] }
 tokio-stream = "0.1"
 tokio-tungstenite = "=0.26"
-tokio-util = { version = "0.7", features = ["codec"] }
+tokio-util = { version = "0.7", features = ["codec", "rt"] }
 
 [target.'cfg(not(any(target_env = "musl", target_os = "freebsd", target_os = "openbsd", target_os = "windows")))'.dependencies]
 tikv-jemallocator = { version = "0.6.0", default-features = false, features = ["disable_initial_exec_tls"] }

Cargo.toml

@@ -631,6 +631,7 @@ macro_rules! serve_mtr {
                 ret
             });
             let rth = rt.handler();
+            let tasks = tokio_util::task::TaskTracker::new();
             let mut srx = signal.get().rx.lock().unwrap().take().unwrap();
 
             let main_loop = crate::runtime::run_until_complete(rt, event_loop.clone(), async move {
@@ -642,7 +643,7 @@ macro_rules! serve_mtr {
                     backpressure,
                     rth,
                     callback_wrapper,
-                    tokio::spawn,
+                    |task| tasks.spawn(task),
                     hyper_util::rt::TokioExecutor::new,
                     http1_opts,
                     http2_opts,
@@ -652,6 +653,9 @@ macro_rules! serve_mtr {
 
                 log::info!("Stopping worker-{}", worker_id);
 
+                tasks.close();
+                tasks.wait().await;
+
                 Python::with_gil(|_| drop(callback_wrapper));
                 Ok(())
             });
@@ -699,6 +703,7 @@ macro_rules! serve_mtr_ssl {
                 ret
             });
             let rth = rt.handler();
+            let tasks = tokio_util::task::TaskTracker::new();
             let mut srx = signal.get().rx.lock().unwrap().take().unwrap();
 
             let main_loop = crate::runtime::run_until_complete(rt, event_loop.clone(), async move {
@@ -711,7 +716,7 @@ macro_rules! serve_mtr_ssl {
                     backpressure,
                     rth,
                     callback_wrapper,
-                    tokio::spawn,
+                    |task| tasks.spawn(task),
                     hyper_util::rt::TokioExecutor::new,
                     http1_opts,
                     http2_opts,
@@ -721,6 +726,9 @@ macro_rules! serve_mtr_ssl {
 
                 log::info!("Stopping worker-{}", worker_id);
 
+                tasks.close();
+                tasks.wait().await;
+
                 Python::with_gil(|_| drop(callback_wrapper));
                 Ok(())
             });
@@ -759,6 +767,7 @@ macro_rules! serve_str_inner {
                     crate::runtime::init_runtime_st(blocking_threads, py_threads, py_threads_idle_timeout, py_loop);
                 let rth = rt.handler();
                 let local = tokio::task::LocalSet::new();
+                let tasks = tokio_util::task::TaskTracker::new();
 
                 crate::runtime::block_on_local(&rt, local, async move {
                     crate::workers::loop_match!(
@@ -769,7 +778,7 @@ macro_rules! serve_str_inner {
                         backpressure,
                         rth,
                         callback_wrapper,
-                        tokio::task::spawn_local,
+                        |task| tasks.spawn_local(task),
                         crate::workers::WorkerExecutor::new,
                         http1_opts,
                         http2_opts,
@@ -779,6 +788,9 @@ macro_rules! serve_str_inner {
 
                     log::info!("Stopping worker-{} runtime-{}", $wid, thread_id + 1);
 
+                    tasks.close();
+                    tasks.wait().await;
+
                     Python::with_gil(|_| drop(callback_wrapper));
                 });
 
@@ -852,6 +864,7 @@ macro_rules! serve_str_ssl_inner {
                     crate::runtime::init_runtime_st(blocking_threads, py_threads, py_threads_idle_timeout, py_loop);
                 let rth = rt.handler();
                 let local = tokio::task::LocalSet::new();
+                let tasks = tokio_util::task::TaskTracker::new();
 
                 crate::runtime::block_on_local(&rt, local, async move {
                     crate::workers::loop_match_tls!(
@@ -863,7 +876,7 @@ macro_rules! serve_str_ssl_inner {
                         backpressure,
                         rth,
                         callback_wrapper,
-                        tokio::task::spawn_local,
+                        |task| tasks.spawn_local(task),
                         crate::workers::WorkerExecutor::new,
                         http1_opts,
                         http2_opts,
@@ -873,6 +886,9 @@ macro_rules! serve_str_ssl_inner {
 
                     log::info!("Stopping worker-{} runtime-{}", $wid, thread_id + 1);
 
+                    tasks.close();
+                    tasks.wait().await;
+
                     Python::with_gil(|_| drop(callback_wrapper));
                 });
 

src/workers.rs

@@ -43,6 +43,7 @@ impl WSGIWorker {
             )
         });
         let rth = rt.handler();
+        let tasks = tokio_util::task::TaskTracker::new();
 
         let (stx, mut srx) = tokio::sync::watch::channel(false);
         let main_loop = rt.inner.spawn(async move {
@@ -54,7 +55,7 @@ impl WSGIWorker {
                 backpressure,
                 rth,
                 callback_wrapper,
-                tokio::spawn,
+                |task| tasks.spawn(task),
                 hyper_util::rt::TokioExecutor::new,
                 http1_opts,
                 http2_opts,
@@ -64,6 +65,9 @@ impl WSGIWorker {
 
             log::info!("Stopping worker-{}", worker_id);
 
+            tasks.close();
+            tasks.wait().await;
+
             Python::with_gil(|_| drop(callback_wrapper));
         });
 
@@ -153,6 +157,7 @@ impl WSGIWorker {
             )
         });
         let rth = rt.handler();
+        let tasks = tokio_util::task::TaskTracker::new();
 
         let (stx, mut srx) = tokio::sync::watch::channel(false);
         rt.inner.spawn(async move {
@@ -165,7 +170,7 @@ impl WSGIWorker {
                 backpressure,
                 rth,
                 callback_wrapper,
-                tokio::spawn,
+                |task| tasks.spawn(task),
                 hyper_util::rt::TokioExecutor::new,
                 http1_opts,
                 http2_opts,
@@ -175,6 +180,9 @@ impl WSGIWorker {
 
             log::info!("Stopping worker-{}", worker_id);
 
+            tasks.close();
+            tasks.wait().await;
+
             Python::with_gil(|_| drop(callback_wrapper));
         });