atuin_server/
lib.rs

1#![forbid(unsafe_code)]
2
3use std::future::Future;
4use std::net::SocketAddr;
5
6use atuin_server_database::Database;
7use axum::{Router, serve};
8use eyre::{Context, Result};
9
10mod handlers;
11mod metrics;
12mod router;
13
14pub use settings::Settings;
15pub use settings::example_config;
16
17pub mod settings;
18
19use tokio::net::TcpListener;
20use tokio::signal;
21
22#[cfg(target_family = "unix")]
23async fn shutdown_signal() {
24    let mut term = signal::unix::signal(signal::unix::SignalKind::terminate())
25        .expect("failed to register signal handler");
26    let mut interrupt = signal::unix::signal(signal::unix::SignalKind::interrupt())
27        .expect("failed to register signal handler");
28
29    tokio::select! {
30        _ = term.recv() => {},
31        _ = interrupt.recv() => {},
32    };
33    eprintln!("Shutting down gracefully...");
34}
35
36#[cfg(target_family = "windows")]
37async fn shutdown_signal() {
38    signal::windows::ctrl_c()
39        .expect("failed to register signal handler")
40        .recv()
41        .await;
42    eprintln!("Shutting down gracefully...");
43}
44
45pub async fn launch<Db: Database>(settings: Settings, addr: SocketAddr) -> Result<()> {
46    launch_with_tcp_listener::<Db>(
47        settings,
48        TcpListener::bind(addr)
49            .await
50            .context("could not connect to socket")?,
51        shutdown_signal(),
52    )
53    .await
54}
55
56pub async fn launch_with_tcp_listener<Db: Database>(
57    settings: Settings,
58    listener: TcpListener,
59    shutdown: impl Future<Output = ()> + Send + 'static,
60) -> Result<()> {
61    let r = make_router::<Db>(settings).await?;
62
63    serve(listener, r.into_make_service())
64        .with_graceful_shutdown(shutdown)
65        .await?;
66
67    Ok(())
68}
69
70// The separate listener means it's much easier to ensure metrics are not accidentally exposed to
71// the public.
72pub async fn launch_metrics_server(host: String, port: u16) -> Result<()> {
73    let listener = TcpListener::bind((host, port))
74        .await
75        .context("failed to bind metrics tcp")?;
76
77    let recorder_handle = metrics::setup_metrics_recorder();
78
79    let router = Router::new().route(
80        "/metrics",
81        axum::routing::get(move || std::future::ready(recorder_handle.render())),
82    );
83
84    serve(listener, router.into_make_service())
85        .with_graceful_shutdown(shutdown_signal())
86        .await?;
87
88    Ok(())
89}
90
91async fn make_router<Db: Database>(settings: Settings) -> Result<Router, eyre::Error> {
92    let db = Db::new(&settings.db_settings)
93        .await
94        .wrap_err_with(|| format!("failed to connect to db: {:?}", settings.db_settings))?;
95    let r = router::router(db, settings);
96    Ok(r)
97}
atuin_server/lib.rs

atuin_server/
lib.rs
