double up the sample types

Author: danielsn

libdd-profiling/src/internal/otel_style_observation.rs

@@ -5,22 +5,23 @@ use crate::{
     api::SampleType,
     internal::{Sample, Timestamp},
 };
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 
 #[derive(Default)]
 pub struct Observations {
     sample_types: Box<[SampleType]>,
     pub aggregated: HashMap<SampleType, HashMap<Sample, i64>>,
+    pub aggregated2: HashMap<(SampleType, SampleType), HashMap<Sample, (i64, i64)>>,
     pub timestamped: HashMap<SampleType, HashMap<Sample, Vec<(i64, Timestamp)>>>,
+    pub timestamped2:
+        HashMap<(SampleType, SampleType), HashMap<Sample, Vec<(i64, i64, Timestamp)>>>,
 }
 
 impl Observations {
     pub fn new(sample_types: Box<[SampleType]>) -> Self {
-        let len = sample_types.len();
         Self {
             sample_types,
-            aggregated: HashMap::with_capacity(len),
-            timestamped: HashMap::with_capacity(len),
+            ..Default::default()
         }
     }
 
@@ -30,27 +31,64 @@ impl Observations {
         timestamp: Option<Timestamp>,
         values: &[i64],
     ) -> anyhow::Result<()> {
+        let mut first = None;
+        let mut second = None;
+
         for (idx, v) in values.iter().enumerate() {
             if *v != 0 {
                 let sample_type = self.sample_types[idx];
+                if first.is_none() {
+                    first = Some((sample_type, *v));
+                } else if second.is_none() {
+                    second = Some((sample_type, *v));
+                } else {
+                    anyhow::bail!("too many values");
+                }
+            }
+        }
+
+        match (first, second) {
+            (None, None) => {}
+            (Some((st, val)), None) => {
                 if let Some(ts) = timestamp {
                     self.timestamped
-                        .entry(sample_type)
+                        .entry(st)
                         .or_default()
                         .entry(sample)
                         .or_default()
-                        .push((*v, ts));
+                        .push((val, ts));
                 } else {
-                    let val = self
+                    let entry = self
                         .aggregated
-                        .entry(sample_type)
+                        .entry(st)
                         .or_default()
                         .entry(sample)
                         .or_insert(0);
-                    *val = val.saturating_add(*v);
+                    *entry = entry.saturating_add(val);
                 }
             }
+            (Some((st1, val1)), Some((st2, val2))) => {
+                if let Some(ts) = timestamp {
+                    self.timestamped2
+                        .entry((st1, st2))
+                        .or_default()
+                        .entry(sample)
+                        .or_default()
+                        .push((val1, val2, ts));
+                } else {
+                    let entry = self
+                        .aggregated2
+                        .entry((st1, st2))
+                        .or_default()
+                        .entry(sample)
+                        .or_insert((0, 0));
+                    entry.0 = entry.0.saturating_add(val1);
+                    entry.1 = entry.1.saturating_add(val2);
+                }
+            }
+            (None, Some(_)) => unreachable!("second set implies first set"),
         }
+
         Ok(())
     }
 
@@ -59,21 +97,27 @@ impl Observations {
     }
 
     pub fn aggregated_samples_count(&self) -> usize {
-        // TODO: this is the actual sample count, but doesn't reflect aggregated samples being
-        // overlapping self.aggregated.iter().map(|(_, v)| v.len()).sum()
-        let samples: HashSet<&Sample> = self
-            .aggregated
-            .iter()
-            .flat_map(|(_, samples)| samples.keys())
-            .collect();
-        samples.len()
+        self.aggregated
+            .values()
+            .map(|samples| samples.len())
+            .sum::<usize>()
+            + self
+                .aggregated2
+                .values()
+                .map(|samples| samples.len())
+                .sum::<usize>()
     }
 
     pub fn timestamped_samples_count(&self) -> usize {
         self.timestamped
             .iter()
             .flat_map(|(_, v)| v.values().map(|v| v.len()))
-            .sum()
+            .sum::<usize>()
+            + self
+                .timestamped2
+                .iter()
+                .flat_map(|(_, v)| v.values().map(|v| v.len()))
+                .sum::<usize>()
     }
 
     pub fn try_into_iter(
@@ -94,6 +138,8 @@ impl Observations {
 
         let len: usize = self.sample_types.len();
         let index_map_ts = index_map.clone();
+        let index_map_ts2 = index_map.clone();
+        let index_map_accum2 = index_map.clone();
         let accum_iter = self
             .aggregated
             .into_iter()
@@ -107,6 +153,22 @@ impl Observations {
                 })
             });
 
