Skip to content

refactor: extract sampling primitives into shared module#418

Merged
rcoh merged 2 commits into
mainfrom
extract-sampler
May 18, 2026
Merged

refactor: extract sampling primitives into shared module#418
rcoh merged 2 commits into
mainfrom
extract-sampler

Conversation

@rcoh
Copy link
Copy Markdown
Contributor

@rcoh rcoh commented May 18, 2026

Move SplitMix64 and draw_exponential from task_dumped.rs into a new pub(crate) sampling module. Rename draw_exponential_ns → draw_exponential since the math is unit-agnostic (the caller decides whether 'mean' is nanoseconds, bytes, or anything else).

This prepares for the memory profiler (commit 2+), which will reuse the same geometric/Poisson sampling primitive to sample on bytes allocated rather than nanoseconds idle.

Pure refactor — no behavior change. Existing task-dump tests pass unmodified (only the function name changed).

Move SplitMix64 and draw_exponential from task_dumped.rs into a new
pub(crate) sampling module. Rename draw_exponential_ns → draw_exponential
since the math is unit-agnostic (the caller decides whether 'mean' is
nanoseconds, bytes, or anything else).

This prepares for the memory profiler (commit 2+), which will reuse the
same geometric/Poisson sampling primitive to sample on bytes allocated
rather than nanoseconds idle.

Pure refactor — no behavior change. Existing task-dump tests pass
unmodified (only the function name changed).
@rcoh rcoh requested a review from yulnr May 18, 2026 15:03
@rcoh rcoh enabled auto-merge May 18, 2026 15:03
Comment on lines +40 to +90
mod tests {
use super::SplitMix64;

#[test]
fn splitmix_deterministic_with_fixed_seed() {
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_returns_at_least_1() {
let mut rng = SplitMix64::new(0);
for _ in 0..1000 {
assert!(rng.draw_exponential(1) >= 1);
}
}

#[test]
fn draw_exponential_mean_approximates_target() {
let mut rng = SplitMix64::new(123);
let mean: u64 = 1024;
let n = 100_000;
let sum: f64 = (0..n).map(|_| rng.draw_exponential(mean) as f64).sum();
let observed_mean = sum / n as f64;
// Within ±5% of the configured mean.
assert!(
(observed_mean - mean as f64).abs() < mean as f64 * 0.05,
"observed mean {observed_mean} too far from expected {mean}"
);
}

#[test]
fn draw_exponential_handles_large_mean() {
let mut rng = SplitMix64::new(999);
let mean: u64 = 1_000_000_000;
let mut saw_large = false;
for _ in 0..1000 {
let v = rng.draw_exponential(mean);
assert!(v >= 1);
if v > 1_000_000 {
saw_large = true;
}
}
assert!(saw_large, "expected some values much larger than 1");
}
}
Copy link
Copy Markdown
Contributor

@Fluzko Fluzko May 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These are duplicated, they still live at task_dumped.rs. They're not exactly the same, but pretty much they are, probably worth consolidating them in one place

Comment thread dial9-tokio-telemetry/src/sampling.rs Outdated
///
/// The unit is whatever the caller treats `mean` as (nanoseconds,
/// bytes, etc.).
pub(crate) fn draw_exponential(&mut self, mean: u64) -> i64 {
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
pub(crate) fn draw_exponential(&mut self, mean: u64) -> i64 {
pub(crate) fn draw_exponential(&mut self, mean: u64) -> u64 {

this is guaranteed to return >1 so I would change it to a narrower and more correct type

@rcoh rcoh added this pull request to the merge queue May 18, 2026
@rcoh rcoh removed this pull request from the merge queue due to a manual request May 18, 2026
…ted tests

- Change draw_exponential return type from i64 to u64 since the value
  is guaranteed >= 1
- Remove duplicated sampling tests from task_dumped.rs (already covered
  by sampling.rs with stricter tolerances)
@rcoh rcoh added this pull request to the merge queue May 18, 2026
Merged via the queue into main with commit a9b1a3f May 18, 2026
21 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants