use std::sync::LazyLock;
use majra::ratelimit::RateLimiter;
pub struct NetworkPool {
http: RateLimiter,
dns: RateLimiter,
port: RateLimiter,
}
impl NetworkPool {
#[must_use]
pub fn new() -> Self {
Self {
http: RateLimiter::new(10.0, 50),
dns: RateLimiter::new(100.0, 200),
port: RateLimiter::new(50.0, 100),
}
}
#[inline]
pub fn check_http(&self, host: &str) -> bool {
self.http.check(host)
}
#[inline]
pub fn check_dns(&self, domain: &str) -> bool {
self.dns.check(domain)
}
#[inline]
pub fn check_port(&self, key: &str) -> bool {
self.port.check(key)
}
}
impl Default for NetworkPool {
fn default() -> Self {
Self::new()
}
}
pub static NETWORK_POOL: LazyLock<NetworkPool> = LazyLock::new(NetworkPool::new);
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn http_rate_limit_exhausts_burst() {
let pool = NetworkPool::new();
let mut allowed = 0;
for _ in 0..60 {
if pool.check_http("example.com") {
allowed += 1;
}
}
assert_eq!(allowed, 50, "expected exactly 50 allowed from burst");
}
#[test]
fn http_rate_limit_eventually_rejects() {
let pool = NetworkPool::new();
for _ in 0..50 {
pool.check_http("test.com");
}
assert!(
!pool.check_http("test.com"),
"expected rejection after burst exhausted"
);
}
#[test]
fn dns_rate_limit_exhausts_burst() {
let pool = NetworkPool::new();
let mut allowed = 0;
for _ in 0..210 {
if pool.check_dns("example.com") {
allowed += 1;
}
}
assert_eq!(allowed, 200, "expected exactly 200 allowed from DNS burst");
}
#[test]
fn port_rate_limit_exhausts_burst() {
let pool = NetworkPool::new();
let mut allowed = 0;
for _ in 0..110 {
if pool.check_port("host:8080") {
allowed += 1;
}
}
assert_eq!(allowed, 100, "expected exactly 100 allowed from port burst");
}
#[test]
fn separate_hosts_have_independent_limits() {
let pool = NetworkPool::new();
for _ in 0..50 {
pool.check_http("host-a.com");
}
assert!(!pool.check_http("host-a.com"));
assert!(pool.check_http("host-b.com"));
}
#[test]
fn global_pool_accessible() {
assert!(NETWORK_POOL.check_dns("global-test.com"));
}
#[test]
fn refills_after_time() {
let pool = NetworkPool::new();
for _ in 0..50 {
pool.check_http("refill.com");
}
assert!(!pool.check_http("refill.com"));
std::thread::sleep(std::time::Duration::from_millis(150));
assert!(
pool.check_http("refill.com"),
"expected token refill after waiting"
);
}
}