shared/extensions/
shutdown_handlers.rs

1use crate::State;
2use futures_util::FutureExt;
3use std::{
4    borrow::Cow,
5    collections::{HashMap, HashSet},
6    panic::AssertUnwindSafe,
7    pin::Pin,
8    sync::Arc,
9};
10use tokio::sync::{OwnedRwLockReadGuard, RwLock};
11
12pub type ShutdownFunc =
13    dyn Fn(State) -> Pin<Box<dyn Future<Output = Result<(), anyhow::Error>> + Send>> + Send + Sync;
14
15pub struct ShutdownHandler {
16    pub name: &'static str,
17    pub task: Box<ShutdownFunc>,
18}
19
20pub struct ShutdownHandlerBuilder {
21    state: State,
22    tasks: Arc<RwLock<HashMap<&'static str, ShutdownHandler>>>,
23}
24
25impl ShutdownHandlerBuilder {
26    pub fn new(state: State) -> Self {
27        Self {
28            state,
29            tasks: Arc::new(RwLock::new(HashMap::new())),
30        }
31    }
32
33    pub(crate) fn share(&self) -> Self {
34        Self {
35            state: self.state.clone(),
36            tasks: Arc::clone(&self.tasks),
37        }
38    }
39
40    pub(crate) async fn handler_names(&self) -> HashSet<&'static str> {
41        self.tasks.read().await.keys().copied().collect()
42    }
43
44    /// Removes every handler not in `keep`.
45    pub(crate) async fn rollback(&self, keep: &HashSet<&'static str>) {
46        self.tasks
47            .write()
48            .await
49            .retain(|name, _| keep.contains(name));
50    }
51
52    /// Adds a shutdown handler that will be executed when a gradual shutdown is initiated.
53    /// This will run on primary and on non-primary instances, so be aware of that when implementing your handler and
54    /// perhaps check `state.env.app_primary`.
55    pub async fn add_handler<
56        F: Fn(State) -> Fut + Send + Sync + 'static,
57        Fut: Future<Output = Result<(), anyhow::Error>> + Send + 'static,
58    >(
59        &self,
60        name: &'static str,
61        shutdown_fn: F,
62    ) {
63        let state = self.state.clone();
64        let tasks = Arc::clone(&self.tasks);
65
66        let previous = self.tasks.write().await.insert(
67            name,
68            ShutdownHandler {
69                name,
70                task: Box::new(move |state: State| Box::pin(shutdown_fn(state))),
71            },
72        );
73
74        if previous.is_some() {
75            tracing::warn!(
76                name,
77                "shutdown handler registered twice, replacing the previous one"
78            );
79        }
80    }
81}
82
83#[derive(Default)]
84pub struct ShutdownHandlerManager {
85    builder: RwLock<Option<ShutdownHandlerBuilder>>,
86}
87
88impl ShutdownHandlerManager {
89    pub async fn merge_builder(&self, builder: ShutdownHandlerBuilder) {
90        self.builder.write().await.replace(builder);
91    }
92
93    pub async fn get_handlers(
94        &self,
95    ) -> OwnedRwLockReadGuard<HashMap<&'static str, ShutdownHandler>> {
96        let inner = self.builder.read().await;
97
98        inner.as_ref().unwrap().tasks.clone().read_owned().await
99    }
100
101    pub async fn handle_shutdown(&self) {
102        let handlers = self.get_handlers().await;
103
104        for (name, handler) in handlers.iter() {
105            tracing::info!(name, "running shutdown task");
106            let result = AssertUnwindSafe((handler.task)(
107                self.builder.read().await.as_ref().unwrap().state.clone(),
108            ))
109            .catch_unwind()
110            .await;
111
112            match result {
113                Ok(result) => {
114                    if let Err(err) = result {
115                        tracing::error!(name, %err, "shutdown task failed");
116                        sentry_anyhow::capture_anyhow(&err);
117                    }
118                }
119                Err(err) => {
120                    let err_msg: Cow<'_, str> = if let Some(s) = err.downcast_ref::<&str>() {
121                        (*s).into()
122                    } else if let Some(s) = err.downcast_ref::<String>() {
123                        s.as_str().into()
124                    } else {
125                        "unknown panic".into()
126                    };
127
128                    tracing::error!(name, %err_msg, "shutdown task panicked");
129                    sentry::capture_message(
130                        &format!("Shutdown handler '{}' panicked: {}", name, err_msg),
131                        sentry::Level::Error,
132                    );
133                }
134            }
135        }
136    }
137}
shared/extensions/shutdown_handlers.rs

shared/extensions/
shutdown_handlers.rs
