cuprate_p2p_core/
handles.rs

1//! Connection Handles.
2//!
3//! This module contains the [`ConnectionHandle`] which allows banning a peer, disconnecting a peer and
4//! checking if the peer is still connected.
5use std::{
6    sync::{Arc, OnceLock},
7    time::Duration,
8};
9
10use tokio::sync::OwnedSemaphorePermit;
11use tokio_util::sync::{CancellationToken, WaitForCancellationFutureOwned};
12
13/// A [`ConnectionHandle`] builder.
14#[derive(Default, Debug)]
15pub struct HandleBuilder {
16    permit: Option<OwnedSemaphorePermit>,
17}
18
19impl HandleBuilder {
20    /// Create a new builder.
21    pub const fn new() -> Self {
22        Self { permit: None }
23    }
24
25    /// Sets the permit for this connection.
26    #[must_use]
27    pub fn with_permit(mut self, permit: Option<OwnedSemaphorePermit>) -> Self {
28        self.permit = permit;
29        self
30    }
31
32    /// Builds the [`ConnectionGuard`] which should be handed to the connection task and the [`ConnectionHandle`].
33    ///
34    /// This will panic if a permit was not set [`HandleBuilder::with_permit`]
35    pub fn build(self) -> (ConnectionGuard, ConnectionHandle) {
36        let token = CancellationToken::new();
37
38        (
39            ConnectionGuard {
40                token: token.clone(),
41                on_close: None,
42                _permit: self.permit,
43            },
44            ConnectionHandle {
45                token,
46                ban: Arc::new(OnceLock::new()),
47            },
48        )
49    }
50}
51
52/// A struct representing the time a peer should be banned for.
53#[derive(Debug, Copy, Clone)]
54pub struct BanPeer(pub Duration);
55
56/// A struct given to the connection task.
57pub struct ConnectionGuard {
58    token: CancellationToken,
59    on_close: Option<Box<dyn FnOnce() + Send>>,
60    _permit: Option<OwnedSemaphorePermit>,
61}
62
63impl ConnectionGuard {
64    /// Checks if we should close the connection.
65    pub fn should_shutdown(&self) -> WaitForCancellationFutureOwned {
66        self.token.clone().cancelled_owned()
67    }
68    /// Tell the corresponding [`ConnectionHandle`]s that this connection is closed.
69    ///
70    /// This will be called on [`Drop::drop`].
71    pub fn connection_closed(&self) {
72        self.token.cancel();
73    }
74    /// Sets a callback to run when the connection closes.
75    #[must_use]
76    pub(crate) fn with_on_close(mut self, callback: impl FnOnce() + Send + 'static) -> Self {
77        self.on_close = Some(Box::new(callback));
78        self
79    }
80}
81
82impl Drop for ConnectionGuard {
83    fn drop(&mut self) {
84        self.token.cancel();
85
86        if let Some(callback) = self.on_close.take() {
87            callback();
88        }
89    }
90}
91
92/// A handle given to a task that needs to ban, disconnect, check if the peer should be banned or check
93/// the peer is still connected.
94#[derive(Debug, Clone)]
95pub struct ConnectionHandle {
96    token: CancellationToken,
97    ban: Arc<OnceLock<BanPeer>>,
98}
99
100impl ConnectionHandle {
101    pub fn closed(&self) -> WaitForCancellationFutureOwned {
102        self.token.clone().cancelled_owned()
103    }
104    /// Bans the peer for the given `duration`.
105    pub fn ban_peer(&self, duration: Duration) {
106        #[expect(
107            clippy::let_underscore_must_use,
108            reason = "error means peer is already banned; fine to ignore"
109        )]
110        let _ = self.ban.set(BanPeer(duration));
111        self.token.cancel();
112    }
113    /// Checks if this connection is closed.
114    pub fn is_closed(&self) -> bool {
115        self.token.is_cancelled()
116    }
117    /// Returns if this peer has been banned and the [`Duration`] of that ban.
118    pub fn check_should_ban(&mut self) -> Option<BanPeer> {
119        self.ban.get().copied()
120    }
121    /// Sends the signal to the connection task to disconnect.
122    pub fn send_close_signal(&self) {
123        self.token.cancel();
124    }
125}
cuprate_p2p_core/handles.rs

cuprate_p2p_core/
handles.rs
