[profiling] Only convert labels once per labelset (#1215)

Co-authored-by: realFlowControl <[email protected]>

Author: danielsn

datadog-profiling/src/internal/label.rs

@@ -76,6 +76,12 @@ impl From<&Label> for prost_impls::Label {
 
 impl From<Label> for datadog_profiling_protobuf::Label {
     fn from(label: Label) -> Self {
+        Self::from(&label)
+    }
+}
+
+impl From<&Label> for datadog_profiling_protobuf::Label {
+    fn from(label: &Label) -> Self {
         let (str, num, num_unit) = match label.value {
             LabelValue::Str(str) => (str, 0, StringOffset::ZERO),
             LabelValue::Num { num, num_unit } => (StringOffset::ZERO, num, num_unit),
@@ -125,10 +131,22 @@ pub struct LabelSet {
 }
 
 impl LabelSet {
+    pub fn is_empty(&self) -> bool {
+        self.labels.is_empty()
+    }
+
     pub fn iter(&self) -> core::slice::Iter<'_, LabelId> {
         self.labels.iter()
     }
 
+    pub fn labels(&self) -> &[LabelId] {
+        &self.labels
+    }
+
+    pub fn len(&self) -> usize {
+        self.labels.len()
+    }
+
     pub fn new(labels: Box<[LabelId]>) -> Self {
         // Once upon a time label ids were guaranteed to be sorted. However,
         // this makes testing difficult because the order of input labels and

datadog-profiling/src/internal/profile/mod.rs

@@ -371,25 +371,50 @@ impl Profile {
             .as_nanos()
             .min(i64::MAX as u128) as i64;
 
+        let mut extended_label_sets: Vec<Vec<Label>> = Vec::with_capacity(self.label_sets.len());
+
+        for label_set in std::mem::take(&mut self.label_sets) {
+            let endpoint_label = self.get_endpoint_for_label_set(&label_set)?;
+            // Leave one space for the timestamp if needed
+            let mut labels = Vec::with_capacity(
+                label_set.len() + 1 + if endpoint_label.is_some() { 1 } else { 0 },
+            );
+            for l in label_set.iter() {
+                labels.push(*self.get_label(*l)?);
+            }
+            if let Some(endpoint_label) = endpoint_label {
+                labels.push(endpoint_label);
+            }
+            extended_label_sets.push(labels);
+        }
+
         for (sample, timestamp, mut values) in std::mem::take(&mut self.observations).into_iter() {
-            let labels = self.enrich_sample_labels(sample, timestamp)?;
+            let labels = &mut extended_label_sets[sample.labels.to_raw_id()];
             let location_ids: Vec<_> = self
                 .get_stacktrace(sample.stacktrace)?
                 .locations
                 .iter()
                 .map(Id::to_raw_id)
                 .collect();
             self.check_location_ids_are_valid(&location_ids, self.locations.len())?;
-            self.upscaling_rules.upscale_values(&mut values, &labels)?;
+            self.upscaling_rules.upscale_values(&mut values, labels)?;
+
+            // Use the extra slot in the labels vector to store the timestamp without any reallocs.
+            if let Some(ts) = timestamp {
+                labels.push(Label::num(self.timestamp_key, ts.get(), StringId::ZERO))
+            }
+            let pprof_labels: Vec<_> = labels.iter().map(protobuf::Label::from).collect();
+            if timestamp.is_some() {
+                labels.pop();
+            }
 
-            let labels: Vec<_> = labels.into_iter().map(protobuf::Label::from).collect();
             let item = protobuf::Sample {
                 location_ids: Record::from(location_ids.as_slice()),
                 values: Record::from(values.as_slice()),
                 // SAFETY: converting &[Label] to &[Field<Label,..>] which is
                 // safe, because Field is repr(transparent).
                 labels: unsafe {
-                    &*(labels.as_slice() as *const [protobuf::Label]
+                    &*(pprof_labels.as_slice() as *const [protobuf::Label]
                         as *const [Record<protobuf::Label, 3, NO_OPT_ZERO>])
                 },
             };
@@ -679,16 +704,15 @@ impl Profile {
             .map(|v| Label::str(self.endpoints.endpoint_label, *v)))
     }
 
-    fn get_endpoint_for_labels(&self, label_set_id: LabelSetId) -> anyhow::Result<Option<Label>> {
-        let label = self.get_label_set(label_set_id)?.iter().find_map(|id| {
+    fn get_endpoint_for_label_set(&self, label_set: &LabelSet) -> anyhow::Result<Option<Label>> {
+        if let Some(label) = label_set.iter().find_map(|id| {
             if let Ok(label) = self.get_label(*id) {
                 if label.get_key() == self.endpoints.local_root_span_id_label {
                     return Some(label);
                 }
             }
             None
-        });
-        if let Some(label) = label {
+        }) {
             self.get_endpoint_for_label(label)
         } else {
             Ok(None)
@@ -701,6 +725,7 @@ impl Profile {
             .context("LabelId to have a valid interned index")
     }
 
+    #[allow(dead_code)]
     fn get_label_set(&self, id: LabelSetId) -> anyhow::Result<&LabelSet> {
         self.label_sets
             .get_index(id.to_offset())
@@ -800,19 +825,6 @@ impl Profile {
         profile
     }
 
-    fn enrich_sample_labels(
-        &self,
-        sample: Sample,
-        timestamp: Option<Timestamp>,
-    ) -> anyhow::Result<Vec<Label>> {
-        self.get_label_set(sample.labels)?
-            .iter()
-            .map(|l| self.get_label(*l).copied())
-            .chain(self.get_endpoint_for_labels(sample.labels).transpose())
-            .chain(timestamp.map(|ts| Ok(Label::num(self.timestamp_key, ts.get(), StringId::ZERO))))
-            .collect()
-    }
-
     #[cfg(debug_assertions)]
     fn validate_sample_labels(&mut self, sample: &api::Sample) -> anyhow::Result<()> {
         let mut seen: HashMap<&str, &api::Label> = HashMap::new();