Skip to main content

faucet_cli/serve/
state.rs

1//! Shared, cheaply-cloneable server state handed to every handler via
2//! `axum::extract::State`. Holds auth, the Prometheus render handle, the
3//! server-wide shutdown token, the run registry, the execution semaphore, the
4//! run-history backend, and the `--default-config` merge base.
5
6use crate::serve::config::{AuthMode, ServeConfig};
7use crate::serve::history::RunHistory;
8use crate::serve::logs::LogHub;
9use crate::serve::registry::Registry;
10use metrics_exporter_prometheus::PrometheusHandle;
11use serde_json::Value;
12use std::sync::Arc;
13use std::time::Duration;
14use tokio::sync::Semaphore;
15use tokio_util::sync::CancellationToken;
16
17#[derive(Clone)]
18pub struct ServerState {
19    inner: Arc<Inner>,
20}
21
22struct Inner {
23    auth: AuthMode,
24    prometheus: Option<PrometheusHandle>,
25    shutdown: CancellationToken,
26    registry: Registry,
27    semaphore: Arc<Semaphore>,
28    history: Arc<dyn RunHistory>,
29    log_hub: LogHub,
30    default_base: Option<Value>,
31    idempotency_retention: Duration,
32    probe_timeout: Duration,
33}
34
35impl ServerState {
36    #[allow(clippy::too_many_arguments)]
37    pub fn new(
38        config: &ServeConfig,
39        prometheus: Option<PrometheusHandle>,
40        shutdown: CancellationToken,
41        history: Arc<dyn RunHistory>,
42        log_hub: LogHub,
43        default_base: Option<Value>,
44    ) -> Self {
45        Self {
46            inner: Arc::new(Inner {
47                auth: config.auth.clone(),
48                prometheus,
49                shutdown,
50                registry: Registry::new(config.max_queued_runs),
51                semaphore: Arc::new(Semaphore::new(config.max_concurrent_runs)),
52                history,
53                log_hub,
54                default_base,
55                idempotency_retention: config.idempotency_retention,
56                probe_timeout: config.probe_timeout,
57            }),
58        }
59    }
60
61    pub fn auth_token(&self) -> Option<&str> {
62        match &self.inner.auth {
63            AuthMode::Token(t) => Some(t),
64            AuthMode::None => None,
65        }
66    }
67
68    pub fn render_metrics(&self) -> Option<String> {
69        self.inner.prometheus.as_ref().map(|h| h.render())
70    }
71
72    pub fn shutdown_token(&self) -> CancellationToken {
73        self.inner.shutdown.clone()
74    }
75
76    pub fn registry(&self) -> &Registry {
77        &self.inner.registry
78    }
79
80    pub fn semaphore(&self) -> Arc<Semaphore> {
81        Arc::clone(&self.inner.semaphore)
82    }
83
84    pub fn history(&self) -> Arc<dyn RunHistory> {
85        Arc::clone(&self.inner.history)
86    }
87
88    /// The per-run log buffer registry shared with the tracing [`LogHub`] layer.
89    pub fn log_hub(&self) -> &LogHub {
90        &self.inner.log_hub
91    }
92
93    pub fn default_base(&self) -> Option<&Value> {
94        self.inner.default_base.as_ref()
95    }
96
97    pub fn idempotency_retention(&self) -> Duration {
98        self.inner.idempotency_retention
99    }
100
101    pub fn probe_timeout(&self) -> Duration {
102        self.inner.probe_timeout
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109    use crate::serve::config::HistoryBackendSpec;
110    use crate::serve::history::memory::MemoryHistory;
111
112    fn cfg(auth: AuthMode) -> ServeConfig {
113        ServeConfig {
114            listen: "127.0.0.1:0".parse().unwrap(),
115            auth,
116            max_concurrent_runs: 4,
117            max_queued_runs: 32,
118            default_config_path: None,
119            history: HistoryBackendSpec::Memory,
120            cors_origins: vec![],
121            body_limit_bytes: 1_048_576,
122            shutdown_grace: Duration::from_secs(60),
123            retain_terminal_runs: Duration::from_secs(60),
124            idempotency_retention: Duration::from_secs(60),
125            lease_ttl: Duration::from_secs(30),
126            probe_timeout: Duration::from_secs(10),
127            env_file: None,
128            no_env_file: false,
129            log_level: "info".into(),
130        }
131    }
132
133    fn state(auth: AuthMode) -> ServerState {
134        use crate::serve::logs::LogHub;
135        let history = Arc::new(MemoryHistory::new(Duration::from_secs(60))) as Arc<dyn RunHistory>;
136        ServerState::new(
137            &cfg(auth),
138            None,
139            CancellationToken::new(),
140            history,
141            LogHub::new(),
142            None,
143        )
144    }
145
146    #[test]
147    fn auth_token_reflects_mode() {
148        assert_eq!(state(AuthMode::Token("x".into())).auth_token(), Some("x"));
149        assert_eq!(state(AuthMode::None).auth_token(), None);
150    }
151
152    #[test]
153    fn render_metrics_none_without_handle() {
154        assert!(state(AuthMode::None).render_metrics().is_none());
155    }
156
157    #[test]
158    fn registry_starts_empty() {
159        let s = state(AuthMode::None);
160        assert_eq!(s.registry().queued(), 0);
161        assert_eq!(s.registry().in_flight(), 0);
162    }
163}