Unit tests for location_table access pattern

src/collector/otlp/service.rs

@@ -70,7 +70,13 @@ impl pb_collector::profiles_service_server::ProfilesService for ProfilesService
 
                     let st = &profile.sample_type[0];
                     for sample in &profile.sample {
-                        process_sample(dict, &loc_mapping, &st, sample, &profile.location_indices)?;
+                        let frame_list = collect_frame_list(
+                            sample.locations_start_index as usize,
+                            sample.locations_length as usize,
+                            &loc_mapping,
+                            &profile.location_indices,
+                        )?;
+                        process_sample(dict, &st, sample, frame_list)?;
                     }
                 }
             }
@@ -277,31 +283,10 @@ fn ingest_locations(dic: &ProfilesDictionary) -> Result<Vec<Frame>, Status> {
 
 fn process_sample(
     dict: &ProfilesDictionary,
-    loc_mapping: &Vec<Frame>,
     sample_type: &ValueType,
     sample: &Sample,
-    profile_location_indices: &Vec<i32>,
+    frame_list: Vec<Frame>,
 ) -> Result<(), Status> {
-    let loc_start = sample.locations_start_index as usize;
-    let loc_len = sample.locations_length as usize;
-    let loc_rng = loc_start..loc_start.saturating_add(loc_len);
-
-    // Collect frame list.
-    let mut frame_list = Vec::with_capacity(loc_rng.len().min(128));
-    for loc_index in loc_rng {
-        let Some(location_table_idx) = profile_location_indices.get(loc_index as usize) else {
-            return Err(Status::invalid_argument(
-                "location_indices: index is out of bounds",
-            ));
-        };
-        let Some(frame) = loc_mapping.get(*location_table_idx as usize) else {
-            return Err(Status::invalid_argument(
-                "location_table: index is out of bounds",
-            ));
-        };
-        frame_list.push(*frame);
-    }
-
     // Insert frame list.
     let mut hasher = xxh3::Xxh3::new();
     frame_list.hash(&mut hasher);
@@ -376,3 +361,100 @@ fn process_sample(
 
     Ok(())
 }
+
+fn collect_frame_list<V>(
+    loc_start: usize,
+    loc_len: usize,
+    loc_mapping: &Vec<V>,
+    profile_location_indices: &Vec<i32>,
+) -> Result<Vec<V>, Status>
+where
+    V: Copy,
+{
+    let loc_rng = loc_start..loc_start.saturating_add(loc_len);
+
+    // Collect frame list.
+    let mut frame_list = Vec::with_capacity(loc_rng.len().min(128));
+    for loc_index in loc_rng {
+        let Some(location_table_idx) = profile_location_indices.get(loc_index as usize) else {
+            return Err(Status::invalid_argument(
+                "location_indices: index is out of bounds",
+            ));
+        };
+        let Some(frame) = loc_mapping.get(*location_table_idx as usize) else {
+            return Err(Status::invalid_argument(
+                "location_table: index is out of bounds",
+            ));
+        };
+        frame_list.push(*frame);
+    }
+
+    return Ok(frame_list);
+}
+
+#[cfg(test)]
+mod tests {
+    use itertools::Itertools;
+
+    use super::*;
+
+    #[test]
+    fn sample_frame_list() -> Result<(), Status> {
+        let loc_mapping = (0..11).collect_vec();
+        let location_indices = vec![4, 9, 6, 2, 7, 4, 4, 2, 0, 1, 2, 3, 5];
+
+        assert_eq!(
+            collect_frame_list(0, 2, &loc_mapping, &location_indices)?,
+            vec![4, 9],
+            "location_indices: {{0,1}}"
+        );
+        assert_eq!(
+            collect_frame_list(1, 0, &loc_mapping, &location_indices)?,
+            Vec::<i32>::new(),
+            "zero-length trace"
+        );
+        assert_eq!(
+            collect_frame_list(0, location_indices.len(), &loc_mapping, &location_indices)?,
+            location_indices,
+            "trace takes all indices in location_indices"
+        );
+        assert_eq!(
+            collect_frame_list(2, 0, &loc_mapping, &vec![0i32, 1i32])?,
+            Vec::<i32>::new(),
+            "zero-length trace with loc_start out-of-bounds"
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn sample_frame_list_err() -> Result<(), Status> {
+        let loc_mapping = (0..11).collect_vec();
+
+        assert_eq!(
+            collect_frame_list(0, 3, &loc_mapping, &vec![0i32, 1i32])
+                .unwrap_err()
+                .message(),
+            "location_indices: index is out of bounds",
+            "sample trace size: 3, len(location_indices): 2"
+        );
+
+        assert_eq!(
+            collect_frame_list(1, 2, &loc_mapping, &vec![0i32, 1i32])
+                .unwrap_err()
+                .message(),
+            "location_indices: index is out of bounds",
+            "sample trace index start: 1, sample trace length: 2, len(location_indices): 2"
+        );
+
+        assert_eq!(
+            collect_frame_list(0, 2, &loc_mapping, &vec![1i32, 13i32])
+                .unwrap_err()
+                .message(),
+            "location_table: index is out of bounds",
+            "trace location indices: {{1,13}}, len(location_table): 2"
+        );
+
+        Ok(())
+    }
+}