Browse Source

fix: ctrl+c not exit sometimes

pull/20/head
sigoden 3 years ago
parent
commit
882a9ae716
  1. 15
      src/main.rs
  2. 100
      src/server.rs

15
src/main.rs

@ -14,10 +14,23 @@ async fn main() {
async fn run() -> BoxResult<()> { async fn run() -> BoxResult<()> {
let args = Args::parse(matches())?; let args = Args::parse(matches())?;
serve(args).await tokio::select! {
ret = serve(args) => {
ret
},
_ = shutdown_signal() => {
Ok(())
},
}
} }
fn handle_err<T>(err: Box<dyn std::error::Error>) -> T { fn handle_err<T>(err: Box<dyn std::error::Error>) -> T {
eprintln!("error: {}", err); eprintln!("error: {}", err);
std::process::exit(1); std::process::exit(1);
} }
async fn shutdown_signal() {
tokio::signal::ctrl_c()
.await
.expect("Failed to install CTRL+C signal handler")
}

100
src/server.rs

@ -54,65 +54,55 @@ macro_rules! status {
} }
pub async fn serve(args: Args) -> BoxResult<()> { pub async fn serve(args: Args) -> BoxResult<()> {
match args.tls.as_ref() {
Some(_) => serve_https(args).await,
None => serve_http(args).await,
}
}
pub async fn serve_https(args: Args) -> BoxResult<()> {
let args = Arc::new(args); let args = Arc::new(args);
let socket_addr = args.address()?; let socket_addr = args.address()?;
let (certs, key) = args.tls.clone().unwrap();
let inner = Arc::new(InnerService::new(args.clone())); let inner = Arc::new(InnerService::new(args.clone()));
let config = ServerConfig::builder() match args.tls.clone() {
.with_safe_defaults() Some((certs, key)) => {
.with_no_client_auth() let config = ServerConfig::builder()
.with_single_cert(certs, key)?; .with_safe_defaults()
let tls_acceptor = TlsAcceptor::from(Arc::new(config)); .with_no_client_auth()
let arc_acceptor = Arc::new(tls_acceptor); .with_single_cert(certs, key)?;
let listener = TcpListener::bind(&socket_addr).await?; let tls_acceptor = TlsAcceptor::from(Arc::new(config));
let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener); let arc_acceptor = Arc::new(tls_acceptor);
let incoming = hyper::server::accept::from_stream(incoming.filter_map(|socket| async { let listener = TcpListener::bind(&socket_addr).await?;
match socket { let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener);
Ok(stream) => match arc_acceptor.clone().accept(stream).await { let incoming =
Ok(val) => Some(Ok::<_, Infallible>(val)), hyper::server::accept::from_stream(incoming.filter_map(|socket| async {
Err(_) => None, match socket {
}, Ok(stream) => match arc_acceptor.clone().accept(stream).await {
Err(_) => None, Ok(val) => Some(Ok::<_, Infallible>(val)),
} Err(_) => None,
})); },
let server = hyper::Server::builder(incoming).serve(make_service_fn(move |_| { Err(_) => None,
let inner = inner.clone(); }
async move { }));
Ok::<_, Infallible>(service_fn(move |req| { let server = hyper::Server::builder(incoming).serve(make_service_fn(move |_| {
let inner = inner.clone(); let inner = inner.clone();
inner.call(req) async move {
})) Ok::<_, Infallible>(service_fn(move |req| {
let inner = inner.clone();
inner.call(req)
}))
}
}));
print_listening(args.address.as_str(), args.port, &args.uri_prefix, true);
server.await?;
} }
})); None => {
print_listening(args.address.as_str(), args.port, &args.uri_prefix, true); let server = hyper::Server::try_bind(&socket_addr)?.serve(make_service_fn(move |_| {
let graceful = server.with_graceful_shutdown(shutdown_signal());
graceful.await?;
Ok(())
}
pub async fn serve_http(args: Args) -> BoxResult<()> {
let args = Arc::new(args);
let socket_addr = args.address()?;
let inner = Arc::new(InnerService::new(args.clone()));
let server = hyper::Server::try_bind(&socket_addr)?.serve(make_service_fn(move |_| {
let inner = inner.clone();
async move {
Ok::<_, Infallible>(service_fn(move |req| {
let inner = inner.clone(); let inner = inner.clone();
inner.call(req) async move {
})) Ok::<_, Infallible>(service_fn(move |req| {
let inner = inner.clone();
inner.call(req)
}))
}
}));
print_listening(args.address.as_str(), args.port, &args.uri_prefix, false);
server.await?;
} }
})); }
print_listening(args.address.as_str(), args.port, &args.uri_prefix, false);
let graceful = server.with_graceful_shutdown(shutdown_signal());
graceful.await?;
Ok(()) Ok(())
} }
@ -1012,9 +1002,3 @@ fn retrieve_listening_addrs(address: &str) -> Vec<String> {
} }
vec![address.to_owned()] vec![address.to_owned()]
} }
async fn shutdown_signal() {
tokio::signal::ctrl_c()
.await
.expect("Failed to install CTRL+C signal handler")
}

Loading…
Cancel
Save