diff --git a/profiling/src/api.rs b/profiling/src/api.rs index e90c5039f..17ea9481b 100644 --- a/profiling/src/api.rs +++ b/profiling/src/api.rs @@ -196,6 +196,12 @@ pub enum UpscalingInfo { count_value_offset: usize, sampling_distance: u64, }, + PoissonNonSampleTypeCount { + // sum_value_offset is an offset in the profile values type array + sum_value_offset: usize, + count_value: u64, + sampling_distance: u64, + }, Proportional { scale: f64, }, @@ -213,6 +219,15 @@ impl std::fmt::Display for UpscalingInfo { "Poisson = sum_value_offset: {}, count_value_offset: {}, sampling_distance: {}", sum_value_offset, count_value_offset, sampling_distance ), + UpscalingInfo::PoissonNonSampleTypeCount { + sum_value_offset, + count_value, + sampling_distance, + } => write!( + f, + "Poisson = sum_value_offset: {}, count_value: {}, sampling_distance: {}", + sum_value_offset, count_value, sampling_distance + ), UpscalingInfo::Proportional { scale } => { write!(f, "Proportional = scale: {}", scale) } @@ -241,6 +256,28 @@ impl UpscalingInfo { sampling_distance ) } + UpscalingInfo::PoissonNonSampleTypeCount { + sum_value_offset, + count_value, + sampling_distance, + } => { + anyhow::ensure!( + sum_value_offset < &number_of_values, + "sum_value_offset {} must be strictly less than {}", + sum_value_offset, + number_of_values + ); + anyhow::ensure!( + count_value != &0, + "count_value {} must be greater than 0", + count_value + ); + anyhow::ensure!( + sampling_distance != &0, + "sampling_distance {} must be greater than 0", + sampling_distance + ) + } UpscalingInfo::Proportional { scale: _ } => (), } anyhow::Ok(()) diff --git a/profiling/src/internal/profile/mod.rs b/profiling/src/internal/profile/mod.rs index 109b4e323..670f5024c 100644 --- a/profiling/src/internal/profile/mod.rs +++ b/profiling/src/internal/profile/mod.rs @@ -1428,6 +1428,38 @@ mod api_tests { assert_eq!(first.values, vec![1, 298, 29]); } + #[test] + fn test_upscaling_by_value_on_one_value_with_poisson_count() { + let sample_types = create_samples_types(); + + let mut profile = Profile::new(SystemTime::now(), &sample_types, None); + + let sample1 = api::Sample { + locations: vec![], + values: vec![1, 16, 29], + labels: vec![], + }; + + profile.add_sample(sample1, None).expect("add to success"); + + let upscaling_info = UpscalingInfo::PoissonNonSampleTypeCount { + sum_value_offset: 1, + count_value: 29, + sampling_distance: 10, + }; + let values_offset: Vec = vec![1]; + profile + .add_upscaling_rule(values_offset.as_slice(), "", "", upscaling_info) + .expect("Rule added"); + + let serialized_profile = pprof::roundtrip_to_pprof(profile).unwrap(); + + assert_eq!(serialized_profile.samples.len(), 1); + let first = serialized_profile.samples.first().expect("one sample"); + + assert_eq!(first.values, vec![1, 298, 29]); + } + #[test] fn test_upscaling_by_value_on_zero_value_with_poisson() { let sample_types = create_samples_types(); diff --git a/profiling/src/internal/upscaling.rs b/profiling/src/internal/upscaling.rs index 5227d9eb1..2142fef49 100644 --- a/profiling/src/internal/upscaling.rs +++ b/profiling/src/internal/upscaling.rs @@ -28,6 +28,20 @@ impl UpscalingRule { let avg = values[sum_value_offset] as f64 / values[count_value_offset] as f64; 1_f64 / (1_f64 - (-avg / sampling_distance as f64).exp()) } + UpscalingInfo::PoissonNonSampleTypeCount { + sum_value_offset, + count_value, + sampling_distance, + } => { + // This should not happen, but if it happens, + // do not upscale + if values[sum_value_offset] == 0 || count_value == 0 { + return 1_f64; + } + + let avg = values[sum_value_offset] as f64 / count_value as f64; + 1_f64 / (1_f64 - (-avg / sampling_distance as f64).exp()) + } UpscalingInfo::Proportional { scale } => scale, } }