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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
use std::fmt;

#[cfg(unix)]
use std::collections::HashSet;

use futures::stream::Stream;
use serde::{Deserialize, Serialize};

/// A Unix signal for triggering graceful shutdown.
///
/// Each variant corresponds to a Unix process signal which can be used to
/// trigger a graceful shutdown. See [`Shutdown`] for details.
///
/// ## (De)serialization
///
/// A `Sig` variant serializes and deserializes as a lowercase string equal to
/// the name of the variant: `"alrm"` for [`Sig::Alrm`], `"chld"` for
/// [`Sig::Chld`], and so on.
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
#[cfg_attr(nightly, doc(cfg(unix)))]
pub enum Sig {
    /// The `SIGALRM` Unix signal.
    Alrm,
    /// The `SIGCHLD` Unix signal.
    Chld,
    /// The `SIGHUP` Unix signal.
    Hup,
    /// The `SIGINT` Unix signal.
    Int,
    /// The `SIGIO` Unix signal.
    Io,
    /// The `SIGPIPE` Unix signal.
    Pipe,
    /// The `SIGQUIT` Unix signal.
    Quit,
    /// The `SIGTERM` Unix signal.
    Term,
    /// The `SIGUSR1` Unix signal.
    Usr1,
    /// The `SIGUSR2` Unix signal.
    Usr2
}

impl fmt::Display for Sig {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let s = match self {
            Sig::Alrm => "SIGALRM",
            Sig::Chld => "SIGCHLD",
            Sig::Hup => "SIGHUP",
            Sig::Int => "SIGINT",
            Sig::Io => "SIGIO",
            Sig::Pipe => "SIGPIPE",
            Sig::Quit => "SIGQUIT",
            Sig::Term => "SIGTERM",
            Sig::Usr1 => "SIGUSR1",
            Sig::Usr2 => "SIGUSR2",
        };

        s.fmt(f)
    }
}

