From 3f1c660927acacf5eeb69d28ece4791955c53d5c Mon Sep 17 00:00:00 2001
From: Yanhao Mo <yanhaocs@gmail.com>
Date: Mon, 18 Sep 2023 22:44:30 +0800
Subject: [PATCH] support shutdown hooks

---
 volo-thrift/src/server.rs | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/volo-thrift/src/server.rs b/volo-thrift/src/server.rs
index 865f3e12..a4aeb8e0 100644
--- a/volo-thrift/src/server.rs
+++ b/volo-thrift/src/server.rs
@@ -1,9 +1,11 @@
 use std::{
+    collections::VecDeque,
     marker::PhantomData,
     sync::{atomic::Ordering, Arc},
     time::Duration,
 };
 
+use futures::future::BoxFuture;
 use motore::{
     layer::{Identity, Layer, Stack},
     service::Service,
@@ -42,6 +44,9 @@ pub struct Server<S, L, Req, MkC, SP> {
     #[cfg(feature = "multiplex")]
     multiplex: bool,
     span_provider: SP,
+
+    shutdown_hooks: VecDeque<Box<dyn FnOnce() -> BoxFuture<'static, ()>>>,
+
     _marker: PhantomData<Req>,
 }
 
@@ -66,12 +71,25 @@ impl<S, Req>
             #[cfg(feature = "multiplex")]
             multiplex: false,
             span_provider: DefaultProvider {},
+            shutdown_hooks: VecDeque::new(),
             _marker: PhantomData,
         }
     }
 }
 
 impl<S, L, Req, MkC, SP> Server<S, L, Req, MkC, SP> {
+    /// Register shutdown hook.
+    ///
+    /// Hook functions will be called just before volo's own gracefull existing code starts,
+    /// in reverse order of registration.
+    pub fn register_shutdown_hook(
+        mut self,
+        hook: impl FnOnce() -> BoxFuture<'static, ()> + 'static,
+    ) -> Self {
+        self.shutdown_hooks.push_front(Box::new(hook));
+        self
+    }
+
     /// Adds a new inner layer to the server.
     ///
     /// The layer's `Service` should be `Send + Sync + Clone + 'static`.
@@ -92,6 +110,7 @@ impl<S, L, Req, MkC, SP> Server<S, L, Req, MkC, SP> {
             #[cfg(feature = "multiplex")]
             multiplex: self.multiplex,
             span_provider: self.span_provider,
+            shutdown_hooks: self.shutdown_hooks,
             _marker: PhantomData,
         }
     }
@@ -116,6 +135,7 @@ impl<S, L, Req, MkC, SP> Server<S, L, Req, MkC, SP> {
             #[cfg(feature = "multiplex")]
             multiplex: self.multiplex,
             span_provider: self.span_provider,
+            shutdown_hooks: self.shutdown_hooks,
             _marker: PhantomData,
         }
     }
@@ -144,6 +164,7 @@ impl<S, L, Req, MkC, SP> Server<S, L, Req, MkC, SP> {
             #[cfg(feature = "multiplex")]
             multiplex: self.multiplex,
             span_provider: self.span_provider,
+            shutdown_hooks: self.shutdown_hooks,
             _marker: PhantomData,
         }
     }
@@ -291,6 +312,14 @@ impl<S, L, Req, MkC, SP> Server<S, L, Req, MkC, SP> {
             }
         }
 
+        if !self.shutdown_hooks.is_empty() {
+            info!("[VOLO] call shutdown hooks");
+
+            for hook in self.shutdown_hooks {
+                (hook)().await;
+            }
+        }
+
         // received signal, graceful shutdown now
         info!("[VOLO] received signal, gracefully exiting now");
         *exit_flag.write() = true;
@@ -343,6 +372,7 @@ impl<S, L, Req, MkC, SP> Server<S, L, Req, MkC, SP> {
             #[cfg(feature = "multiplex")]
             multiplex: self.multiplex,
             span_provider: provider,
+            shutdown_hooks: self.shutdown_hooks,
             _marker: PhantomData,
         }
     }