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

cuprate_p2p/peer_set/
client_wrappers.rs