/// Graceful shutdown configuration.
///
/// # Summary
///
/// This structure configures when and how graceful shutdown occurs. The `ctrlc`
/// and `signals` properties control _when_ and the `grace` and `mercy`
/// properties control _how_.
///
/// When a shutdown is triggered by an externally or internally initiated
/// [`Shutdown::notify()`], Rocket allows application I/O to make progress for
/// at most `grace` seconds before initiating connection-level shutdown.
/// Connection shutdown forcibly terminates _application_ I/O, but connections
/// are allowed an additional `mercy` seconds to shutdown before being
/// forcefully terminated. This implies that a _cooperating_ and active remote
/// client maintaining an open connection can stall shutdown for at most `grace`
/// seconds, while an _uncooperative_ remote client can stall shutdown for at
/// most `grace + mercy` seconds.
///
/// # Triggers
///
/// _All_ graceful shutdowns are initiated via [`Shutdown::notify()`]. Rocket
/// can be configured to call [`Shutdown::notify()`] automatically on certain
/// conditions, specified via the `ctrlc` and `signals` properties of this
/// structure. More specifically, if `ctrlc` is `true` (the default), `ctrl-c`
/// (`SIGINT`) initiates a server shutdown, and on Unix, `signals` specifies a
/// list of IPC signals that trigger a shutdown (`["term"]` by default).
///
/// [`Shutdown::notify()`]: crate::Shutdown::notify()
///
/// # Grace Period
///
/// Once a shutdown is triggered, Rocket stops accepting new connections and
/// waits at most `grace` seconds before initiating connection shutdown.
/// Applications can `await` the [`Shutdown`](crate::Shutdown) future to detect
/// a shutdown and cancel any server-initiated I/O, such as from [infinite
/// responders](crate::response::stream#graceful-shutdown), to avoid abrupt I/O
/// cancellation.
///
/// # Mercy Period
///
/// After the grace period has elapsed, Rocket initiates connection shutdown,
/// allowing connection-level I/O termination such as TLS's `close_notify` to
/// proceed nominally. Rocket waits at most `mercy` seconds for connections to
/// shutdown before forcefully terminating all connections.
///
/// # Runaway I/O
///
/// If tasks are _still_ executing after both periods _and_ a Rocket configured
/// async runtime is in use, Rocket waits an unspecified amount of time (not to
/// exceed 1s) and forcefully terminates the asynchronous runtime. This
/// guarantees that the server process terminates, prohibiting uncooperative,
/// runaway I/O from preventing shutdown altogether.
///
/// A "Rocket configured runtime" is one started by the `#[rocket::main]` and
/// `#[launch]` attributes. Rocket _never_ forcefully terminates a custom
/// runtime. A server that creates its own async runtime must take care to
/// terminate itself if tasks it spawns fail to cooperate.
///
/// Under normal circumstances, forced termination should never occur. No use of
/// "normal" cooperative I/O (that is, via `.await` or `task::spawn()`) should
/// trigger abrupt termination. Instead, forced cancellation is intended to
/// prevent _buggy_ code, such as an unintended infinite loop or unknown use of
/// blocking I/O, from preventing shutdown.
///
/// This behavior can be disabled by setting [`Shutdown::force`] to `false`.
///
/// # Example
///
/// As with all Rocket configuration options, when using the default
/// [`Config::figment()`](crate::Config::figment()), `Shutdown` can be
/// configured via a `Rocket.toml` file. As always, defaults are provided
/// (documented below), and thus configuration only needs to provided to change
/// defaults.
///
/// ```rust
/// # use rocket::figment::{Figment, providers::{Format, Toml}};
/// use rocket::Config;
///
/// // If these are the contents of `Rocket.toml`...
/// # let toml = Toml::string(r#"
/// [default.shutdown]
/// ctrlc = false
/// signals = ["term", "hup"]
/// grace = 10
/// mercy = 5
/// # force = false
/// # "#).nested();
///
/// // The config parses as follows:
/// # let config = Config::from(Figment::from(Config::debug_default()).merge(toml));
/// assert_eq!(config.shutdown.ctrlc, false);
/// assert_eq!(config.shutdown.grace, 10);
/// assert_eq!(config.shutdown.mercy, 5);
/// # assert_eq!(config.shutdown.force, false);
///
/// # #[cfg(unix)] {
/// use rocket::config::Sig;
///
/// assert_eq!(config.shutdown.signals.len(), 2);
/// assert!(config.shutdown.signals.contains(&Sig::Term));
/// assert!(config.shutdown.signals.contains(&Sig::Hup));
/// # }
/// ```
///
/// Or, as with all configuration options, programmatically:
///
/// ```rust
/// # use rocket::figment::{Figment, providers::{Format, Toml}};
/// use rocket::config::{Config, Shutdown};
///
/// #[cfg(unix)]
/// use rocket::config::Sig;
///
/// let config = Config {
///     shutdown: Shutdown {
///         ctrlc: false,
///         #[cfg(unix)]
///         signals: {
///             let mut set = std::collections::HashSet::new();
///             set.insert(Sig::Term);
///             set.insert(Sig::Hup);
///             set
///         },
///         grace: 10,
///         mercy: 5,
///         force: true,
///         ..Default::default()
///     },
///     ..Config::default()
/// };
///
/// assert_eq!(config.shutdown.ctrlc, false);
/// assert_eq!(config.shutdown.grace, 10);
/// assert_eq!(config.shutdown.mercy, 5);
/// assert_eq!(config.shutdown.force, true);
///
/// #[cfg(unix)] {
///     assert_eq!(config.shutdown.signals.len(), 2);
///     assert!(config.shutdown.signals.contains(&Sig::Term));
///     assert!(config.shutdown.signals.contains(&Sig::Hup));
/// }
/// ```
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Shutdown {
    /// Whether `ctrl-c` (`SIGINT`) initiates a server shutdown.
    ///
    /// **default: `true`**
    #[serde(deserialize_with = "figment::util::bool_from_str_or_int")]
    pub ctrlc: bool,
    /// On Unix, a set of signal which trigger a shutdown. On non-Unix, this
    /// option is unavailable and silently ignored.
    ///
    /// **default: { [`Sig::Term`] }**
    #[cfg(unix)]
    #[cfg_attr(nightly, doc(cfg(unix)))]
    pub signals: HashSet<Sig>,
    /// The grace period: number of seconds to continue to try to finish
    /// outstanding _server_ I/O for before forcibly terminating it.
    ///
    /// **default: `2`**
    pub grace: u32,
    /// The mercy period: number of seconds to continue to try to finish
    /// outstanding _connection_ I/O for before forcibly terminating it.
    ///
    /// **default: `3`**
    pub mercy: u32,
    /// Whether to force termination of an async runtime that refuses to
    /// cooperatively shutdown.
    ///
    /// Rocket _never_ forcefully terminates a custom runtime, irrespective of
    /// this value. A server that creates its own async runtime must take care
    /// to terminate itself if it fails to cooperate.
    ///
    /// _**Note:** Rocket only reads this value from sources in the [default
    /// provider](crate::Config::figment())._
    ///
    /// **default: `true`**
    #[serde(deserialize_with = "figment::util::bool_from_str_or_int")]
    pub force: bool,
    /// PRIVATE: This structure may grow (but never change otherwise) in a
    /// non-breaking release. As such, constructing this structure should
    /// _always_ be done using a public constructor or update syntax:
    ///
    /// ```rust
    /// use rocket::config::Shutdown;
    ///
    /// let config = Shutdown {
    ///     grace: 5,
    ///     mercy: 10,
    ///     ..Default::default()
    /// };
    /// ```
    #[doc(hidden)]
    #[serde(skip)]
    pub __non_exhaustive: (),
}

