From 3bb6c1994c39b68d9f03a2dca60002159135ada4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= <joakim.frostegard@gmail.com>
Date: Fri, 15 Oct 2021 02:53:33 +0200
Subject: [PATCH] aquatic_udp: move TorrentMap cleaning logic to TorrentMap
 impl

---
 aquatic_udp/src/lib/common.rs | 62 +++++++++++++++++++++++++++++++++++
 aquatic_udp/src/lib/tasks.rs  | 59 ++++-----------------------------
 2 files changed, 69 insertions(+), 52 deletions(-)

diff --git a/aquatic_udp/src/lib/common.rs b/aquatic_udp/src/lib/common.rs
index b7eba08..0fc0ebe 100644
--- a/aquatic_udp/src/lib/common.rs
+++ b/aquatic_udp/src/lib/common.rs
@@ -1,7 +1,9 @@
 use std::hash::Hash;
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
 use std::sync::{atomic::AtomicUsize, Arc};
+use std::time::Instant;
 
+use aquatic_common::AccessListType;
 use hashbrown::HashMap;
 use indexmap::IndexMap;
 use parking_lot::Mutex;
@@ -108,6 +110,66 @@ pub struct TorrentMaps {
     pub ipv6: TorrentMap<Ipv6Addr>,
 }
 
+impl TorrentMaps {
+    /// Remove inactive torrents
+    pub fn clean(&mut self, now: Instant) {
+        self.ipv4
+            .retain(|_, torrent| Self::clean_torrent_and_peers(now, torrent));
+        self.ipv4.shrink_to_fit();
+
+        self.ipv6
+            .retain(|_, torrent| Self::clean_torrent_and_peers(now, torrent));
+        self.ipv6.shrink_to_fit();
+    }
+
+    /// Remove disallowed and inactive torrents
+    pub fn clean_with_access_list(
+        &mut self,
+        access_list_type: AccessListType,
+        access_list: &AccessList,
+        now: Instant,
+    ) {
+        self.ipv4.retain(|info_hash, torrent| {
+            access_list.allows(access_list_type, &info_hash.0)
+                && Self::clean_torrent_and_peers(now, torrent)
+        });
+        self.ipv4.shrink_to_fit();
+
+        self.ipv6.retain(|info_hash, torrent| {
+            access_list.allows(access_list_type, &info_hash.0)
+                && Self::clean_torrent_and_peers(now, torrent)
+        });
+        self.ipv6.shrink_to_fit();
+    }
+
+    /// Returns true if torrent is to be kept
+    #[inline]
+    fn clean_torrent_and_peers<I: Ip>(now: Instant, torrent: &mut TorrentData<I>) -> bool {
+        let num_seeders = &mut torrent.num_seeders;
+        let num_leechers = &mut torrent.num_leechers;
+
+        torrent.peers.retain(|_, peer| {
+            let keep = peer.valid_until.0 > now;
+
+            if !keep {
+                match peer.status {
+                    PeerStatus::Seeding => {
+                        *num_seeders -= 1;
+                    }
+                    PeerStatus::Leeching => {
+                        *num_leechers -= 1;
+                    }
+                    _ => (),
+                };
+            }
+
+            keep
+        });
+
+        !torrent.peers.is_empty()
+    }
+}
+
 #[derive(Default)]
 pub struct Statistics {
     pub requests_received: AtomicUsize,
diff --git a/aquatic_udp/src/lib/tasks.rs b/aquatic_udp/src/lib/tasks.rs
index 391038d..628a5ab 100644
--- a/aquatic_udp/src/lib/tasks.rs
+++ b/aquatic_udp/src/lib/tasks.rs
@@ -19,70 +19,25 @@ pub fn clean_connections_and_torrents(config: &Config, state: &State) {
     }
 
     match config.access_list.list_type {
-        access_list_type @ (AccessListType::Allow | AccessListType::Deny) => {
+        AccessListType::Allow | AccessListType::Deny => {
             let mut access_list = state.access_list.lock();
 
             if let Err(err) = access_list.update_from_path(&config.access_list.path) {
                 ::log::error!("Update access list from path: {:?}", err);
             }
 
-            let mut torrents = state.torrents.lock();
-
-            torrents.ipv4.retain(|info_hash, torrent| {
-                access_list.allows(access_list_type, &info_hash.0)
-                    && clean_torrent_and_peers(now, torrent)
-            });
-            torrents.ipv4.shrink_to_fit();
-
-            torrents.ipv6.retain(|info_hash, torrent| {
-                access_list.allows(access_list_type, &info_hash.0)
-                    && clean_torrent_and_peers(now, torrent)
-            });
-            torrents.ipv6.shrink_to_fit();
+            state.torrents.lock().clean_with_access_list(
+                config.access_list.list_type,
+                &access_list,
+                now,
+            );
         }
         AccessListType::Ignore => {
-            let mut torrents = state.torrents.lock();
-
-            torrents
-                .ipv4
-                .retain(|_, torrent| clean_torrent_and_peers(now, torrent));
-            torrents.ipv4.shrink_to_fit();
-
-            torrents
-                .ipv6
-                .retain(|_, torrent| clean_torrent_and_peers(now, torrent));
-            torrents.ipv6.shrink_to_fit();
+            state.torrents.lock().clean(now);
         }
     }
 }
 
-/// Returns true if torrent is to be kept
-#[inline]
-fn clean_torrent_and_peers<I: Ip>(now: Instant, torrent: &mut TorrentData<I>) -> bool {
-    let num_seeders = &mut torrent.num_seeders;
-    let num_leechers = &mut torrent.num_leechers;
-
-    torrent.peers.retain(|_, peer| {
-        let keep = peer.valid_until.0 > now;
-
-        if !keep {
-            match peer.status {
-                PeerStatus::Seeding => {
-                    *num_seeders -= 1;
-                }
-                PeerStatus::Leeching => {
-                    *num_leechers -= 1;
-                }
-                _ => (),
-            };
-        }
-
-        keep
-    });
-
-    !torrent.peers.is_empty()
-}
-
 pub fn gather_and_print_statistics(state: &State, config: &Config) {
     let interval = config.statistics.interval;