Cleanups

Author: rcoh

dial9-tokio-telemetry/src/task_dumped.rs

@@ -1,5 +1,5 @@
 //! `TaskDumped<F>` wraps a future and captures async backtraces at yield
-//! points using geometric (Poisson) sampling keyed on idle duration.
+//! points using Poisson sampling keyed on idle duration.
 //!
 //! This wrapper is intentionally separate from [`crate::traced::Traced`]: the
 //! wake-event capture in `Traced` runs on every instrumented spawn regardless
@@ -83,7 +83,7 @@ impl SplitMix64 {
 
 pin_project! {
     /// Future wrapper that captures async backtraces at yield points using
-    /// geometric (Poisson) sampling keyed on idle duration.
+    /// Poisson sampling keyed on idle duration.
     pub(crate) struct TaskDumped<F> {
         #[pin]
         inner: F,
@@ -93,7 +93,7 @@ pin_project! {
         // Monotonic nanoseconds when the frames in `frames` were captured.
         // Only meaningful when `frames.has_data()`.
         pending_capture_ts: Option<NonZeroU64>,
-        // Geometric sampling state: remaining nanoseconds of idle time before
+        // Sampling state: remaining nanoseconds of idle time before
         // the next sample triggers. Signed so subtracting a large idle from a
         // small remaining value goes negative rather than wrapping.
         next_sample_ns: i64,
@@ -107,15 +107,14 @@ pin_project! {
 impl<F> TaskDumped<F> {
     pub(crate) fn new(inner: F, shared: Arc<SharedState>, task_id: TaskId) -> Self {
         let sample_mean_ns = shared.task_dump_idle_threshold_ns.load(Ordering::Relaxed);
-        let base_seed = shared.task_dump_rng_seed.load(Ordering::Relaxed);
-        // When a fixed seed is configured (non-zero), use it directly for
-        // deterministic tests. Otherwise use task_id + timestamp for
-        // production uniqueness across tasks.
-        let seed = if base_seed != 0 {
-            base_seed
-        } else {
-            (task_id.to_u64()).wrapping_mul(0x517cc1b727220a95)
-                ^ crate::telemetry::events::clock_monotonic_ns()
+        // When a fixed seed is configured, use it directly for deterministic
+        // tests. Otherwise use task_id + timestamp for production uniqueness.
+        let seed = match shared.task_dump_rng_seed {
+            Some(s) => s,
+            None => {
+                (task_id.to_u64()).wrapping_mul(0x517cc1b727220a95)
+                    ^ crate::telemetry::events::clock_monotonic_ns()
+            }
         };
         let mut rng = SplitMix64::new(seed);
         let next_sample_ns = rng.draw_exponential_ns(sample_mean_ns);
@@ -148,7 +147,7 @@ impl<F: Future> Future for TaskDumped<F> {
             return this.inner.poll(cx);
         }
 
-        // Geometric sampling: subtract the idle duration from the counter.
+        // Poisson sampling: subtract the idle duration from the counter.
         // If it goes to zero or below, we should emit.
         let poll_start = crate::telemetry::recorder::poll_start_ts_or_now();
         let should_emit = match *this.pending_capture_ts {
@@ -289,3 +288,43 @@ impl Encodable for TaskDumpData<'_> {
         });
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::SplitMix64;
+
+    #[test]
+    fn splitmix64_deterministic() {
+        let mut rng = SplitMix64::new(42);
+        let a = rng.next_u64();
+        let b = rng.next_u64();
+
+        let mut rng2 = SplitMix64::new(42);
+        assert_eq!(a, rng2.next_u64());
+        assert_eq!(b, rng2.next_u64());
+    }
+
+    #[test]
+    fn draw_exponential_ns_mean_is_reasonable() {
+        let mut rng = SplitMix64::new(123);
+        let mean_ns: u64 = 10_000_000; // 10ms
+        let n = 10_000;
+        let sum: f64 = (0..n)
+            .map(|_| rng.draw_exponential_ns(mean_ns) as f64)
+            .sum();
+        let observed_mean = sum / n as f64;
+        // Within 10% of the configured mean.
+        assert!(
+            (observed_mean - mean_ns as f64).abs() < mean_ns as f64 * 0.1,
+            "observed mean {observed_mean} too far from expected {mean_ns}"
+        );
+    }
+
+    #[test]
+    fn draw_exponential_ns_always_positive() {
+        let mut rng = SplitMix64::new(0);
+        for _ in 0..10_000 {
+            assert!(rng.draw_exponential_ns(1_000_000) >= 1);
+        }
+    }
+}

dial9-tokio-telemetry/src/telemetry/recorder/mod.rs

@@ -1473,15 +1473,13 @@ impl TelemetryCore {
         worker_metrics_sink: Option<metrique_writer::BoxEntrySink>,
     ) -> std::io::Result<TelemetryGuard> {
         let start_mono_ns = crate::telemetry::events::clock_monotonic_ns();
-        let shared = Arc::new(SharedState::new(start_mono_ns));
+        let rng_seed = task_dump_config.as_ref().and_then(|cfg| cfg.rng_seed());
+        let shared = Arc::new(SharedState::new(start_mono_ns, rng_seed));
         if let Some(cfg) = task_dump_config.as_ref() {
             shared.task_dumps_enabled.store(true, Ordering::Relaxed);
             shared
                 .task_dump_idle_threshold_ns
                 .store(cfg.idle_threshold().as_nanos() as u64, Ordering::Relaxed);
-            if let Some(seed) = cfg.rng_seed() {
-                shared.task_dump_rng_seed.store(seed, Ordering::Relaxed);
-            }
         }
         #[allow(unused_mut)]
         let mut event_writer = EventWriter::new(Box::new(writer));
@@ -2130,7 +2128,7 @@ mod tests {
 
     #[test]
     fn test_shared_state_no_spawn_location_fields() {
-        let _shared = SharedState::new(crate::telemetry::events::clock_monotonic_ns());
+        let _shared = SharedState::new(crate::telemetry::events::clock_monotonic_ns(), None);
     }
 
     #[test]

dial9-tokio-telemetry/src/telemetry/recorder/shared_state.rs

@@ -32,8 +32,9 @@ pub(crate) struct SharedState {
     /// each `TaskDumped` instance at construction time so the hot poll path
     /// does not need an atomic load.
     pub(crate) task_dump_idle_threshold_ns: AtomicU64,
-    /// Fixed RNG seed for deterministic task dump sampling (0 = use timestamp).
-    pub(crate) task_dump_rng_seed: AtomicU64,
+    /// Fixed RNG seed for deterministic task dump sampling. Set once at
+    /// construction before the `Arc` is shared; read-only thereafter.
+    pub(crate) task_dump_rng_seed: Option<u64>,
     pub(crate) collector: Arc<CentralCollector>,
     /// Absolute `CLOCK_MONOTONIC` nanosecond timestamp captured at trace start.
     pub(crate) start_time_ns: u64,
@@ -60,12 +61,12 @@ pub(crate) struct SharedState {
 }
 
 impl SharedState {
-    pub(super) fn new(start_time_ns: u64) -> Self {
+    pub(super) fn new(start_time_ns: u64, task_dump_rng_seed: Option<u64>) -> Self {
         Self {
             enabled: AtomicBool::new(false),
             task_dumps_enabled: AtomicBool::new(false),
             task_dump_idle_threshold_ns: AtomicU64::new(0),
-            task_dump_rng_seed: AtomicU64::new(0),
+            task_dump_rng_seed,
             collector: Arc::new(CentralCollector::new()),
             start_time_ns,
             next_worker_id: AtomicU64::new(0),
@@ -252,7 +253,7 @@ mod tests {
 
     /// Helper: create a SharedState with recording enabled.
     fn enabled_shared_state() -> SharedState {
-        let ss = SharedState::new(0);
+        let ss = SharedState::new(0, None);
         ss.enabled.store(true, Ordering::Relaxed);
         ss
     }

dial9-tokio-telemetry/src/telemetry/task_dump_config.rs

@@ -1,7 +1,7 @@
 //! Configuration for task dump capture.
 //!
 //! Task dumps capture async backtraces at yield points for tasks that have
-//! been idle, using geometric (Poisson) sampling keyed on idle duration.
+//! been idle, using Poisson sampling keyed on idle duration.
 //! Use [`TaskDumpConfig`] with
 //! [`TracedRuntimeBuilder::with_task_dumps`](crate::telemetry::TracedRuntimeBuilder::with_task_dumps)
 //! or [`TelemetryCoreBuilder::task_dump_config`](crate::telemetry::TelemetryCoreBuilder::task_dump_config).
@@ -12,13 +12,13 @@
 
 use std::time::Duration;
 
-/// Default mean idle duration for geometric sampling.
+/// Default mean idle duration for Poisson sampling.
 const DEFAULT_IDLE_THRESHOLD: Duration = Duration::from_millis(10);
 
 /// Configuration for task dump capture.
 #[derive(Debug, Clone, bon::Builder)]
 pub struct TaskDumpConfig {
-    /// Mean idle duration for geometric (Poisson) sampling. On average, one
+    /// Mean idle duration for Poisson sampling. On average, one
     /// task dump is emitted per this amount of cumulative idle time. Shorter
     /// idles have a lower (but non-zero) probability of triggering a dump;
     /// longer idles are very likely to trigger. Defaults to 10ms.
@@ -39,7 +39,7 @@ impl Default for TaskDumpConfig {
 }
 
 impl TaskDumpConfig {
-    /// Mean idle duration for geometric sampling.
+    /// Mean idle duration for Poisson sampling.
     pub fn idle_threshold(&self) -> Duration {
         self.idle_threshold
     }

dial9-tokio-telemetry/src/unwind.rs

@@ -68,6 +68,6 @@ pub(crate) fn collect_frames(
         leaf_addr,
     };
     unsafe {
-        _Unwind_Backtrace(trace_fn, core::ptr::addr_of_mut!(data).cast());
+        _Unwind_Backtrace(trace_fn, (&raw mut data).cast());
     }
 }