Skip to content

Commit

Permalink
Update fault injector so that rate limiting works without std.
Browse files Browse the repository at this point in the history
@whitequark
whitequark committed Jul 23, 2017
1 parent c799bfc commit d73c2bd
Showing 2 changed files with 29 additions and 62 deletions.
28 changes: 14 additions & 14 deletions examples/utils.rs
View file
Original file line number Diff line number Diff line change
@@ -3,7 +3,7 @@ use std::str::{self, FromStr};
use std::rc::Rc;
use std::io;
use std::fs::File;
use std::time::{Instant, Duration, SystemTime, UNIX_EPOCH};
use std::time::{Instant, SystemTime, UNIX_EPOCH};
use std::env;
use std::process;
use log::{LogLevel, LogLevelFilter, LogRecord};
@@ -106,18 +106,18 @@ pub fn add_middleware_options(opts: &mut Options, _free: &mut Vec<&str>) {

pub fn parse_middleware_options<D: Device>(matches: &mut Matches, device: D, loopback: bool)
-> FaultInjector<EthernetTracer<PcapWriter<D, Rc<PcapSink>>>> {
let drop_chance = u8::from_str(&matches.opt_str("drop-chance")
.unwrap_or("0".to_string())).unwrap();
let corrupt_chance = u8::from_str(&matches.opt_str("corrupt-chance")
.unwrap_or("0".to_string())).unwrap();
let size_limit = usize::from_str(&matches.opt_str("size-limit")
.unwrap_or("0".to_string())).unwrap();
let tx_rate_limit = u64::from_str(&matches.opt_str("tx-rate-limit")
.unwrap_or("0".to_string())).unwrap();
let rx_rate_limit = u64::from_str(&matches.opt_str("rx-rate-limit")
.unwrap_or("0".to_string())).unwrap();
let shaping_interval = u32::from_str(&matches.opt_str("shaping-interval")
.unwrap_or("0".to_string())).unwrap();
let drop_chance = matches.opt_str("drop-chance").map(|s| u8::from_str(&s).unwrap())
.unwrap_or(0);
let corrupt_chance = matches.opt_str("corrupt-chance").map(|s| u8::from_str(&s).unwrap())
.unwrap_or(0);
let size_limit = matches.opt_str("size-limit").map(|s| usize::from_str(&s).unwrap())
.unwrap_or(0);
let tx_rate_limit = matches.opt_str("tx-rate-limit").map(|s| u64::from_str(&s).unwrap())
.unwrap_or(0);
let rx_rate_limit = matches.opt_str("rx-rate-limit").map(|s| u64::from_str(&s).unwrap())
.unwrap_or(0);
let shaping_interval = matches.opt_str("shaping-interval").map(|s| u64::from_str(&s).unwrap())
.unwrap_or(0);

let pcap_writer: Box<io::Write>;
if let Some(pcap_filename) = matches.opt_str("pcap") {
@@ -138,6 +138,6 @@ pub fn parse_middleware_options<D: Device>(matches: &mut Matches, device: D, loo
device.set_max_packet_size(size_limit);
device.set_max_tx_rate(tx_rate_limit);
device.set_max_rx_rate(rx_rate_limit);
device.set_bucket_interval(Duration::from_millis(shaping_interval as u64));
device.set_bucket_interval(shaping_interval);
device
}
63 changes: 15 additions & 48 deletions src/phy/fault_injector.rs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,3 @@
#[cfg(feature = "std")]
use std::time::{Instant, Duration};

use Error;
use super::{DeviceLimits, Device};

@@ -24,22 +21,16 @@ struct Config {
drop_pct: u8,
reorder_pct: u8,
max_size: usize,
#[cfg(feature = "std")]
max_tx_rate: u64,
#[cfg(feature = "std")]
max_rx_rate: u64,
#[cfg(feature = "std")]
interval: Duration,
interval: u64,
}

#[derive(Debug, Clone, Copy)]
struct State {
rng_seed: u32,
#[cfg(feature = "std")]
refilled_at: Instant,
#[cfg(feature = "std")]
refilled_at: u64,
tx_bucket: u64,
#[cfg(feature = "std")]
rx_bucket: u64,
}

@@ -56,20 +47,18 @@ impl State {
buffer[index] ^= bit;
}

#[cfg(feature = "std")]
fn refill(&mut self, config: &Config) {
if self.refilled_at.elapsed() > config.interval {
fn refill(&mut self, config: &Config, timestamp: u64) {
if self.refilled_at - timestamp > config.interval {
self.tx_bucket = config.max_tx_rate;
self.rx_bucket = config.max_rx_rate;
self.refilled_at = Instant::now();
self.refilled_at = timestamp;
}
}

#[cfg(feature = "std")]
fn maybe_transmit(&mut self, config: &Config) -> bool {
fn maybe_transmit(&mut self, config: &Config, timestamp: u64) -> bool {
if config.max_tx_rate == 0 { return true }

self.refill(config);
self.refill(config, timestamp);
if self.tx_bucket > 0 {
self.tx_bucket -= 1;
true
@@ -78,28 +67,17 @@ impl State {
}
}

#[cfg(not(feature = "std"))]
fn maybe_transmit(&mut self, _config: &Config) -> bool {
true
}

#[cfg(feature = "std")]
fn maybe_receive(&mut self, config: &Config) -> bool {
fn maybe_receive(&mut self, config: &Config, timestamp: u64) -> bool {
if config.max_rx_rate == 0 { return true }

self.refill(config);
self.refill(config, timestamp);
if self.rx_bucket > 0 {
self.rx_bucket -= 1;
true
} else {
false
}
}

#[cfg(not(feature = "std"))]
fn maybe_receive(&mut self, _config: &Config) -> bool {
true
}
}

/// A fault injector device.
@@ -117,17 +95,12 @@ pub struct FaultInjector<D: Device> {
impl<D: Device> FaultInjector<D> {
/// Create a fault injector device, using the given random number generator seed.
pub fn new(inner: D, seed: u32) -> FaultInjector<D> {
#[cfg(feature = "std")]
let state = State {
rng_seed: seed,
refilled_at: Instant::now(),
refilled_at: 0,
tx_bucket: 0,
rx_bucket: 0,
};
#[cfg(not(feature = "std"))]
let state = State {
rng_seed: seed,
};
FaultInjector {
inner: inner,
state: state,
@@ -156,20 +129,17 @@ impl<D: Device> FaultInjector<D> {
}

/// Return the maximum packet transmission rate, in packets per second.
#[cfg(feature = "std")]
pub fn max_tx_rate(&self) -> u64 {
self.config.max_rx_rate
}

/// Return the maximum packet reception rate, in packets per second.
#[cfg(feature = "std")]
pub fn max_rx_rate(&self) -> u64 {
self.config.max_tx_rate
}

/// Return the interval for packet rate limiting, in milliseconds.
#[cfg(feature = "std")]
pub fn bucket_interval(&self) -> Duration {
pub fn bucket_interval(&self) -> u64 {
self.config.interval
}

@@ -197,21 +167,18 @@ impl<D: Device> FaultInjector<D> {
}

/// Set the maximum packet transmission rate, in packets per interval.
#[cfg(feature = "std")]
pub fn set_max_tx_rate(&mut self, rate: u64) {
self.config.max_tx_rate = rate
}

/// Set the maximum packet reception rate, in packets per interval.
#[cfg(feature = "std")]
pub fn set_max_rx_rate(&mut self, rate: u64) {
self.config.max_rx_rate = rate
}

/// Set the interval for packet rate limiting, in milliseconds.
#[cfg(feature = "std")]
pub fn set_bucket_interval(&mut self, interval: Duration) {
self.state.refilled_at = Instant::now() - self.config.interval;
pub fn set_bucket_interval(&mut self, interval: u64) {
self.state.refilled_at = 0;
self.config.interval = interval
}
}
@@ -243,7 +210,7 @@ impl<D: Device> Device for FaultInjector<D>
net_trace!("rx: dropping a packet that is too large");
return Err(Error::Exhausted)
}
if !self.state.maybe_receive(&self.config) {
if !self.state.maybe_receive(&self.config, timestamp) {
net_trace!("rx: dropping a packet because of rate limiting");
return Err(Error::Exhausted)
}
@@ -258,7 +225,7 @@ impl<D: Device> Device for FaultInjector<D>
} else if self.config.max_size > 0 && length > self.config.max_size {
net_trace!("tx: dropping a packet that is too large");
buffer = None;
} else if !self.state.maybe_transmit(&self.config) {
} else if !self.state.maybe_transmit(&self.config, timestamp) {
net_trace!("tx: dropping a packet because of rate limiting");
buffer = None;
} else {

0 comments on commit d73c2bd

Please sign in to comment.