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