fix(profiling): IO/allocation sampling intervals

profiling/src/allocation/mod.rs

@@ -6,12 +6,11 @@ use crate::profiling::bindings::{self as zend};
 use crate::profiling::config::SystemSettings;
 use crate::profiling::module_globals;
 use crate::profiling::profiler::Profiler;
-use crate::profiling::{RefCellExt, REQUEST_LOCALS};
+use crate::profiling::{sample_exponential_interval, RefCellExt, REQUEST_LOCALS};
 use core::cell::Cell;
 use core::ptr;
 use libc::size_t;
 use log::{debug, trace};
-use rand_distr::{Distribution, Poisson};
 use std::ffi::c_void;
 use std::num::{NonZero, NonZeroU32, NonZeroU64};
 use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
@@ -150,7 +149,7 @@ unsafe extern "C" fn _zend_mm_realloc(
 /// Default sampling interval in bytes (4 MiB).
 pub const DEFAULT_ALLOCATION_SAMPLING_INTERVAL: NonZeroU32 = NonZero::new(1024 * 4096).unwrap();
 
-/// Sampling distance feed into poison sampling algo. This must be > 0.
+/// Mean distance between allocation samples in bytes. This must be > 0.
 pub static ALLOCATION_PROFILING_INTERVAL: AtomicU64 =
     AtomicU64::new(DEFAULT_ALLOCATION_SAMPLING_INTERVAL.get() as u64);
 
@@ -169,7 +168,7 @@ pub static ALLOCATION_PROFILING_SIZE: AtomicU64 = AtomicU64::new(0);
 pub struct AllocationProfilingStats {
     /// Number of bytes remaining until the next sample collection.
     next_sample: i64,
-    poisson: Poisson<f64>,
+    mean: f64,
     #[cfg(php_zts)]
     rng: ThreadRng,
     #[cfg(not(php_zts))]
@@ -178,11 +177,9 @@ pub struct AllocationProfilingStats {
 
 impl AllocationProfilingStats {
     fn new(sampling_distance: NonZeroU64) -> AllocationProfilingStats {
-        // SAFETY: this will only error if lambda <= 0, and it's NonZeroU64.
-        let poisson = unsafe { Poisson::new(sampling_distance.get() as f64).unwrap_unchecked() };
         let mut stats = AllocationProfilingStats {
             next_sample: 0,
-            poisson,
+            mean: sampling_distance.get() as f64,
             #[cfg(php_zts)]
             rng: rand::rng(),
             #[cfg(not(php_zts))]
@@ -193,7 +190,7 @@ impl AllocationProfilingStats {
     }
 
     fn next_sampling_interval(&mut self) {
-        self.next_sample = self.poisson.sample(&mut self.rng) as i64;
+        self.next_sample = sample_exponential_interval(&mut self.rng, self.mean) as i64;
     }
 
     fn should_collect_allocation(&mut self, len: size_t) -> bool {
@@ -342,6 +339,39 @@ pub fn alloc_prof_rshutdown() {
     allocation_ge84::alloc_prof_rshutdown(heap_live_enabled);
 }
 
+#[cfg(all(test, not(php_zts)))]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn allocation_sampling_matches_upscaling_probability() {
+        let default_mean = DEFAULT_ALLOCATION_SAMPLING_INTERVAL.get() as f64;
+        let trials = 100_000;
+        for (mean, size) in [
+            (default_mean, (0.1 * default_mean) as usize),
+            (default_mean, (1.1 * default_mean) as usize),
+            (default_mean, (3.0 * default_mean) as usize),
+            (1.0, 1),
+            (1.0, 4),
+            (4.0, 4),
+        ] {
+            let mut stats = AllocationProfilingStats::new(NonZeroU64::new(mean as u64).unwrap());
+            stats.rng = StdRng::seed_from_u64(42);
+            stats.next_sampling_interval();
+            let sampled = (0..trials)
+                .filter(|_| stats.should_collect_allocation(size))
+                .count();
+            let probability = 1.0 - (-(size as f64) / mean).exp();
+            let expected = trials as f64 * probability;
+            let sigma = (expected * (1.0 - probability)).sqrt();
+            assert!(
+                (sampled as f64 - expected).abs() < 8.0 * sigma,
+                "mean={mean}, size={size}: sampled {sampled}, expected {expected}"
+            );
+        }
+    }
+}
+
 #[cfg(php_zend_mm_set_custom_handlers_ex)]
 #[track_caller]
 fn initialization_panic() -> ! {

profiling/src/io/mod.rs

@@ -4,10 +4,9 @@ pub mod got_elf64;
 pub mod got_macho;
 
 use crate::profiling::profiler::Profiler;
-use crate::profiling::{zend, RefCellExt, REQUEST_LOCALS};
+use crate::profiling::{sample_exponential_interval, zend, RefCellExt, REQUEST_LOCALS};
 use libc::{c_int, c_void, fstat, stat, S_IFMT, S_IFSOCK};
 use rand::rngs::ThreadRng;
-use rand_distr::{Distribution, Poisson};
 use rustc_hash::FxHashMap;
 use std::cell::RefCell;
 use std::mem::MaybeUninit;
@@ -603,25 +602,24 @@ fn collect_file_write_size(value: u64) {
 
 pub struct IOProfilingStats {
     next_sample: u64,
-    poisson: Poisson<f64>,
+    mean: f64,
     rng: ThreadRng,
 }
 
 impl IOProfilingStats {
-    fn new(lambda: f64) -> Self {
-        // Safety: this will only error if lambda <= 0
-        let poisson = Poisson::new(lambda).unwrap();
+    fn new(mean: u64) -> Self {
+        assert!(mean > 0);
         let mut stats = IOProfilingStats {
+            mean: mean as f64,
             next_sample: 0,
-            poisson,
             rng: rand::rng(),
         };
         stats.next_sampling_interval();
         stats
     }
 
     fn next_sampling_interval(&mut self) {
-        self.next_sample = self.poisson.sample(&mut self.rng) as u64;
+        self.next_sample = sample_exponential_interval(&mut self.rng, self.mean) as u64;
     }
 
     fn should_collect(&mut self, value: u64) -> bool {
@@ -634,8 +632,8 @@ impl IOProfilingStats {
             // (or risking a crash) we refrain from collection I/O.
             return false;
         }
-        if let Some(next_sample) = self.next_sample.checked_sub(value) {
-            self.next_sample = next_sample;
+        if self.next_sample > value {
+            self.next_sample -= value;
             return false;
         }
         self.next_sampling_interval();
@@ -646,42 +644,42 @@ impl IOProfilingStats {
 thread_local! {
     static SOCKET_READ_TIME_PROFILING_STATS: RefCell<IOProfilingStats> = RefCell::new(
         IOProfilingStats::new(
-            SOCKET_READ_TIME_PROFILING_INTERVAL.load(Ordering::Relaxed) as f64,
+            SOCKET_READ_TIME_PROFILING_INTERVAL.load(Ordering::Relaxed),
         )
     );
     static SOCKET_WRITE_TIME_PROFILING_STATS: RefCell<IOProfilingStats> = RefCell::new(
         IOProfilingStats::new(
-            SOCKET_WRITE_TIME_PROFILING_INTERVAL.load(Ordering::Relaxed) as f64,
+            SOCKET_WRITE_TIME_PROFILING_INTERVAL.load(Ordering::Relaxed),
         )
     );
     static FILE_READ_TIME_PROFILING_STATS: RefCell<IOProfilingStats> = RefCell::new(
         IOProfilingStats::new(
-            FILE_READ_TIME_PROFILING_INTERVAL.load(Ordering::Relaxed) as f64,
+            FILE_READ_TIME_PROFILING_INTERVAL.load(Ordering::Relaxed),
         )
     );
     static FILE_WRITE_TIME_PROFILING_STATS: RefCell<IOProfilingStats> = RefCell::new(
         IOProfilingStats::new(
-            FILE_WRITE_TIME_PROFILING_INTERVAL.load(Ordering::Relaxed) as f64,
+            FILE_WRITE_TIME_PROFILING_INTERVAL.load(Ordering::Relaxed),
         )
     );
     static SOCKET_READ_SIZE_PROFILING_STATS: RefCell<IOProfilingStats> = RefCell::new(
         IOProfilingStats::new(
-            SOCKET_READ_SIZE_PROFILING_INTERVAL.load(Ordering::Relaxed) as f64,
+            SOCKET_READ_SIZE_PROFILING_INTERVAL.load(Ordering::Relaxed),
         )
     );
     static SOCKET_WRITE_SIZE_PROFILING_STATS: RefCell<IOProfilingStats> = RefCell::new(
         IOProfilingStats::new(
-            SOCKET_WRITE_SIZE_PROFILING_INTERVAL.load(Ordering::Relaxed) as f64,
+            SOCKET_WRITE_SIZE_PROFILING_INTERVAL.load(Ordering::Relaxed),
         )
     );
     static FILE_READ_SIZE_PROFILING_STATS: RefCell<IOProfilingStats> = RefCell::new(
         IOProfilingStats::new(
-            FILE_READ_SIZE_PROFILING_INTERVAL.load(Ordering::Relaxed) as f64,
+            FILE_READ_SIZE_PROFILING_INTERVAL.load(Ordering::Relaxed),
         )
     );
     static FILE_WRITE_SIZE_PROFILING_STATS: RefCell<IOProfilingStats> = RefCell::new(
         IOProfilingStats::new(
-            FILE_WRITE_SIZE_PROFILING_INTERVAL.load(Ordering::Relaxed) as f64,
+            FILE_WRITE_SIZE_PROFILING_INTERVAL.load(Ordering::Relaxed),
         )
     );
 }
@@ -809,6 +807,20 @@ mod tests {
         assert!(!slot_fits_range(usize::MAX, 0x1000, 0x1000));
     }
 
+    #[test]
+    fn sampling_collects_at_interval_boundary() {
+        let vm_interrupt = std::sync::atomic::AtomicBool::new(false);
+        let previous = super::REQUEST_LOCALS.with_borrow_mut(|locals| {
+            std::mem::replace(&mut locals.vm_interrupt_addr, &vm_interrupt)
+        });
+        let mut stats = super::IOProfilingStats::new(100);
+        stats.next_sample = 8;
+        assert!(!stats.should_collect(0));
+        assert!(!stats.should_collect(4));
+        assert!(stats.should_collect(4));
+        super::REQUEST_LOCALS.with_borrow_mut(|locals| locals.vm_interrupt_addr = previous);
+    }
+
     #[test]
     fn no_hooks_are_safe_to_unload() {
         let mut restores = Vec::new();

profiling/src/lib.rs

@@ -42,6 +42,7 @@ use libdd_common::cstr;
 use log::{debug, error, info, trace, warn};
 use profile_tags::{ProfileTagSegment, UnifiedServiceTagSegment};
 use profiler::{LocalRootSpanResourceMessage, Profiler, VmInterrupt};
+use rand::Rng;
 use sapi::Sapi;
 use std::borrow::Cow;
 use std::cell::{BorrowError, BorrowMutError, RefCell};
@@ -56,6 +57,14 @@ use uuid::Uuid;
 /// interior null bytes and must be null terminated.
 static PROFILER_NAME: &CStr = c"datadog-profiling";
 
+/// Draws the exponential sampling distance assumed by libdatadog's upscaler,
+/// rounded up to whole units and clamped to `[1, 20 * mean]`.
+fn sample_exponential_interval(rng: &mut impl Rng, mean: f64) -> f64 {
+    let sample: f64 = rng.random();
+    let sample = if sample <= 0.0 { 1e-10 } else { sample };
+    (-sample.ln() * mean).ceil().clamp(1.0, 20.0 * mean)
+}
+
 // SAFETY: PROFILER_NAME is a valid utf8 string.
 static PROFILER_NAME_STR: &str = match PROFILER_NAME.to_str() {
     Ok(s) => s,

profiling/tests/correctness/allocation_time_combined.json

@@ -1,19 +1,20 @@
 {
   "scale_by_duration": true,
   "test_name": "php_allocation_time_combined",
+  "note": "Each iteration allocates three equal 10 MB strings: two in str_replace and one in str_repeat. With at least 128 iterations and p = 1 - exp(-10000000 / 4194304) = 0.9078, the binomial model gives an approximate share standard deviation of at most 0.77 percentage points. The 6-point margins allow over eight standard deviations, including the analyzer's integer truncation of the exact 2/3 and 1/3 shares.",
   "stacks": [
     {
       "profile-type": "alloc-size",
       "stack-content": [
         {
           "regular_expression": "<?php;main;standard\\|str_replace$",
           "percent": 66,
-          "error_margin": 1
+          "error_margin": 6
         },
         {
           "regular_expression": "<?php;main;standard\\|str_repeat$",
           "percent": 33,
-          "error_margin": 1
+          "error_margin": 6
         }
       ]
     },
@@ -23,12 +24,12 @@
         {
           "regular_expression": "<?php;main;standard\\|str_replace$",
           "percent": 66,
-          "error_margin": 3
+          "error_margin": 6
         },
         {
           "regular_expression": "<?php;main;standard\\|str_repeat$",
           "percent": 33,
-          "error_margin": 3
+          "error_margin": 6
         }
       ]
     },

profiling/tests/correctness/allocation_time_combined.php

@@ -13,7 +13,8 @@ function main() {
     $duration = $_ENV["EXECUTION_TIME"] ?? 10;
     $end = microtime(true) + $duration;
 
-    while (microtime(true) < $end) {
+    // Keep enough allocations for stable shares, even on slower CI workers.
+    for ($i = 0; $i < 128 || microtime(true) < $end; $i++) {
         // str_replace is frameless in PHP 8.4+ and allocates a new string
         $xs = str_repeat("x", 10_000_000); // 10MB source
         $ys = str_replace("x", "y", $xs); // 10MB allocation in frameless function

profiling/tests/correctness/allocations.json

@@ -1,54 +1,59 @@
 {
-  "scale_by_duration": true,
+  "scale_by_duration": false,
   "test_name": "php_allocations",
+  "note": "512 iterations allocate 18874368000 payload bytes in 2048 strings; string headers add less than 0.001%. At the 4 MiB interval, p = 1 - exp(-size / interval) is 0.9466 or 0.7689. Binomial sampling gives relative standard deviations of 0.76% for total bytes and 0.93% for total count; the 6% margins exceed six standard deviations. Approximate stack-share standard deviations are at most 0.36 percentage points for bytes and 0.49 for count; the 3-point margins exceed six standard deviations. The analyzer truncates shares to integer percentages, hence 33/16 for the exact 1/3 and 1/6 byte shares.",
   "stacks": [
     {
       "profile-type": "alloc-size",
+      "value-matching-sum": 18874368000,
+      "error-margin": 6,
       "stack-content": [
         {
           "regular_expression": "<?php;main;a;standard\\|str_repeat$",
           "percent": 33,
-          "error_margin": 5
+          "error_margin": 3
         },
         {
           "regular_expression": "<?php;main;a;standard\\|str_replace$",
           "percent": 33,
-          "error_margin": 5
+          "error_margin": 3
         },
         {
           "regular_expression": "<?php;main;b;standard\\|str_repeat$",
           "percent": 16,
-          "error_margin": 5
+          "error_margin": 3
         },
         {
           "regular_expression": "<?php;main;b;standard\\|str_replace$",
           "percent": 16,
-          "error_margin": 5
+          "error_margin": 3
         }
       ]
     },
     {
       "profile-type": "alloc-samples",
+      "value-matching-sum": 2048,
+      "error-margin": 6,
       "stack-content": [
         {
           "regular_expression": "<?php;main;a;standard\\|str_repeat$",
           "percent": 25,
-          "error_margin": 5
+          "error_margin": 3
         },
         {
           "regular_expression": "<?php;main;a;standard\\|str_replace$",
           "percent": 25,
-          "error_margin": 5
+          "error_margin": 3
         },
         {
           "regular_expression": "<?php;main;b;standard\\|str_repeat$",
           "percent": 25,
-          "error_margin": 5
+          "error_margin": 3
         },
         {
           "regular_expression": "<?php;main;b;standard\\|str_replace$",
           "percent": 25,
-          "error_margin": 5
+          "error_margin": 3
         }
       ]
     }

profiling/tests/correctness/allocations.php

@@ -3,30 +3,24 @@
 function a()
 {
     $a = str_repeat("a", 1024 * 12_000);
-    str_replace('a', 'b', $a);
+    // One replacement still allocates a full copy, without per-byte work.
+    $a[0] = 'b';
+    str_replace('b', 'c', $a);
 }
 
 function b()
 {
     $a = str_repeat("a", 1024 * 6_000);
-    str_replace('a', 'b', $a);
+    $a[0] = 'b';
+    str_replace('b', 'c', $a);
 }
 
 function main()
 {
-    $duration = $_ENV["EXECUTION_TIME"] ?? 10;
-    $end = microtime(true) + $duration;
-    while (microtime(true) < $end) {
-        $start = microtime(true);
+    // Fixed work makes allocation totals independent of machine speed.
+    for ($i = 0; $i < 512; $i++) {
         a();
         b();
-        $elapsed = microtime(true) - $start;
-        // sleep for the remainder to 100 ms
-        // so we end up doing 10 iterations per second
-        $sleep = (0.1 - $elapsed);
-        if ($sleep > 0.0) {
-            usleep((int) ($sleep * 1_000_000));
-        }
     }
 }
 main();