use std::collections::VecDeque;
use std::io::{stdout, Write};
use std::thread;
use std::time::{Duration, Instant};
use clap::{Arg, Command};
use nenya::pid_controller::PIDController;
use nenya::RateLimiter;
const LINE_LENGTH: usize = 80;
fn main() {
let matches = Command::new("Rate Limiter Simulation")
.about("Simulates a rate limiter using a PID controller")
.arg(
Arg::new("base_tps")
.short('b')
.long("base_tps")
.value_parser(clap::value_parser!(f64))
.default_value("50.0")
.help("Base TPS for the request generator"),
)
.arg(
Arg::new("min_tps")
.short('m')
.long("min_tps")
.value_parser(clap::value_parser!(f64))
.default_value("1.0")
.help("Lower bound of TPS for the rate limiter"),
)
.arg(
Arg::new("max_tps")
.short('x')
.long("max_tps")
.value_parser(clap::value_parser!(f64))
.default_value("60.0")
.help("Upper bound of TPS for the rate limiter"),
)
.arg(
Arg::new("target_tps")
.short('t')
.long("target_tps")
.value_parser(clap::value_parser!(f64))
.default_value("40.0")
.help("Target TPS for the rate limiter"),
)
.arg(
Arg::new("trailing_window")
.short('w')
.long("trailing_window")
.value_parser(clap::value_parser!(u64))
.default_value("5")
.help("Trailing window for TPS calculation (seconds)"),
)
.arg(
Arg::new("duration")
.short('d')
.long("duration")
.value_parser(clap::value_parser!(u64))
.default_value("60")
.help("Duration of the simulation (seconds)"),
)
.arg(
Arg::new("amplitudes")
.short('a')
.long("amplitudes")
.value_parser(clap::value_parser!(f64))
.num_args(1..)
.use_value_delimiter(true)
.default_value("20.0,10.0")
.help("Amplitudes for the sine waves"),
)
.arg(
Arg::new("frequencies")
.short('f')
.long("frequencies")
.value_parser(clap::value_parser!(f64))
.num_args(1..)
.use_value_delimiter(true)
.default_value("0.1,0.5")
.help("Frequencies for the sine waves"),
)
.arg(
Arg::new("kp")
.long("kp")
.value_parser(clap::value_parser!(f64))
.default_value("0.5")
.help("Proportional gain for the PID controller"),
)
.arg(
Arg::new("ki")
.long("ki")
.value_parser(clap::value_parser!(f64))
.default_value("0.1")
.help("Integral gain for the PID controller"),
)
.arg(
Arg::new("kd")
.long("kd")
.value_parser(clap::value_parser!(f64))
.default_value("0.05")
.help("Derivative gain for the PID controller"),
)
.arg(
Arg::new("error_bias")
.long("error_bias")
.value_parser(clap::value_parser!(f64))
.default_value("1.5")
.help("Bias factor for the integral term"),
)
.arg(
Arg::new("error_limit")
.long("error_limit")
.value_parser(clap::value_parser!(f64))
.help("Error limit for the PID controller"),
)
.arg(
Arg::new("output_limit")
.long("output_limit")
.value_parser(clap::value_parser!(f64))
.help("Output limit for the PID controller"),
)
.arg(
Arg::new("update_interval")
.long("update_interval")
.value_parser(clap::value_parser!(u64))
.default_value("1000")
.help("Update interval for the PID controller (milliseconds)"),
)
.get_matches();
let base_tps = *matches.get_one::<f64>("base_tps").unwrap();
let target_tps = *matches.get_one::<f32>("target_tps").unwrap();
let trailing_window = Duration::from_secs(*matches.get_one::<u64>("trailing_window").unwrap());
let duration = Duration::from_secs(*matches.get_one::<u64>("duration").unwrap());
let amplitudes: Vec<f64> = matches
.get_many::<f64>("amplitudes")
.unwrap()
.copied()
.collect();
let frequencies: Vec<f64> = matches
.get_many::<f64>("frequencies")
.unwrap()
.copied()
.collect();
let min_tps = *matches.get_one::<f32>("min_tps").unwrap();
let max_tps = *matches.get_one::<f32>("max_tps").unwrap();
let kp = *matches.get_one::<f32>("kp").unwrap();
let ki = *matches.get_one::<f32>("ki").unwrap();
let kd = *matches.get_one::<f32>("kd").unwrap();
let error_bias = *matches.get_one::<f32>("error_bias").unwrap();
let error_limit = matches.get_one::<f32>("error_limit").copied();
let output_limit = matches.get_one::<f32>("output_limit").copied();
let update_interval =
Duration::from_millis(*matches.get_one::<u64>("update_interval").unwrap());
let pid_controller = PIDController::new(
target_tps,
kp,
ki,
kd,
error_bias,
error_limit,
output_limit,
);
let mut rate_limiter = RateLimiter::new(
target_tps,
min_tps,
max_tps,
pid_controller,
update_interval,
);
let generator = RequestGenerator::new(base_tps, amplitudes, frequencies);
generate_requests(&mut rate_limiter, &generator, trailing_window, duration);
}
fn generate_requests(
rate_limiter: &mut RateLimiter<f32>,
generator: &RequestGenerator,
trailing_window: Duration,
duration: Duration,
) {
let start = Instant::now();
let mut accepted_requests = 0;
let mut total_requests = 0;
let mut total_tps = 0.0;
let mut accepted_tps = 0.0;
let mut trailing_tps = 0.0;
print!("\x1B[2J");
print!("\x1B[0;0H");
println!("Rate Limiter Target");
println!("-------------------");
let mut output_buffer = vec![' '; LINE_LENGTH];
let mut request_times = VecDeque::new();
while Instant::now().duration_since(start) < duration {
let elapsed_seconds = Instant::now().duration_since(start).as_secs_f64();
let generated_tps = generator.generate_request_rate(elapsed_seconds);
let inter_request_delay = if generated_tps > 0.0 {
(1000.0 / generated_tps) as u64
} else {
1000
};
let should_accept_request = rate_limiter.should_throttle();
total_requests += 1;
let now = Instant::now();
for i in 1..LINE_LENGTH {
output_buffer[i - 1] = output_buffer[i];
}
if should_accept_request {
accepted_requests += 1;
output_buffer[LINE_LENGTH - 1] = '.';
request_times.push_back(now);
} else {
output_buffer[LINE_LENGTH - 1] = '!';
}
while let Some(&time) = request_times.front() {
if now.duration_since(time) > trailing_window {
request_times.pop_front();
} else {
break;
}
}
trailing_tps = request_times.len() as f64 / trailing_window.as_secs_f64();
print!("\x1B7");
print!("\x1B[0J");
print!("\r[{}]\n", output_buffer.iter().collect::<String>());
let elapsed = Instant::now().duration_since(start).as_secs_f64();
accepted_tps = accepted_requests as f64 / elapsed;
total_tps = total_requests as f64 / elapsed;
print_metrics(
&total_tps,
&accepted_tps,
&trailing_tps,
rate_limiter,
generated_tps,
);
println!();
print!("\x1B8");
stdout().flush().unwrap();
thread::sleep(Duration::from_millis(inter_request_delay));
}
let elapsed = Instant::now().duration_since(start).as_secs_f64();
print!("\x1B[4;0H");
print_metrics(&total_tps, &accepted_tps, &trailing_tps, rate_limiter, 0.0);
println!("\rElapsed Time (s): {:.2}", elapsed);
println!("\rAccepted Requests: {}", accepted_requests);
}
fn print_metrics(
total_tps: &f64,
accepted_tps: &f64,
trailing_tps: &f64,
rate_limiter: &RateLimiter<f32>,
generated_tps: f64,
) {
println!("\rTotal TPS: {:.2}", total_tps);
println!("\rAccepted TPS: {:.2}", accepted_tps);
println!("\rTrailing Accepted TPS: {:.2}", trailing_tps);
println!("\rGenerated TPS: {:.2}", generated_tps);
println!("\rTarget TPS: {:.2}", rate_limiter.target_rate());
println!("\rMeasured TPS: {:.2}", rate_limiter.request_rate());
}
pub struct RequestGenerator {
pub base_tps: f64,
pub amplitudes: Vec<f64>,
pub frequencies: Vec<f64>,
}
impl RequestGenerator {
pub fn new(base_tps: f64, amplitudes: Vec<f64>, frequencies: Vec<f64>) -> Self {
RequestGenerator {
base_tps,
amplitudes,
frequencies,
}
}
pub fn generate_request_rate(&self, elapsed_seconds: f64) -> f64 {
let mut rate = self.base_tps;
for (amplitude, frequency) in self.amplitudes.iter().zip(self.frequencies.iter()) {
rate += amplitude * (2.0 * std::f64::consts::PI * frequency * elapsed_seconds).sin();
}
rate
}
}