udp load test: move network.rs and handler.rs into new worker module

aquatic_udp_load_test/src/worker/mod.rs

@@ -0,0 +1,206 @@
+mod request_gen;
+
+use std::io::Cursor;
+use std::net::SocketAddr;
+use std::sync::atomic::Ordering;
+use std::time::Duration;
+
+use mio::{net::UdpSocket, Events, Interest, Poll, Token};
+use rand::Rng;
+use rand::{prelude::SmallRng, thread_rng, SeedableRng};
+use rand_distr::Pareto;
+use socket2::{Domain, Protocol, Socket, Type};
+
+use aquatic_udp_protocol::*;
+
+use crate::config::Config;
+use crate::{common::*, utils::*};
+use request_gen::process_response;
+
+const MAX_PACKET_SIZE: usize = 4096;
+
+pub fn run_worker_thread(
+    state: LoadTestState,
+    pareto: Pareto<f64>,
+    config: &Config,
+    addr: SocketAddr,
+    thread_id: ThreadId,
+) {
+    let mut socket = UdpSocket::from_std(create_socket(config, addr));
+    let mut buffer = [0u8; MAX_PACKET_SIZE];
+
+    let mut rng = SmallRng::from_rng(thread_rng()).expect("create SmallRng from thread_rng()");
+    let mut torrent_peers = TorrentPeerMap::default();
+
+    let token = Token(thread_id.0 as usize);
+    let interests = Interest::READABLE;
+    let timeout = Duration::from_micros(config.network.poll_timeout);
+
+    let mut poll = Poll::new().expect("create poll");
+
+    poll.registry()
+        .register(&mut socket, token, interests)
+        .unwrap();
+
+    let mut events = Events::with_capacity(config.network.poll_event_capacity);
+
+    let mut statistics = SocketWorkerLocalStatistics::default();
+
+    // Bootstrap request cycle
+    let initial_request = create_connect_request(generate_transaction_id(&mut thread_rng()));
+    send_request(&mut socket, &mut buffer, &mut statistics, initial_request);
+
+    loop {
+        poll.poll(&mut events, Some(timeout))
+            .expect("failed polling");
+
+        for event in events.iter() {
+            if (event.token() == token) & event.is_readable() {
+                while let Ok(amt) = socket.recv(&mut buffer) {
+                    match Response::from_bytes(&buffer[0..amt]) {
+                        Ok(response) => {
+                            match response {
+                                Response::AnnounceIpv4(ref r) => {
+                                    statistics.responses_announce += 1;
+                                    statistics.response_peers += r.peers.len();
+                                }
+                                Response::AnnounceIpv6(ref r) => {
+                                    statistics.responses_announce += 1;
+                                    statistics.response_peers += r.peers.len();
+                                }
+                                Response::Scrape(_) => {
+                                    statistics.responses_scrape += 1;
+                                }
+                                Response::Connect(_) => {
+                                    statistics.responses_connect += 1;
+                                }
+                                Response::Error(_) => {
+                                    statistics.responses_error += 1;
+                                }
+                            }
+
+                            let opt_request = process_response(
+                                &mut rng,
+                                pareto,
+                                &state.info_hashes,
+                                &config,
+                                &mut torrent_peers,
+                                response,
+                            );
+
+                            if let Some(request) = opt_request {
+                                send_request(&mut socket, &mut buffer, &mut statistics, request);
+                            }
+                        }
+                        Err(err) => {
+                            eprintln!("Received invalid response: {:#?}", err);
+                        }
+                    }
+                }
+
+                if rng.gen::<f32>() <= config.requests.additional_request_probability {
+                    let additional_request =
+                        create_connect_request(generate_transaction_id(&mut rng));
+
+                    send_request(
+                        &mut socket,
+                        &mut buffer,
+                        &mut statistics,
+                        additional_request,
+                    );
+                }
+
+                update_shared_statistics(&state, &mut statistics);
+            }
+        }
+    }
+}
+
+fn send_request(
+    socket: &mut UdpSocket,
+    buffer: &mut [u8],
+    statistics: &mut SocketWorkerLocalStatistics,
+    request: Request,
+) {
+    let mut cursor = Cursor::new(buffer);
+
+    match request.write(&mut cursor) {
+        Ok(()) => {
+            let position = cursor.position() as usize;
+            let inner = cursor.get_ref();
+
+            match socket.send(&inner[..position]) {
+                Ok(_) => {
+                    statistics.requests += 1;
+                }
+                Err(err) => {
+                    eprintln!("Couldn't send packet: {:?}", err);
+                }
+            }
+        }
+        Err(err) => {
+            eprintln!("request_to_bytes err: {}", err);
+        }
+    }
+}
+
+fn update_shared_statistics(state: &LoadTestState, statistics: &mut SocketWorkerLocalStatistics) {
+    state
+        .statistics
+        .requests
+        .fetch_add(statistics.requests, Ordering::Relaxed);
+    state
+        .statistics
+        .responses_connect
+        .fetch_add(statistics.responses_connect, Ordering::Relaxed);
+    state
+        .statistics
+        .responses_announce
+        .fetch_add(statistics.responses_announce, Ordering::Relaxed);
+    state
+        .statistics
+        .responses_scrape
+        .fetch_add(statistics.responses_scrape, Ordering::Relaxed);
+    state
+        .statistics
+        .responses_error
+        .fetch_add(statistics.responses_error, Ordering::Relaxed);
+    state
+        .statistics
+        .response_peers
+        .fetch_add(statistics.response_peers, Ordering::Relaxed);
+
+    *statistics = SocketWorkerLocalStatistics::default();
+}
+
+fn create_socket(config: &Config, addr: SocketAddr) -> ::std::net::UdpSocket {
+    let socket = if addr.is_ipv4() {
+        Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::UDP))
+    } else {
+        Socket::new(Domain::IPV6, Type::DGRAM, Some(Protocol::UDP))
+    }
+    .expect("create socket");
+
+    socket
+        .set_nonblocking(true)
+        .expect("socket: set nonblocking");
+
+    if config.network.recv_buffer != 0 {
+        if let Err(err) = socket.set_recv_buffer_size(config.network.recv_buffer) {
+            eprintln!(
+                "socket: failed setting recv buffer to {}: {:?}",
+                config.network.recv_buffer, err
+            );
+        }
+    }
+
+    socket
+        .bind(&addr.into())
+        .unwrap_or_else(|err| panic!("socket: bind to {}: {:?}", addr, err));
+
+    socket
+        .connect(&config.server_address.into())
+        .expect("socket: connect to server");
+
+    socket.into()
+}

