Skip to main content

libdd_profiling_protobuf/
sample.rs

1// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::{Label, Record, Value, WireType, NO_OPT_ZERO};
5
6#[cfg(feature = "prost_impls")]
7use crate::prost_impls;
8
9use std::io::{self, Write};
10
11/// Each Sample records values encountered in some program context. The
12/// program context is typically a stack trace, perhaps augmented with
13/// auxiliary information like the thread-id, some indicator of a higher level
14/// request being handled, etc.
15///
16/// It borrows its data but requires it to be a slice. An iterator wouldn't
17/// work well because we have to walk over the fields twice: one to calculate
18/// the length, and one to encode it.
19#[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
20pub struct Sample<'a> {
21    /// The ids recorded here correspond to a Profile.location.id.
22    /// The leaf is at location_id\[0\].
23    pub location_ids: Record<&'a [u64], 1, NO_OPT_ZERO>,
24    /// The type and unit of each value is defined by the corresponding entry
25    /// in Profile.sample_type. All samples must have the same number of
26    /// values, the same as the length of Profile.sample_type. When
27    /// aggregating multiple samples into a single sample, the result has a
28    /// list of values that is the element-wise sum of the original lists.
29    pub values: Record<&'a [i64], 2, NO_OPT_ZERO>,
30    /// NOTE: While possible, having multiple values for the same label key is
31    /// strongly discouraged and should never be used. Most tools (e.g. pprof)
32    /// do not have good (or any) support for multi-value labels. And an even
33    /// more discouraged case is having a string label and a numeric label of
34    /// the same name on a sample. Again, possible to express, but should not
35    /// be used.
36    pub labels: &'a [Record<Label, 3, NO_OPT_ZERO>],
37}
38
39/// # Safety
40/// The Default implementation will return all zero-representations.
41unsafe impl Value for Sample<'_> {
42    const WIRE_TYPE: WireType = WireType::LengthDelimited;
43
44    fn proto_len(&self) -> u64 {
45        self.location_ids.proto_len()
46            + self.values.proto_len()
47            + self.labels.iter().map(Record::proto_len).sum::<u64>()
48    }
49
50    fn encode<W: Write>(&self, writer: &mut W) -> io::Result<()> {
51        self.location_ids.encode(writer)?;
52        self.values.encode(writer)?;
53        for label in self.labels {
54            label.encode(writer)?;
55        }
56        Ok(())
57    }
58}
59
60#[cfg(feature = "prost_impls")]
61impl From<Sample<'_>> for prost_impls::Sample {
62    fn from(sample: Sample) -> Self {
63        // If the prost file is regenerated, this may pick up new members.
64        #[allow(clippy::needless_update)]
65        Self {
66            location_ids: Vec::from_iter(sample.location_ids.value.iter().copied()),
67            values: Vec::from_iter(sample.values.value.iter().copied()),
68            labels: sample
69                .labels
70                .iter()
71                .map(|field| field.value)
72                .map(prost_impls::Label::from)
73                .collect(),
74            ..Self::default()
75        }
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82    use bolero::generator::TypeGenerator;
83    use prost::Message;
84
85    #[test]
86    fn empty() {
87        let sample = Sample {
88            location_ids: [].as_slice().into(),
89            values: [].as_slice().into(),
90            labels: &[],
91        };
92        let prost_sample = prost_impls::Sample {
93            location_ids: vec![],
94            values: vec![],
95            labels: vec![],
96        };
97
98        use prost::Message;
99        let len = sample.proto_len() as usize;
100        let mut buffer = Vec::with_capacity(len);
101        sample.encode(&mut buffer).unwrap();
102        let roundtrip = prost_impls::Sample::decode(buffer.as_slice()).unwrap();
103        assert_eq!(prost_sample, roundtrip);
104    }
105
106    #[test]
107    fn roundtrip() {
108        let locations = Vec::<u64>::produce();
109        let values = Vec::<i64>::produce();
110        let labels = Vec::<Label>::produce();
111
112        bolero::check!()
113            .with_generator((locations, values, labels))
114            .for_each(|(location_ids, values, labels)| {
115                let labels = labels
116                    .iter()
117                    .map(|l| Record::<_, 3, NO_OPT_ZERO>::from(*l))
118                    .collect::<Vec<_>>();
119                let sample = Sample {
120                    location_ids: Record::from(location_ids.as_slice()),
121                    values: Record::from(values.as_slice()),
122                    labels: labels.as_slice(),
123                };
124
125                let prost_sample = prost_impls::Sample::from(sample);
126
127                let mut buffer = Vec::with_capacity(sample.proto_len() as usize);
128                sample.encode(&mut buffer).unwrap();
129                let roundtrip = prost_impls::Sample::decode(buffer.as_slice()).unwrap();
130                assert_eq!(prost_sample, roundtrip);
131
132                let mut buffer2 = Vec::with_capacity(sample.proto_len() as usize);
133                prost_sample.encode(&mut buffer2).unwrap();
134                let roundtrip2 = prost_impls::Sample::decode(buffer2.as_slice()).unwrap();
135                assert_eq!(roundtrip, roundtrip2);
136            });
137    }
138}