impl fmt::Display for Shutdown {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "ctrlc = {}, force = {}, ", self.ctrlc, self.force)?;

        #[cfg(unix)] {
            write!(f, "signals = [")?;
            for (i, sig) in self.signals.iter().enumerate() {
                if i != 0 { write!(f, ", ")?; }
                write!(f, "{}", sig)?;
            }
            write!(f, "], ")?;
        }

        write!(f, "grace = {}s, mercy = {}s", self.grace, self.mercy)?;
        Ok(())
    }
}

impl Default for Shutdown {
    fn default() -> Self {
        Shutdown {
            ctrlc: true,
            #[cfg(unix)]
            signals: { let mut set = HashSet::new(); set.insert(Sig::Term); set },
            grace: 2,
            mercy: 3,
            force: true,
            __non_exhaustive: (),
        }
    }
}

impl Shutdown {
    #[cfg(unix)]
    pub(crate) fn signal_stream(&self) -> Option<impl Stream<Item = Sig>> {
        use tokio_stream::{StreamExt, StreamMap, wrappers::SignalStream};
        use tokio::signal::unix::{signal, SignalKind};

        if !self.ctrlc && self.signals.is_empty() {
            return None;
        }

        let mut signals = self.signals.clone();
        if self.ctrlc {
            signals.insert(Sig::Int);
        }

        let mut map = StreamMap::new();
        for sig in signals {
            let sigkind = match sig {
                Sig::Alrm => SignalKind::alarm(),
                Sig::Chld => SignalKind::child(),
                Sig::Hup => SignalKind::hangup(),
                Sig::Int => SignalKind::interrupt(),
                Sig::Io => SignalKind::io(),
                Sig::Pipe => SignalKind::pipe(),
                Sig::Quit => SignalKind::quit(),
                Sig::Term => SignalKind::terminate(),
                Sig::Usr1 => SignalKind::user_defined1(),
                Sig::Usr2 => SignalKind::user_defined2()
            };

            match signal(sigkind) {
                Ok(signal) => { map.insert(sig, SignalStream::new(signal)); },
                Err(e) => warn!("Failed to enable `{}` shutdown signal: {}", sig, e),
            }
        }

        Some(map.map(|(k, _)| k))
    }

    #[cfg(not(unix))]
    pub(crate) fn signal_stream(&self) -> Option<impl Stream<Item = Sig>> {
        use tokio_stream::StreamExt;
        use futures::stream::once;

        self.ctrlc.then(|| tokio::signal::ctrl_c())
            .map(|signal| once(Box::pin(signal)))
            .map(|stream| stream.filter_map(|result| {
                result.map(|_| Sig::Int)
                    .map_err(|e| warn!("Failed to enable `ctrl-c` shutdown signal: {}", e))
                    .ok()
            }))
    }
}