Skip to main content

camel_prometheus/
service.rs

1use std::net::SocketAddr;
2use std::sync::Arc;
3use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU16, Ordering};
4
5use crate::PrometheusMetrics;
6use async_trait::async_trait;
7use camel_api::{CamelError, HealthSource, Lifecycle, MetricsCollector, ServiceStatus};
8use tokio::sync::oneshot;
9use tokio::task::JoinHandle;
10use tokio::time::{Duration, timeout};
11use tracing::{debug, info, warn};
12
13pub struct PrometheusService {
14    addr: SocketAddr,
15    metrics: Arc<PrometheusMetrics>,
16    server_handle: Option<JoinHandle<()>>,
17    shutdown_tx: Option<oneshot::Sender<()>>,
18    /// The actual bound port (set after start(), 0 before)
19    bound_port: Arc<AtomicU16>,
20    /// Current service status (Stopped=0, Started=1, Failed=2)
21    status: Arc<AtomicU8>,
22    health_source: Option<Arc<dyn HealthSource>>,
23    /// Guard to prevent double-start
24    started: AtomicBool,
25}
26
27impl PrometheusService {
28    pub fn new(addr: SocketAddr) -> Self {
29        Self {
30            addr,
31            metrics: Arc::new(PrometheusMetrics::new()),
32            server_handle: None,
33            shutdown_tx: None,
34            bound_port: Arc::new(AtomicU16::new(0)),
35            status: Arc::new(AtomicU8::new(0)),
36            health_source: None,
37            started: AtomicBool::new(false),
38        }
39    }
40
41    /// Returns the actual port the server is listening on (after start())
42    ///
43    /// Returns 0 if the service hasn't been started yet.
44    pub fn port(&self) -> u16 {
45        self.bound_port.load(Ordering::SeqCst)
46    }
47
48    /// Returns a cloneable accessor for the bound port.
49    ///
50    /// Use this when you need to read the port after the service has been
51    /// moved into a CamelContext or other container.
52    pub fn port_accessor(&self) -> Arc<AtomicU16> {
53        Arc::clone(&self.bound_port)
54    }
55
56    pub fn status_arc(&self) -> Arc<AtomicU8> {
57        Arc::clone(&self.status)
58    }
59
60    pub fn set_health_source(&mut self, source: Arc<dyn HealthSource>) {
61        self.health_source = Some(source);
62    }
63
64    pub fn health_source(&self) -> Option<Arc<dyn HealthSource>> {
65        self.health_source.clone()
66    }
67}
68
69#[async_trait]
70impl Lifecycle for PrometheusService {
71    fn name(&self) -> &str {
72        "prometheus"
73    }
74
75    fn as_metrics_collector(&self) -> Option<Arc<dyn MetricsCollector>> {
76        Some(Arc::clone(&self.metrics) as Arc<dyn MetricsCollector>)
77    }
78
79    fn status(&self) -> ServiceStatus {
80        match self.status.load(Ordering::SeqCst) {
81            0 => ServiceStatus::Stopped,
82            1 => ServiceStatus::Started,
83            2 => ServiceStatus::Failed,
84            _ => ServiceStatus::Failed,
85        }
86    }
87
88    async fn start(&mut self) -> Result<(), CamelError> {
89        use tokio::net::TcpListener;
90
91        if self
92            .started
93            .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
94            .is_err()
95        {
96            return Err(CamelError::Config(
97                "PrometheusService already started".to_string(),
98            ));
99        }
100
101        let listener = TcpListener::bind(self.addr).await.map_err(|e| {
102            self.status.store(2, Ordering::SeqCst);
103            self.started.store(false, Ordering::SeqCst);
104            CamelError::Io(e.to_string())
105        })?;
106
107        let actual_port = listener.local_addr().map(|addr| addr.port()).map_err(|e| {
108            self.status.store(2, Ordering::SeqCst);
109            self.started.store(false, Ordering::SeqCst);
110            CamelError::Io(e.to_string())
111        })?;
112
113        self.bound_port.store(actual_port, Ordering::SeqCst);
114
115        if !self.addr.ip().is_loopback() {
116            warn!(
117                addr = %self.addr,
118                "prometheus metrics endpoint bound to non-loopback address; endpoint is reachable from all interfaces without application-layer restriction (ADR-0052)"
119            );
120        }
121
122        self.status.store(1, Ordering::SeqCst);
123
124        let metrics = Arc::clone(&self.metrics);
125        let health_source = self.health_source.clone();
126        let status = Arc::clone(&self.status);
127        let (shutdown_tx, shutdown_rx) = oneshot::channel();
128
129        let handle = tokio::spawn(async move {
130            if let Err(err) =
131                crate::MetricsServer::run_with_listener_and_health_source_with_shutdown(
132                    listener,
133                    metrics,
134                    health_source,
135                    shutdown_rx,
136                )
137                .await
138            {
139                status.store(2, Ordering::SeqCst);
140                warn!("prometheus metrics server exited with error: {err}");
141            }
142        });
143
144        self.shutdown_tx = Some(shutdown_tx);
145        self.server_handle = Some(handle);
146        info!(port = %actual_port, "prometheus metrics service started");
147        Ok(())
148    }
149
150    async fn stop(&mut self) -> Result<(), CamelError> {
151        if let Some(tx) = self.shutdown_tx.take() {
152            let _ = tx.send(());
153        }
154
155        if let Some(handle) = self.server_handle.take() {
156            let mut handle = handle;
157            match timeout(Duration::from_secs(5), &mut handle).await {
158                Ok(join_result) => {
159                    if let Err(e) = join_result {
160                        return Err(CamelError::Io(format!(
161                            "prometheus server task join failed: {e}"
162                        )));
163                    }
164                }
165                Err(_) => {
166                    debug!("prometheus server shutdown timed out; aborting task");
167                    handle.abort();
168                    let _ = handle.await;
169                }
170            }
171        }
172        self.status.store(0, Ordering::SeqCst);
173        self.started.store(false, Ordering::SeqCst);
174        debug!("prometheus metrics service stopped");
175        Ok(())
176    }
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182    use camel_api::HealthStatus;
183    use std::net::{IpAddr, Ipv4Addr};
184    use std::sync::atomic::Ordering;
185
186    struct MockHealthSource {
187        readiness: HealthStatus,
188    }
189
190    #[async_trait]
191    impl HealthSource for MockHealthSource {
192        async fn liveness(&self) -> HealthStatus {
193            HealthStatus::Healthy
194        }
195
196        async fn readiness(&self) -> HealthStatus {
197            self.readiness
198        }
199
200        async fn startup(&self) -> HealthStatus {
201            HealthStatus::Healthy
202        }
203    }
204
205    #[test]
206    fn test_create_prometheus_service() {
207        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9090);
208        let service = PrometheusService::new(addr);
209        assert_eq!(service.name(), "prometheus");
210    }
211
212    #[tokio::test]
213    async fn test_prometheus_service_status_transitions() {
214        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0);
215        let mut service = PrometheusService::new(addr);
216
217        assert_eq!(service.status(), ServiceStatus::Stopped);
218
219        service.start().await.unwrap();
220        assert_eq!(service.status(), ServiceStatus::Started);
221
222        service.stop().await.unwrap();
223        assert_eq!(service.status(), ServiceStatus::Stopped);
224    }
225
226    #[tokio::test]
227    async fn test_stop_uses_graceful_shutdown_signal() {
228        use std::sync::atomic::Ordering as AtomicOrdering;
229
230        crate::server::test_reset_graceful_shutdown_observability();
231
232        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0);
233        let mut service = PrometheusService::new(addr);
234        service.start().await.unwrap();
235
236        service.stop().await.unwrap();
237
238        let signal_count = crate::server::test_graceful_shutdown_signal_count();
239        assert!(
240            signal_count.load(AtomicOrdering::SeqCst) >= 1,
241            "expected at least one graceful shutdown signal"
242        );
243    }
244
245    #[tokio::test]
246    async fn test_port_and_port_accessor_after_start() {
247        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0);
248        let mut service = PrometheusService::new(addr);
249
250        assert_eq!(service.port(), 0);
251        let accessor = service.port_accessor();
252        assert_eq!(accessor.load(Ordering::SeqCst), 0);
253
254        service.start().await.unwrap();
255        let port = service.port();
256        assert!(port > 0);
257        assert_eq!(accessor.load(Ordering::SeqCst), port);
258
259        service.stop().await.unwrap();
260    }
261
262    #[test]
263    fn test_as_metrics_collector_returns_metrics_instance() {
264        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 9090);
265        let service = PrometheusService::new(addr);
266        let collector = service.as_metrics_collector().unwrap();
267        collector.increment_exchanges("route-a");
268        let output = service.metrics.gather();
269        assert!(output.contains("camel_exchanges_total"));
270        assert!(output.contains("route-a"));
271    }
272
273    #[test]
274    fn test_unknown_internal_status_maps_to_failed() {
275        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 9090);
276        let service = PrometheusService::new(addr);
277        service.status_arc().store(9, Ordering::SeqCst);
278        assert_eq!(service.status(), ServiceStatus::Failed);
279    }
280
281    #[tokio::test]
282    async fn test_health_source_injection() {
283        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9090);
284        let mut service = PrometheusService::new(addr);
285
286        assert!(service.health_source().is_none());
287
288        let source = Arc::new(MockHealthSource {
289            readiness: HealthStatus::Healthy,
290        });
291
292        service.set_health_source(source);
293        assert!(service.health_source().is_some());
294
295        let status = service.health_source().unwrap().readiness().await;
296        assert_eq!(status, HealthStatus::Healthy);
297    }
298
299    #[test]
300    fn test_prometheus_service_with_socket_addr() {
301        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 9091);
302        let service = PrometheusService::new(addr);
303        assert_eq!(service.name(), "prometheus");
304    }
305
306    #[tokio::test]
307    async fn test_double_start_returns_error() {
308        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0);
309        let mut service = PrometheusService::new(addr);
310
311        service.start().await.unwrap();
312        assert_eq!(service.status(), ServiceStatus::Started);
313
314        let result = service.start().await;
315        assert!(result.is_err(), "second start() should return an error");
316        let err_msg = result.unwrap_err().to_string();
317        assert!(
318            err_msg.contains("already started"),
319            "error should mention 'already started', got: {err_msg}"
320        );
321
322        service.stop().await.unwrap();
323    }
324
325    #[tokio::test]
326    async fn test_start_allowed_again_after_stop() {
327        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0);
328        let mut service = PrometheusService::new(addr);
329
330        service.start().await.unwrap();
331        service.stop().await.unwrap();
332
333        // After stop, start should succeed again
334        service.start().await.unwrap();
335        assert_eq!(service.status(), ServiceStatus::Started);
336
337        service.stop().await.unwrap();
338    }
339
340    #[tracing_test::traced_test]
341    #[tokio::test]
342    async fn loopback_bind_emits_no_warning() {
343        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 0);
344        let mut service = PrometheusService::new(addr);
345
346        assert!(service.start().await.is_ok());
347        assert!(!logs_contain("non-loopback"));
348
349        service.stop().await.unwrap();
350    }
351
352    #[tracing_test::traced_test]
353    #[tokio::test]
354    async fn non_loopback_bind_emits_warning() {
355        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0);
356        let mut service = PrometheusService::new(addr);
357
358        assert!(service.start().await.is_ok());
359        assert!(logs_contain("non-loopback"));
360
361        service.stop().await.unwrap();
362    }
363
364    #[test]
365    fn server_task_error_sets_status_failed() {
366        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0);
367        let service = PrometheusService::new(addr);
368
369        // Simulate server-task failure: the spawned closure would store
370        // status=2 (Failed) on Err. We reproduce that via the public
371        // status_arc() accessor.
372        let status = service.status_arc();
373        status.store(2, Ordering::SeqCst);
374
375        assert_eq!(service.status(), ServiceStatus::Failed);
376    }
377
378    #[tokio::test]
379    async fn clean_shutdown_does_not_set_failed() {
380        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0);
381        let mut service = PrometheusService::new(addr);
382
383        service.start().await.unwrap();
384        service.stop().await.unwrap();
385
386        assert_ne!(service.status(), ServiceStatus::Failed);
387        assert_eq!(service.status(), ServiceStatus::Stopped);
388    }
389
390    #[tokio::test]
391    async fn status_started_before_spawn() {
392        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0);
393        let mut service = PrometheusService::new(addr);
394
395        service.start().await.unwrap();
396        assert_eq!(service.status(), ServiceStatus::Started);
397
398        service.stop().await.unwrap();
399    }
400}