diff --git a/pbs-api-types/src/traffic_control.rs b/pbs-api-types/src/traffic_control.rs index 210f53ac..a1fcb7b5 100644 --- a/pbs-api-types/src/traffic_control.rs +++ b/pbs-api-types/src/traffic_control.rs @@ -30,13 +30,6 @@ pub const TRAFFIC_CONTROL_BURST_SCHEMA: Schema = IntegerSchema::new( #[api( properties: { - name: { - schema: TRAFFIC_CONTROL_ID_SCHEMA, - }, - comment: { - optional: true, - schema: SINGLE_LINE_COMMENT_SCHEMA, - }, "rate-in": { type: HumanByte, optional: true, @@ -53,6 +46,45 @@ pub const TRAFFIC_CONTROL_BURST_SCHEMA: Schema = IntegerSchema::new( type: HumanByte, optional: true, }, + }, +)] +#[derive(Serialize,Deserialize,Default,Clone,Updater)] +#[serde(rename_all = "kebab-case")] +/// Rate Limit Configuration +pub struct RateLimitConfig { + #[serde(skip_serializing_if="Option::is_none")] + pub rate_in: Option, + #[serde(skip_serializing_if="Option::is_none")] + pub burst_in: Option, + #[serde(skip_serializing_if="Option::is_none")] + pub rate_out: Option, + #[serde(skip_serializing_if="Option::is_none")] + pub burst_out: Option, +} + +impl RateLimitConfig { + pub fn with_same_inout(rate: Option, burst: Option) -> Self { + Self { + rate_in: rate, + burst_in: burst, + rate_out: rate, + burst_out: burst, + } + } +} + +#[api( + properties: { + name: { + schema: TRAFFIC_CONTROL_ID_SCHEMA, + }, + comment: { + optional: true, + schema: SINGLE_LINE_COMMENT_SCHEMA, + }, + limit: { + type: RateLimitConfig, + }, network: { type: Array, items: { @@ -78,14 +110,8 @@ pub struct TrafficControlRule { pub comment: Option, /// Rule applies to Source IPs within this networks pub network: Vec, - #[serde(skip_serializing_if="Option::is_none")] - pub rate_in: Option, - #[serde(skip_serializing_if="Option::is_none")] - pub burst_in: Option, - #[serde(skip_serializing_if="Option::is_none")] - pub rate_out: Option, - #[serde(skip_serializing_if="Option::is_none")] - pub burst_out: Option, + #[serde(flatten)] + pub limit: RateLimitConfig, // fixme: expose this? // /// Bandwidth is shared accross all connections // #[serde(skip_serializing_if="Option::is_none")] diff --git a/pbs-config/src/traffic_control.rs b/pbs-config/src/traffic_control.rs index ec0f5dc4..48e2a514 100644 --- a/pbs-config/src/traffic_control.rs +++ b/pbs-config/src/traffic_control.rs @@ -22,7 +22,7 @@ fn init() -> SectionConfig { let mut config = SectionConfig::new(&TRAFFIC_CONTROL_ID_SCHEMA); let obj_schema = match TrafficControlRule::API_SCHEMA { - Schema::Object(ref obj_schema) => obj_schema, + Schema::AllOf(ref allof_schema) => allof_schema, _ => unreachable!(), }; let plugin = SectionConfigPlugin::new("rule".to_string(), Some("name".to_string()), obj_schema); diff --git a/src/api2/config/traffic_control.rs b/src/api2/config/traffic_control.rs index dec0e3b2..aca3b5c4 100644 --- a/src/api2/config/traffic_control.rs +++ b/src/api2/config/traffic_control.rs @@ -168,10 +168,10 @@ pub fn update_traffic_control( if let Some(delete) = delete { for delete_prop in delete { match delete_prop { - DeletableProperty::rate_in => { data.rate_in = None; }, - DeletableProperty::rate_out => { data.rate_out = None; }, - DeletableProperty::burst_in => { data.burst_in = None; }, - DeletableProperty::burst_out => { data.burst_out = None; }, + DeletableProperty::rate_in => { data.limit.rate_in = None; }, + DeletableProperty::rate_out => { data.limit.rate_out = None; }, + DeletableProperty::burst_in => { data.limit.burst_in = None; }, + DeletableProperty::burst_out => { data.limit.burst_out = None; }, DeletableProperty::comment => { data.comment = None; }, DeletableProperty::timeframe => { data.timeframe = None; }, } @@ -187,12 +187,22 @@ pub fn update_traffic_control( } } - if update.rate_in.is_some() { data.rate_in = update.rate_in; } - if update.rate_out.is_some() { data.rate_out = update.rate_out; } + if update.limit.rate_in.is_some() { + data.limit.rate_in = update.limit.rate_in; + } + + if update.limit.rate_out.is_some() { + data.limit.rate_out = update.limit.rate_out; + } + + if update.limit.burst_in.is_some() { + data.limit.burst_in = update.limit.burst_in; + } + + if update.limit.burst_out.is_some() { + data.limit.burst_out = update.limit.burst_out; + } - if update.burst_in.is_some() { data.burst_in = update.burst_in; } - if update.burst_out.is_some() { data.burst_out = update.burst_out; } - if let Some(network) = update.network { data.network = network; } if update.timeframe.is_some() { data.timeframe = update.timeframe; } diff --git a/src/cached_traffic_control.rs b/src/cached_traffic_control.rs index b45e564c..ba552215 100644 --- a/src/cached_traffic_control.rs +++ b/src/cached_traffic_control.rs @@ -219,27 +219,28 @@ impl TrafficControlCache { for rule in rules { let entry = self.limiter_map.entry(rule.name.clone()).or_insert((None, None)); + let limit = &rule.limit; match entry.0 { Some(ref read_limiter) => { - match rule.rate_in { + match limit.rate_in { Some(rate_in) => { read_limiter.update_rate( rate_in.as_u64(), - rule.burst_in.unwrap_or(rate_in).as_u64(), + limit.burst_in.unwrap_or(rate_in).as_u64(), ); } None => entry.0 = None, } } None => { - if let Some(rate_in) = rule.rate_in { + if let Some(rate_in) = limit.rate_in { let name = format!("{}.in", rule.name); let limiter = create_limiter( self.use_shared_memory, &name, rate_in.as_u64(), - rule.burst_in.unwrap_or(rate_in).as_u64(), + limit.burst_in.unwrap_or(rate_in).as_u64(), )?; entry.0 = Some(limiter); } @@ -248,24 +249,24 @@ impl TrafficControlCache { match entry.1 { Some(ref write_limiter) => { - match rule.rate_out { + match limit.rate_out { Some(rate_out) => { write_limiter.update_rate( rate_out.as_u64(), - rule.burst_out.unwrap_or(rate_out).as_u64(), + limit.burst_out.unwrap_or(rate_out).as_u64(), ); } None => entry.1 = None, } } None => { - if let Some(rate_out) = rule.rate_out { + if let Some(rate_out) = limit.rate_out { let name = format!("{}.out", rule.name); let limiter = create_limiter( self.use_shared_memory, &name, rate_out.as_u64(), - rule.burst_out.unwrap_or(rate_out).as_u64(), + limit.burst_out.unwrap_or(rate_out).as_u64(), )?; entry.1 = Some(limiter); }