4676fe79e2
- ml-observability (1.2K lines): alerts, dashboards, metrics modules. Depends on ml-core + common (ModelType). 4 tests passing. - ml-stress-testing (1.3K lines): load_generator, market_simulator, performance_analyzer modules. Depends on ml-core + common + config. 5 tests passing. - ml-security (1.4K lines): anomaly_detector, prediction_validator modules. Depends on ml-core + ml-ensemble (EnsembleDecision, ModelVote, TradingAction). 17 tests passing. Total: 18 sub-crates extracted from ml monolith. Workspace: 0 errors, ml tests 876 + 26 in new sub-crates. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
108 lines
3.1 KiB
Rust
108 lines
3.1 KiB
Rust
//! Load generation for ML model stress testing
|
|
|
|
use anyhow::Result;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::sync::Arc;
|
|
use std::time::{Duration, Instant};
|
|
use tokio::sync::Semaphore;
|
|
|
|
/// Load generation profiles
|
|
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
|
|
pub enum LoadProfile {
|
|
Constant,
|
|
Ramp,
|
|
Spike,
|
|
Burst,
|
|
Sine,
|
|
}
|
|
|
|
/// Traffic pattern configuration
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct TrafficPattern {
|
|
pub profile: LoadProfile,
|
|
pub base_rps: u32,
|
|
pub peak_rps: u32,
|
|
pub pattern_duration_seconds: u64,
|
|
}
|
|
|
|
/// Load generator for stress testing
|
|
#[derive(Debug)]
|
|
pub struct LoadGenerator {
|
|
target_rps: u32,
|
|
concurrent_limit: Arc<Semaphore>,
|
|
load_multiplier: f64,
|
|
}
|
|
|
|
impl LoadGenerator {
|
|
pub fn new(target_rps: u32, concurrent_connections: u32) -> Result<Self> {
|
|
Ok(Self {
|
|
target_rps,
|
|
concurrent_limit: Arc::new(Semaphore::new(concurrent_connections as usize)),
|
|
load_multiplier: 1.0,
|
|
})
|
|
}
|
|
|
|
pub fn set_load_multiplier(&mut self, multiplier: f64) {
|
|
self.load_multiplier = multiplier;
|
|
}
|
|
|
|
pub fn get_current_rps(&self) -> u32 {
|
|
(self.target_rps as f64 * self.load_multiplier) as u32
|
|
}
|
|
|
|
/// Generate load according to traffic pattern
|
|
pub async fn generate_load<F>(&self, pattern: TrafficPattern, mut operation: F) -> Result<()>
|
|
where
|
|
F: FnMut() -> Result<()> + Send,
|
|
{
|
|
let start_time = Instant::now();
|
|
let pattern_duration = Duration::from_secs(pattern.pattern_duration_seconds);
|
|
|
|
while start_time.elapsed() < pattern_duration {
|
|
let elapsed_ratio = start_time.elapsed().as_secs_f64() / pattern_duration.as_secs_f64();
|
|
let current_rps = self.calculate_rps_for_pattern(&pattern, elapsed_ratio);
|
|
|
|
let interval = Duration::from_secs_f64(1.0 / current_rps as f64);
|
|
|
|
// Acquire semaphore permit for concurrency control
|
|
let _permit = self.concurrent_limit.acquire().await?;
|
|
|
|
// Execute operation
|
|
operation()?;
|
|
|
|
tokio::time::sleep(interval).await;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn calculate_rps_for_pattern(&self, pattern: &TrafficPattern, elapsed_ratio: f64) -> u32 {
|
|
let base = pattern.base_rps as f64;
|
|
let peak = pattern.peak_rps as f64;
|
|
|
|
let current_rps = match pattern.profile {
|
|
LoadProfile::Constant => base,
|
|
LoadProfile::Ramp => base + (peak - base) * elapsed_ratio,
|
|
LoadProfile::Spike => {
|
|
if elapsed_ratio > 0.8 && elapsed_ratio < 0.9 {
|
|
peak
|
|
} else {
|
|
base
|
|
}
|
|
},
|
|
LoadProfile::Burst => {
|
|
if elapsed_ratio.rem_euclid(0.2) < 0.1 {
|
|
peak
|
|
} else {
|
|
base
|
|
}
|
|
},
|
|
LoadProfile::Sine => {
|
|
base + (peak - base) * (std::f64::consts::PI * elapsed_ratio * 2.0).sin().abs()
|
|
},
|
|
};
|
|
|
|
(current_rps * self.load_multiplier) as u32
|
|
}
|
|
}
|