volans_ping/
inbound.rs

1use std::{
2    collections::VecDeque,
3    convert::Infallible,
4    io,
5    task::{Context, Poll, Waker},
6    time::{Duration, Instant},
7};
8
9use futures::{FutureExt, future::BoxFuture};
10use futures_timer::Delay;
11use volans_core::{PeerId, Multiaddr, upgrade::ReadyUpgrade};
12use volans_swarm::{
13    BehaviorEvent, ConnectionDenied, ConnectionHandler, ConnectionHandlerEvent, ConnectionId,
14    InboundStreamHandler, InboundUpgradeSend, NetworkBehavior, NetworkIncomingBehavior,
15    StreamProtocol, Substream, SubstreamProtocol, THandlerAction, THandlerEvent,
16};
17
18use crate::{Config, Event, Failure, protocol};
19
20type PongFuture = BoxFuture<'static, Result<Substream, io::Error>>;
21
22pub struct Handler {
23    interval: Delay,
24    config: Config,
25    last_ping: Instant,
26    failed: bool,
27    inbound: Option<PongFuture>,
28    pending_errors: VecDeque<Failure>,
29}
30
31impl Handler {
32    pub fn new(config: Config) -> Self {
33        Self {
34            interval: Delay::new(config.interval * config.failures),
35            config,
36            last_ping: Instant::now(),
37            failed: false,
38            inbound: None,
39            pending_errors: VecDeque::new(),
40        }
41    }
42}
43
44impl ConnectionHandler for Handler {
45    type Action = Infallible;
46
47    type Event = Result<Duration, Failure>;
48
49    fn handle_action(&mut self, _action: Self::Action) {
50        unreachable!("Ping handler does not support actions");
51    }
52
53    fn poll_close(&mut self, _: &mut Context<'_>) -> Poll<Option<Self::Event>> {
54        if let Some(error) = self.pending_errors.pop_back() {
55            return Poll::Ready(Some(Err(error)));
56        }
57        Poll::Ready(None)
58    }
59
60    fn poll(&mut self, cx: &mut Context<'_>) -> Poll<ConnectionHandlerEvent<Self::Event>> {
61        loop {
62            if let Some(error) = self.pending_errors.pop_back() {
63                return Poll::Ready(ConnectionHandlerEvent::Notify(Err(error)));
64            }
65
66            if self.failed {
67                return Poll::Ready(ConnectionHandlerEvent::CloseConnection);
68            }
69
70            if let Some(fut) = self.inbound.as_mut() {
71                match fut.poll_unpin(cx) {
72                    Poll::Pending => {}
73                    Poll::Ready(Ok(substream)) => {
74                        //重新开始新的延迟
75                        self.inbound = Some(protocol::recv_ping(substream).boxed());
76                        // 重置为新的周期间隔
77                        self.interval
78                            .reset(self.config.interval * self.config.failures);
79
80                        let elapsed = self.last_ping.elapsed();
81                        self.last_ping = Instant::now();
82
83                        return Poll::Ready(ConnectionHandlerEvent::Notify(Ok(elapsed)));
84                    }
85                    Poll::Ready(Err(err)) => {
86                        self.inbound = None;
87                        self.failed = true;
88                        self.pending_errors.push_back(Failure::other(err));
89                        continue;
90                    }
91                }
92            }
93
94            match self.interval.poll_unpin(cx) {
95                Poll::Pending => {}
96                Poll::Ready(()) => {
97                    // 重置为新的周期间隔
98                    tracing::debug!("Ping timeout, sending ping");
99                    self.interval
100                        .reset(self.config.interval * self.config.failures);
101                    self.inbound = None;
102                    self.failed = true;
103                    self.pending_errors.push_back(Failure::Timeout);
104                    continue;
105                }
106            }
107            return Poll::Pending;
108        }
109    }
110}
111
112impl InboundStreamHandler for Handler {
113    type InboundUpgrade = ReadyUpgrade<StreamProtocol>;
114
115    type InboundUserData = ();
116
117    fn listen_protocol(&self) -> SubstreamProtocol<Self::InboundUpgrade, Self::InboundUserData> {
118        SubstreamProtocol::new(ReadyUpgrade::new(protocol::PROTOCOL_NAME), ())
119    }
120
121    fn on_fully_negotiated(
122        &mut self,
123        _user_data: Self::InboundUserData,
124        protocol: <Self::InboundUpgrade as InboundUpgradeSend>::Output,
125    ) {
126        self.inbound = Some(protocol::recv_ping(protocol).boxed());
127        self.last_ping = Instant::now();
128    }
129
130    fn on_upgrade_error(
131        &mut self,
132        _user_data: Self::InboundUserData,
133        error: <Self::InboundUpgrade as InboundUpgradeSend>::Error,
134    ) {
135        tracing::debug!("Ping protocol upgrade error: {}", error);
136        self.inbound = None;
137        self.interval.reset(Duration::new(0, 0));
138    }
139}
140
141pub struct Behavior {
142    config: Config,
143    events: VecDeque<Event>,
144    none_event_waker: Option<Waker>,
145}
146
147impl Behavior {
148    pub fn new(config: Config) -> Self {
149        Self {
150            config,
151            events: VecDeque::new(),
152            none_event_waker: None,
153        }
154    }
155}
156
157impl Default for Behavior {
158    fn default() -> Self {
159        Self::new(Config::default())
160    }
161}
162
163impl NetworkBehavior for Behavior {
164    type ConnectionHandler = Handler;
165    type Event = Event;
166
167    fn on_connection_handler_event(
168        &mut self,
169        id: ConnectionId,
170        peer_id: PeerId,
171        event: THandlerEvent<Self>,
172    ) {
173        self.events.push_front(Event {
174            peer_id,
175            connection: id,
176            result: event,
177        });
178        if let Some(waker) = self.none_event_waker.take() {
179            waker.wake();
180        }
181    }
182
183    fn poll(
184        &mut self,
185        _cx: &mut Context<'_>,
186    ) -> Poll<BehaviorEvent<Self::Event, THandlerAction<Self>>> {
187        if let Some(event) = self.events.pop_back() {
188            return Poll::Ready(BehaviorEvent::Behavior(event));
189        }
190        self.none_event_waker = Some(_cx.waker().clone());
191        Poll::Pending
192    }
193}
194
195impl NetworkIncomingBehavior for Behavior {
196    /// 处理已建立的连接
197    fn handle_established_connection(
198        &mut self,
199        _id: ConnectionId,
200        peer_id: PeerId,
201        _local_addr: &Multiaddr,
202        _remote_addr: &Multiaddr,
203    ) -> Result<Self::ConnectionHandler, ConnectionDenied> {
204        tracing::trace!("Ping handler established for peer: {}", peer_id);
205        Ok(Handler::new(self.config.clone()))
206    }
207}
volans_ping/inbound.rs

volans_ping/
inbound.rs
