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" ] } | pxar = { version = "0.10.1", features = [ "tokio-io" ] } | ||||||
|  |  | ||||||
| proxmox = { version = "0.15.0", features = [ "sortable-macro" ] } | 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-io = "1" | ||||||
| proxmox-lang = "1" | proxmox-lang = "1" | ||||||
| proxmox-router = { version = "1.1", features = [ "cli" ] } | proxmox-router = { version = "1.1", features = [ "cli" ] } | ||||||
|  | |||||||
| @ -30,7 +30,7 @@ xdg = "2.2" | |||||||
| pathpatterns = "0.1.2" | pathpatterns = "0.1.2" | ||||||
| proxmox = "0.15.0" | proxmox = "0.15.0" | ||||||
| proxmox-fuse = "0.1.1" | 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-io = { version = "1", features = [ "tokio" ] } | ||||||
| proxmox-lang = "1" | proxmox-lang = "1" | ||||||
| proxmox-router = { version = "1.1", features = [ "cli" ] } | 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::tools::fs::CreateOptions; | ||||||
| use proxmox_lang::try_block; | use proxmox_lang::try_block; | ||||||
| use proxmox_router::{RpcEnvironment, RpcEnvironmentType, UserInformation}; | 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_tools::{task_log, task_warn}; | ||||||
| use pbs_datastore::DataStore; | use pbs_datastore::DataStore; | ||||||
| @ -1093,7 +1093,7 @@ lazy_static::lazy_static!{ | |||||||
|  |  | ||||||
| fn lookup_rate_limiter( | fn lookup_rate_limiter( | ||||||
|     peer: Option<std::net::SocketAddr>, |     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 mut cache = TRAFFIC_CONTROL_CACHE.lock().unwrap(); | ||||||
|  |  | ||||||
|     let now = proxmox_time::epoch_i64(); |     let now = proxmox_time::epoch_i64(); | ||||||
|  | |||||||
| @ -6,7 +6,7 @@ use std::net::{IpAddr, Ipv4Addr, SocketAddr}; | |||||||
| use anyhow::Error; | use anyhow::Error; | ||||||
| use cidr::IpInet; | use cidr::IpInet; | ||||||
|  |  | ||||||
| use proxmox_http::client::RateLimiter; | use proxmox_http::client::{ShareableRateLimit, RateLimiter}; | ||||||
| use proxmox_section_config::SectionConfigData; | use proxmox_section_config::SectionConfigData; | ||||||
|  |  | ||||||
| use proxmox_systemd::daily_duration::{parse_daily_duration, DailyDuration}; | use proxmox_systemd::daily_duration::{parse_daily_duration, DailyDuration}; | ||||||
| @ -26,7 +26,7 @@ pub struct TrafficControlCache { | |||||||
|     last_update: i64, |     last_update: i64, | ||||||
|     last_traffic_control_generation: usize, |     last_traffic_control_generation: usize, | ||||||
|     rules: Vec<ParsedTcRule>, |     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 |     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 { | impl TrafficControlCache { | ||||||
|  |  | ||||||
|     pub fn new() -> Self { |     pub fn new() -> Self { | ||||||
| @ -130,6 +138,7 @@ impl TrafficControlCache { | |||||||
|         self.update_config(&config) |         self.update_config(&config) | ||||||
|     } |     } | ||||||
|  |  | ||||||
|  |  | ||||||
|     fn update_config(&mut self, config: &SectionConfigData) -> Result<(), Error> { |     fn update_config(&mut self, config: &SectionConfigData) -> Result<(), Error> { | ||||||
|         self.limiter_map.retain(|key, _value| config.sections.contains_key(key)); |         self.limiter_map.retain(|key, _value| config.sections.contains_key(key)); | ||||||
|  |  | ||||||
| @ -146,16 +155,15 @@ impl TrafficControlCache { | |||||||
|                 Some(ref read_limiter) => { |                 Some(ref read_limiter) => { | ||||||
|                     match rule.rate_in { |                     match rule.rate_in { | ||||||
|                         Some(rate_in) => { |                         Some(rate_in) => { | ||||||
|                             read_limiter.lock().unwrap(). |                             read_limiter.update_rate(rate_in, rule.burst_in.unwrap_or(rate_in)); | ||||||
|                                 update_rate(rate_in, rule.burst_in.unwrap_or(rate_in)); |  | ||||||
|                         } |                         } | ||||||
|                         None => entry.0 = None, |                         None => entry.0 = None, | ||||||
|                     } |                     } | ||||||
|                 } |                 } | ||||||
|                 None => { |                 None => { | ||||||
|                     if let Some(rate_in) = rule.rate_in { |                     if let Some(rate_in) = rule.rate_in { | ||||||
|                         let limiter = RateLimiter::new(rate_in, rule.burst_in.unwrap_or(rate_in)); |                         let limiter = create_limiter(rate_in, rule.burst_in.unwrap_or(rate_in), false)?; | ||||||
|                         entry.0 = Some(Arc::new(Mutex::new(limiter))); |                         entry.0 = Some(limiter); | ||||||
|                     } |                     } | ||||||
|                 } |                 } | ||||||
|             } |             } | ||||||
| @ -164,16 +172,15 @@ impl TrafficControlCache { | |||||||
|                 Some(ref write_limiter) => { |                 Some(ref write_limiter) => { | ||||||
|                     match rule.rate_out { |                     match rule.rate_out { | ||||||
|                         Some(rate_out) => { |                         Some(rate_out) => { | ||||||
|                             write_limiter.lock().unwrap(). |                             write_limiter.update_rate(rate_out, rule.burst_out.unwrap_or(rate_out)); | ||||||
|                                 update_rate(rate_out, rule.burst_out.unwrap_or(rate_out)); |  | ||||||
|                         } |                         } | ||||||
|                         None => entry.1 = None, |                         None => entry.1 = None, | ||||||
|                     } |                     } | ||||||
|                 } |                 } | ||||||
|                 None => { |                 None => { | ||||||
|                     if let Some(rate_out) = rule.rate_out { |                     if let Some(rate_out) = rule.rate_out { | ||||||
|                         let limiter = RateLimiter::new(rate_out, rule.burst_out.unwrap_or(rate_out)); |                         let limiter = create_limiter(rate_out, rule.burst_out.unwrap_or(rate_out), true)?; | ||||||
|                         entry.1 = Some(Arc::new(Mutex::new(limiter))); |                         entry.1 = Some(limiter); | ||||||
|                     } |                     } | ||||||
|                 } |                 } | ||||||
|             } |             } | ||||||
| @ -212,7 +219,7 @@ impl TrafficControlCache { | |||||||
|         &self, |         &self, | ||||||
|         peer: Option<SocketAddr>, |         peer: Option<SocketAddr>, | ||||||
|         now: i64, |         now: i64, | ||||||
|     ) -> (&str, Option<Arc<Mutex<RateLimiter>>>, Option<Arc<Mutex<RateLimiter>>>) { |     ) -> (&str, Option<Arc<dyn ShareableRateLimit>>, Option<Arc<dyn ShareableRateLimit>>) { | ||||||
|  |  | ||||||
|         let peer = match peer { |         let peer = match peer { | ||||||
|             None => return ("", None, None), |             None => return ("", None, None), | ||||||
|  | |||||||
		Reference in New Issue
	
	Block a user