cuprate_p2p/peer_set/
client_wrappers.rs

1use std::{
2    ops::{Deref, DerefMut},
3    sync::{
4        atomic::{AtomicBool, Ordering},
5        Arc,
6    },
7};
8
9use cuprate_p2p_core::{client::Client, NetworkZone};
10
11/// A client stored in the peer-set.
12pub(super) struct StoredClient<N: NetworkZone> {
13    pub client: Client<N>,
14    /// An [`AtomicBool`] for if the peer is currently downloading blocks.
15    downloading_blocks: Arc<AtomicBool>,
16    /// An [`AtomicBool`] for if the peer is currently being used to stem txs.
17    stem_peer: Arc<AtomicBool>,
18}
19
20impl<N: NetworkZone> StoredClient<N> {
21    pub(super) fn new(client: Client<N>) -> Self {
22        Self {
23            client,
24            downloading_blocks: Arc::new(AtomicBool::new(false)),
25            stem_peer: Arc::new(AtomicBool::new(false)),
26        }
27    }
28
29    /// Returns [`true`] if the [`StoredClient`] is currently downloading blocks.
30    pub(super) fn is_downloading_blocks(&self) -> bool {
31        self.downloading_blocks.load(Ordering::Relaxed)
32    }
33
34    /// Returns [`true`] if the [`StoredClient`] is currently being used to stem txs.
35    pub(super) fn is_a_stem_peer(&self) -> bool {
36        self.stem_peer.load(Ordering::Relaxed)
37    }
38
39    /// Returns a [`ClientDropGuard`] that while it is alive keeps the [`StoredClient`] in the downloading blocks state.
40    pub(super) fn downloading_blocks_guard(&self) -> ClientDropGuard<N> {
41        self.downloading_blocks.store(true, Ordering::Relaxed);
42
43        ClientDropGuard {
44            client: self.client.clone(),
45            bool: Arc::clone(&self.downloading_blocks),
46        }
47    }
48
49    /// Returns a [`ClientDropGuard`] that while it is alive keeps the [`StoredClient`] in the stemming peers state.
50    pub(super) fn stem_peer_guard(&self) -> ClientDropGuard<N> {
51        self.stem_peer.store(true, Ordering::Relaxed);
52
53        ClientDropGuard {
54            client: self.client.clone(),
55            bool: Arc::clone(&self.stem_peer),
56        }
57    }
58}
59
60/// A [`Drop`] guard for a client returned from the peer-set.
61pub struct ClientDropGuard<N: NetworkZone> {
62    client: Client<N>,
63    bool: Arc<AtomicBool>,
64}
65
66impl<N: NetworkZone> Deref for ClientDropGuard<N> {
67    type Target = Client<N>;
68    fn deref(&self) -> &Self::Target {
69        &self.client
70    }
71}
72
73impl<N: NetworkZone> DerefMut for ClientDropGuard<N> {
74    fn deref_mut(&mut self) -> &mut Self::Target {
75        &mut self.client
76    }
77}
78
79impl<N: NetworkZone> Drop for ClientDropGuard<N> {
80    fn drop(&mut self) {
81        self.bool.store(false, Ordering::Relaxed);
82    }
83}
cuprate_p2p/peer_set/client_wrappers.rs

cuprate_p2p/peer_set/
client_wrappers.rs
