From 2af98936a98477a5229d5c9552ba8a0dd2e141ea Mon Sep 17 00:00:00 2001
From: Vitali Lovich <vlovich@openinfer.io>
Date: Fri, 31 Jan 2025 18:10:01 -0800
Subject: [PATCH] Expose on_worker_stop method

If setting up TLS variables in on_worker_start, they need to be torn
down in on_worker_stop or you might end up accidentally accessing
destroyed TLS variables.
---
 server/src/builder.rs    | 21 +++++++++++++++++++++
 server/src/server/mod.rs |  7 +++++++
 web/src/server.rs        | 10 ++++++++++
 3 files changed, 38 insertions(+)

diff --git a/server/src/builder.rs b/server/src/builder.rs
index 911b14d65..a703dc4bb 100644
--- a/server/src/builder.rs
+++ b/server/src/builder.rs
@@ -19,6 +19,7 @@ pub struct Builder {
     pub(crate) enable_signal: bool,
     pub(crate) shutdown_timeout: Duration,
     pub(crate) on_worker_start: Box<dyn Fn() -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>,
+    pub(crate) on_worker_stop: Box<dyn Fn() -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>,
     backlog: u32,
 }
 
@@ -40,6 +41,7 @@ impl Builder {
             enable_signal: true,
             shutdown_timeout: Duration::from_secs(30),
             on_worker_start: Box::new(|| Box::pin(async {})),
+            on_worker_stop: Box::new(|| Box::pin(async {})),
             backlog: 2048,
         }
     }
@@ -131,6 +133,25 @@ impl Builder {
         self
     }
 
+    #[doc(hidden)]
+    /// Async callback called when worker thread is being shut down.
+    ///
+    /// *. This API is subject to change with no stable guarantee.
+    pub fn on_worker_stop<F, Fut>(mut self, on_stop: F) -> Self
+    where
+        F: Fn() -> Fut + Send + Sync + 'static,
+        Fut: Future + Send + 'static,
+    {
+        self.on_worker_stop = Box::new(move || {
+            let fut = on_stop();
+            Box::pin(async {
+                fut.await;
+            })
+        });
+
+        self
+    }
+
     pub fn listen<N, L, F, St>(mut self, name: N, listener: L, service: F) -> Self
     where
         N: AsRef<str>,
diff --git a/server/src/server/mod.rs b/server/src/server/mod.rs
index 09aab92e5..d7e7900b6 100644
--- a/server/src/server/mod.rs
+++ b/server/src/server/mod.rs
@@ -38,6 +38,7 @@ impl Server {
             factories,
             shutdown_timeout,
             on_worker_start,
+            on_worker_stop,
             ..
         } = builder;
 
@@ -60,6 +61,7 @@ impl Server {
         let is_graceful_shutdown = Arc::new(AtomicBool::new(false));
 
         let on_start_fut = on_worker_start();
+        let on_stop_fut = on_worker_stop();
 
         let fut = async {
             on_start_fut.await;
@@ -78,6 +80,8 @@ impl Server {
 
             worker::wait_for_stop(handles, services, shutdown_timeout, &is_graceful_shutdown).await;
 
+            on_stop_fut.await;
+
             Ok::<_, io::Error>(())
         };
 
@@ -96,6 +100,7 @@ impl Server {
             factories,
             shutdown_timeout,
             on_worker_start,
+            on_worker_stop,
             ..
         } = builder;
 
@@ -153,6 +158,8 @@ impl Server {
                             }
 
                             worker::wait_for_stop(handles, services, shutdown_timeout, &is_graceful_shutdown).await;
+
+                            on_worker_stop().await;
                         };
 
                         #[cfg(not(feature = "io-uring"))]
diff --git a/web/src/server.rs b/web/src/server.rs
index 3314fa2f0..71ddbefff 100644
--- a/web/src/server.rs
+++ b/web/src/server.rs
@@ -166,6 +166,16 @@ where
         self
     }
 
+    #[doc(hidden)]
+    pub fn on_worker_stop<FS, Fut>(mut self, on_stop: FS) -> Self
+    where
+        FS: Fn() -> Fut + Send + Sync + 'static,
+        Fut: Future + Send + 'static,
+    {
+        self.builder = self.builder.on_worker_stop(on_stop);
+        self
+    }
+
     #[cfg(not(target_family = "wasm"))]
     pub fn bind<A, ResB, BE>(mut self, addr: A) -> std::io::Result<Self>
     where