diff --git a/aquatic_udp/src/workers/swarm/mod.rs b/aquatic_udp/src/workers/swarm/mod.rs index dbf2b03..2c054b3 100644 --- a/aquatic_udp/src/workers/swarm/mod.rs +++ b/aquatic_udp/src/workers/swarm/mod.rs @@ -17,7 +17,7 @@ use aquatic_udp_protocol::*; use crate::common::*; use crate::config::Config; -use storage::{Peer, TorrentMap, TorrentMaps}; +use storage::{TorrentMap, TorrentMaps}; pub fn run_swarm_worker( _sentinel: PanicSentinel, @@ -145,16 +145,17 @@ fn handle_announce_request( ) }; - let peer = Peer { - ip_address: peer_ip, - port: request.port, - status: PeerStatus::from_event_and_bytes_left(request.event, request.bytes_left), - valid_until: peer_valid_until, - }; - let torrent_data = torrents.0.entry(request.info_hash).or_default(); - torrent_data.update_peer(request.peer_id, peer); + let peer_status = PeerStatus::from_event_and_bytes_left(request.event, request.bytes_left); + + torrent_data.update_peer( + request.peer_id, + peer_ip, + request.port, + peer_status, + peer_valid_until, + ); let response_peers = torrent_data.extract_response_peers(rng, request.peer_id, max_num_peers_to_take); diff --git a/aquatic_udp/src/workers/swarm/storage.rs b/aquatic_udp/src/workers/swarm/storage.rs index 6523d25..cebe1df 100644 --- a/aquatic_udp/src/workers/swarm/storage.rs +++ b/aquatic_udp/src/workers/swarm/storage.rs @@ -20,15 +20,15 @@ use crate::config::Config; use super::create_torrent_scrape_statistics; #[derive(Clone, Debug)] -pub struct Peer { - pub ip_address: I, - pub port: Port, - pub status: PeerStatus, - pub valid_until: ValidUntil, +struct Peer { + ip_address: I, + port: Port, + is_seeder: bool, + valid_until: ValidUntil, } impl Peer { - pub fn to_response_peer(&self) -> ResponsePeer { + fn to_response_peer(&self) -> ResponsePeer { ResponsePeer { ip_address: self.ip_address, port: self.port, @@ -36,7 +36,7 @@ impl Peer { } } -pub type PeerMap = IndexMap>; +type PeerMap = IndexMap>; pub struct TorrentData { peers: PeerMap, @@ -45,14 +45,35 @@ pub struct TorrentData { } impl TorrentData { - pub fn update_peer(&mut self, peer_id: PeerId, peer: Peer) { - let opt_removed_peer = match peer.status { + pub fn update_peer( + &mut self, + peer_id: PeerId, + ip_address: I, + port: Port, + status: PeerStatus, + valid_until: ValidUntil, + ) { + let opt_removed_peer = match status { PeerStatus::Leeching => { + let peer = Peer { + ip_address, + port, + is_seeder: false, + valid_until, + }; + self.num_leechers += 1; self.peers.insert(peer_id, peer) } PeerStatus::Seeding => { + let peer = Peer { + ip_address, + port, + is_seeder: true, + valid_until, + }; + self.num_seeders += 1; self.peers.insert(peer_id, peer) @@ -60,14 +81,14 @@ impl TorrentData { PeerStatus::Stopped => self.peers.remove(&peer_id), }; - match opt_removed_peer.map(|peer| peer.status) { - Some(PeerStatus::Leeching) => { + match opt_removed_peer.map(|peer| peer.is_seeder) { + Some(true) => { self.num_leechers -= 1; } - Some(PeerStatus::Seeding) => { + Some(false) => { self.num_seeders -= 1; } - _ => {} + None => {} } } @@ -107,15 +128,11 @@ impl TorrentData { if peer.valid_until.valid(now) { true } else { - match peer.status { - PeerStatus::Seeding => { - self.num_seeders -= 1; - } - PeerStatus::Leeching => { - self.num_leechers -= 1; - } - _ => (), - }; + if peer.is_seeder { + self.num_seeders -= 1; + } else { + self.num_leechers -= 1; + } false } @@ -265,7 +282,7 @@ mod tests { Peer { ip_address: Ipv4Addr::from(i.to_be_bytes()), port: Port(1), - status: PeerStatus::Leeching, + is_seeder: false, valid_until: ValidUntil::new(ServerStartInstant::new(), 0), } }