mtorrent_core/data/
piece_requests.rs

1use std::collections::{HashMap, HashSet};
2use std::net::SocketAddr;
3
4/// Keeps track of outstanding requests with piece-level granularity (as opposed to blocks).
5#[derive(Default, Debug)]
6pub struct PendingRequests {
7    piece_requested_from: HashMap<usize, HashSet<SocketAddr>>,
8}
9
10impl PendingRequests {
11    /// Add a record of a new request sent to `peer` asking for the piece index `piece`.
12    pub fn add(&mut self, piece: usize, peer: &SocketAddr) {
13        self.piece_requested_from.entry(piece).or_default().insert(*peer);
14    }
15
16    /// Forget all pending requests asking peers for the piece index `piece`.
17    pub fn clear_requests_of(&mut self, piece: usize) {
18        self.piece_requested_from.remove(&piece);
19    }
20
21    /// Forget all pending requests sent to `peer`.
22    pub fn clear_requests_to(&mut self, peer: &SocketAddr) {
23        for peers in self.piece_requested_from.values_mut() {
24            peers.remove(peer);
25        }
26    }
27
28    /// Check presense of any pending requests asking for the piece index `piece`.
29    pub fn is_piece_requested(&self, piece: usize) -> bool {
30        self.piece_requested_from.get(&piece).is_some_and(|peers| !peers.is_empty())
31    }
32
33    /// Check if a request asking for the piece index `piece` has been sent to `peer`.
34    pub fn is_piece_requested_from(&self, peer: &SocketAddr, piece: usize) -> bool {
35        self.piece_requested_from.get(&piece).is_some_and(|peers| peers.contains(peer))
36    }
37
38    /// Count of all requests currently in-flight.
39    pub fn requests_in_flight(&self) -> usize {
40        self.piece_requested_from.values().flatten().count()
41    }
42
43    /// Count of all pieces currently requested from at least one peer.
44    pub fn pieces_requested(&self) -> usize {
45        self.piece_requested_from.values().filter(|peers| !peers.is_empty()).count()
46    }
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52    use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
53
54    #[test]
55    fn test_pending_requests_from_single_peer() {
56        let peer = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 6666));
57        let mut pr = PendingRequests::default();
58
59        pr.add(42, &peer);
60        pr.add(43, &peer);
61        pr.add(44, &peer);
62        assert!(pr.is_piece_requested(42));
63        assert!(pr.is_piece_requested(43));
64        assert!(pr.is_piece_requested(44));
65
66        pr.clear_requests_of(43);
67        assert!(pr.is_piece_requested(42));
68        assert!(!pr.is_piece_requested(43));
69        assert!(pr.is_piece_requested(44));
70
71        pr.clear_requests_to(&peer);
72        assert!(!pr.is_piece_requested(42));
73        assert!(!pr.is_piece_requested(43));
74        assert!(!pr.is_piece_requested(44));
75    }
76
77    #[test]
78    fn test_pending_requests_from_multiple_peers() {
79        let peer1 = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 6666));
80        let peer2 = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 6667));
81        let peer3 = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 6668));
82        let mut pr = PendingRequests::default();
83
84        pr.add(42, &peer1);
85        pr.add(42, &peer2);
86        pr.add(42, &peer3);
87        assert!(pr.is_piece_requested(42));
88
89        pr.clear_requests_to(&peer2);
90        assert!(pr.is_piece_requested(42));
91
92        pr.clear_requests_of(42);
93        assert!(!pr.is_piece_requested(42));
94    }
95}
mtorrent_core/data/piece_requests.rs

mtorrent_core/data/
piece_requests.rs
