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 trace;
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(
65        listener,
66        r.into_make_service_with_connect_info::<SocketAddr>(),
67    )
68    .with_graceful_shutdown(shutdown)
69    .await?;
70
71    Ok(())
72}
73
74// The separate listener means it's much easier to ensure metrics are not accidentally exposed to
75// the public.
76pub async fn launch_metrics_server(host: String, port: u16) -> Result<()> {
77    let listener = TcpListener::bind((host, port))
78        .await
79        .context("failed to bind metrics tcp")?;
80
81    let recorder_handle = metrics::setup_metrics_recorder();
82
83    let router = Router::new().route(
84        "/metrics",
85        axum::routing::get(move || std::future::ready(recorder_handle.render())),
86    );
87
88    serve(listener, router.into_make_service())
89        .with_graceful_shutdown(shutdown_signal())
90        .await?;
91
92    Ok(())
93}
94
95async fn make_router<Db: Database>(settings: Settings) -> Result<Router, eyre::Error> {
96    let db = Db::new(&settings.db_settings)
97        .await
98        .wrap_err_with(|| format!("failed to connect to db: {:?}", settings.db_settings))?;
99    let r = router::router(db, settings);
100    Ok(r)
101}
atuin_server/lib.rs

atuin_server/
lib.rs
