From 0b6a02e1a772e485590d92417f926a9bdcf51e60 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= <joakim.frostegard@gmail.com>
Date: Tue, 6 Feb 2024 18:32:26 +0100
Subject: [PATCH] udp load test: use connection IDs from responses in requests

---
 crates/udp_load_test/src/config.rs |  2 +-
 crates/udp_load_test/src/worker.rs | 15 ++++++++++++++-
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/crates/udp_load_test/src/config.rs b/crates/udp_load_test/src/config.rs
index 5dbbd33..5f87318 100644
--- a/crates/udp_load_test/src/config.rs
+++ b/crates/udp_load_test/src/config.rs
@@ -124,7 +124,7 @@ impl Default for RequestConfig {
             number_of_peers: 2_000_000,
             scrape_max_torrents: 10,
             announce_peers_wanted: 30,
-            weight_connect: 0,
+            weight_connect: 1,
             weight_announce: 100,
             weight_scrape: 1,
             peer_seeder_probability: 0.75,
diff --git a/crates/udp_load_test/src/worker.rs b/crates/udp_load_test/src/worker.rs
index a090a18..0d8037c 100644
--- a/crates/udp_load_test/src/worker.rs
+++ b/crates/udp_load_test/src/worker.rs
@@ -89,7 +89,10 @@ impl Worker {
                 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);
+                            self.send_connect_request(
+                                connect_socket_index,
+                                connect_socket_index.into(),
+                            );
 
                             connect_socket_index = connect_socket_index.wrapping_add(1)
                                 % self.config.network.sockets_per_worker;
@@ -117,6 +120,16 @@ impl Worker {
                 match socket.recv(&mut self.buffer[..]) {
                     Ok(amt) => {
                         match Response::parse_bytes(&self.buffer[0..amt], self.addr.is_ipv4()) {
+                            Ok(Response::Connect(r)) => {
+                                // If we're sending connect requests, we might
+                                // as well keep connection IDs valid
+                                let connection_id_index =
+                                    u32::from_ne_bytes(r.transaction_id.0.get().to_ne_bytes())
+                                        as usize;
+                                connection_ids[connection_id_index] = r.connection_id;
+
+                                self.handle_response(Response::Connect(r));
+                            }
                             Ok(response) => {
                                 self.handle_response(response);
                             }