Skip to content

Commit

Permalink
Add global default wait timeout
Browse files Browse the repository at this point in the history
  • Loading branch information
@notriddle
notriddle committed Aug 27, 2020
1 parent a4a1047 commit c01b864
Show file tree
Hide file tree
Showing 3 changed files with 46 additions and 18 deletions.
23 changes: 23 additions & 0 deletions core/lib/src/config/config.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,8 @@ pub struct Config {
pub(crate) config_file_path: Option<PathBuf>,
/// The path root-relative files will be rooted from.
pub(crate) root_path: Option<PathBuf>,
/// Default duraion to hold open a connection on shutting down.
pub wait_on_shutdown: u32,
}

macro_rules! config_from_raw {
Expand Down Expand Up @@ -283,6 +285,7 @@ impl Config {
extras: HashMap::new(),
config_file_path: None,
root_path: None,
wait_on_shutdown: 0,
}
}
Staging => {
Expand All @@ -299,6 +302,7 @@ impl Config {
extras: HashMap::new(),
config_file_path: None,
root_path: None,
wait_on_shutdown: 1,
}
}
Production => {
Expand All @@ -315,6 +319,7 @@ impl Config {
extras: HashMap::new(),
config_file_path: None,
root_path: None,
wait_on_shutdown: 1,
}
}
})
Expand Down Expand Up @@ -358,6 +363,7 @@ impl Config {
secret_key => (str, set_secret_key, id),
tls => (tls_config, set_raw_tls, id),
limits => (limits, set_limits, ok),
wait_on_shutdown => (u32, set_wait_on_shutdown, ok),
| _ => {
self.extras.insert(name.into(), val.clone());
Ok(())
Expand Down Expand Up @@ -426,6 +432,22 @@ impl Config {
self.port = port;
}

/// Sets the amount of time to wait for a response to finish on graceful shutdown (in seconds).
///
/// # Example
///
/// ```rust
/// use rocket::config::{Config, Environment};
///
/// let mut config = Config::new(Environment::Staging);
/// config.set_wait_on_shutdown(2);
/// assert_eq!(config.wait_on_shutdown, 2);
/// ```
#[inline]
pub fn set_wait_on_shutdown(&mut self, wait_on_shutdown: u32) {
self.wait_on_shutdown = wait_on_shutdown;
}

/// Sets the number of `workers` in `self` to `workers`.
///
/// # Example
Expand Down Expand Up @@ -972,6 +994,7 @@ impl fmt::Debug for Config {
s.field("workers", &self.workers);
s.field("keep_alive", &self.keep_alive);
s.field("log_level", &self.log_level);
s.field("wait_on_shutdown", &self.wait_on_shutdown);

for (key, value) in self.extras() {
s.field(key, &value);
Expand Down
25 changes: 13 additions & 12 deletions core/lib/src/config/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,18 +31,19 @@
//! not used by Rocket itself but can be used by external libraries. The
//! standard configuration parameters are:
//!
//! | name | type | description | examples |
//! |------------|----------------|-------------------------------------------------------------|----------------------------|
//! | address | string | ip address or host to listen on | `"localhost"`, `"1.2.3.4"` |
//! | port | integer | port number to listen on | `8000`, `80` |
//! | keep_alive | integer | keep-alive timeout in seconds | `0` (disable), `10` |
//! | workers | integer | number of concurrent thread workers | `36`, `512` |
//! | log | string | max log level: `"off"`, `"normal"`, `"debug"`, `"critical"` | `"off"`, `"normal"` |
//! | secret_key | 256-bit base64 | secret key for private cookies | `"8Xui8SI..."` (44 chars) |
//! | tls | table | tls config table with two keys (`certs`, `key`) | _see below_ |
//! | tls.certs | string | path to certificate chain in PEM format | `"private/cert.pem"` |
//! | tls.key | string | path to private key for `tls.certs` in PEM format | `"private/key.pem"` |
//! | limits | table | map from data type (string) to data limit (integer: bytes) | `{ forms = 65536 }` |
//! | name | type | description | examples |
//! |------------------|----------------|-------------------------------------------------------------|----------------------------|
//! | address | string | ip address or host to listen on | `"localhost"`, `"1.2.3.4"` |
//! | port | integer | port number to listen on | `8000`, `80` |
//! | keep_alive | integer | keep-alive timeout in seconds | `0` (disable), `10` |
//! | wait_on_shutdown | integer | timeout in seconds when stopping the server | `0` (disable), `10` |
//! | workers | integer | number of concurrent thread workers | `36`, `512` |
//! | log | string | max log level: `"off"`, `"normal"`, `"debug"`, `"critical"` | `"off"`, `"normal"` |
//! | secret_key | 256-bit base64 | secret key for private cookies | `"8Xui8SI..."` (44 chars) |
//! | tls | table | tls config table with two keys (`certs`, `key`) | _see below_ |
//! | tls.certs | string | path to certificate chain in PEM format | `"private/cert.pem"` |
//! | tls.key | string | path to private key for `tls.certs` in PEM format | `"private/key.pem"` |
//! | limits | table | map from data type (string) to data limit (integer: bytes) | `{ forms = 65536 }` |
//!
//! ### Rocket.toml
//!
Expand Down
16 changes: 10 additions & 6 deletions core/lib/src/rocket.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -236,21 +236,25 @@ impl Rocket {
tx: oneshot::Sender<hyper::Response<hyper::Body>>,
mut shutdown_receiver: broadcast::Receiver<()>,
) {
let wait_on_shutdown = response.wait_on_shutdown();
let result = self.write_response(response, tx);
let mut shutdown_receiver = if wait_on_shutdown != Duration::from_millis(0) {
let wait_on_shutdown = {
let wait_on_shutdown = response.wait_on_shutdown();
if wait_on_shutdown != Duration::from_millis(0) {
wait_on_shutdown
} else {
Duration::from_millis((self.config.wait_on_shutdown as u64) * 1_000)
}
};
let mut shutdown_receiver = {
let (tx, rx) = broadcast::channel(1);
tokio::spawn(async move {
let _ = shutdown_receiver.recv().await;
tokio::time::delay_for(wait_on_shutdown).await;
tx.send(()).expect("there should always be at least one shutdown listener");
});
rx
} else {
shutdown_receiver
};
tokio::select!{
result = result => {
result = self.write_response(response, tx) => {
match result {
Ok(()) => {
info_!("{}", Paint::green("Response succeeded."));
Expand Down

0 comments on commit c01b864

Please sign in to comment.