From 36954e5f487c39b9b5b715049dfc15083764c38f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= <joakim.frostegard@gmail.com>
Date: Mon, 8 Jan 2024 18:50:17 +0100
Subject: [PATCH] ws: SwarmControlMessage::ConnectionClosed: use Vec for info
 hashes

---
 crates/ws/src/common.rs                    |  5 ++---
 crates/ws/src/workers/socket/connection.rs | 24 +++++++++++++---------
 crates/ws/src/workers/swarm/mod.rs         | 11 +++++-----
 3 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/crates/ws/src/common.rs b/crates/ws/src/common.rs
index dae682f..8c41fb1 100644
--- a/crates/ws/src/common.rs
+++ b/crates/ws/src/common.rs
@@ -67,11 +67,10 @@ impl Into<OutMessageMeta> for InMessageMeta {
     }
 }
 
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Debug)]
 pub enum SwarmControlMessage {
     ConnectionClosed {
-        info_hash: InfoHash,
-        peer_id: PeerId,
         ip_version: IpVersion,
+        announced_info_hashes: Vec<(InfoHash, PeerId)>,
     },
 }
diff --git a/crates/ws/src/workers/socket/connection.rs b/crates/ws/src/workers/socket/connection.rs
index 74908e9..bf24186 100644
--- a/crates/ws/src/workers/socket/connection.rs
+++ b/crates/ws/src/workers/socket/connection.rs
@@ -599,19 +599,23 @@ impl ConnectionCleanupData {
         config: &Config,
         control_message_senders: Rc<Senders<SwarmControlMessage>>,
     ) {
-        // Use RefCell::take to avoid issues with Rc borrow across await
-        let announced_info_hashes = self.announced_info_hashes.take();
-
-        // Tell swarm workers to remove peer
-        for (info_hash, peer_id) in announced_info_hashes.into_iter() {
-            let message = SwarmControlMessage::ConnectionClosed {
-                info_hash,
-                peer_id,
-                ip_version: self.ip_version,
-            };
+        let mut announced_info_hashes = HashMap::new();
 
+        for (info_hash, peer_id) in self.announced_info_hashes.take().into_iter() {
             let consumer_index = calculate_in_message_consumer_index(&config, info_hash);
 
+            announced_info_hashes
+                .entry(consumer_index)
+                .or_insert(Vec::new())
+                .push((info_hash, peer_id));
+        }
+
+        for (consumer_index, announced_info_hashes) in announced_info_hashes.into_iter() {
+            let message = SwarmControlMessage::ConnectionClosed {
+                ip_version: self.ip_version,
+                announced_info_hashes,
+            };
+
             control_message_senders
                 .send_to(consumer_index, message)
                 .await
diff --git a/crates/ws/src/workers/swarm/mod.rs b/crates/ws/src/workers/swarm/mod.rs
index 8fff337..599ad67 100644
--- a/crates/ws/src/workers/swarm/mod.rs
+++ b/crates/ws/src/workers/swarm/mod.rs
@@ -102,13 +102,14 @@ where
     while let Some(message) = stream.next().await {
         match message {
             SwarmControlMessage::ConnectionClosed {
-                info_hash,
-                peer_id,
                 ip_version,
+                announced_info_hashes,
             } => {
-                torrents
-                    .borrow_mut()
-                    .handle_connection_closed(info_hash, peer_id, ip_version);
+                let mut torrents = torrents.borrow_mut();
+
+                for (info_hash, peer_id) in announced_info_hashes {
+                    torrents.handle_connection_closed(info_hash, peer_id, ip_version);
+                }
             }
         }
     }