cached_traffic_control: use ShareableRateLimit trait object
This commit is contained in:
		| @ -98,7 +98,7 @@ pathpatterns = "0.1.2" | ||||
| pxar = { version = "0.10.1", features = [ "tokio-io" ] } | ||||
|  | ||||
| proxmox = { version = "0.15.0", features = [ "sortable-macro" ] } | ||||
| proxmox-http = { version = "0.5.2", features = [ "client", "http-helpers", "websocket" ] } | ||||
| proxmox-http = { version = "0.5.3", features = [ "client", "http-helpers", "websocket" ] } | ||||
| proxmox-io = "1" | ||||
| proxmox-lang = "1" | ||||
| proxmox-router = { version = "1.1", features = [ "cli" ] } | ||||
|  | ||||
| @ -30,7 +30,7 @@ xdg = "2.2" | ||||
| pathpatterns = "0.1.2" | ||||
| proxmox = "0.15.0" | ||||
| proxmox-fuse = "0.1.1" | ||||
| proxmox-http = { version = "0.5.2", features = [ "client", "http-helpers", "websocket" ] } | ||||
| proxmox-http = { version = "0.5.3", features = [ "client", "http-helpers", "websocket" ] } | ||||
| proxmox-io = { version = "1", features = [ "tokio" ] } | ||||
| proxmox-lang = "1" | ||||
| proxmox-router = { version = "1.1", features = [ "cli" ] } | ||||
|  | ||||
| @ -21,7 +21,7 @@ use proxmox::sys::linux::socket::set_tcp_keepalive; | ||||
| use proxmox::tools::fs::CreateOptions; | ||||
| use proxmox_lang::try_block; | ||||
| use proxmox_router::{RpcEnvironment, RpcEnvironmentType, UserInformation}; | ||||
| use proxmox_http::client::{RateLimiter, RateLimitedStream}; | ||||
| use proxmox_http::client::{RateLimitedStream, ShareableRateLimit}; | ||||
|  | ||||
| use pbs_tools::{task_log, task_warn}; | ||||
| use pbs_datastore::DataStore; | ||||
| @ -1093,7 +1093,7 @@ lazy_static::lazy_static!{ | ||||
|  | ||||
| fn lookup_rate_limiter( | ||||
|     peer: Option<std::net::SocketAddr>, | ||||
| ) -> (Option<Arc<Mutex<RateLimiter>>>, Option<Arc<Mutex<RateLimiter>>>) { | ||||
| ) -> (Option<Arc<dyn ShareableRateLimit>>, Option<Arc<dyn ShareableRateLimit>>) { | ||||
|     let mut cache = TRAFFIC_CONTROL_CACHE.lock().unwrap(); | ||||
|  | ||||
|     let now = proxmox_time::epoch_i64(); | ||||
|  | ||||
| @ -6,7 +6,7 @@ use std::net::{IpAddr, Ipv4Addr, SocketAddr}; | ||||
| use anyhow::Error; | ||||
| use cidr::IpInet; | ||||
|  | ||||
| use proxmox_http::client::RateLimiter; | ||||
| use proxmox_http::client::{ShareableRateLimit, RateLimiter}; | ||||
| use proxmox_section_config::SectionConfigData; | ||||
|  | ||||
| use proxmox_systemd::daily_duration::{parse_daily_duration, DailyDuration}; | ||||
| @ -26,7 +26,7 @@ pub struct TrafficControlCache { | ||||
|     last_update: i64, | ||||
|     last_traffic_control_generation: usize, | ||||
|     rules: Vec<ParsedTcRule>, | ||||
|     limiter_map: HashMap<String, (Option<Arc<Mutex<RateLimiter>>>, Option<Arc<Mutex<RateLimiter>>>)>, | ||||
|     limiter_map: HashMap<String, (Option<Arc<dyn ShareableRateLimit>>, Option<Arc<dyn ShareableRateLimit>>)>, | ||||
|     use_utc: bool, // currently only used for testing | ||||
| } | ||||
|  | ||||
| @ -84,6 +84,14 @@ fn cannonical_ip(ip: IpAddr) -> IpAddr { | ||||
|     } | ||||
| } | ||||
|  | ||||
| fn create_limiter( | ||||
|     rate: u64, | ||||
|     burst: u64, | ||||
|     _direction: bool, // false => in, true => out | ||||
| ) -> Result<Arc<dyn ShareableRateLimit>, Error> { | ||||
|     Ok(Arc::new(Mutex::new(RateLimiter::new(rate, burst)))) | ||||
| } | ||||
|  | ||||
| impl TrafficControlCache { | ||||
|  | ||||
|     pub fn new() -> Self { | ||||
| @ -130,6 +138,7 @@ impl TrafficControlCache { | ||||
|         self.update_config(&config) | ||||
|     } | ||||
|  | ||||
|  | ||||
|     fn update_config(&mut self, config: &SectionConfigData) -> Result<(), Error> { | ||||
|         self.limiter_map.retain(|key, _value| config.sections.contains_key(key)); | ||||
|  | ||||
| @ -146,16 +155,15 @@ impl TrafficControlCache { | ||||
|                 Some(ref read_limiter) => { | ||||
|                     match rule.rate_in { | ||||
|                         Some(rate_in) => { | ||||
|                             read_limiter.lock().unwrap(). | ||||
|                                 update_rate(rate_in, rule.burst_in.unwrap_or(rate_in)); | ||||
|                             read_limiter.update_rate(rate_in, rule.burst_in.unwrap_or(rate_in)); | ||||
|                         } | ||||
|                         None => entry.0 = None, | ||||
|                     } | ||||
|                 } | ||||
|                 None => { | ||||
|                     if let Some(rate_in) = rule.rate_in { | ||||
|                         let limiter = RateLimiter::new(rate_in, rule.burst_in.unwrap_or(rate_in)); | ||||
|                         entry.0 = Some(Arc::new(Mutex::new(limiter))); | ||||
|                         let limiter = create_limiter(rate_in, rule.burst_in.unwrap_or(rate_in), false)?; | ||||
|                         entry.0 = Some(limiter); | ||||
|                     } | ||||
|                 } | ||||
|             } | ||||
| @ -164,16 +172,15 @@ impl TrafficControlCache { | ||||
|                 Some(ref write_limiter) => { | ||||
|                     match rule.rate_out { | ||||
|                         Some(rate_out) => { | ||||
|                             write_limiter.lock().unwrap(). | ||||
|                                 update_rate(rate_out, rule.burst_out.unwrap_or(rate_out)); | ||||
|                             write_limiter.update_rate(rate_out, rule.burst_out.unwrap_or(rate_out)); | ||||
|                         } | ||||
|                         None => entry.1 = None, | ||||
|                     } | ||||
|                 } | ||||
|                 None => { | ||||
|                     if let Some(rate_out) = rule.rate_out { | ||||
|                         let limiter = RateLimiter::new(rate_out, rule.burst_out.unwrap_or(rate_out)); | ||||
|                         entry.1 = Some(Arc::new(Mutex::new(limiter))); | ||||
|                         let limiter = create_limiter(rate_out, rule.burst_out.unwrap_or(rate_out), true)?; | ||||
|                         entry.1 = Some(limiter); | ||||
|                     } | ||||
|                 } | ||||
|             } | ||||
| @ -212,7 +219,7 @@ impl TrafficControlCache { | ||||
|         &self, | ||||
|         peer: Option<SocketAddr>, | ||||
|         now: i64, | ||||
|     ) -> (&str, Option<Arc<Mutex<RateLimiter>>>, Option<Arc<Mutex<RateLimiter>>>) { | ||||
|     ) -> (&str, Option<Arc<dyn ShareableRateLimit>>, Option<Arc<dyn ShareableRateLimit>>) { | ||||
|  | ||||
|         let peer = match peer { | ||||
|             None => return ("", None, None), | ||||
|  | ||||
		Reference in New Issue
	
	Block a user