Skip to main content

espresso_telemetry/
lib.rs

1mod lifecycle;
2mod push_task;
3mod rate_limit;
4pub mod remote_write;
5pub mod token;
6
7use std::time::{SystemTime, UNIX_EPOCH};
8
9pub use lifecycle::{TelemetryHandle, TelemetryOptions, init, registry, set_registry};
10use prometheus::proto::{LabelPair, MetricFamily, MetricType};
11use prost::Message;
12use remote_write::{Label, Sample, TimeSeries, WriteRequest};
13pub use token::{
14    Token, TokenParseError, TokenVerifyError, UnauthenticatedToken, load_bls_signing_key,
15    parse_bls_signing_key,
16};
17
18const NAME_LABEL: &str = "__name__";
19
20pub fn build_write_request(families: &[MetricFamily]) -> anyhow::Result<WriteRequest> {
21    let timestamp_ms = now_unix_millis()?;
22    let mut series: Vec<TimeSeries> = Vec::new();
23
24    for family in families {
25        let name = family.get_name();
26        match family.get_field_type() {
27            MetricType::COUNTER | MetricType::GAUGE => {
28                for metric in family.get_metric() {
29                    let value = if family.get_field_type() == MetricType::COUNTER {
30                        metric.get_counter().get_value()
31                    } else {
32                        metric.get_gauge().get_value()
33                    };
34                    series.push(make_series(name, metric.get_label(), value, timestamp_ms));
35                }
36            },
37            MetricType::HISTOGRAM => {
38                let bucket_name = format!("{name}_bucket");
39                let sum_name = format!("{name}_sum");
40                let count_name = format!("{name}_count");
41                for metric in family.get_metric() {
42                    let h = metric.get_histogram();
43                    for bucket in h.get_bucket() {
44                        // Skip any user-configured +Inf bucket. The unconditional
45                        // emit below covers it via sample_count, so emitting both
46                        // would collide on `(name, le="+Inf")` at the receiver.
47                        if bucket.get_upper_bound().is_infinite() {
48                            continue;
49                        }
50                        series.push(make_bucket_series(
51                            &bucket_name,
52                            metric.get_label(),
53                            &format_float(bucket.get_upper_bound()),
54                            bucket.get_cumulative_count() as f64,
55                            timestamp_ms,
56                        ));
57                    }
58                    series.push(make_bucket_series(
59                        &bucket_name,
60                        metric.get_label(),
61                        "+Inf",
62                        h.get_sample_count() as f64,
63                        timestamp_ms,
64                    ));
65                    series.push(make_series(
66                        &sum_name,
67                        metric.get_label(),
68                        h.get_sample_sum(),
69                        timestamp_ms,
70                    ));
71                    series.push(make_series(
72                        &count_name,
73                        metric.get_label(),
74                        h.get_sample_count() as f64,
75                        timestamp_ms,
76                    ));
77                }
78            },
79            other => {
80                tracing::warn!(
81                    "telemetry: skipping unsupported metric type {other:?} for family {name}"
82                );
83            },
84        }
85    }
86
87    Ok(WriteRequest {
88        timeseries: series,
89        ..Default::default()
90    })
91}
92
93pub fn encode_to_snappy(req: &WriteRequest) -> anyhow::Result<Vec<u8>> {
94    let buf = req.encode_to_vec();
95    snap::raw::Encoder::new()
96        .compress_vec(&buf)
97        .map_err(|e| anyhow::anyhow!("snappy compress: {e}"))
98}
99
100fn make_series(name: &str, label_pairs: &[LabelPair], value: f64, timestamp_ms: i64) -> TimeSeries {
101    let labels = label_pairs_to_labels(label_pairs);
102    named_series(name, labels, value, timestamp_ms)
103}
104
105fn make_bucket_series(
106    name: &str,
107    label_pairs: &[LabelPair],
108    le: &str,
109    value: f64,
110    timestamp_ms: i64,
111) -> TimeSeries {
112    let mut labels = label_pairs_to_labels(label_pairs);
113    labels.push(Label {
114        name: "le".into(),
115        value: le.into(),
116    });
117    named_series(name, labels, value, timestamp_ms)
118}
119
120fn named_series(name: &str, mut labels: Vec<Label>, value: f64, timestamp_ms: i64) -> TimeSeries {
121    labels.push(Label {
122        name: NAME_LABEL.into(),
123        value: name.into(),
124    });
125    // Prometheus remote-write 1.0 requires labels sorted by name within each
126    // TimeSeries. Receivers are typically lenient but the spec is explicit.
127    labels.sort_by(|a, b| a.name.cmp(&b.name));
128    TimeSeries {
129        labels,
130        samples: vec![Sample {
131            value,
132            timestamp: timestamp_ms,
133        }],
134        ..Default::default()
135    }
136}
137
138fn label_pairs_to_labels(pairs: &[LabelPair]) -> Vec<Label> {
139    pairs
140        .iter()
141        .map(|p| Label {
142            name: p.get_name().to_owned(),
143            value: p.get_value().to_owned(),
144        })
145        .collect()
146}
147
148fn format_float(v: f64) -> String {
149    if v.is_infinite() {
150        if v.is_sign_positive() {
151            "+Inf".to_string()
152        } else {
153            "-Inf".to_string()
154        }
155    } else {
156        format!("{v}")
157    }
158}
159
160fn now_unix_millis() -> anyhow::Result<i64> {
161    let d = SystemTime::now()
162        .duration_since(UNIX_EPOCH)
163        .map_err(|e| anyhow::anyhow!("system clock is before UNIX_EPOCH: {e}"))?;
164    Ok(d.as_millis() as i64)
165}