ws load test: refactor network module

crates/ws_load_test/src/network.rs

@@ -7,7 +7,9 @@ use std::{
 };
 
 use aquatic_ws_protocol::{
-    InMessage, InfoHash, OfferId, OutMessage, PeerId, RtcAnswer, RtcAnswerType,
+    AnnounceAction, AnnounceEvent, AnnounceRequest, AnnounceRequestOffer, InMessage, InfoHash,
+    OfferId, OutMessage, PeerId, RtcAnswer, RtcAnswerType, RtcOffer, RtcOfferType, ScrapeAction,
+    ScrapeRequest, ScrapeRequestInfoHashes,
 };
 use async_tungstenite::{client_async, WebSocketStream};
 use futures::{SinkExt, StreamExt};
@@ -15,8 +17,13 @@ use futures_rustls::{client::TlsStream, TlsConnector};
 use glommio::net::TcpStream;
 use glommio::{prelude::*, timer::TimerActionRepeat};
 use rand::{prelude::SmallRng, Rng, SeedableRng};
+use rand_distr::{Distribution, WeightedIndex};
 
-use crate::{common::LoadTestState, config::Config, utils::create_random_request};
+use crate::{
+    common::{LoadTestState, RequestType},
+    config::Config,
+    utils::select_info_hash_index,
+};
 
 pub async fn run_socket_thread(
     config: Config,
@@ -24,6 +31,7 @@ pub async fn run_socket_thread(
     load_test_state: LoadTestState,
 ) -> anyhow::Result<()> {
     let config = Rc::new(config);
+    let rng = Rc::new(RefCell::new(SmallRng::from_entropy()));
     let num_active_connections = Rc::new(RefCell::new(0usize));
     let connection_creation_interval =
         Duration::from_millis(config.connection_creation_interval_ms);
@@ -34,12 +42,14 @@ pub async fn run_socket_thread(
             tls_config.clone(),
             load_test_state.clone(),
             num_active_connections.clone(),
+            rng.clone(),
             connection_creation_interval,
         )
     })
     .join()