+        let accum2_iter = self
+            .aggregated2
+            .into_iter()
+            .flat_map(move |((st1, st2), inner)| {
+                #[allow(clippy::unwrap_used)]
+                let index1 = *index_map_accum2.get(&st1).unwrap();
+                #[allow(clippy::unwrap_used)]
+                let index2 = *index_map_accum2.get(&st2).unwrap();
+                inner.into_iter().map(move |(sample, (val1, val2))| {
+                    let mut vals = vec![0; len];
+                    vals[index1] = val1;
+                    vals[index2] = val2;
+                    (sample, None, vals)
+                })
+            });
+
         // let mut accum: HashMap<Sample, Vec<i64>> = HashMap::new();
         // for (sample_type, samples) in self.aggregated.into_iter() {
         //     let Some(idx) = index_map.get(&sample_type) else {
@@ -133,6 +195,25 @@ impl Observations {
                     })
                 })
             });
-        accum_iter.chain(ts_iter)
+
+        let ts2_iter = self
+            .timestamped2
+            .into_iter()
+            .flat_map(move |((st1, st2), inner)| {
+                #[allow(clippy::unwrap_used)]
+                let index1 = *index_map_ts2.get(&st1).unwrap();
+                #[allow(clippy::unwrap_used)]
+                let index2 = *index_map_ts2.get(&st2).unwrap();
+                inner.into_iter().flat_map(move |(sample, ts_vals)| {
+                    ts_vals.into_iter().map(move |(val1, val2, ts)| {
+                        let mut vals = vec![0; len];
+                        vals[index1] = val1;
+                        vals[index2] = val2;
+                        (sample, Some(ts), vals)
+                    })
+                })
+            });
+
+        accum_iter.chain(accum2_iter).chain(ts_iter).chain(ts2_iter)
     }
 }

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

@@ -462,6 +462,7 @@ impl Profile {
     ///   duration based on the end time minus the start time, but under anomalous conditions this
     ///   may fail as system clocks can be adjusted. The programmer may also accidentally pass an
     ///   earlier time. The duration will be set to zero these cases.
+    #[cfg(feature = "otel")]
     pub fn serialize_into_compressed_otel(
         self,
         end_time: Option<SystemTime>,
@@ -783,7 +784,7 @@ impl Profile {
             end,
             buffer,
             endpoints_stats,
-        })    
+        })
     }
 
     /// Encodes the profile. Note that the buffer will be empty. The caller
@@ -1694,7 +1695,6 @@ mod api_tests {
     }
 
     #[test]
-    #[cfg_attr(feature = "otel", ignore)]
     fn lazy_endpoints() -> anyhow::Result<()> {
         let sample_types = [api::SampleType::CpuSamples, api::SampleType::WallTime];
 
@@ -1858,8 +1858,6 @@ mod api_tests {
     }
 
     #[test]
-    // This works if we merge accumulated samples
-    #[cfg_attr(feature = "otel", ignore)]
     fn test_no_upscaling_if_no_rules() {
         let sample_types = vec![api::SampleType::CpuSamples, api::SampleType::WallTime];
 
@@ -1909,8 +1907,6 @@ mod api_tests {
     }
 
     #[test]
-    // This works if we accumulate
-    #[cfg_attr(feature = "otel", ignore)]
     fn test_upscaling_by_value_a_zero_value() {
         let sample_types = create_samples_types();
 
@@ -2048,9 +2044,6 @@ mod api_tests {
     }
 
     #[test]
-    // This works if we accumulate
-    #[cfg_attr(feature = "otel", ignore)]
-
     fn test_upscaling_by_value_on_zero_value_with_poisson() {
         let sample_types = create_samples_types();
 
@@ -2839,6 +2832,7 @@ mod api_tests {
     }
 
     #[test]
+    #[cfg_attr(feature = "otel", ignore)] // too many non-zero values
     fn test_fails_when_adding_byvalue_rule_colliding_on_offset_with_existing_bylabel_rule() {
         let sample_types = create_samples_types();
 
@@ -2878,8 +2872,6 @@ mod api_tests {
     }
 
     #[test]
-    // This works if we merge accumulated samples
-    #[cfg_attr(feature = "otel", ignore)]
     fn local_root_span_id_label_as_i64() -> anyhow::Result<()> {
         let sample_types = vec![api::SampleType::CpuSamples, api::SampleType::WallTime];