pub use prometheus_client::{self, *}; // std use std::fmt::{Debug, Error, Formatter}; use std::sync::{Arc, Mutex}; // crates use futures::StreamExt; use overwatch_rs::services::life_cycle::LifecycleMessage; use overwatch_rs::services::{ handle::ServiceStateHandle, relay::RelayMessage, state::{NoOperator, NoState}, ServiceCore, ServiceData, }; use prometheus_client::encoding::text::encode; use prometheus_client::registry::Registry; use tokio::sync::oneshot::Sender; use tracing::error; // internal // A wrapper for prometheus_client Registry. // Lock is only used during services initialization and prometheus pull query. pub type NomosRegistry = Arc>; pub struct Metrics { service_state: ServiceStateHandle, registry: NomosRegistry, } #[derive(Clone, Debug)] pub struct MetricsSettings { pub registry: Option, } pub enum MetricsMsg { Gather { reply_channel: Sender }, } impl RelayMessage for MetricsMsg {} impl Debug for MetricsMsg { fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> { match self { Self::Gather { .. } => { write!(f, "MetricsMsg::Gather") } } } } impl ServiceData for Metrics { const SERVICE_ID: &'static str = "Metrics"; type Settings = MetricsSettings; type State = NoState; type StateOperator = NoOperator; type Message = MetricsMsg; } #[async_trait::async_trait] impl ServiceCore for Metrics { fn init(service_state: ServiceStateHandle) -> Result { let config = service_state.settings_reader.get_updated_settings(); Ok(Self { service_state, registry: config.registry.ok_or("No registry provided")?, }) } async fn run(self) -> Result<(), overwatch_rs::DynError> { let Self { mut service_state, registry, } = self; let mut lifecycle_stream = service_state.lifecycle_handle.message_stream(); loop { tokio::select! { Some(msg) = service_state.inbound_relay.recv() => { let MetricsMsg::Gather{reply_channel} = msg; let mut buf = String::new(); { let reg = registry.lock().unwrap(); // If encoding fails, we need to stop trying process subsequent metrics gather // requests. If it succeds, encode method returns empty unit type. _ = encode(&mut buf, ®).map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?; } reply_channel .send(buf) .unwrap_or_else(|_| tracing::debug!("could not send back metrics")); } Some(msg) = lifecycle_stream.next() => { if Self::should_stop_service(msg).await { break; } } } } Ok(()) } } impl Metrics { async fn should_stop_service(message: LifecycleMessage) -> bool { match message { LifecycleMessage::Shutdown(sender) => { if sender.send(()).is_err() { error!( "Error sending successful shutdown signal from service {}", Self::SERVICE_ID ); } true } LifecycleMessage::Kill => true, } } }