consume observations on drop of iterator (#1362)

consume observations on drop of iterator

add test for anyhow bailout

`upscale_values` can not fail

Add missing comment

Co-authored-by: florian.engelhardt <florian.engelhardt@datadoghq.com>

Author: realFlowControl

libdd-profiling/src/internal/observation/observations.rs

@@ -115,9 +115,10 @@ impl Observations {
                     .try_into_iter()?
                     .map(|(s, t, o)| (s, Some(t), o));
 
-                let agg_it = std::mem::take(&mut aggregated_data.data)
-                    .into_iter()
-                    .map(move |(s, o)| (s, None, unsafe { o.into_vec(obs_len) }));
+                let agg_it = AggregatedObservationsIter {
+                    iter: std::mem::take(&mut aggregated_data.data).into_iter(),
+                    obs_len,
+                };
 
                 Ok(ObservationsIntoIter {
                     it: Box::new(ts_it.chain(agg_it)),
@@ -195,6 +196,35 @@ impl Drop for AggregatedObservations {
     }
 }
 
+/// This iterator does only exist to make the drop work as in: consume the rest of the iterator on
+/// drop to clean up and not leak memory
+struct AggregatedObservationsIter {
+    iter: std::collections::hash_map::IntoIter<Sample, TrimmedObservation>,
+    obs_len: ObservationLength,
+}
+
+impl Iterator for AggregatedObservationsIter {
+    type Item = (Sample, Option<Timestamp>, Vec<i64>);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let (sample, observation) = self.iter.next()?;
+        // SAFETY: The only way to build one of these is through
+        // [Observations::add], which already checked that the length was correct.
+        let vec = unsafe { observation.into_vec(self.obs_len) };
+        Some((sample, None, vec))
+    }
+}
+
+impl Drop for AggregatedObservationsIter {
+    fn drop(&mut self) {
+        for (_, observation) in &mut self.iter {
+            // SAFETY: The only way to build one of these is through
+            // [Observations::add], which already checked that the length was correct.
+            unsafe { observation.consume(self.obs_len) };
+        }
+    }
+}
+
 pub struct ObservationsIntoIter {
     it: Box<dyn Iterator<Item = <ObservationsIntoIter as Iterator>::Item>>,
 }

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

@@ -424,7 +424,7 @@ impl Profile {
                 .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 {
@@ -2605,4 +2605,61 @@ mod api_tests {
             .iter()
             .any(|s| s == "world"));
     }
+
+    #[test]
+    fn reproduce_crash_with_anyhow_bailout() {
+        let sample_types = [api::ValueType::new("samples", "count")];
+        let mapping = api::Mapping {
+            filename: "test.php",
+            ..Default::default()
+        };
+
+        let mut profile = Profile::new(&sample_types, None);
+
+        let locations = vec![api::Location {
+            mapping,
+            function: api::Function {
+                name: "test_function",
+                system_name: "test_function",
+                filename: "test.php",
+            },
+            line: 0,
+            ..Default::default()
+        }];
+
+        let sample = api::Sample {
+            locations,
+            values: &[1],
+            labels: vec![api::Label {
+                key: "iteration",
+                num: 1,
+                ..Default::default()
+            }],
+        };
+
+        profile
+            .try_add_sample(sample, None) // No timestamp = aggregated
+            .expect("profile to not be full");
+
+        // We want to trigger an error inside the loop over observations.
+        // By clearing the locations, the location IDs referenced by the samples/stacktraces
+        // will become invalid (checking against len=0), causing check_location_ids_are_valid to
+        // fail.
+        profile.locations.clear();
+
+        let result = profile.serialize_into_compressed_pprof(None, None);
+
+        match result {
+            Ok(_) => panic!(
+                "Expected serialization to fail due to invalid location IDs, but it succeeded"
+            ),
+            Err(err) => {
+                assert!(
+                    err.to_string().contains("invalid location id"),
+                    "Expected error about invalid location id, got: {}",
+                    err
+                );
+            }
+        }
+    }
 }

libdd-profiling/src/internal/upscaling.rs

@@ -165,37 +165,36 @@ impl UpscalingRules {
         self.rules.is_empty()
     }
 
-    pub fn upscale_values(&self, values: &mut [i64], labels: &[Label]) -> anyhow::Result<()> {
-        if !self.is_empty() {
-            // get bylabel rules first (if any)
-            let mut group_of_rules = labels
-                .iter()
-                .filter_map(|label| {
-                    self.get(&(
-                        label.get_key(),
-                        match label.get_value() {
-                            LabelValue::Str(str) => *str,
-                            LabelValue::Num { .. } => StringId::ZERO,
-                        },
-                    ))
-                })
-                .collect::<Vec<&Vec<UpscalingRule>>>();
-
-            // get byvalue rules if any
-            if let Some(byvalue_rules) = self.get(&(StringId::ZERO, StringId::ZERO)) {
-                group_of_rules.push(byvalue_rules);
-            }
+    pub fn upscale_values(&self, values: &mut [i64], labels: &[Label]) {
+        if self.is_empty() {
+            return;
+        }
+        // get bylabel rules first (if any)
+        let mut group_of_rules = labels
+            .iter()
+            .filter_map(|label| {
+                self.get(&(
+                    label.get_key(),
+                    match label.get_value() {
+                        LabelValue::Str(str) => *str,
+                        LabelValue::Num { .. } => StringId::ZERO,
+                    },
+                ))
+            })
+            .collect::<Vec<&Vec<UpscalingRule>>>();
 
-            group_of_rules.iter().for_each(|rules| {
-                rules.iter().for_each(|rule| {
-                    let scale = rule.compute_scale(values);
-                    rule.values_offset.iter().for_each(|offset| {
-                        values[*offset] = (values[*offset] as f64 * scale).round() as i64
-                    })
-                })
-            });
+        // get byvalue rules if any
+        if let Some(byvalue_rules) = self.get(&(StringId::ZERO, StringId::ZERO)) {
+            group_of_rules.push(byvalue_rules);
         }
 
-        Ok(())
+        group_of_rules.iter().for_each(|rules| {
+            rules.iter().for_each(|rule| {
+                let scale = rule.compute_scale(values);
+                rule.values_offset.iter().for_each(|offset| {
+                    values[*offset] = (values[*offset] as f64 * scale).round() as i64
+                })
+            })
+        });
     }
 }