🐛 Add forced shutdown to ptth_file_server.

main
_ 2020-11-07 01:40:29 +00:00
parent 75177cec80
commit f02e12aecc
3 changed files with 11 additions and 42 deletions

View File

@ -45,8 +45,6 @@ pub struct Config {
struct ServerState <'a> { struct ServerState <'a> {
config: Config, config: Config,
handlebars: handlebars::Handlebars <'a>, handlebars: handlebars::Handlebars <'a>,
shutdown_watch_rx: watch::Receiver <bool>,
} }
fn status_reply <B: Into <Body>> (status: StatusCode, b: B) fn status_reply <B: Into <Body>> (status: StatusCode, b: B)
@ -76,19 +74,12 @@ async fn handle_all (req: Request <Body>, state: Arc <ServerState <'static>>)
.as_ref () .as_ref ()
.unwrap_or (&default_root); .unwrap_or (&default_root);
let mut shutdown_watch_rx = state.shutdown_watch_rx.clone ();
if shutdown_watch_rx.recv ().await != Some (false) {
error! ("Can't serve, I'm shutting down");
panic! ("Can't serve, I'm shutting down");
}
let ptth_resp = file_server::serve_all ( let ptth_resp = file_server::serve_all (
&state.handlebars, &state.handlebars,
file_server_root, file_server_root,
ptth_req.method, ptth_req.method,
&ptth_req.uri, &ptth_req.uri,
&ptth_req.headers, &ptth_req.headers
Some (shutdown_watch_rx)
).await; ).await;
let mut resp = Response::builder () let mut resp = Response::builder ()
@ -129,15 +120,11 @@ async fn main () -> Result <(), Box <dyn Error>> {
let handlebars = file_server::load_templates ()?; let handlebars = file_server::load_templates ()?;
let (shutdown_watch_tx, shutdown_watch_rx) = watch::channel (false);
let state = Arc::new (ServerState { let state = Arc::new (ServerState {
config: Config { config: Config {
file_server_root: config_file.file_server_root, file_server_root: config_file.file_server_root,
}, },
handlebars, handlebars,
shutdown_watch_rx,
}); });
let make_svc = make_service_fn (|_conn| { let make_svc = make_service_fn (|_conn| {
@ -161,13 +148,15 @@ async fn main () -> Result <(), Box <dyn Error>> {
shutdown_oneshot.await.ok (); shutdown_oneshot.await.ok ();
info! ("Received graceful shutdown"); info! ("Received graceful shutdown");
shutdown_watch_tx.broadcast (true).unwrap (); // Kick off a timer for a forced shutdown
force_shutdown_tx.send (()).unwrap (); force_shutdown_tx.send (()).unwrap ();
}); });
let force_shutdown_fut = async move { let force_shutdown_fut = async move {
force_shutdown_rx.await.unwrap (); force_shutdown_rx.await.unwrap ();
delay_for (Duration::from_secs (5)).await; let timeout = 10;
info! ("Forced shutdown in {} seconds", timeout);
delay_for (Duration::from_secs (timeout)).await;
error! ("Forcing shutdown"); error! ("Forcing shutdown");
}; };

View File

@ -141,13 +141,12 @@ async fn serve_dir (
resp resp
} }
#[instrument (level = "debug", skip (f, cancel_rx))] #[instrument (level = "debug", skip (f))]
async fn serve_file ( async fn serve_file (
mut f: File, mut f: File,
should_send_body: bool, should_send_body: bool,
range_start: Option <u64>, range_start: Option <u64>,
range_end: Option <u64>, range_end: Option <u64>
mut cancel_rx: Option <watch::Receiver <bool>>
) -> http_serde::Response { ) -> http_serde::Response {
let (tx, rx) = channel (1); let (tx, rx) = channel (1);
let body = if should_send_body { let body = if should_send_body {
@ -172,7 +171,6 @@ async fn serve_file (
if should_send_body { if should_send_body {
tokio::spawn (async move { tokio::spawn (async move {
{
//println! ("Opening file {:?}", path); //println! ("Opening file {:?}", path);
let mut tx = tx; let mut tx = tx;
@ -191,20 +189,7 @@ async fn serve_file (
break; break;
} }
let send_fut = tx.send (Ok::<_, Infallible> (buffer)); if tx.send (Ok::<_, Infallible> (buffer)).await.is_err () {
let send_result = match &mut cancel_rx {
Some (cancel_rx) => futures::select! {
x = send_fut.fuse () => x,
_ = cancel_rx.recv ().fuse () => {
error! ("Cancelled");
break;
},
},
None => send_fut.await,
};
if send_result.is_err () {
warn! ("Cancelling file stream (Sent {} out of {} bytes)", bytes_sent, end - start); warn! ("Cancelling file stream (Sent {} out of {} bytes)", bytes_sent, end - start);
break; break;
} }
@ -220,8 +205,6 @@ async fn serve_file (
//delay_for (Duration::from_millis (50)).await; //delay_for (Duration::from_millis (50)).await;
} }
}
debug! ("Exited stream scope");
}); });
} }
@ -267,8 +250,7 @@ pub async fn serve_all (
root: &Path, root: &Path,
method: http_serde::Method, method: http_serde::Method,
uri: &str, uri: &str,
headers: &HashMap <String, Vec <u8>>, headers: &HashMap <String, Vec <u8>>
cancel_rx: Option <watch::Receiver <bool>>
) )
-> http_serde::Response -> http_serde::Response
{ {
@ -317,8 +299,7 @@ pub async fn serve_all (
file, file,
should_send_body, should_send_body,
range_start, range_start,
range_end, range_end
cancel_rx
).await ).await
} }
else { else {

View File

@ -76,8 +76,7 @@ async fn handle_req_resp <'a> (
file_server_root, file_server_root,
parts.method, parts.method,
uri, uri,
&parts.headers, &parts.headers
None
).await ).await
} }
else { else {