rpki/src/main_rtr.rs
2026-07-24 11:10:32 +08:00

134 lines
4.4 KiB
Rust

use anyhow::Result;
use chrono::Utc;
use tokio::sync::mpsc;
use tracing::info;
use rpki::rtr::admin::{
AdminState, LogTailConfig, RuntimeConfigHandle, SourceReloadHandle, spawn_admin_config_server,
};
use rpki::rtr::bootstrap::{init_shared_cache, open_store, start_servers};
use rpki::rtr::config::{AppConfig, log_startup_config};
use rpki::rtr::report::{ReportConfiguration, ReportContext};
use rpki::rtr::runtime::spawn_refresh_task;
use rpki::rtr::server::{RtrService, RtrShutdownReason};
use rpki::slurm::admin::SlurmAdmin;
#[tokio::main]
async fn main() -> Result<()> {
let config = AppConfig::from_env()?;
init_tracing(config.timezone);
log_startup_config(&config);
let report_context = ReportContext::new(ReportConfiguration::new(
config.source_refresh_interval.as_secs(),
config.runtime_report_interval.as_secs(),
config.report_history_limit,
config.max_delta,
config.prune_delta_by_snapshot_size,
config.timezone,
(
config.timing.refresh,
config.timing.retry,
config.timing.expire,
),
));
let store = open_store(&config)?;
let shared_cache = init_shared_cache(&config, &store, &report_context)?;
let runtime_config = RuntimeConfigHandle::new(config.runtime_config());
let service = RtrService::with_config(shared_cache.clone(), config.service_config.clone());
let notifier = service.notifier();
let service_stats = service.stats();
let shutdown_handle = service.shutdown_handle();
let (source_reload_tx, source_reload_rx) = mpsc::channel(8);
let source_reload = SourceReloadHandle::new(source_reload_tx);
let (process_shutdown_tx, process_shutdown_rx) = mpsc::channel(1);
let admin_task = config.admin_addr.map(|addr| {
let slurm_admin = config.slurm_dir.as_ref().map(SlurmAdmin::new);
let admin_state = AdminState::new(
runtime_config.clone(),
Some(source_reload.clone()),
Some(shutdown_handle.clone()),
Some(process_shutdown_tx.clone()),
slurm_admin,
LogTailConfig::from_env(),
);
spawn_admin_config_server(addr, config.admin_token.clone(), admin_state)
});
let running = start_servers(&config, &service);
let refresh_task = spawn_refresh_task(
&config,
runtime_config,
source_reload_rx,
shared_cache.clone(),
store.clone(),
notifier,
service_stats,
report_context,
);
let stop_reason = wait_for_shutdown(process_shutdown_rx).await?;
info!("stopping RTR service: reason={}", stop_reason.as_str());
running.shutdown(stop_reason);
running.wait().await;
refresh_task.abort();
let _ = refresh_task.await;
if let Some(admin_task) = admin_task {
admin_task.abort();
let _ = admin_task.await;
}
info!("RTR service stopped");
Ok(())
}
async fn wait_for_shutdown(
mut process_shutdown_rx: mpsc::Receiver<RtrShutdownReason>,
) -> Result<RtrShutdownReason> {
tokio::select! {
signal = tokio::signal::ctrl_c() => {
signal?;
let reason = RtrShutdownReason::from_env();
info!("shutdown signal received: reason={}", reason.as_str());
Ok(reason)
}
reason = process_shutdown_rx.recv() => {
let reason = reason.unwrap_or_else(RtrShutdownReason::from_env);
info!("admin process lifecycle request received: reason={}", reason.as_str());
Ok(reason)
}
}
}
fn init_tracing(timezone: chrono_tz::Tz) {
let filter = tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("warn"));
struct LocalTimer {
timezone: chrono_tz::Tz,
}
impl tracing_subscriber::fmt::time::FormatTime for LocalTimer {
fn format_time(
&self,
w: &mut tracing_subscriber::fmt::format::Writer<'_>,
) -> std::fmt::Result {
let now = Utc::now().with_timezone(&self.timezone);
write!(w, "{}", now.format("%Y-%m-%d %H:%M:%S%.3f %:z"))
}
}
if let Err(err) = tracing_subscriber::fmt()
.with_timer(LocalTimer { timezone })
.with_env_filter(filter)
.with_target(true)
.with_thread_ids(true)
.with_level(true)
.try_init()
{
eprintln!("failed to initialize tracing subscriber: {err}");
}
}