use std::{net::TcpListener, sync::Arc};
use axum::{
response::{Html, IntoResponse},
routing::{get, post},
Extension, Router, Server,
};
use tokio::sync::{
oneshot::{channel, Sender},
Mutex,
};
struct State {
shutdown: Option<Sender<()>>,
}
async fn index() -> impl IntoResponse {
Html(
r#"Hello World. <form method="post" action="/shutdown"><input type="submit" value="Shut Down"></form>"#,
)
}
async fn shutdown(Extension(state): Extension<Arc<Mutex<State>>>) -> impl IntoResponse {
if let Some(shutdown) = state.lock().await.shutdown.take() {
let _ = shutdown.send(());
"Shutting down"
} else {
"Shutdown already initiated"
}
}
fn main() {
let (shutdown_tx, shutdown_rx) = channel::<()>();
let state = Arc::new(Mutex::new(State {
shutdown: Some(shutdown_tx),
}));
let app = Router::new()
.route("/", get(index))
.route("/shutdown", post(shutdown))
.layer(Extension(state));
let ear = TcpListener::bind("127.0.0.1:0").expect("bind port");
let addr = ear.local_addr().expect("retrieve port");
drteeth::launch("Axum Demo", addr, async move {
Server::from_tcp(ear)
.unwrap()
.serve(app.into_make_service())
.with_graceful_shutdown(async {
shutdown_rx.await.ok();
})
.await
.unwrap();
})
.unwrap();
}