use std::ops::Range;
use std::path::PathBuf;
use std::time::Duration;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) enum SimTransport {
Tcp,
Grpc,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) enum SimDataFormat {
HyperBinary,
ArrowStream,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct DurationRangeMs {
pub min_ms: u64,
pub max_ms: u64,
}
impl DurationRangeMs {
pub(crate) fn new(min: Duration, max: Duration) -> Self {
Self {
min_ms: min.as_millis() as u64,
max_ms: max.as_millis() as u64,
}
}
pub(crate) fn to_range(&self) -> Range<Duration> {
Duration::from_millis(self.min_ms)..Duration::from_millis(self.max_ms)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct SimulationConfig {
pub duration_secs: u64,
pub transport: SimTransport,
pub data_format: SimDataFormat,
pub num_databases: usize,
pub max_total_disk_bytes: u64,
pub inserter_users: usize,
pub query_users: usize,
pub mixed_users: usize,
pub max_memory_mb: u64,
pub max_cpu_percent: f64,
pub log_resource_interval_secs: u64,
pub log_resource_mode: u16,
pub memory_limit: String,
pub seed: Option<u64>,
pub think_time: DurationRangeMs,
pub batch_size_min: usize,
pub batch_size_max: usize,
pub query_complexity_min: usize,
pub query_complexity_max: usize,
pub output_dir: Option<PathBuf>,
}
impl SimulationConfig {
pub(crate) fn total_users(&self) -> usize {
self.inserter_users + self.query_users + self.mixed_users
}
pub(crate) fn duration(&self) -> Duration {
Duration::from_secs(self.duration_secs)
}
pub(crate) fn from_env() -> Self {
fn env_u64(key: &str, default: u64) -> u64 {
std::env::var(key)
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(default)
}
fn env_usize(key: &str, default: usize) -> usize {
std::env::var(key)
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(default)
}
fn env_f64(key: &str, default: f64) -> f64 {
std::env::var(key)
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(default)
}
fn env_str(key: &str, default: &str) -> String {
std::env::var(key).unwrap_or_else(|_| default.to_string())
}
let transport = match env_str("STRESS_TRANSPORT", "tcp").to_lowercase().as_str() {
"grpc" => SimTransport::Grpc,
_ => SimTransport::Tcp,
};
let data_format = match env_str("STRESS_DATA_FORMAT", "hyperbinary")
.to_lowercase()
.as_str()
{
"arrow" | "arrowstream" => SimDataFormat::ArrowStream,
_ => SimDataFormat::HyperBinary,
};
SimulationConfig {
duration_secs: env_u64("STRESS_DURATION", 300), transport,
data_format,
num_databases: env_usize("STRESS_DATABASES", 3),
max_total_disk_bytes: env_u64("STRESS_MAX_DISK_MB", 2048) * 1024 * 1024,
inserter_users: env_usize("STRESS_INSERTER_USERS", 4),
query_users: env_usize("STRESS_QUERY_USERS", 3),
mixed_users: env_usize("STRESS_MIXED_USERS", 2),
max_memory_mb: env_u64("STRESS_MAX_MEMORY_MB", 4096),
max_cpu_percent: env_f64("STRESS_MAX_CPU_PERCENT", 90.0),
log_resource_interval_secs: env_u64("STRESS_LOG_INTERVAL", 5),
log_resource_mode: env_u64("STRESS_LOG_MODE", 2047) as u16,
memory_limit: env_str("STRESS_MEMORY_LIMIT", "80%"),
seed: std::env::var("STRESS_SEED")
.ok()
.and_then(|v| v.parse().ok()),
think_time: DurationRangeMs::new(
Duration::from_millis(env_u64("STRESS_THINK_MIN_MS", 10)),
Duration::from_millis(env_u64("STRESS_THINK_MAX_MS", 200)),
),
batch_size_min: env_usize("STRESS_BATCH_MIN", 100),
batch_size_max: env_usize("STRESS_BATCH_MAX", 10_000),
query_complexity_min: 1,
query_complexity_max: env_usize("STRESS_QUERY_COMPLEXITY_MAX", 3),
output_dir: std::env::var("STRESS_OUTPUT_DIR").ok().map(PathBuf::from),
}
}
}
impl Default for SimulationConfig {
fn default() -> Self {
Self::from_env()
}
}