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 { #[cfg(feature = "multiplex")] multiplex: bool, span_provider: SP, + + shutdown_hooks: VecDeque BoxFuture<'static, ()>>>, + _marker: PhantomData, } @@ -66,12 +71,25 @@ impl #[cfg(feature = "multiplex")] multiplex: false, span_provider: DefaultProvider {}, + shutdown_hooks: VecDeque::new(), _marker: PhantomData, } } } impl Server { + /// 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 Server { #[cfg(feature = "multiplex")] multiplex: self.multiplex, span_provider: self.span_provider, + shutdown_hooks: self.shutdown_hooks, _marker: PhantomData, } } @@ -116,6 +135,7 @@ impl Server { #[cfg(feature = "multiplex")] multiplex: self.multiplex, span_provider: self.span_provider, + shutdown_hooks: self.shutdown_hooks, _marker: PhantomData, } } @@ -144,6 +164,7 @@ impl Server { #[cfg(feature = "multiplex")] multiplex: self.multiplex, span_provider: self.span_provider, + shutdown_hooks: self.shutdown_hooks, _marker: PhantomData, } } @@ -291,6 +312,14 @@ impl Server { } } + 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 Server { #[cfg(feature = "multiplex")] multiplex: self.multiplex, span_provider: provider, + shutdown_hooks: self.shutdown_hooks, _marker: PhantomData, } }