use crate::NetworkConfig;
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::{Duration, Instant};
use thiserror::Error;
#[derive(Debug, Error)]
pub enum AutoTunerError {
#[error("Failed to detect system resources: {0}")]
ResourceDetectionFailed(String),
#[error("Invalid configuration: {0}")]
InvalidConfig(String),
#[error("Tuning not initialized")]
NotInitialized,
#[error("Monitoring already running")]
MonitoringActive,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystemResources {
pub total_memory: u64,
pub available_memory: u64,
pub cpu_cores: usize,
pub network_bandwidth: u64,
pub is_battery_powered: bool,
}
impl SystemResources {
pub fn detect() -> Result<Self, AutoTunerError> {
Ok(Self {
total_memory: 4 * 1024 * 1024 * 1024, available_memory: 2 * 1024 * 1024 * 1024, cpu_cores: num_cpus::get(),
network_bandwidth: 0, is_battery_powered: false,
})
}
pub fn memory_category(&self) -> &'static str {
match self.total_memory {
0..134_217_728 => "very_low", 134_217_728..536_870_912 => "low", 536_870_912..2_147_483_648 => "medium", 2_147_483_648..8_589_934_592 => "high", _ => "very_high", }
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkloadProfile {
pub avg_connections: f64,
pub avg_query_rate: f64,
pub avg_bandwidth_usage: f64,
pub peak_memory_usage: u64,
pub cpu_bound: bool,
pub bandwidth_bound: bool,
pub memory_bound: bool,
}
impl Default for WorkloadProfile {
fn default() -> Self {
Self {
avg_connections: 0.0,
avg_query_rate: 0.0,
avg_bandwidth_usage: 0.0,
peak_memory_usage: 0,
cpu_bound: false,
bandwidth_bound: false,
memory_bound: false,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AutoTunerConfig {
pub enable_auto_adjust: bool,
pub adjustment_interval: Duration,
pub stabilization_period: Duration,
pub safety_margin: f64,
pub aggressive_mode: bool,
}
impl Default for AutoTunerConfig {
fn default() -> Self {
Self {
enable_auto_adjust: true,
adjustment_interval: Duration::from_secs(300), stabilization_period: Duration::from_secs(60), safety_margin: 0.2,
aggressive_mode: false,
}
}
}
impl AutoTunerConfig {
pub fn conservative() -> Self {
Self {
enable_auto_adjust: true,
adjustment_interval: Duration::from_secs(600), stabilization_period: Duration::from_secs(120), safety_margin: 0.3,
aggressive_mode: false,
}
}
pub fn aggressive() -> Self {
Self {
enable_auto_adjust: true,
adjustment_interval: Duration::from_secs(60), stabilization_period: Duration::from_secs(30), safety_margin: 0.1,
aggressive_mode: true,
}
}
}
#[derive(Debug, Clone, Default)]
pub struct AutoTunerStats {
pub adjustments_made: u64,
pub resource_checks: u64,
pub workload_checks: u64,
pub last_adjustment: Option<Instant>,
pub optimization_score: f64,
}
pub struct AutoTuner {
config: AutoTunerConfig,
system_resources: Option<SystemResources>,
workload_profile: WorkloadProfile,
stats: Arc<RwLock<AutoTunerStats>>,
monitoring_active: Arc<RwLock<bool>>,
}
impl AutoTuner {
pub fn new() -> Self {
Self::with_config(AutoTunerConfig::default())
}
pub fn with_config(config: AutoTunerConfig) -> Self {
Self {
config,
system_resources: None,
workload_profile: WorkloadProfile::default(),
stats: Arc::new(RwLock::new(AutoTunerStats::default())),
monitoring_active: Arc::new(RwLock::new(false)),
}
}
pub async fn analyze_system(&mut self) -> Result<SystemResources, AutoTunerError> {
let resources = SystemResources::detect()?;
self.system_resources = Some(resources.clone());
let mut stats = self.stats.write();
stats.resource_checks += 1;
Ok(resources)
}
pub async fn generate_config(&mut self) -> Result<NetworkConfig, AutoTunerError> {
if self.system_resources.is_none() {
self.analyze_system().await?;
}
let resources = self
.system_resources
.as_ref()
.expect("just populated by analyze_system() if was None");
let usable_factor = 1.0 - self.config.safety_margin;
let mut config = match resources.memory_category() {
"very_low" => NetworkConfig::low_memory(),
"low" => NetworkConfig::iot(),
"medium" => NetworkConfig::mobile(),
"high" | "very_high" => {
if resources.is_battery_powered {
NetworkConfig::mobile()
} else {
NetworkConfig::high_performance()
}
}
_ => NetworkConfig::default(),
};
let base_connections = resources.cpu_cores * 50;
config.max_connections = Some((base_connections as f64 * usable_factor) as usize);
let memory_mb = resources.available_memory / (1024 * 1024);
if memory_mb < 256 {
config.connection_buffer_size = 8 * 1024; config.max_connections = Some(16);
} else if memory_mb < 512 {
config.connection_buffer_size = 16 * 1024; config.max_connections = Some(32);
}
config.enable_nat_traversal =
resources.memory_category() != "very_high" || resources.is_battery_powered;
let mut stats = self.stats.write();
stats.adjustments_made += 1;
stats.last_adjustment = Some(Instant::now());
stats.optimization_score = self.calculate_optimization_score();
Ok(config)
}
pub fn update_workload(
&mut self,
connections: usize,
query_rate: f64,
bandwidth_usage: f64,
memory_usage: u64,
) {
let alpha = 0.3;
let profile = &mut self.workload_profile;
profile.avg_connections =
profile.avg_connections * (1.0 - alpha) + (connections as f64) * alpha;
profile.avg_query_rate = profile.avg_query_rate * (1.0 - alpha) + query_rate * alpha;
profile.avg_bandwidth_usage =
profile.avg_bandwidth_usage * (1.0 - alpha) + bandwidth_usage * alpha;
profile.peak_memory_usage = profile.peak_memory_usage.max(memory_usage);
if let Some(resources) = &self.system_resources {
let memory_usage_ratio = memory_usage as f64 / resources.available_memory as f64;
profile.memory_bound = memory_usage_ratio > 0.8;
profile.cpu_bound = connections > resources.cpu_cores * 100;
profile.bandwidth_bound = resources.network_bandwidth > 0
&& bandwidth_usage > (resources.network_bandwidth as f64 * 0.8);
}
let mut stats = self.stats.write();
stats.workload_checks += 1;
}
pub async fn start_monitoring(&mut self) -> Result<(), AutoTunerError> {
let mut active = self.monitoring_active.write();
if *active {
return Err(AutoTunerError::MonitoringActive);
}
*active = true;
Ok(())
}
pub fn stop_monitoring(&mut self) {
let mut active = self.monitoring_active.write();
*active = false;
}
pub fn is_monitoring(&self) -> bool {
*self.monitoring_active.read()
}
pub fn workload_profile(&self) -> &WorkloadProfile {
&self.workload_profile
}
pub fn stats(&self) -> AutoTunerStats {
self.stats.read().clone()
}
fn calculate_optimization_score(&self) -> f64 {
if self.system_resources.is_none() {
return 0.0;
}
let resources = self
.system_resources
.as_ref()
.expect("just checked is_some above");
let profile = &self.workload_profile;
let memory_score = if profile.peak_memory_usage > 0 {
1.0 - (profile.peak_memory_usage as f64 / resources.available_memory as f64).min(1.0)
} else {
0.5
};
let cpu_score = if profile.cpu_bound { 0.3 } else { 0.8 };
let bandwidth_score = if profile.bandwidth_bound { 0.3 } else { 0.8 };
(memory_score * 0.4 + cpu_score * 0.3 + bandwidth_score * 0.3).clamp(0.0, 1.0)
}
pub fn recommendations(&self) -> Vec<String> {
let mut recommendations = Vec::new();
if let Some(resources) = &self.system_resources {
let profile = &self.workload_profile;
if profile.memory_bound {
recommendations.push(
"Memory usage is high. Consider reducing max_connections or enabling low_memory_mode.".to_string()
);
}
if profile.cpu_bound {
recommendations.push(
format!("CPU usage is high with {} cores. Consider distributing load across more nodes.",
resources.cpu_cores)
);
}
if profile.bandwidth_bound {
recommendations.push(
"Bandwidth is saturated. Consider enabling bandwidth throttling or upgrading network capacity.".to_string()
);
}
if resources.is_battery_powered && profile.avg_query_rate > 10.0 {
recommendations.push(
"High DHT query rate on battery power. Consider enabling query batching."
.to_string(),
);
}
if resources.memory_category() == "very_low" && !profile.memory_bound {
recommendations.push(
"System resources are underutilized. You can increase max_connections for better performance.".to_string()
);
}
} else {
recommendations.push("Run analyze_system() first to get recommendations.".to_string());
}
recommendations
}
}
impl Default for AutoTuner {
fn default() -> Self {
Self::new()
}
}
mod num_cpus {
pub fn get() -> usize {
std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(4) }
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_auto_tuner_creation() {
let tuner = AutoTuner::new();
assert!(!tuner.is_monitoring());
}
#[tokio::test]
async fn test_system_resource_detection() {
let mut tuner = AutoTuner::new();
let resources = tuner
.analyze_system()
.await
.expect("test: analyze_system should succeed");
assert!(resources.cpu_cores > 0);
assert!(resources.total_memory > 0);
}
#[tokio::test]
async fn test_config_generation() {
let mut tuner = AutoTuner::new();
let config = tuner
.generate_config()
.await
.expect("test: generate_config should succeed");
assert!(config.max_connections.is_some());
}
#[tokio::test]
async fn test_workload_update() {
let mut tuner = AutoTuner::new();
tuner
.analyze_system()
.await
.expect("test: analyze_system should succeed");
tuner.update_workload(10, 5.0, 100_000.0, 50_000_000);
let profile = tuner.workload_profile();
assert!(profile.avg_connections > 0.0);
}
#[tokio::test]
async fn test_monitoring_lifecycle() {
let mut tuner = AutoTuner::new();
assert!(!tuner.is_monitoring());
tuner
.start_monitoring()
.await
.expect("test: start_monitoring should succeed when not yet active");
assert!(tuner.is_monitoring());
tuner.stop_monitoring();
assert!(!tuner.is_monitoring());
}
#[test]
fn test_memory_categories() {
let low = SystemResources {
total_memory: 100 * 1024 * 1024, available_memory: 50 * 1024 * 1024,
cpu_cores: 2,
network_bandwidth: 0,
is_battery_powered: true,
};
assert_eq!(low.memory_category(), "very_low");
let high = SystemResources {
total_memory: 16 * 1024 * 1024 * 1024, available_memory: 8 * 1024 * 1024 * 1024,
cpu_cores: 8,
network_bandwidth: 0,
is_battery_powered: false,
};
assert_eq!(high.memory_category(), "very_high");
}
#[tokio::test]
async fn test_statistics_tracking() {
let mut tuner = AutoTuner::new();
let stats_before = tuner.stats();
assert_eq!(stats_before.adjustments_made, 0);
tuner
.generate_config()
.await
.expect("test: generate_config should succeed");
let stats_after = tuner.stats();
assert_eq!(stats_after.adjustments_made, 1);
assert!(stats_after.last_adjustment.is_some());
}
#[tokio::test]
async fn test_recommendations() {
let mut tuner = AutoTuner::new();
tuner
.analyze_system()
.await
.expect("test: analyze_system should succeed");
if let Some(resources) = &tuner.system_resources {
let high_memory = (resources.available_memory as f64 * 0.85) as u64;
tuner.update_workload(50, 20.0, 1_000_000.0, high_memory);
}
let recommendations = tuner.recommendations();
assert!(!recommendations.is_empty());
}
#[tokio::test]
async fn test_config_presets() {
let conservative = AutoTunerConfig::conservative();
assert!(!conservative.aggressive_mode);
assert!(conservative.safety_margin > 0.2);
let aggressive = AutoTunerConfig::aggressive();
assert!(aggressive.aggressive_mode);
assert!(aggressive.safety_margin < 0.2);
}
#[tokio::test]
async fn test_optimization_score() {
let mut tuner = AutoTuner::new();
tuner
.analyze_system()
.await
.expect("test: analyze_system should succeed");
let stats = tuner.stats();
assert!(stats.optimization_score >= 0.0 && stats.optimization_score <= 1.0);
}
}