Skip to main content

espresso_telemetry/
lib.rs

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