Skip to main content

krabka_client_core/telemetry/
otlp.rs

1//! The part of the OpenTelemetry metrics protocol (OTLP) that a KIP-714
2//! client pushes.
3//!
4//! The `metrics` field of a `PushTelemetry` request holds a serialized
5//! `opentelemetry.proto.metrics.v1.MetricsData`. Kafka's client builds it
6//! with the generated protobuf classes. A client only needs number gauges and
7//! number sums with string attributes, so this module encodes and decodes that
8//! subset of the protobuf wire format by hand. Fields outside the subset are
9//! skipped on decode.
10
11use thiserror::Error;
12
13/// `opentelemetry.proto.metrics.v1.MetricsData`.
14#[derive(Clone, Debug, Default, PartialEq)]
15pub struct MetricsData {
16    /// Field 1.
17    pub resource_metrics: Vec<ResourceMetrics>,
18}
19
20/// `opentelemetry.proto.metrics.v1.ResourceMetrics`.
21#[derive(Clone, Debug, Default, PartialEq)]
22pub struct ResourceMetrics {
23    /// Field 1. Kafka always sets it, and an empty resource encodes as an
24    /// empty message.
25    pub resource: Option<Resource>,
26    /// Field 2.
27    pub scope_metrics: Vec<ScopeMetrics>,
28}
29
30/// `opentelemetry.proto.resource.v1.Resource`.
31#[derive(Clone, Debug, Default, PartialEq, Eq)]
32pub struct Resource {
33    /// Field 1.
34    pub attributes: Vec<KeyValue>,
35}
36
37/// `opentelemetry.proto.metrics.v1.ScopeMetrics`, without the
38/// instrumentation scope, which Kafka does not set.
39#[derive(Clone, Debug, Default, PartialEq)]
40pub struct ScopeMetrics {
41    /// Field 2.
42    pub metrics: Vec<Metric>,
43}
44
45/// `opentelemetry.proto.metrics.v1.Metric` with a gauge or a sum.
46#[derive(Clone, Debug, PartialEq)]
47pub struct Metric {
48    /// Field 1.
49    pub name: String,
50    /// Field 5 (`gauge`) or field 7 (`sum`). `None` for a metric of another
51    /// kind, which this module does not decode.
52    pub data: Option<MetricData>,
53}
54
55/// The data of a [`Metric`].
56#[derive(Clone, Debug, PartialEq)]
57pub enum MetricData {
58    /// `opentelemetry.proto.metrics.v1.Gauge`.
59    Gauge {
60        /// Field 1.
61        data_points: Vec<NumberDataPoint>,
62    },
63    /// `opentelemetry.proto.metrics.v1.Sum`.
64    Sum {
65        /// Field 1.
66        data_points: Vec<NumberDataPoint>,
67        /// Field 2.
68        aggregation_temporality: AggregationTemporality,
69        /// Field 3.
70        is_monotonic: bool,
71    },
72}
73
74/// `opentelemetry.proto.metrics.v1.AggregationTemporality`.
75#[derive(Clone, Copy, Debug, PartialEq, Eq)]
76pub enum AggregationTemporality {
77    /// `AGGREGATION_TEMPORALITY_UNSPECIFIED` (0).
78    Unspecified,
79    /// `AGGREGATION_TEMPORALITY_DELTA` (1).
80    Delta,
81    /// `AGGREGATION_TEMPORALITY_CUMULATIVE` (2).
82    Cumulative,
83}
84
85impl AggregationTemporality {
86    const fn number(self) -> u64 {
87        match self {
88            Self::Unspecified => 0,
89            Self::Delta => 1,
90            Self::Cumulative => 2,
91        }
92    }
93
94    const fn from_number(number: u64) -> Option<Self> {
95        match number {
96            0 => Some(Self::Unspecified),
97            1 => Some(Self::Delta),
98            2 => Some(Self::Cumulative),
99            _ => None,
100        }
101    }
102}
103
104/// `opentelemetry.proto.metrics.v1.NumberDataPoint`.
105#[derive(Clone, Debug, PartialEq)]
106pub struct NumberDataPoint {
107    /// Field 7.
108    pub attributes: Vec<KeyValue>,
109    /// Field 2. Zero when unset.
110    pub start_time_unix_nano: u64,
111    /// Field 3.
112    pub time_unix_nano: u64,
113    /// Field 4 (`as_double`) or field 6 (`as_int`).
114    pub value: Option<NumberValue>,
115}
116
117/// The `value` oneof of a [`NumberDataPoint`].
118#[derive(Clone, Copy, Debug, PartialEq)]
119pub enum NumberValue {
120    /// `as_double`.
121    Double(f64),
122    /// `as_int`.
123    Int(i64),
124}
125
126/// `opentelemetry.proto.common.v1.KeyValue` with a string value, the only
127/// kind of attribute that Kafka writes.
128#[derive(Clone, Debug, PartialEq, Eq)]
129pub struct KeyValue {
130    /// Field 1.
131    pub key: String,
132    /// The `string_value` (field 1) of the `AnyValue` in field 2.
133    pub value: String,
134}
135
136impl KeyValue {
137    /// A key and a string value.
138    #[must_use]
139    pub fn new(key: impl Into<String>, value: impl Into<String>) -> Self {
140        Self {
141            key: key.into(),
142            value: value.into(),
143        }
144    }
145}
146
147/// A payload that is not a well-formed `MetricsData`.
148#[derive(Clone, Debug, Error, PartialEq, Eq)]
149#[error("malformed OTLP metrics payload: {0}")]
150pub struct OtlpDecodeError(&'static str);
151
152const WIRE_VARINT: u8 = 0;
153const WIRE_FIXED64: u8 = 1;
154const WIRE_LEN: u8 = 2;
155const WIRE_FIXED32: u8 = 5;
156
157fn put_varint(buf: &mut Vec<u8>, mut value: u64) {
158    while value >= 0x80 {
159        buf.push(u8::try_from(value & 0x7f).unwrap_or(0) | 0x80);
160        value >>= 7;
161    }
162    buf.push(u8::try_from(value).unwrap_or(0));
163}
164
165fn put_key(buf: &mut Vec<u8>, field: u32, wire: u8) {
166    put_varint(buf, (u64::from(field) << 3) | u64::from(wire));
167}
168
169fn put_bytes(buf: &mut Vec<u8>, field: u32, bytes: &[u8]) {
170    put_key(buf, field, WIRE_LEN);
171    put_varint(buf, u64::try_from(bytes.len()).unwrap_or(u64::MAX));
172    buf.extend_from_slice(bytes);
173}
174
175fn put_message(buf: &mut Vec<u8>, field: u32, encode: impl FnOnce(&mut Vec<u8>)) {
176    let mut inner = Vec::new();
177    encode(&mut inner);
178    put_bytes(buf, field, &inner);
179}
180
181/// A proto3 string: an empty one is the default and is not written.
182fn put_string(buf: &mut Vec<u8>, field: u32, value: &str) {
183    if !value.is_empty() {
184        put_bytes(buf, field, value.as_bytes());
185    }
186}
187
188fn put_fixed64(buf: &mut Vec<u8>, field: u32, value: u64) {
189    put_key(buf, field, WIRE_FIXED64);
190    buf.extend_from_slice(&value.to_le_bytes());
191}
192
193impl MetricsData {
194    /// The protobuf encoding of this message.
195    #[must_use]
196    pub fn encode(&self) -> Vec<u8> {
197        let mut buf = Vec::new();
198        for resource_metrics in &self.resource_metrics {
199            put_message(&mut buf, 1, |buf| resource_metrics.encode(buf));
200        }
201        buf
202    }
203
204    /// Decode a protobuf `MetricsData`.
205    ///
206    /// # Errors
207    ///
208    /// Returns [`OtlpDecodeError`] for a truncated message, a bad wire type,
209    /// a string that is not UTF-8, or an unknown aggregation temporality.
210    pub fn decode(bytes: &[u8]) -> Result<Self, OtlpDecodeError> {
211        let mut out = Self::default();
212        for field in Fields(bytes) {
213            if let (1, Value::Len(bytes)) = field? {
214                out.resource_metrics.push(ResourceMetrics::decode(bytes)?);
215            }
216        }
217        Ok(out)
218    }
219}
220
221impl ResourceMetrics {
222    fn encode(&self, buf: &mut Vec<u8>) {
223        if let Some(resource) = &self.resource {
224            put_message(buf, 1, |buf| {
225                for attribute in &resource.attributes {
226                    put_message(buf, 1, |buf| attribute.encode(buf));
227                }
228            });
229        }
230        for scope_metrics in &self.scope_metrics {
231            put_message(buf, 2, |buf| {
232                for metric in &scope_metrics.metrics {
233                    put_message(buf, 2, |buf| metric.encode(buf));
234                }
235            });
236        }
237    }
238
239    fn decode(bytes: &[u8]) -> Result<Self, OtlpDecodeError> {
240        let mut out = Self::default();
241        for field in Fields(bytes) {
242            match field? {
243                (1, Value::Len(bytes)) => {
244                    let mut resource = Resource::default();
245                    for field in Fields(bytes) {
246                        if let (1, Value::Len(bytes)) = field? {
247                            resource.attributes.push(KeyValue::decode(bytes)?);
248                        }
249                    }
250                    out.resource = Some(resource);
251                }
252                (2, Value::Len(bytes)) => {
253                    let mut scope_metrics = ScopeMetrics::default();
254                    for field in Fields(bytes) {
255                        if let (2, Value::Len(bytes)) = field? {
256                            scope_metrics.metrics.push(Metric::decode(bytes)?);
257                        }
258                    }
259                    out.scope_metrics.push(scope_metrics);
260                }
261                _ => {}
262            }
263        }
264        Ok(out)
265    }
266}
267
268impl Metric {
269    fn encode(&self, buf: &mut Vec<u8>) {
270        put_string(buf, 1, &self.name);
271        match &self.data {
272            Some(MetricData::Gauge { data_points }) => put_message(buf, 5, |buf| {
273                for point in data_points {
274                    put_message(buf, 1, |buf| point.encode(buf));
275                }
276            }),
277            Some(MetricData::Sum {
278                data_points,
279                aggregation_temporality,
280                is_monotonic,
281            }) => put_message(buf, 7, |buf| {
282                for point in data_points {
283                    put_message(buf, 1, |buf| point.encode(buf));
284                }
285                if *aggregation_temporality != AggregationTemporality::Unspecified {
286                    put_key(buf, 2, WIRE_VARINT);
287                    put_varint(buf, aggregation_temporality.number());
288                }
289                if *is_monotonic {
290                    put_key(buf, 3, WIRE_VARINT);
291                    put_varint(buf, 1);
292                }
293            }),
294            None => {}
295        }
296    }
297
298    fn decode(bytes: &[u8]) -> Result<Self, OtlpDecodeError> {
299        let mut name = String::new();
300        let mut data = None;
301        for field in Fields(bytes) {
302            match field? {
303                (1, Value::Len(bytes)) => name = utf8(bytes)?,
304                (5, Value::Len(bytes)) => {
305                    let mut data_points = Vec::new();
306                    for field in Fields(bytes) {
307                        if let (1, Value::Len(bytes)) = field? {
308                            data_points.push(NumberDataPoint::decode(bytes)?);
309                        }
310                    }
311                    data = Some(MetricData::Gauge { data_points });
312                }
313                (7, Value::Len(bytes)) => {
314                    let mut data_points = Vec::new();
315                    let mut aggregation_temporality = AggregationTemporality::Unspecified;
316                    let mut is_monotonic = false;
317                    for field in Fields(bytes) {
318                        match field? {
319                            (1, Value::Len(bytes)) => {
320                                data_points.push(NumberDataPoint::decode(bytes)?);
321                            }
322                            (2, Value::Varint(number)) => {
323                                aggregation_temporality =
324                                    AggregationTemporality::from_number(number)
325                                        .ok_or(OtlpDecodeError("unknown temporality"))?;
326                            }
327                            (3, Value::Varint(flag)) => is_monotonic = flag != 0,
328                            _ => {}
329                        }
330                    }
331                    data = Some(MetricData::Sum {
332                        data_points,
333                        aggregation_temporality,
334                        is_monotonic,
335                    });
336                }
337                _ => {}
338            }
339        }
340        Ok(Self { name, data })
341    }
342}
343
344impl NumberDataPoint {
345    /// Protobuf-java writes the fields in field-number order.
346    fn encode(&self, buf: &mut Vec<u8>) {
347        if self.start_time_unix_nano != 0 {
348            put_fixed64(buf, 2, self.start_time_unix_nano);
349        }
350        if self.time_unix_nano != 0 {
351            put_fixed64(buf, 3, self.time_unix_nano);
352        }
353        // A oneof member is written even when it holds the default value.
354        match self.value {
355            Some(NumberValue::Double(value)) => put_fixed64(buf, 4, value.to_bits()),
356            Some(NumberValue::Int(value)) => {
357                put_fixed64(buf, 6, u64::from_le_bytes(value.to_le_bytes()));
358            }
359            None => {}
360        }
361        for attribute in &self.attributes {
362            put_message(buf, 7, |buf| attribute.encode(buf));
363        }
364    }
365
366    fn decode(bytes: &[u8]) -> Result<Self, OtlpDecodeError> {
367        let mut out = Self {
368            attributes: Vec::new(),
369            start_time_unix_nano: 0,
370            time_unix_nano: 0,
371            value: None,
372        };
373        for field in Fields(bytes) {
374            match field? {
375                (2, Value::Fixed64(time)) => out.start_time_unix_nano = time,
376                (3, Value::Fixed64(time)) => out.time_unix_nano = time,
377                (4, Value::Fixed64(bits)) => {
378                    out.value = Some(NumberValue::Double(f64::from_bits(bits)));
379                }
380                (6, Value::Fixed64(bits)) => {
381                    out.value = Some(NumberValue::Int(i64::from_le_bytes(bits.to_le_bytes())));
382                }
383                (7, Value::Len(bytes)) => out.attributes.push(KeyValue::decode(bytes)?),
384                _ => {}
385            }
386        }
387        Ok(out)
388    }
389}
390
391impl KeyValue {
392    fn encode(&self, buf: &mut Vec<u8>) {
393        put_string(buf, 1, &self.key);
394        // `string_value` is a oneof member of `AnyValue`: written even when
395        // empty.
396        put_message(buf, 2, |buf| put_bytes(buf, 1, self.value.as_bytes()));
397    }
398
399    fn decode(bytes: &[u8]) -> Result<Self, OtlpDecodeError> {
400        let mut out = Self::new("", "");
401        for field in Fields(bytes) {
402            match field? {
403                (1, Value::Len(bytes)) => out.key = utf8(bytes)?,
404                (2, Value::Len(bytes)) => {
405                    for field in Fields(bytes) {
406                        if let (1, Value::Len(bytes)) = field? {
407                            out.value = utf8(bytes)?;
408                        }
409                    }
410                }
411                _ => {}
412            }
413        }
414        Ok(out)
415    }
416}
417
418fn utf8(bytes: &[u8]) -> Result<String, OtlpDecodeError> {
419    String::from_utf8(bytes.to_vec()).map_err(|_| OtlpDecodeError("string is not UTF-8"))
420}
421
422/// The value of one protobuf field.
423enum Value<'a> {
424    Varint(u64),
425    Fixed64(u64),
426    Fixed32,
427    Len(&'a [u8]),
428}
429
430/// The fields of one protobuf message, in wire order.
431struct Fields<'a>(&'a [u8]);
432
433impl<'a> Fields<'a> {
434    fn varint(&mut self) -> Result<u64, OtlpDecodeError> {
435        let mut value = 0u64;
436        for shift in (0..64).step_by(7) {
437            let (&byte, rest) = self
438                .0
439                .split_first()
440                .ok_or(OtlpDecodeError("truncated varint"))?;
441            self.0 = rest;
442            value |= u64::from(byte & 0x7f) << shift;
443            if byte & 0x80 == 0 {
444                return Ok(value);
445            }
446        }
447        Err(OtlpDecodeError("varint longer than 10 bytes"))
448    }
449
450    fn take_bytes(&mut self, len: usize) -> Result<&'a [u8], OtlpDecodeError> {
451        if self.0.len() < len {
452            return Err(OtlpDecodeError("truncated field"));
453        }
454        let (head, rest) = self.0.split_at(len);
455        self.0 = rest;
456        Ok(head)
457    }
458
459    fn field(&mut self) -> Result<(u32, Value<'a>), OtlpDecodeError> {
460        let key = self.varint()?;
461        let field = u32::try_from(key >> 3).map_err(|_| OtlpDecodeError("field number"))?;
462        let value = match u8::try_from(key & 0x7).unwrap_or(u8::MAX) {
463            WIRE_VARINT => Value::Varint(self.varint()?),
464            WIRE_FIXED64 => {
465                let bytes: [u8; 8] = self
466                    .take_bytes(8)?
467                    .try_into()
468                    .map_err(|_| OtlpDecodeError("truncated fixed64"))?;
469                Value::Fixed64(u64::from_le_bytes(bytes))
470            }
471            WIRE_LEN => {
472                let len = usize::try_from(self.varint()?)
473                    .map_err(|_| OtlpDecodeError("length does not fit"))?;
474                Value::Len(self.take_bytes(len)?)
475            }
476            WIRE_FIXED32 => {
477                self.take_bytes(4)?;
478                Value::Fixed32
479            }
480            _ => return Err(OtlpDecodeError("unsupported wire type")),
481        };
482        Ok((field, value))
483    }
484}
485
486impl<'a> Iterator for Fields<'a> {
487    type Item = Result<(u32, Value<'a>), OtlpDecodeError>;
488
489    fn next(&mut self) -> Option<Self::Item> {
490        if self.0.is_empty() {
491            return None;
492        }
493        let field = self.field();
494        if field.is_err() {
495            // Stop after the first error.
496            self.0 = &[];
497        }
498        Some(field)
499    }
500}
501
502#[cfg(test)]
503mod tests {
504    use assert2::assert;
505
506    use super::*;
507
508    fn point(value: NumberValue, attributes: Vec<KeyValue>) -> NumberDataPoint {
509        NumberDataPoint {
510            attributes,
511            start_time_unix_nano: 1_700_000_000_000_000_000,
512            time_unix_nano: 1_700_000_001_000_000_000,
513            value: Some(value),
514        }
515    }
516
517    #[test]
518    fn metrics_data_round_trips() {
519        let cases = [
520            ("empty", MetricsData::default()),
521            (
522                "gauge with an empty resource",
523                MetricsData {
524                    resource_metrics: vec![ResourceMetrics {
525                        resource: Some(Resource::default()),
526                        scope_metrics: vec![ScopeMetrics {
527                            metrics: vec![Metric {
528                                name: "org.apache.kafka.producer.connection.count".into(),
529                                data: Some(MetricData::Gauge {
530                                    data_points: vec![point(NumberValue::Double(0.0), vec![])],
531                                }),
532                            }],
533                        }],
534                    }],
535                },
536            ),
537            (
538                "sums with attributes",
539                MetricsData {
540                    resource_metrics: vec![ResourceMetrics {
541                        resource: Some(Resource {
542                            attributes: vec![KeyValue::new("transactional_id", "tx")],
543                        }),
544                        scope_metrics: vec![ScopeMetrics {
545                            metrics: vec![
546                                Metric {
547                                    name: "org.apache.kafka.producer.request.total".into(),
548                                    data: Some(MetricData::Sum {
549                                        data_points: vec![point(
550                                            NumberValue::Double(42.5),
551                                            vec![KeyValue::new("node_id", "node-1")],
552                                        )],
553                                        aggregation_temporality: AggregationTemporality::Delta,
554                                        is_monotonic: true,
555                                    }),
556                                },
557                                Metric {
558                                    name: "negative".into(),
559                                    data: Some(MetricData::Sum {
560                                        data_points: vec![point(NumberValue::Int(-7), vec![])],
561                                        aggregation_temporality: AggregationTemporality::Cumulative,
562                                        is_monotonic: false,
563                                    }),
564                                },
565                            ],
566                        }],
567                    }],
568                },
569            ),
570        ];
571        for (name, data) in cases {
572            let decoded = MetricsData::decode(&data.encode());
573            assert!(decoded == Ok(data), "{name}");
574        }
575    }
576
577    /// The bytes of a small message, as protobuf-java writes them.
578    #[test]
579    fn encoding_matches_the_protobuf_wire_format() {
580        let data = MetricsData {
581            resource_metrics: vec![ResourceMetrics {
582                resource: Some(Resource::default()),
583                scope_metrics: vec![ScopeMetrics {
584                    metrics: vec![Metric {
585                        name: "m".into(),
586                        data: Some(MetricData::Sum {
587                            data_points: vec![NumberDataPoint {
588                                attributes: vec![KeyValue::new("k", "v")],
589                                start_time_unix_nano: 0,
590                                time_unix_nano: 1,
591                                value: Some(NumberValue::Double(1.0)),
592                            }],
593                            aggregation_temporality: AggregationTemporality::Cumulative,
594                            is_monotonic: true,
595                        }),
596                    }],
597                }],
598            }],
599        };
600        let point = [
601            &[0x19][..],
602            &1u64.to_le_bytes(),
603            &[0x21],
604            &1.0f64.to_bits().to_le_bytes(),
605            &[0x3a, 0x08, 0x0a, 0x01, b'k', 0x12, 0x03, 0x0a, 0x01, b'v'],
606        ]
607        .concat();
608        let sum = [&[0x0a, 0x1c][..], &point, &[0x10, 0x02, 0x18, 0x01]].concat();
609        let metric = [&[0x0a, 0x01, b'm', 0x3a, 0x22][..], &sum].concat();
610        let scope = [&[0x12, 0x27][..], &metric].concat();
611        let resource_metrics = [&[0x0a, 0x00, 0x12, 0x29][..], &scope].concat();
612        let expected = [&[0x0a, 0x2d][..], &resource_metrics].concat();
613        assert!(data.encode() == expected);
614    }
615
616    #[test]
617    fn decode_skips_unknown_fields_and_rejects_malformed_input() {
618        // A `MetricsData` with an unknown varint field 9 and an unknown
619        // fixed32 field 10 around an empty resource metrics.
620        let unknown = [0x48, 0x01, 0x0a, 0x00, 0x55, 1, 2, 3, 4];
621        assert!(
622            MetricsData::decode(&unknown)
623                == Ok(MetricsData {
624                    resource_metrics: vec![ResourceMetrics::default()],
625                })
626        );
627        let cases: [(&str, &[u8], &str); 4] = [
628            ("truncated length", &[0x0a, 0x05, 0x00], "truncated field"),
629            ("truncated varint", &[0x08, 0x80], "truncated varint"),
630            ("group wire type", &[0x0b], "unsupported wire type"),
631            (
632                "unknown temporality",
633                &[0x0a, 0x08, 0x12, 0x06, 0x12, 0x04, 0x3a, 0x02, 0x10, 0x07],
634                "unknown temporality",
635            ),
636        ];
637        for (name, bytes, reason) in cases {
638            assert!(
639                MetricsData::decode(bytes) == Err(OtlpDecodeError(reason)),
640                "{name}"
641            );
642        }
643    }
644}