Skip to main content

aion_server/observability/
metrics.rs

1//! Prometheus metrics registry and recording helpers.
2
3use std::sync::Arc;
4use std::time::Duration;
5
6use axum::http::header::CONTENT_TYPE;
7use axum::http::{HeaderValue, StatusCode};
8use axum::response::{IntoResponse, Response};
9use prometheus::{
10    Encoder, HistogramOpts, HistogramVec, IntCounterVec, IntGaugeVec, Opts, Registry, TextEncoder,
11};
12use thiserror::Error;
13
14const TEXT_FORMAT: &str = "text/plain; version=0.0.4; charset=utf-8";
15
16/// Prometheus registry construction or exposition error.
17#[derive(Debug, Error)]
18pub enum MetricsError {
19    /// A metric failed to register in the prometheus registry.
20    #[error("failed to register prometheus metric: {0}")]
21    Register(#[from] prometheus::Error),
22    /// The prometheus text encoder failed to encode gathered metric families.
23    #[error("failed to encode prometheus metrics: {0}")]
24    Encode(String),
25}
26
27/// Cloneable server metrics handle backed by a prometheus registry.
28#[derive(Clone, Debug)]
29pub struct Metrics {
30    inner: Arc<MetricsInner>,
31}
32
33#[derive(Debug)]
34struct MetricsInner {
35    registry: Registry,
36    workflows_started: IntCounterVec,
37    workflows_completed: IntCounterVec,
38    activities_dispatched: IntCounterVec,
39    activities_completed: IntCounterVec,
40    activity_duration: HistogramVec,
41    store_operation_duration: HistogramVec,
42    connected_workers: IntGaugeVec,
43    inflight_activities: IntGaugeVec,
44    signals_delivered: IntCounterVec,
45    schedules_fired: IntCounterVec,
46    deploy_operations: IntCounterVec,
47    deploy_denied: IntCounterVec,
48    loaded_workflow_versions: IntGaugeVec,
49}
50
51impl MetricsInner {
52    fn register_collectors(&self) -> Result<(), prometheus::Error> {
53        self.registry
54            .register(Box::new(self.workflows_started.clone()))?;
55        self.registry
56            .register(Box::new(self.workflows_completed.clone()))?;
57        self.registry
58            .register(Box::new(self.activities_dispatched.clone()))?;
59        self.registry
60            .register(Box::new(self.activities_completed.clone()))?;
61        self.registry
62            .register(Box::new(self.activity_duration.clone()))?;
63        self.registry
64            .register(Box::new(self.store_operation_duration.clone()))?;
65        self.registry
66            .register(Box::new(self.connected_workers.clone()))?;
67        self.registry
68            .register(Box::new(self.inflight_activities.clone()))?;
69        self.registry
70            .register(Box::new(self.signals_delivered.clone()))?;
71        self.registry
72            .register(Box::new(self.schedules_fired.clone()))?;
73        self.registry
74            .register(Box::new(self.deploy_operations.clone()))?;
75        self.registry
76            .register(Box::new(self.deploy_denied.clone()))?;
77        self.registry
78            .register(Box::new(self.loaded_workflow_versions.clone()))?;
79        Ok(())
80    }
81}
82
83impl Metrics {
84    /// Construct the server metrics registry and register all exported metrics.
85    ///
86    /// # Errors
87    ///
88    /// Returns [`MetricsError::Register`] if prometheus rejects a metric descriptor.
89    pub fn new() -> Result<Self, MetricsError> {
90        let inner = build_metrics_inner()?;
91        inner.register_collectors()?;
92        initialize_default_label_sets(&inner);
93        Ok(Self {
94            inner: Arc::new(inner),
95        })
96    }
97
98    /// Encode all currently gathered metrics in prometheus text exposition format.
99    ///
100    /// # Errors
101    ///
102    /// Returns [`MetricsError::Encode`] if prometheus cannot encode gathered metrics.
103    pub fn encode(&self) -> Result<Vec<u8>, MetricsError> {
104        let encoder = TextEncoder::new();
105        let families = self.inner.registry.gather();
106        let mut buffer = Vec::new();
107        encoder
108            .encode(&families, &mut buffer)
109            .map_err(|error| MetricsError::Encode(error.to_string()))?;
110        Ok(buffer)
111    }
112
113    /// Increment the workflow-start counter.
114    pub fn workflow_started(&self, namespace: &str, workflow_type: &str) {
115        self.inner
116            .workflows_started
117            .with_label_values(&[namespace, workflow_type])
118            .inc();
119    }
120
121    /// Increment the workflow-terminal counter.
122    pub fn workflow_completed(&self, namespace: &str, status: &str) {
123        self.inner
124            .workflows_completed
125            .with_label_values(&[namespace, status])
126            .inc();
127    }
128
129    /// Increment the activity-dispatch counter and in-flight gauge.
130    pub fn activity_dispatched(&self, namespace: &str, activity_type: &str) {
131        self.inner
132            .activities_dispatched
133            .with_label_values(&[namespace, activity_type])
134            .inc();
135        self.inner
136            .inflight_activities
137            .with_label_values(&[namespace])
138            .inc();
139    }
140
141    /// Increment the activity-completion counter, observe duration, and decrement in-flight gauge.
142    pub fn activity_completed(
143        &self,
144        namespace: &str,
145        activity_type: &str,
146        outcome: &str,
147        duration: Duration,
148    ) {
149        self.inner
150            .activities_completed
151            .with_label_values(&[namespace, outcome])
152            .inc();
153        self.inner
154            .activity_duration
155            .with_label_values(&[namespace, activity_type])
156            .observe(duration.as_secs_f64());
157        self.inner
158            .inflight_activities
159            .with_label_values(&[namespace])
160            .dec();
161    }
162
163    /// Decrement in-flight activity gauge when dispatch fails before a result can arrive.
164    pub fn activity_abandoned(&self, namespace: &str) {
165        self.inner
166            .inflight_activities
167            .with_label_values(&[namespace])
168            .dec();
169    }
170
171    /// Observe a store operation duration.
172    pub fn store_operation(&self, operation: &str, duration: Duration) {
173        self.inner
174            .store_operation_duration
175            .with_label_values(&[operation])
176            .observe(duration.as_secs_f64());
177    }
178
179    /// Increment connected worker gauge for a namespace.
180    pub fn worker_connected(&self, namespace: &str) {
181        self.inner
182            .connected_workers
183            .with_label_values(&[namespace])
184            .inc();
185    }
186
187    /// Decrement connected worker gauge for a namespace.
188    pub fn worker_disconnected(&self, namespace: &str) {
189        self.inner
190            .connected_workers
191            .with_label_values(&[namespace])
192            .dec();
193    }
194
195    /// Increment signal delivery counter.
196    pub fn signal_delivered(&self, namespace: &str, residency: &str) {
197        self.inner
198            .signals_delivered
199            .with_label_values(&[namespace, residency])
200            .inc();
201    }
202
203    /// Increment schedule-fired counter.
204    pub fn schedule_fired(&self, namespace: &str) {
205        self.inner
206            .schedules_fired
207            .with_label_values(&[namespace])
208            .inc();
209    }
210
211    /// Increment the deploy-operation counter for one mutation outcome.
212    pub fn deploy_operation(&self, operation: &str, outcome: &str) {
213        self.inner
214            .deploy_operations
215            .with_label_values(&[operation, outcome])
216            .inc();
217    }
218
219    /// Increment the deploy-denied counter for a transport.
220    pub fn deploy_denied(&self, transport: &str) {
221        self.inner
222            .deploy_denied
223            .with_label_values(&[transport])
224            .inc();
225    }
226
227    /// Set the loaded-version gauge for one workflow type from the
228    /// post-operation listing.
229    pub fn set_loaded_workflow_versions(&self, workflow_type: &str, count: i64) {
230        self.inner
231            .loaded_workflow_versions
232            .with_label_values(&[workflow_type])
233            .set(count);
234    }
235}
236
237fn build_metrics_inner() -> Result<MetricsInner, MetricsError> {
238    let registry = Registry::new();
239    let workflows_started = IntCounterVec::new(
240        Opts::new(
241            "aion_workflows_started_total",
242            "Total workflow executions started by namespace and workflow type.",
243        ),
244        &["namespace", "workflow_type"],
245    )?;
246    let workflows_completed = IntCounterVec::new(
247        Opts::new(
248            "aion_workflows_completed_total",
249            "Total workflow executions that reached a terminal status by namespace and status.",
250        ),
251        &["namespace", "status"],
252    )?;
253    let activities_dispatched = IntCounterVec::new(
254        Opts::new(
255            "aion_activities_dispatched_total",
256            "Total activities dispatched to workers by namespace and activity type.",
257        ),
258        &["namespace", "activity_type"],
259    )?;
260    let activities_completed = IntCounterVec::new(
261        Opts::new(
262            "aion_activities_completed_total",
263            "Total activity results received by namespace and outcome.",
264        ),
265        &["namespace", "outcome"],
266    )?;
267    let activity_duration = HistogramVec::new(
268        HistogramOpts::new(
269            "aion_activity_duration_seconds",
270            "Wall-clock activity execution latency from dispatch to result by namespace and activity type.",
271        )
272        .buckets(vec![
273            0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0,
274        ]),
275        &["namespace", "activity_type"],
276    )?;
277    let store_operation_duration = HistogramVec::new(
278        HistogramOpts::new(
279            "aion_store_operation_duration_seconds",
280            "Store operation latency by operation.",
281        )
282        .buckets(vec![
283            0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0,
284        ]),
285        &["operation"],
286    )?;
287    let connected_workers = IntGaugeVec::new(
288        Opts::new(
289            "aion_connected_workers",
290            "Current connected worker streams by namespace.",
291        ),
292        &["namespace"],
293    )?;
294    let inflight_activities = IntGaugeVec::new(
295        Opts::new(
296            "aion_inflight_activities",
297            "Current dispatched activities awaiting worker completion by namespace.",
298        ),
299        &["namespace"],
300    )?;
301    let signals_delivered = IntCounterVec::new(
302        Opts::new(
303            "aion_signals_delivered_total",
304            "Total signals delivered by namespace and residency classification.",
305        ),
306        &["namespace", "residency"],
307    )?;
308    let schedules_fired = IntCounterVec::new(
309        Opts::new(
310            "aion_schedules_fired_total",
311            "Total schedule timer evaluations that started a workflow by namespace.",
312        ),
313        &["namespace"],
314    )?;
315    let (deploy_operations, deploy_denied, loaded_workflow_versions) = build_deploy_metrics()?;
316
317    Ok(MetricsInner {
318        registry,
319        workflows_started,
320        workflows_completed,
321        activities_dispatched,
322        activities_completed,
323        activity_duration,
324        store_operation_duration,
325        connected_workers,
326        inflight_activities,
327        signals_delivered,
328        schedules_fired,
329        deploy_operations,
330        deploy_denied,
331        loaded_workflow_versions,
332    })
333}
334
335/// Deploy API collectors: mutation counter, denial counter, and the
336/// loaded-version gauge fed from the post-operation listing.
337fn build_deploy_metrics() -> Result<(IntCounterVec, IntCounterVec, IntGaugeVec), MetricsError> {
338    let deploy_operations = IntCounterVec::new(
339        Opts::new(
340            "aion_deploy_operations_total",
341            "Total deploy API mutations by operation and outcome class.",
342        ),
343        &["operation", "outcome"],
344    )?;
345    let deploy_denied = IntCounterVec::new(
346        Opts::new(
347            "aion_deploy_denied_total",
348            "Total deploy API authorization denials by transport.",
349        ),
350        &["transport"],
351    )?;
352    let loaded_workflow_versions = IntGaugeVec::new(
353        Opts::new(
354            "aion_loaded_workflow_versions",
355            "Currently loaded package versions per workflow type.",
356        ),
357        &["workflow_type"],
358    )?;
359    Ok((deploy_operations, deploy_denied, loaded_workflow_versions))
360}
361
362/// Pre-initialize known label sets so all metric families appear in the
363/// prometheus text output before any workflow or activity traffic occurs.
364fn initialize_default_label_sets(inner: &MetricsInner) {
365    for operation in ["append", "read_history", "list_active", "list_workflow_ids"] {
366        inner
367            .store_operation_duration
368            .with_label_values(&[operation]);
369    }
370    inner
371        .activity_duration
372        .with_label_values(&["default", "default"]);
373}
374
375/// Axum handler for `/metrics`.
376pub async fn metrics_handler(
377    axum::extract::State(metrics): axum::extract::State<Metrics>,
378) -> Response {
379    match metrics.encode() {
380        Ok(body) => {
381            let mut response = body.into_response();
382            response
383                .headers_mut()
384                .insert(CONTENT_TYPE, HeaderValue::from_static(TEXT_FORMAT));
385            response
386        }
387        Err(error) => (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()).into_response(),
388    }
389}