From 910accad31bac17b9ab810ea00ecf2d5d678e2ed Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= <joakim.frostegard@gmail.com>
Date: Wed, 8 Apr 2020 13:32:18 +0200
Subject: [PATCH] announce handler: keep mutable access to torrent map less
 time

---
 TODO.md                                      |  3 --
 aquatic/src/lib/handlers.rs                  | 48 +++++++++++---------
 aquatic_bench/src/bin/bench_handlers/main.rs |  6 +--
 3 files changed, 29 insertions(+), 28 deletions(-)

diff --git a/TODO.md b/TODO.md
index f2732d9..ddaff43 100644
--- a/TODO.md
+++ b/TODO.md
@@ -17,9 +17,6 @@
 * Performance
     * cpu-target=native good?
     * mialloc good?
-    * TorrentMap: mutable access only to insert peer, then drop reference
-      and get read-only access to gather peers. This could speed up
-      multi-threaded performance a lot
     * Use less bytes from PeerId for hashing? Would need to implement
       "faulty" PartialEq too (on PeerMapKey, which would be OK)
 * bittorrent_udp
diff --git a/aquatic/src/lib/handlers.rs b/aquatic/src/lib/handlers.rs
index 4317d46..26f2b6b 100644
--- a/aquatic/src/lib/handlers.rs
+++ b/aquatic/src/lib/handlers.rs
@@ -59,10 +59,6 @@ pub fn handle_announce_requests(
             return None;
         }
 
-        let mut torrent_data = state.torrents
-            .entry(request.info_hash)
-            .or_insert_with(|| TorrentData::default());
-
         let peer_key = PeerMapKey {
             ip: src.ip(),
             peer_id: request.peer_id,
@@ -70,12 +66,25 @@ pub fn handle_announce_requests(
 
         let peer = Peer::from_announce_and_ip(&request, src.ip());
         let peer_status = peer.status;
-        
-        let opt_removed_peer = if peer.status == PeerStatus::Stopped {
-            torrent_data.peers.remove(&peer_key)
-        } else {
-            torrent_data.peers.insert(peer_key, peer)
+
+        let opt_removed_peer_status = {
+            let mut torrent_data = state.torrents
+                .entry(request.info_hash)
+                .or_insert_with(|| TorrentData::default());
+            
+            if peer_status == PeerStatus::Stopped {
+                torrent_data.peers.remove(&peer_key)
+                    .map(|peer| peer.status)
+            } else {
+                torrent_data.peers.insert(peer_key, peer)
+                    .map(|peer| peer.status)
+            }
         };
+
+        let max_num_peers_to_take = (request.peers_wanted.0.max(0) as usize)
+            .min(config.max_response_peers);
+
+        let torrent_data = state.torrents.get(&request.info_hash).unwrap();
         
         match peer_status {
             PeerStatus::Leeching => {
@@ -87,21 +96,16 @@ pub fn handle_announce_requests(
             PeerStatus::Stopped => {}
         };
 
-        if let Some(removed_peer) = opt_removed_peer {
-            match removed_peer.status {
-                PeerStatus::Leeching => {
-                    torrent_data.num_leechers.fetch_sub(1, Ordering::SeqCst);
-                },
-                PeerStatus::Seeding => {
-                    torrent_data.num_seeders.fetch_sub(1, Ordering::SeqCst);
-                },
-                PeerStatus::Stopped => {}
-            }
+        match opt_removed_peer_status {
+            Some(PeerStatus::Leeching) => {
+                torrent_data.num_leechers.fetch_sub(1, Ordering::SeqCst);
+            },
+            Some(PeerStatus::Seeding) => {
+                torrent_data.num_seeders.fetch_sub(1, Ordering::SeqCst);
+            },
+            _ => {}
         }
 
-        let max_num_peers_to_take = (request.peers_wanted.0.max(0) as usize)
-            .min(config.max_response_peers);
-
         let response_peers = extract_response_peers(
             &mut rng,
             &torrent_data.peers,
diff --git a/aquatic_bench/src/bin/bench_handlers/main.rs b/aquatic_bench/src/bin/bench_handlers/main.rs
index 03fa21e..9eb2ba5 100644
--- a/aquatic_bench/src/bin/bench_handlers/main.rs
+++ b/aquatic_bench/src/bin/bench_handlers/main.rs
@@ -4,9 +4,9 @@
 //! ```
 //! ## Average results over 10 rounds
 //!
-//! Connect handler:   2 351 042 requests/second,   425.85 ns/request
-//! Announce handler:    266 609 requests/second,  3754.52 ns/request
-//! Scrape handler:      463 222 requests/second,  2161.85 ns/request
+//! Connect handler:   2 330 506 requests/second,   429.30 ns/request
+//! Announce handler:    264 596 requests/second,  3782.88 ns/request
+//! Scrape handler:      461 530 requests/second,  2169.41 ns/request
 //! ```
 
 use std::time::{Duration, Instant};