Test that str and num labels work with endpoint profiling

profiling/src/profile/mod.rs

@@ -942,13 +942,13 @@ mod api_test {
         let string_table = &serialized_profile.string_table;
 
         let samples_labels: [&[(&str, &str)]; 2] = [
-            // The trace endpoint label should be added to the first sample
+            // The trace endpoint label should be added to the first sample.
             &[
                 ("other", "test"),
                 ("local root span id", "10"),
                 ("trace endpoint", "my endpoint"),
             ],
-            // The trace endpoint label shouldn't be added to second sample (span id doesn't match)
+            // The trace endpoint label shouldn't be added to second sample (span id doesn't match).
             &[("local root span id", "11"), ("other", "test")],
         ];
         assert_eq!(serialized_profile.samples.len(), samples_labels.len());
@@ -962,6 +962,82 @@ mod api_test {
         }
     }
 
+    #[test]
+    fn test_local_root_span_id_can_be_str_or_num_and_endpoint_label_works_either_way() {
+        let sample_types = vec![api::ValueType {
+            r#type: "wall-time",
+            unit: "nanoseconds",
+        }];
+
+        let mut profile: Profile = Profile::builder().sample_types(sample_types).build();
+
+        let id_str = api::Label {
+            key: "local root span id",
+            str: Some("10"),
+            num: 0,
+            num_unit: None,
+        };
+
+        let id_num = api::Label {
+            key: "local root span id",
+            str: None,
+            num: 10,
+            num_unit: None,
+        };
+
+        let other_label = api::Label {
+            key: "other",
+            str: Some("test"),
+            num: 0,
+            num_unit: None,
+        };
+
+        let sample1 = api::Sample {
+            locations: vec![],
+            values: vec![10000],
+            labels: vec![other_label, id_str],
+        };
+
+        let sample2 = api::Sample {
+            locations: vec![],
+            values: vec![10000],
+            labels: vec![id_num, other_label],
+        };
+
+        // Sample 3 shouldn't get a trace endpoint label.
+        let sample3 = api::Sample {
+            locations: vec![],
+            values: vec![10000],
+            labels: vec![other_label],
+        };
+
+        profile.add(sample1).expect("add to succeed");
+        profile.add(sample2).expect("add to succeed");
+        profile.add(sample3).expect("add to succeed");
+
+        profile.add_endpoint(10, Cow::from("my endpoint"));
+
+        let serialized_profile = pprof::Profile::from(&profile);
+        assert_eq!(serialized_profile.samples.len(), 3);
+
+        let string_table = &serialized_profile.string_table;
+        let has_trace_endpoint_my_endpoint = |label: &pprof::Label| -> bool {
+            let key = string_table.get(label.key as usize).unwrap();
+            let value = string_table.get(label.str as usize).unwrap();
+            key == "trace endpoint" && value == "my endpoint"
+        };
+
+        // The first two samples should have a "trace endpoint" label with a "my endpoint" value.
+        for sample in serialized_profile.samples[0..2].iter() {
+            assert_eq!(sample.labels.len(), 3);
+            assert!(sample.labels.iter().any(has_trace_endpoint_my_endpoint));
+        }
+
+        // The last sample should not have this label.
+        let sample = serialized_profile.samples.last().unwrap();
+        assert!(!sample.labels.iter().any(has_trace_endpoint_my_endpoint));
+    }
+
     #[test]
     fn endpoint_counts_empty_test() {
         let sample_types = vec![