1use core::time::Duration;
4
5use metrics_exporter_prometheus::BuildError;
6use tokio::net::TcpListener;
7
8#[derive(Debug, thiserror::Error)]
10#[non_exhaustive]
11pub enum ServerError {
12 #[error(transparent)]
14 Listen(#[from] std::io::Error),
15
16 #[error("failed to initialize Prometheus metrics recorder: {0}")]
18 MetricsInit(#[source] BuildError),
19
20 #[error("failed to build tokio runtime: {0}")]
22 RuntimeBuild(#[source] std::io::Error),
23}
24
25pub struct ServerConfig {
27 pub host: String,
30 pub port: u16,
32}
33
34impl ServerConfig {
35 #[inline]
37 #[must_use]
38 pub fn new<Host>(host: Host, port: u16) -> Self
39 where
40 Host: Into<String>,
41 {
42 Self {
43 host: host.into(),
44 port,
45 }
46 }
47}
48
49#[inline]
63pub async fn bind(config: &ServerConfig) -> std::io::Result<TcpListener> {
64 TcpListener::bind((config.host.as_str(), config.port)).await
65}
66
67#[inline]
80pub async fn serve(config: ServerConfig) -> Result<(), ServerError> {
81 let listener = bind(&config).await?;
82 let bound_addr = listener.local_addr()?;
83
84 let handle = crate::metrics::install_global().map_err(ServerError::MetricsInit)?;
85 let upkeep_handle = handle.clone();
86 let upkeep_task = tokio::spawn(async move {
87 let mut interval = tokio::time::interval(Duration::from_secs(5));
88 interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
89 loop {
90 interval.tick().await;
91 upkeep_handle.run_upkeep();
92 }
93 });
94
95 tracing::info!(addr = %bound_addr, "docspec-http listening");
96
97 let server_result = axum::serve(listener, crate::router::router_with_metrics(handle))
98 .with_graceful_shutdown(shutdown_signal())
99 .await;
100
101 upkeep_task.abort();
102 match upkeep_task.await {
103 Err(error) if error.is_cancelled() => {}
104 Err(error) => tracing::error!(%error, "metrics upkeep task failed during shutdown"),
105 Ok(()) => {}
106 }
107 server_result?;
108
109 Ok(())
110}
111
112#[allow(clippy::single_call_fn)]
116#[inline]
117async fn shutdown_signal() {
118 use core::future;
119
120 use tokio::signal;
121
122 let ctrl_c = async {
123 if let Err(error) = signal::ctrl_c().await {
124 tracing::error!(%error, "failed to install Ctrl+C handler");
125 future::pending::<()>().await;
126 }
127 };
128
129 #[cfg(unix)]
130 let terminate = async {
131 match signal::unix::signal(signal::unix::SignalKind::terminate()) {
132 Ok(mut stream) => {
133 stream.recv().await;
134 }
135 Err(error) => {
136 tracing::error!(%error, "failed to install SIGTERM handler");
137 future::pending::<()>().await;
138 }
139 }
140 };
141
142 #[cfg(not(unix))]
143 let terminate = future::pending::<()>();
144
145 tokio::select! {
146 () = ctrl_c => {
147 tracing::info!("shutdown signal received, draining");
148 },
149 () = terminate => {
150 tracing::info!("shutdown signal received, draining");
151 },
152 }
153}