aquatic_udp_load_test/src/worker/request_gen.rs

@@ -0,0 +1,218 @@
+use std::sync::Arc;
+
+use rand::distributions::WeightedIndex;
+use rand::prelude::*;
+use rand_distr::Pareto;
+
+use aquatic_udp_protocol::*;
+
+use crate::common::*;
+use crate::config::Config;
+use crate::utils::*;
+
+pub fn process_response(
+    rng: &mut impl Rng,
+    pareto: Pareto<f64>,
+    info_hashes: &Arc<Vec<InfoHash>>,
+    config: &Config,
+    torrent_peers: &mut TorrentPeerMap,
+    response: Response,
+) -> Option<Request> {
+    match response {
+        Response::Connect(r) => {
+            // Fetch the torrent peer or create it if is doesn't exists. Update
+            // the connection id if fetched. Create a request and move the
+            // torrent peer appropriately.
+
+            let torrent_peer = torrent_peers
+                .remove(&r.transaction_id)
+                .map(|mut torrent_peer| {
+                    torrent_peer.connection_id = r.connection_id;
+
+                    torrent_peer
+                })
+                .unwrap_or_else(|| {
+                    create_torrent_peer(config, rng, pareto, info_hashes, r.connection_id)
+                });
+
+            let new_transaction_id = generate_transaction_id(rng);
+
+            let request =
+                create_random_request(config, rng, info_hashes, new_transaction_id, &torrent_peer);
+
+            torrent_peers.insert(new_transaction_id, torrent_peer);
+
+            Some(request)
+        }
+        Response::AnnounceIpv4(r) => if_torrent_peer_move_and_create_random_request(
+            config,
+            rng,
+            info_hashes,
+            torrent_peers,
+            r.transaction_id,
+        ),
+        Response::AnnounceIpv6(r) => if_torrent_peer_move_and_create_random_request(
+            config,
+            rng,
+            info_hashes,
+            torrent_peers,
+            r.transaction_id,
+        ),
+        Response::Scrape(r) => if_torrent_peer_move_and_create_random_request(
+            config,
+            rng,
+            info_hashes,
+            torrent_peers,
+            r.transaction_id,
+        ),
+        Response::Error(r) => {
+            if !r.message.to_lowercase().contains("connection") {
+                eprintln!(
+                    "Received error response which didn't contain the word 'connection': {}",
+                    r.message
+                );
+            }
+
+            if let Some(torrent_peer) = torrent_peers.remove(&r.transaction_id) {
+                let new_transaction_id = generate_transaction_id(rng);
+
+                torrent_peers.insert(new_transaction_id, torrent_peer);
+
+                Some(create_connect_request(new_transaction_id))
+            } else {
+                Some(create_connect_request(generate_transaction_id(rng)))
+            }
+        }
+    }
+}
+
+fn if_torrent_peer_move_and_create_random_request(
+    config: &Config,
+    rng: &mut impl Rng,
+    info_hashes: &Arc<Vec<InfoHash>>,
+    torrent_peers: &mut TorrentPeerMap,
+    transaction_id: TransactionId,
+) -> Option<Request> {
+    if let Some(torrent_peer) = torrent_peers.remove(&transaction_id) {
+        let new_transaction_id = generate_transaction_id(rng);
+
+        let request =
+            create_random_request(config, rng, info_hashes, new_transaction_id, &torrent_peer);
+
+        torrent_peers.insert(new_transaction_id, torrent_peer);
+
+        Some(request)
+    } else {
+        None
+    }
+}
+
+fn create_random_request(
+    config: &Config,
+    rng: &mut impl Rng,
+    info_hashes: &Arc<Vec<InfoHash>>,
+    transaction_id: TransactionId,
+    torrent_peer: &TorrentPeer,
+) -> Request {
+    let weights = vec![
+        config.requests.weight_announce as u32,
+        config.requests.weight_connect as u32,
+        config.requests.weight_scrape as u32,
+    ];
+
+    let items = vec![
+        RequestType::Announce,
+        RequestType::Connect,
+        RequestType::Scrape,
+    ];
+
+    let dist = WeightedIndex::new(&weights).expect("random request weighted index");
+
+    match items[dist.sample(rng)] {
+        RequestType::Announce => create_announce_request(config, rng, torrent_peer, transaction_id),
+        RequestType::Connect => create_connect_request(transaction_id),
+        RequestType::Scrape => create_scrape_request(&info_hashes, torrent_peer, transaction_id),
+    }
+}
+
+fn create_announce_request(
+    config: &Config,
+    rng: &mut impl Rng,
+    torrent_peer: &TorrentPeer,
+    transaction_id: TransactionId,
+) -> Request {
+    let (event, bytes_left) = {
+        if rng.gen_bool(config.requests.peer_seeder_probability) {
+            (AnnounceEvent::Completed, NumberOfBytes(0))
+        } else {
+            (AnnounceEvent::Started, NumberOfBytes(50))
+        }
+    };
+
+    (AnnounceRequest {
+        connection_id: torrent_peer.connection_id,
+        transaction_id,
+        info_hash: torrent_peer.info_hash,
+        peer_id: torrent_peer.peer_id,
+        bytes_downloaded: NumberOfBytes(50),
+        bytes_uploaded: NumberOfBytes(50),
+        bytes_left,
+        event,
+        ip_address: None,
+        key: PeerKey(12345),
+        peers_wanted: NumberOfPeers(100),
+        port: torrent_peer.port,
+    })
+    .into()
+}
+
+fn create_scrape_request(
+    info_hashes: &Arc<Vec<InfoHash>>,
+    torrent_peer: &TorrentPeer,
+    transaction_id: TransactionId,
+) -> Request {
+    let indeces = &torrent_peer.scrape_hash_indeces;
+
+    let mut scape_hashes = Vec::with_capacity(indeces.len());
+
+    for i in indeces {
+        scape_hashes.push(info_hashes[*i].to_owned())
+    }
+
+    (ScrapeRequest {
+        connection_id: torrent_peer.connection_id,
+        transaction_id,
+        info_hashes: scape_hashes,
+    })
+    .into()
+}
+
+fn create_torrent_peer(
+    config: &Config,
+    rng: &mut impl Rng,
+    pareto: Pareto<f64>,
+    info_hashes: &Arc<Vec<InfoHash>>,
+    connection_id: ConnectionId,
+) -> TorrentPeer {
+    let num_scape_hashes = rng.gen_range(1..config.requests.scrape_max_torrents);
+
+    let mut scrape_hash_indeces = Vec::new();
+
+    for _ in 0..num_scape_hashes {
+        scrape_hash_indeces.push(select_info_hash_index(config, rng, pareto))
+    }
+
+    let info_hash_index = select_info_hash_index(config, rng, pareto);
+
+    TorrentPeer {
+        info_hash: info_hashes[info_hash_index],
+        scrape_hash_indeces,
+        connection_id,
+        peer_id: generate_peer_id(),
+        port: Port(rand::random()),
+    }
+}
+
+fn select_info_hash_index(config: &Config, rng: &mut impl Rng, pareto: Pareto<f64>) -> usize {
+    pareto_usize(rng, pareto, config.requests.number_of_torrents - 1)
+}