poolsim_core/
distribution.rs

1//! Distribution fitting and sampling for workload latency inputs.
2//!
3//! This module converts percentile-based or sample-based latency observations
4//! into a service-time distribution usable by the optimizer and simulation
5//! layers.
6//!
7//! Main entrypoints:
8//!
9//! - [`LatencyDistribution::fit`]
10//! - [`LatencyDistribution::sample_ms`]
11//! - [`LatencyDistribution::percentile_ms`]
12//! - [`LatencyDistribution::mean_ms`]
13//!
14//! Selection rules:
15//!
16//! - percentile-only workloads fit one of the supported parametric models
17//! - `raw_samples_ms` forces empirical fitting regardless of the requested model
18//! - fitted distributions are reused by Monte Carlo and queueing helpers
19
20use rand::Rng;
21use rand_distr::{
22    Distribution as RandDistribution, Exp, Gamma as RandGamma, LogNormal as RandLogNormal,
23};
24use statrs::distribution::{ContinuousCDF, Gamma as StatGamma, LogNormal as StatLogNormal, Normal};
25
26use crate::{
27    error::PoolsimError,
28    types::{DistributionModel, WorkloadConfig},
29};
30
31/// Empirical cumulative distribution built from raw latency samples.
32#[derive(Debug, Clone)]
33pub struct EmpiricalCdf {
34    samples: Vec<f64>,
35}
36
37impl EmpiricalCdf {
38    fn new(mut samples: Vec<f64>) -> Result<Self, PoolsimError> {
39        if samples.is_empty() {
40            return Err(PoolsimError::Distribution(
41                "empirical distribution requires at least one sample".to_string(),
42            ));
43        }
44        samples.sort_by(|a, b| a.total_cmp(b));
45        Ok(Self { samples })
46    }
47
48    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
49        let idx = rng.gen_range(0..self.samples.len());
50        self.samples[idx]
51    }
52
53    fn percentile(&self, p: f64) -> f64 {
54        if self.samples.is_empty() {
55            return 0.0;
56        }
57        let p = p.clamp(0.0, 1.0);
58        let idx = ((self.samples.len() - 1) as f64 * p).round() as usize;
59        self.samples[idx]
60    }
61
62    fn mean(&self) -> f64 {
63        self.samples.iter().sum::<f64>() / self.samples.len() as f64
64    }
65}
66
67/// Service-time distribution used during simulation and queue estimation.
68#[derive(Debug, Clone)]
69pub enum LatencyDistribution {
70    /// Log-normal distribution parameters.
71    LogNormal {
72        /// Log-space mean.
73        mu: f64,
74        /// Log-space standard deviation.
75        sigma: f64,
76    },
77    /// Exponential distribution with the provided mean.
78    Exponential {
79        /// Mean service time in milliseconds.
80        mean_ms: f64,
81    },
82    /// Empirical distribution sampled directly from input samples.
83    Empirical(EmpiricalCdf),
84    /// Gamma distribution parameters.
85    Gamma {
86        /// Shape parameter (k).
87        shape: f64,
88        /// Scale parameter (theta).
89        scale: f64,
90    },
91}
92
93impl LatencyDistribution {
94    /// Fits a latency distribution for a workload and selected model.
95    ///
96    /// If `workload.raw_samples_ms` is present, empirical fitting is used regardless
97    /// of `model`.
98    ///
99    /// # Errors
100    ///
101    /// Returns [`PoolsimError::Distribution`] when distribution parameters cannot
102    /// be derived.
103    pub fn fit(workload: &WorkloadConfig, model: DistributionModel) -> Result<Self, PoolsimError> {
104        if let Some(raw_samples) = &workload.raw_samples_ms {
105            return EmpiricalCdf::new(raw_samples.clone()).map(Self::Empirical);
106        }
107
108        match model {
109            DistributionModel::LogNormal | DistributionModel::Empirical => {
110                let (mu, sigma) = fit_lognormal(workload)?;
111                Ok(Self::LogNormal { mu, sigma })
112            }
113            DistributionModel::Exponential => {
114                let mean_ms = workload.latency_p50_ms / std::f64::consts::LN_2;
115                Ok(Self::Exponential { mean_ms })
116            }
117            DistributionModel::Gamma => {
118                let (shape, scale) = fit_gamma(workload)?;
119                Ok(Self::Gamma { shape, scale })
120            }
121        }
122    }
123
124    /// Draws one latency sample in milliseconds.
125    pub fn sample_ms<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
126        match self {
127            Self::LogNormal { mu, sigma } => {
128                let dist = RandLogNormal::new(*mu, *sigma).expect("valid lognormal params");
129                dist.sample(rng)
130            }
131            Self::Exponential { mean_ms } => {
132                let dist = Exp::new(1.0 / mean_ms).expect("valid exponential rate");
133                dist.sample(rng)
134            }
135            Self::Empirical(empirical) => empirical.sample(rng),
136            Self::Gamma { shape, scale } => {
137                let dist = RandGamma::new(*shape, *scale).expect("valid gamma params");
138                dist.sample(rng)
139            }
140        }
141    }
142
143    /// Returns the requested percentile in milliseconds.
144    ///
145    /// `p` is clamped into `[0, 1]`.
146    ///
147    /// # Errors
148    ///
149    /// Returns [`PoolsimError::Distribution`] when percentile evaluation fails for
150    /// the current parameterization.
151    pub fn percentile_ms(&self, p: f64) -> Result<f64, PoolsimError> {
152        let p = p.clamp(0.0, 1.0);
153        match self {
154            Self::LogNormal { mu, sigma } => {
155                let dist = StatLogNormal::new(*mu, *sigma)
156                    .map_err(|e| PoolsimError::Distribution(e.to_string()))?;
157                Ok(dist.inverse_cdf(p))
158            }
159            Self::Exponential { mean_ms } => Ok(-mean_ms * (1.0 - p).ln()),
160            Self::Empirical(empirical) => Ok(empirical.percentile(p)),
161            Self::Gamma { shape, scale } => {
162                let dist = StatGamma::new(*shape, *scale)
163                    .map_err(|e| PoolsimError::Distribution(e.to_string()))?;
164                Ok(dist.inverse_cdf(p))
165            }
166        }
167    }
168
169    /// Returns the mean service time in milliseconds.
170    pub fn mean_ms(&self) -> f64 {
171        match self {
172            Self::LogNormal { mu, sigma } => (mu + 0.5 * sigma * sigma).exp(),
173            Self::Exponential { mean_ms } => *mean_ms,
174            Self::Empirical(empirical) => empirical.mean(),
175            Self::Gamma { shape, scale } => shape * scale,
176        }
177    }
178}
179
180fn fit_lognormal(workload: &WorkloadConfig) -> Result<(f64, f64), PoolsimError> {
181    let mu = workload.latency_p50_ms.ln();
182    let normal = Normal::new(0.0, 1.0).map_err(|e| PoolsimError::Distribution(e.to_string()))?;
183
184    let mut sigmas = Vec::new();
185    if workload.latency_p95_ms > workload.latency_p50_ms {
186        let z95 = normal.inverse_cdf(0.95);
187        sigmas.push((workload.latency_p95_ms / workload.latency_p50_ms).ln() / z95);
188    }
189    if workload.latency_p99_ms > workload.latency_p50_ms {
190        let z99 = normal.inverse_cdf(0.99);
191        sigmas.push((workload.latency_p99_ms / workload.latency_p50_ms).ln() / z99);
192    }
193
194    let sigma = sigmas
195        .into_iter()
196        .filter(|s| s.is_finite() && *s > 0.0)
197        .sum::<f64>();
198
199    if sigma <= 0.0 {
200        return Err(PoolsimError::Distribution(
201            "unable to derive positive lognormal sigma from percentiles".to_string(),
202        ));
203    }
204
205    let count = if workload.latency_p99_ms > workload.latency_p50_ms {
206        2.0
207    } else {
208        1.0
209    };
210    Ok((mu, sigma / count))
211}
212
213fn fit_gamma(workload: &WorkloadConfig) -> Result<(f64, f64), PoolsimError> {
214    let mean = (workload.latency_p50_ms + workload.latency_p95_ms + workload.latency_p99_ms) / 3.0;
215    let std_est = ((workload.latency_p99_ms - workload.latency_p50_ms) / 2.326_347_874).max(1e-6);
216    let var = std_est * std_est;
217    let shape = (mean * mean / var).max(1e-6);
218    let scale = (var / mean).max(1e-6);
219
220    Ok((shape, scale))
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226
227    #[test]
228    fn empirical_percentile_returns_zero_for_empty_internal_samples() {
229        let empirical = EmpiricalCdf {
230            samples: Vec::new(),
231        };
232        assert_eq!(empirical.percentile(0.5), 0.0);
233    }
234
235    #[test]
236    fn fit_gamma_clamps_non_finite_derived_parameters_to_safe_minimum() {
237        let workload = WorkloadConfig {
238            requests_per_second: 100.0,
239            latency_p50_ms: 10.0,
240            latency_p95_ms: 20.0,
241            latency_p99_ms: f64::INFINITY,
242            raw_samples_ms: None,
243            step_load_profile: None,
244        };
245
246        let (shape, scale) =
247            fit_gamma(&workload).expect("non-finite intermediates should clamp to safe values");
248        assert_eq!(shape, 1e-6);
249        assert_eq!(scale, 1e-6);
250    }
251
252    #[test]
253    fn fit_lognormal_uses_single_tail_when_p99_is_not_above_p50() {
254        let workload = WorkloadConfig {
255            requests_per_second: 100.0,
256            latency_p50_ms: 10.0,
257            latency_p95_ms: 20.0,
258            latency_p99_ms: 10.0,
259            raw_samples_ms: None,
260            step_load_profile: None,
261        };
262
263        let (mu, sigma) =
264            fit_lognormal(&workload).expect("positive p95 spread should fit lognormal");
265        assert!(mu.is_finite());
266        assert!(sigma > 0.0);
267    }
268}
poolsim_core/distribution.rs

poolsim_core/
distribution.rs