-    .await
+    .await;
-    .ok_or_else(|| anyhow::anyhow!("connection opener timer cancelled"))
+
+    Ok(())
 }
 
 async fn periodically_open_connections(
@@ -47,12 +57,19 @@ async fn periodically_open_connections(
     tls_config: Arc<rustls::ClientConfig>,
     load_test_state: LoadTestState,
     num_active_connections: Rc<RefCell<usize>>,
+    rng: Rc<RefCell<SmallRng>>,
     connection_creation_interval: Duration,
 ) -> Option<Duration> {
     if *num_active_connections.borrow() < config.num_connections_per_worker {
         spawn_local(async move {
-            if let Err(err) =
+            if let Err(err) = Connection::run(
-                Connection::run(config, tls_config, load_test_state, num_active_connections).await
+                config,
+                tls_config,
+                load_test_state,
+                num_active_connections,
+                rng,
+            )
+            .await
             {
                 ::log::info!("connection creation error: {:#}", err);
             }
@@ -66,7 +83,7 @@ async fn periodically_open_connections(
 struct Connection {
     config: Rc<Config>,
     load_test_state: LoadTestState,
-    rng: SmallRng,
+    rng: Rc<RefCell<SmallRng>>,
     peer_id: PeerId,
     can_send_answer: Option<(InfoHash, PeerId, OfferId)>,
     stream: WebSocketStream<TlsStream<TcpStream>>,
@@ -78,9 +95,9 @@ impl Connection {
         tls_config: Arc<rustls::ClientConfig>,
         load_test_state: LoadTestState,
         num_active_connections: Rc<RefCell<usize>>,
+        rng: Rc<RefCell<SmallRng>>,
     ) -> anyhow::Result<()> {
-        let mut rng = SmallRng::from_entropy();
+        let peer_id = PeerId(rng.borrow_mut().gen());
-        let peer_id = PeerId(rng.gen());
         let stream = TcpStream::connect(config.server_address)
             .await
             .map_err(|err| anyhow::anyhow!("connect: {:?}", err))?;
@@ -126,36 +143,87 @@ impl Connection {
     }
 
     async fn send_message(&mut self) -> anyhow::Result<()> {
-        let request = create_random_request(
+        let mut rng = self.rng.borrow_mut();
-            &self.config,
+
-            &self.load_test_state,
+        let request = match random_request_type(&self.config, &mut *rng) {
-            &mut self.rng,
+            RequestType::Announce => {
-            self.peer_id,
+                let (event, bytes_left) = {
-            self.can_send_answer.is_none(),
+                    if rng.gen_bool(self.config.torrents.peer_seeder_probability) {
-        );
+                        (AnnounceEvent::Completed, 0)
+                    } else {
+                        (AnnounceEvent::Started, 50)
+                    }
+                };
+
+                const SDP: &str = "abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg";
 
-        // If self.send_answer is set and request is announce request, make
-        // the request an offer answer
-        let request = if let InMessage::AnnounceRequest(mut r) = request {
                 if let Some((info_hash, peer_id, offer_id)) = self.can_send_answer {
-                r.info_hash = info_hash;
+                    InMessage::AnnounceRequest(AnnounceRequest {
-                r.answer_to_peer_id = Some(peer_id);
+                        info_hash,
-                r.answer_offer_id = Some(offer_id);
+                        answer_to_peer_id: Some(peer_id),
-                r.answer = Some(RtcAnswer {
+                        answer_offer_id: Some(offer_id),
+                        answer: Some(RtcAnswer {
                             t: RtcAnswerType::Answer,
-                    sdp: "abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-".into()
+                            sdp: SDP.into(),
-                });
+                        }),
-                r.event = None;
+                        event: None,
-                r.offers = None;
+                        offers: None,
+                        action: aquatic_ws_protocol::AnnounceAction::Announce,
+                        peer_id: self.peer_id,
+                        bytes_left: Some(bytes_left),
+                        numwant: Some(0),
+                    })
+                } else {
+                    let info_hash_index =
+                        select_info_hash_index(&self.config, &self.load_test_state, &mut *rng);
+
+                    let mut offers = Vec::with_capacity(self.config.torrents.offers_per_request);
+
+                    for _ in 0..self.config.torrents.offers_per_request {
+                        offers.push(AnnounceRequestOffer {
+                            offer_id: OfferId(rng.gen()),
+                            offer: RtcOffer {
+                                t: RtcOfferType::Offer,
+                                sdp: SDP.into(),
+                            },
+                        })
                     }
 
-            self.can_send_answer = None;
+                    InMessage::AnnounceRequest(AnnounceRequest {
+                        action: AnnounceAction::Announce,
+                        info_hash: self.load_test_state.info_hashes[info_hash_index],
+                        peer_id: self.peer_id,
+                        bytes_left: Some(bytes_left),
+                        event: Some(event),
+                        numwant: Some(offers.len()),
+                        offers: Some(offers),
+                        answer: None,
+                        answer_to_peer_id: None,
+                        answer_offer_id: None,
+                    })
+                }
+            }
+            RequestType::Scrape => {
+                let mut scrape_hashes = Vec::with_capacity(5);
 
-            InMessage::AnnounceRequest(r)
+                for _ in 0..5 {
-        } else {
+                    let info_hash_index =
-            request
+                        select_info_hash_index(&self.config, &self.load_test_state, &mut *rng);
+
+                    scrape_hashes.push(self.load_test_state.info_hashes[info_hash_index]);
+                }
+
+                InMessage::ScrapeRequest(ScrapeRequest {
+                    action: ScrapeAction::Scrape,
+                    info_hashes: Some(ScrapeRequestInfoHashes::Multiple(scrape_hashes)),
+                })
+            }
         };
 
+        drop(rng);
+
+        self.can_send_answer = None;
+
         self.stream.send(request.to_ws_message()).await?;
 
         self.load_test_state
@@ -229,3 +297,16 @@ impl Connection {
         Ok(())
     }
 }
+
+pub fn random_request_type(config: &Config, rng: &mut impl Rng) -> RequestType {
+    let weights = [
+        config.torrents.weight_announce as u32,
+        config.torrents.weight_scrape as u32,
+    ];
+
+    let items = [RequestType::Announce, RequestType::Scrape];
+
+    let dist = WeightedIndex::new(&weights).expect("random request weighted index");
+
+    items[dist.sample(rng)]
+}

crates/ws_load_test/src/utils.rs

@@ -1,104 +1,13 @@
 use std::sync::Arc;
 
-use rand::distributions::WeightedIndex;
 use rand::prelude::*;
 use rand_distr::Gamma;
 
 use crate::common::*;
 use crate::config::*;
 
-pub fn create_random_request(
-    config: &Config,
-    state: &LoadTestState,
-    rng: &mut impl Rng,
-    peer_id: PeerId,
-    announce_gen_offers: bool,
-) -> InMessage {
-    let weights = [
-        config.torrents.weight_announce as u32,
-        config.torrents.weight_scrape as u32,
-    ];
-
-    let items = [RequestType::Announce, RequestType::Scrape];
-
-    let dist = WeightedIndex::new(&weights).expect("random request weighted index");
-
-    match items[dist.sample(rng)] {
-        RequestType::Announce => {
-            create_announce_request(config, state, rng, peer_id, announce_gen_offers)
-        }
-        RequestType::Scrape => create_scrape_request(config, state, rng),
-    }
-}
-
 #[inline]
-fn create_announce_request(
+pub fn select_info_hash_index(config: &Config, state: &LoadTestState, rng: &mut impl Rng) -> usize {
-    config: &Config,
-    state: &LoadTestState,
-    rng: &mut impl Rng,
-    peer_id: PeerId,
-    gen_offers: bool,
-) -> InMessage {
-    let (event, bytes_left) = {
-        if rng.gen_bool(config.torrents.peer_seeder_probability) {
-            (AnnounceEvent::Completed, 0)
-        } else {
-            (AnnounceEvent::Started, 50)
-        }
-    };
-
-    let info_hash_index = select_info_hash_index(config, &state, rng);
-
-    let offers = if gen_offers {
-        let mut offers = Vec::with_capacity(config.torrents.offers_per_request);
-
-        for _ in 0..config.torrents.offers_per_request {
-            offers.push(AnnounceRequestOffer {
-                offer_id: OfferId(rng.gen()),
-                offer: RtcOffer {
-                    t: RtcOfferType::Offer,
-                    sdp: "abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-".into()
-                },
-            })
-        }
-
-        offers
-    } else {
-        Vec::new()
-    };
-
-    InMessage::AnnounceRequest(AnnounceRequest {
-        action: AnnounceAction::Announce,
-        info_hash: state.info_hashes[info_hash_index],
-        peer_id,
-        bytes_left: Some(bytes_left),
-        event: Some(event),
-        numwant: Some(offers.len()),
-        offers: Some(offers),
-        answer: None,
-        answer_to_peer_id: None,
-        answer_offer_id: None,
-    })
-}
-
-#[inline]
-fn create_scrape_request(config: &Config, state: &LoadTestState, rng: &mut impl Rng) -> InMessage {
-    let mut scrape_hashes = Vec::with_capacity(5);
-
-    for _ in 0..5 {
-        let info_hash_index = select_info_hash_index(config, &state, rng);
-
-        scrape_hashes.push(state.info_hashes[info_hash_index]);
-    }
-
-    InMessage::ScrapeRequest(ScrapeRequest {
-        action: ScrapeAction::Scrape,
-        info_hashes: Some(ScrapeRequestInfoHashes::Multiple(scrape_hashes)),
-    })
-}
-
-#[inline]
-fn select_info_hash_index(config: &Config, state: &LoadTestState, rng: &mut impl Rng) -> usize {
     gamma_usize(rng, &state.gamma, config.torrents.number_of_torrents - 1)
 }