84 lines
2.4 KiB
Rust
84 lines
2.4 KiB
Rust
|
use anyhow::Result;
|
||
|
|
use axum::middleware::map_request;
|
||
|
|
use chrono::{TimeDelta, Utc};
|
||
|
|
use clap::{Parser, Subcommand};
|
||
|
|
use tokio::time::sleep;
|
||
|
|
use tower::ServiceBuilder;
|
||
|
|
use tower_http::{
|
||
|
|
compression::CompressionLayer, normalize_path::NormalizePathLayer, trace::TraceLayer,
|
||
|
|
};
|
||
|
|
|
||
|
|
use crate::{
|
||
|
|
app_state::AppState, middleware::lowercase_uri_path, router::new_router, worker::run_worker,
|
||
|
|
};
|
||
|
|
|
||
|
|
#[derive(Parser)]
|
||
|
|
#[command(version, about, long_about = None)]
|
||
|
|
pub struct Cli {
|
||
|
|
#[command(subcommand)]
|
||
|
|
pub command: Commands,
|
||
|
|
}
|
||
|
|
|
||
|
|
#[derive(Parser)]
|
||
|
|
pub struct WorkerArgs {
|
||
|
|
/// Loop the every n seconds instead of exiting after execution
|
||
|
|
#[arg(long)]
|
||
|
|
auto_loop_seconds: Option<u32>,
|
||
|
|
}
|
||
|
|
|
||
|
|
#[derive(Subcommand)]
|
||
|
|
pub enum Commands {
|
||
|
|
/// Run web server
|
||
|
|
Serve,
|
||
|
|
/// Run background worker
|
||
|
|
Worker(WorkerArgs),
|
||
|
|
// TODO: add a low-frequency worker task exclusively for self-healing
|
||
|
|
// mechanisms like Governor::reset_all()
|
||
|
|
}
|
||
|
|
|
||
|
|
pub async fn serve_command(state: AppState) -> Result<()> {
|
||
|
|
let router = new_router(state.clone()).layer(
|
||
|
|
ServiceBuilder::new()
|
||
|
|
.layer(map_request(lowercase_uri_path))
|
||
|
|
.layer(TraceLayer::new_for_http())
|
||
|
|
.layer(CompressionLayer::new())
|
||
|
|
.layer(NormalizePathLayer::trim_trailing_slash()),
|
||
|
|
);
|
||
|
|
|
||
|
|
let listener =
|
||
|
|
tokio::net::TcpListener::bind((state.settings.host.clone(), state.settings.port))
|
||
|
|
.await
|
||
|
|
.unwrap();
|
||
|
|
tracing::info!(
|
||
|
|
"App running at http://{}:{}{}",
|
||
|
|
state.settings.host,
|
||
|
|
state.settings.port,
|
||
|
|
state.settings.base_path
|
||
|
|
);
|
||
|
|
|
||
|
|
axum::serve(listener, router).await.map_err(Into::into)
|
||
|
|
}
|
||
|
|
|
||
|
|
pub async fn worker_command(args: &WorkerArgs, state: AppState) -> Result<()> {
|
||
|
|
if let Some(loop_seconds) = args.auto_loop_seconds {
|
||
|
|
let loop_delta = TimeDelta::seconds(i64::from(loop_seconds));
|
||
|
|
loop {
|
||
|
|
let t_next_loop = Utc::now() + loop_delta;
|
||
|
|
|
||
|
|
if let Err(err) = run_worker(state.clone()).await {
|
||
|
|
tracing::error!("{}", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
let sleep_delta = t_next_loop - Utc::now();
|
||
|
|
match sleep_delta.to_std() {
|
||
|
|
Ok(duration) => {
|
||
|
|
sleep(duration).await;
|
||
|
|
}
|
||
|
|
Err(_) => { /* sleep_delta was < 0, so don't sleep */ }
|
||
|
|
}
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
run_worker(state).await
|
||
|
|
}
|
||
|
|
}
|