use std::{ net::{IpAddr, Ipv4Addr}, time::Duration, }; use ahash::RandomState; use criterion::{criterion_group, criterion_main, Criterion}; use dashmap::DashMap; use smtp_server::{ config::*, core::{ throttle::{ThrottleKey, ThrottleKeyHasherBuilder}, Envelope, }, }; #[derive(Debug, PartialEq, Eq, Hash)] enum Item { String(String), Uint(u64), Listener(u16), IpAddr(IpAddr), } #[derive(Debug, Clone)] struct TestEnvelope { pub local_ip: IpAddr, pub remote_ip: IpAddr, pub sender_domain: String, pub sender: String, pub rcpt_domain: String, pub rcpt: String, pub helo_domain: String, pub authenticated_as: String, pub mx: String, pub listener_id: u16, pub priority: i16, } impl Envelope for TestEnvelope { fn local_ip(&self) -> IpAddr { self.local_ip } fn remote_ip(&self) -> IpAddr { self.remote_ip } fn sender_domain(&self) -> &str { self.sender_domain.as_str() } fn sender(&self) -> &str { self.sender.as_str() } fn rcpt_domain(&self) -> &str { self.rcpt_domain.as_str() } fn rcpt(&self) -> &str { self.rcpt.as_str() } fn helo_domain(&self) -> &str { self.helo_domain.as_str() } fn authenticated_as(&self) -> &str { self.authenticated_as.as_str() } fn mx(&self) -> &str { self.mx.as_str() } fn listener_id(&self) -> u16 { self.listener_id } fn priority(&self) -> i16 { self.priority } } pub fn criterion_benchmark(c: &mut Criterion) { let envelope = TestEnvelope { local_ip: "127.0.0.1".parse().unwrap(), remote_ip: "A:B::C:D".parse().unwrap(), sender_domain: "domain.org".to_string(), sender: "sender@mydomain.com".to_string(), rcpt_domain: "otherdomain.net".to_string(), rcpt: "rcpt@otherdomain.net".to_string(), authenticated_as: "test@test.org".to_string(), mx: "mx.tester.org".to_string(), listener_id: 0, priority: 1, helo_domain: "ehlo-domain.org".to_string(), }; let throttle = Throttle { conditions: Conditions { conditions: vec![] }, keys: THROTTLE_LISTENER, //u16::MAX, concurrency: 10.into(), rate: Rate { requests: 30, period: Duration::from_secs(60), } .into(), }; let blake3_map: DashMap = DashMap::with_capacity_and_hasher(100, ThrottleKeyHasherBuilder::default()); let string_map: DashMap = DashMap::with_capacity_and_hasher(100, RandomState::default()); let item_map: DashMap, usize, RandomState> = DashMap::with_capacity_and_hasher(100, RandomState::default()); let mut envelope1 = envelope.clone(); let mut envelope2 = envelope.clone(); let mut envelope3 = envelope; let mut remote_ip1: u32 = 0; let mut remote_ip2: u32 = 0; let mut remote_ip3: u32 = 0; c.bench_function("blake3 key", |b| { b.iter(|| { remote_ip1 += 1; let remote_ip = remote_ip1.to_be_bytes(); envelope1.remote_ip = IpAddr::V4(Ipv4Addr::new( remote_ip[0], remote_ip[1], remote_ip[2], remote_ip[3], )); blake3_map.insert(throttle.new_key(&envelope1), 0) }) }); c.bench_function("string key", |b| { b.iter(|| { remote_ip2 += 1; let remote_ip = remote_ip2.to_be_bytes(); envelope2.remote_ip = IpAddr::V4(Ipv4Addr::new( remote_ip[0], remote_ip[1], remote_ip[2], remote_ip[3], )); string_map.insert(to_key(&throttle, &envelope2), 0) }) }); c.bench_function("items key", |b| { b.iter(|| { remote_ip3 += 1; let remote_ip = remote_ip3.to_be_bytes(); envelope3.remote_ip = IpAddr::V4(Ipv4Addr::new( remote_ip[0], remote_ip[1], remote_ip[2], remote_ip[3], )); item_map.insert(to_items(&throttle, &envelope3), 0) }) }); } fn to_key(t: &Throttle, e: &TestEnvelope) -> String { use std::fmt::Write; let mut result = String::with_capacity(32); if (t.keys & THROTTLE_RCPT) != 0 { result.push_str(e.rcpt.as_str()); } if (t.keys & THROTTLE_RCPT_DOMAIN) != 0 { result.push_str(e.rcpt_domain.as_str()); } if (t.keys & THROTTLE_SENDER) != 0 { result.push_str(e.sender.as_str()); } if (t.keys & THROTTLE_SENDER_DOMAIN) != 0 { result.push_str(e.sender_domain.as_str()); } if (t.keys & THROTTLE_AUTH_AS) != 0 { result.push_str(e.authenticated_as.as_str()); } if (t.keys & THROTTLE_LISTENER) != 0 { write!(result, "{}", e.listener_id).ok(); } if (t.keys & THROTTLE_MX) != 0 { result.push_str(e.mx.as_str()); } if (t.keys & THROTTLE_REMOTE_IP) != 0 { write!(result, "{}", e.local_ip).ok(); } if (t.keys & THROTTLE_LOCAL_IP) != 0 { write!(result, "{}", e.remote_ip).ok(); } if let Some(rate_limit) = &t.rate { write!(result, "{}", rate_limit.period.as_secs()).ok(); write!(result, "{}", rate_limit.requests).ok(); } if let Some(concurrency) = &t.concurrency { write!(result, "{}", concurrency).ok(); } result } fn to_items(t: &Throttle, e: &TestEnvelope) -> Vec { let mut result = Vec::new(); if (t.keys & THROTTLE_RCPT) != 0 { result.push(Item::String(e.rcpt.clone())); } if (t.keys & THROTTLE_RCPT_DOMAIN) != 0 { result.push(Item::String(e.rcpt_domain.clone())); } if (t.keys & THROTTLE_SENDER) != 0 { result.push(Item::String(e.sender.clone())); } if (t.keys & THROTTLE_SENDER_DOMAIN) != 0 { result.push(Item::String(e.sender_domain.clone())); } if (t.keys & THROTTLE_AUTH_AS) != 0 { result.push(Item::String(e.authenticated_as.clone())); } if (t.keys & THROTTLE_LISTENER) != 0 { result.push(Item::Listener(e.listener_id)); } if (t.keys & THROTTLE_MX) != 0 { result.push(Item::String(e.mx.clone())); } if (t.keys & THROTTLE_REMOTE_IP) != 0 { result.push(Item::IpAddr(e.local_ip)); } if (t.keys & THROTTLE_LOCAL_IP) != 0 { result.push(Item::IpAddr(e.remote_ip)); } if let Some(rate_limit) = &t.rate { result.push(Item::Uint(rate_limit.period.as_secs())); result.push(Item::Uint(rate_limit.requests)); } if let Some(concurrency) = &t.concurrency { result.push(Item::Uint(*concurrency)); } result } criterion_group!(benches, criterion_benchmark); criterion_main!(benches);