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

atuin_server/
lib.rs
