phonograph/phono-server/src/cli.rs

use std::net::SocketAddr;

use anyhow::Result;
use axum::{
    ServiceExt,
    extract::Request,
    http::{HeaderValue, header::CONTENT_SECURITY_POLICY},
    middleware::map_request,
};
use chrono::{TimeDelta, Utc};
use clap::{Parser, Subcommand};
use tokio::time::sleep;
use tower::ServiceBuilder;
use tower_http::{
    compression::CompressionLayer, set_header::response::SetResponseHeaderLayer, trace::TraceLayer,
};

use crate::{app::App, middleware::lowercase_uri_path, routes::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: App) -> Result<()> {
    let router = ServiceBuilder::new()
        .layer(map_request(lowercase_uri_path))
        .layer(TraceLayer::new_for_http())
        .layer(CompressionLayer::new())
        .layer(SetResponseHeaderLayer::if_not_present(
            CONTENT_SECURITY_POLICY,
            HeaderValue::from_static("frame-ancestors 'none'"),
        ))
        .service(new_router(state.clone()));

    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.root_path
    );

    axum::serve(
        listener,
        ServiceExt::<Request>::into_make_service_with_connect_info::<SocketAddr>(router),
    )
    .await
    .map_err(Into::into)
}

pub async fn worker_command(args: &WorkerArgs, state: App) -> 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
    }
}