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 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
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
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.