Skip to content

Commit

Permalink
udp load tester: open multiple sockets per worker; minor other fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
@greatest-ape
greatest-ape committed Feb 5, 2024
1 parent 6745eba commit d8bdcfc
Show file tree
Hide file tree
Showing 4 changed files with 59 additions and 43 deletions.
1 change: 1 addition & 0 deletions crates/udp_load_test/src/common.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,4 +22,5 @@ pub struct Peer {
pub announce_info_hash: InfoHash,
pub announce_port: Port,
pub scrape_info_hash_indices: Box<[usize]>,
pub socket_index: u8,
}
6 changes: 3 additions & 3 deletions crates/udp_load_test/src/config.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,8 +60,8 @@ pub struct NetworkConfig {
///
/// Setting this to true can cause issues on macOS.
pub multiple_client_ipv4s: bool,
/// Number of first client port
pub first_port: u16,
/// Number of sockets to open per worker
pub sockets_per_worker: u8,
/// Size of socket recv buffer. Use 0 for OS default.
///
/// This setting can have a big impact on dropped packages. It might
Expand All @@ -81,7 +81,7 @@ impl Default for NetworkConfig {
fn default() -> Self {
Self {
multiple_client_ipv4s: true,
first_port: 45_000,
sockets_per_worker: 4,
recv_buffer: 8_000_000,
}
}
Expand Down
5 changes: 2 additions & 3 deletions crates/udp_load_test/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,8 +54,6 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
// Start workers

for (i, peers) in (0..config.workers).zip(peers_by_worker) {
let port = config.network.first_port + (i as u16);

let ip = if config.server_address.is_ipv6() {
Ipv6Addr::LOCALHOST.into()
} else if config.network.multiple_client_ipv4s {
Expand All @@ -64,7 +62,7 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
Ipv4Addr::LOCALHOST.into()
};

let addr = SocketAddr::new(ip, port);
let addr = SocketAddr::new(ip, 0);
let config = config.clone();
let state = state.clone();

Expand Down Expand Up @@ -230,6 +228,7 @@ fn create_peers(config: &Config, info_hash_dist: &InfoHashDist) -> Vec<Box<[Peer
announce_info_hash,
announce_port: Port::new(rng.gen()),
scrape_info_hash_indices,
socket_index: rng.gen_range(0..config.network.sockets_per_worker),
}
})
.take(config.requests.number_of_peers)
Expand Down
90 changes: 53 additions & 37 deletions crates/udp_load_test/src/worker.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,15 +21,20 @@ pub struct Worker {
peers: Box<[Peer]>,
request_type_dist: RequestTypeDist,
addr: SocketAddr,
socket: UdpSocket,
sockets: Vec<UdpSocket>,
buffer: [u8; MAX_PACKET_SIZE],
rng: SmallRng,
statistics: LocalStatistics,
}

impl Worker {
pub fn run(config: Config, shared_state: LoadTestState, peers: Box<[Peer]>, addr: SocketAddr) {
let socket = create_socket(&config, addr);
let mut sockets = Vec::new();

for _ in 0..config.network.sockets_per_worker {
sockets.push(create_socket(&config, addr));
}

let buffer = [0u8; MAX_PACKET_SIZE];
let rng = SmallRng::seed_from_u64(0xc3aa8be617b3acce);
let statistics = LocalStatistics::default();
Expand All @@ -41,7 +46,7 @@ impl Worker {
peers,
request_type_dist,
addr,
socket,
sockets,
buffer,
rng,
statistics,
Expand All @@ -56,48 +61,59 @@ impl Worker {
let mut requests_sent = 0usize;
let mut responses_received = 0usize;

let mut connect_socket_index = 0u8;
let mut peer_index = 0usize;
let mut loop_index = 0usize;

loop {
let response_ratio = responses_received as f64 / requests_sent.max(1) as f64;

if response_ratio >= 0.95 || requests_sent == 0 || self.rng.gen::<u8>() == 0 {
match self.request_type_dist.sample(&mut self.rng) {
RequestType::Connect => {
self.send_connect_request(u32::MAX - 1);
}
RequestType::Announce => {
self.send_announce_request(connection_id, peer_index);
if response_ratio >= 0.90 || requests_sent == 0 || self.rng.gen::<u8>() == 0 {
for _ in 0..self.sockets.len() {
match self.request_type_dist.sample(&mut self.rng) {
RequestType::Connect => {
self.send_connect_request(connect_socket_index, u32::MAX - 1);

peer_index = (peer_index + 1) % self.peers.len();
}
RequestType::Scrape => {
self.send_scrape_request(connection_id, peer_index);
connect_socket_index = connect_socket_index.wrapping_add(1)
% self.config.network.sockets_per_worker;
}
RequestType::Announce => {
self.send_announce_request(connection_id, peer_index);

peer_index = (peer_index + 1) % self.peers.len();
peer_index = (peer_index + 1) % self.peers.len();
}
RequestType::Scrape => {
self.send_scrape_request(connection_id, peer_index);

peer_index = (peer_index + 1) % self.peers.len();
}
}
}

requests_sent += 1;
requests_sent += 1;
}
}

match self.socket.recv(&mut self.buffer[..]) {
Ok(amt) => {
match Response::parse_bytes(&self.buffer[0..amt], self.addr.is_ipv4()) {
Ok(response) => {
self.handle_response(response);
}
Err(err) => {
eprintln!("Received invalid response: {:#?}", err);
for socket_index in 0..self.sockets.len() {
// Do this instead of iterating over Vec to fix borrow checker complaint
let socket = self.sockets.get(socket_index).unwrap();

match socket.recv(&mut self.buffer[..]) {
Ok(amt) => {
match Response::parse_bytes(&self.buffer[0..amt], self.addr.is_ipv4()) {
Ok(response) => {
self.handle_response(response);
}
Err(err) => {
eprintln!("Received invalid response: {:#?}", err);
}
}
}

responses_received += 1;
}
Err(err) if err.kind() == ErrorKind::WouldBlock => (),
Err(err) => {
eprintln!("recv error: {:#}", err);
responses_received += 1;
}
Err(err) if err.kind() == ErrorKind::WouldBlock => (),
Err(err) => {
eprintln!("recv error: {:#}", err);
}
}
}

Expand All @@ -111,10 +127,10 @@ impl Worker {

fn aquire_connection_id(&mut self) -> ConnectionId {
loop {
self.send_connect_request(u32::MAX);
self.send_connect_request(0, u32::MAX);

for _ in 0..100 {
match self.socket.recv(&mut self.buffer[..]) {
match self.sockets[0].recv(&mut self.buffer[..]) {
Ok(amt) => {
match Response::parse_bytes(&self.buffer[0..amt], self.addr.is_ipv4()) {
Ok(Response::Connect(r)) => {
Expand All @@ -139,7 +155,7 @@ impl Worker {
}
}

fn send_connect_request(&mut self, transaction_id: u32) {
fn send_connect_request(&mut self, socket_index: u8, transaction_id: u32) {
let transaction_id = TransactionId::new(i32::from_ne_bytes(transaction_id.to_ne_bytes()));

let request = ConnectRequest { transaction_id };
Expand All @@ -150,7 +166,7 @@ impl Worker {

let position = cursor.position() as usize;

match self.socket.send(&cursor.get_ref()[..position]) {
match self.sockets[socket_index as usize].send(&cursor.get_ref()[..position]) {
Ok(_) => {
self.statistics.requests += 1;
}
Expand Down Expand Up @@ -199,7 +215,7 @@ impl Worker {

let position = cursor.position() as usize;

match self.socket.send(&cursor.get_ref()[..position]) {
match self.sockets[peer.socket_index as usize].send(&cursor.get_ref()[..position]) {
Ok(_) => {
self.statistics.requests += 1;
}
Expand Down Expand Up @@ -233,7 +249,7 @@ impl Worker {

let position = cursor.position() as usize;

match self.socket.send(&cursor.get_ref()[..position]) {
match self.sockets[peer.socket_index as usize].send(&cursor.get_ref()[..position]) {
Ok(_) => {
self.statistics.requests += 1;
}
Expand Down

0 comments on commit d8bdcfc

Please sign in to comment.