scuffle_bootstrap/
signals.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
use std::sync::Arc;

use scuffle_context::ContextFutExt;

use crate::global::Global;
use crate::service::Service;

#[derive(Default, Debug, Clone, Copy)]
pub struct SignalSvc;

pub trait SignalSvcConfig: Global {
	fn signals(&self) -> Vec<tokio::signal::unix::SignalKind> {
		vec![
			tokio::signal::unix::SignalKind::terminate(),
			tokio::signal::unix::SignalKind::interrupt(),
		]
	}

	fn timeout(&self) -> Option<std::time::Duration> {
		Some(std::time::Duration::from_secs(30))
	}

	fn on_shutdown(self: &Arc<Self>) -> impl std::future::Future<Output = anyhow::Result<()>> + Send {
		std::future::ready(Ok(()))
	}

	fn on_force_shutdown(
		&self,
		signal: Option<tokio::signal::unix::SignalKind>,
	) -> impl std::future::Future<Output = anyhow::Result<()>> + Send {
		let err = if let Some(signal) = signal {
			anyhow::anyhow!("received signal, shutting down immediately: {:?}", signal)
		} else {
			anyhow::anyhow!("timeout reached, shutting down immediately")
		};

		std::future::ready(Err(err))
	}
}

impl<Global: SignalSvcConfig> Service<Global> for SignalSvc {
	fn enabled(&self, global: &Arc<Global>) -> impl std::future::Future<Output = anyhow::Result<bool>> + Send {
		std::future::ready(Ok(!global.signals().is_empty()))
	}

	async fn run(self, global: Arc<Global>, ctx: scuffle_context::Context) -> anyhow::Result<()> {
		let timeout = global.timeout();

		let signals = global.signals();
		let mut handler = scuffle_signal::SignalHandler::with_signals(signals);

		// Wait for a signal, or for the context to be done.
		handler.recv().with_context(ctx).await;

		global.on_shutdown().await?;

		tokio::select! {
			signal = handler.recv() => {
				global.on_force_shutdown(Some(signal)).await?;
			},
			_ = scuffle_context::Handler::global().shutdown() => {}
			Some(()) = async {
				if let Some(timeout) = timeout {
					tokio::time::sleep(timeout).await;
					Some(())
				} else {
					None
				}
			} => {
				global.on_force_shutdown(None).await?;
			},
		};

		Ok(())
	}
}