1#[cfg(test)]
9mod tests;
10
11use std::collections::BTreeMap;
12
13use async_channel::{Receiver, Sender, bounded};
14use async_trait::async_trait;
15use fedimint_core::PeerId;
16use fedimint_core::net::peers::{IP2PConnections, Recipient};
17use fedimint_core::task::{TaskGroup, sleep};
18use fedimint_core::util::FmtCompactAnyhow;
19use fedimint_core::util::backoff_util::{FibonacciBackoff, api_networking_backoff};
20use fedimint_logging::{LOG_CONSENSUS, LOG_NET_PEER};
21use fedimint_server_core::dashboard_ui::P2PConnectionStatus;
22use futures::future::select_all;
23use futures::{FutureExt, StreamExt};
24use tokio::sync::watch;
25use tracing::{Instrument, debug, info, info_span, warn};
26
27use crate::metrics::{PEER_CONNECT_COUNT, PEER_DISCONNECT_COUNT, PEER_MESSAGES_COUNT};
28use crate::net::p2p_connection::{DynConnectionStatusUpdates, DynP2PConnection};
29use crate::net::p2p_connector::DynP2PConnector;
30
31#[derive(Clone, Debug, PartialEq, Eq)]
32pub struct P2PConnectionState {
33 pub connected: Option<P2PConnectionStatus>,
35 pub last_error: Option<String>,
37}
38
39pub type P2PStatusSenders = BTreeMap<PeerId, watch::Sender<P2PConnectionState>>;
40pub type P2PStatusReceivers = BTreeMap<PeerId, watch::Receiver<P2PConnectionState>>;
41
42pub fn p2p_status_channels(peers: Vec<PeerId>) -> (P2PStatusSenders, P2PStatusReceivers) {
43 let mut senders = BTreeMap::new();
44 let mut receivers = BTreeMap::new();
45
46 for peer in peers {
47 let (sender, receiver) = watch::channel(P2PConnectionState {
48 connected: None,
49 last_error: None,
50 });
51
52 senders.insert(peer, sender);
53 receivers.insert(peer, receiver);
54 }
55
56 (senders, receivers)
57}
58
59#[derive(Clone)]
60pub struct ReconnectP2PConnections<M> {
61 connections: BTreeMap<PeerId, P2PConnection<M>>,
62}
63
64impl<M: Send + 'static> ReconnectP2PConnections<M> {
65 pub fn new(
66 identity: PeerId,
67 connector: DynP2PConnector<M>,
68 task_group: &TaskGroup,
69 status_senders: P2PStatusSenders,
70 ) -> Self {
71 let mut connection_senders = BTreeMap::new();
72 let mut connections = BTreeMap::new();
73
74 for peer_id in connector.peers() {
75 assert_ne!(peer_id, identity);
76
77 let (connection_sender, connection_receiver) = bounded(4);
78
79 let connection = P2PConnection::new(
80 identity,
81 peer_id,
82 connector.clone(),
83 connection_receiver,
84 status_senders
85 .get(&peer_id)
86 .expect("No p2p status sender for peer")
87 .clone(),
88 task_group,
89 );
90
91 connection_senders.insert(peer_id, connection_sender);
92 connections.insert(peer_id, connection);
93 }
94
95 task_group.spawn_cancellable("handle-incoming-p2p-connections", async move {
96 info!(target: LOG_NET_PEER, "Starting listening task for p2p connections");
97
98 loop {
99 match connector.accept().await {
100 Ok((peer, connection)) => {
101 if connection_senders
102 .get_mut(&peer)
103 .expect("Authenticating connectors dont return unknown peers")
104 .send(connection)
105 .await
106 .is_err()
107 {
108 break;
109 }
110 },
111 Err(err) => {
112 warn!(target: LOG_NET_PEER, our_id = %identity, err = %err.fmt_compact_anyhow(), "Error while opening incoming connection");
113 }
114 }
115 }
116
117 info!(target: LOG_NET_PEER, "Shutting down listening task for p2p connections");
118 });
119
120 ReconnectP2PConnections { connections }
121 }
122}
123
124#[async_trait]
125impl<M: Clone + Send + 'static> IP2PConnections<M> for ReconnectP2PConnections<M> {
126 fn send(&self, recipient: Recipient, message: M) {
127 match recipient {
128 Recipient::Everyone => {
129 for connection in self.connections.values() {
130 connection.try_send(message.clone());
131 }
132 }
133 Recipient::Peer(peer) => match self.connections.get(&peer) {
134 Some(connection) => {
135 connection.try_send(message);
136 }
137 _ => {
138 warn!(target: LOG_NET_PEER, "No connection for peer {peer}");
139 }
140 },
141 }
142 }
143
144 async fn receive(&self) -> Option<(PeerId, M)> {
145 select_all(self.connections.iter().map(|(&peer, connection)| {
146 Box::pin(connection.receive().map(move |m| m.map(|m| (peer, m))))
147 }))
148 .await
149 .0
150 }
151
152 async fn receive_from_peer(&self, peer: PeerId) -> Option<M> {
153 self.connections
154 .get(&peer)
155 .expect("No connection found for peer")
156 .receive()
157 .await
158 }
159}
160
161#[derive(Clone)]
162struct P2PConnection<M> {
163 outgoing_sender: Sender<M>,
164 incoming_receiver: Receiver<M>,
165}
166
167impl<M: Send + 'static> P2PConnection<M> {
168 fn new(
169 our_id: PeerId,
170 peer_id: PeerId,
171 connector: DynP2PConnector<M>,
172 incoming_connections: Receiver<DynP2PConnection<M>>,
173 status_sender: watch::Sender<P2PConnectionState>,
174 task_group: &TaskGroup,
175 ) -> P2PConnection<M> {
176 let (outgoing_sender, outgoing_receiver) = bounded(5);
183 let (incoming_sender, incoming_receiver) = bounded(5);
184
185 task_group.spawn_cancellable(
186 format!("io-state-machine-{peer_id}"),
187 async move {
188 info!(target: LOG_NET_PEER, "Starting peer connection state machine");
189
190 let mut state_machine = P2PConnectionStateMachine {
191 common: P2PConnectionSMCommon {
192 incoming_sender,
193 outgoing_receiver,
194 our_id_str: our_id.to_string(),
195 our_id,
196 peer_id_str: peer_id.to_string(),
197 peer_id,
198 connector,
199 incoming_connections,
200 status_sender,
201 },
202 state: P2PConnectionSMState::Disconnected {
203 backoff: api_networking_backoff(),
204 last_error: None,
205 },
206 };
207
208 while let Some(sm) = state_machine.state_transition().await {
209 state_machine = sm;
210 }
211
212 info!(target: LOG_NET_PEER, "Shutting down peer connection state machine");
213 }
214 .instrument(info_span!("io-state-machine", ?peer_id)),
215 );
216
217 P2PConnection {
218 outgoing_sender,
219 incoming_receiver,
220 }
221 }
222
223 fn try_send(&self, message: M) {
224 if self.outgoing_sender.try_send(message).is_err() {
225 debug!(target: LOG_NET_PEER, "Outgoing message channel is full");
226 }
227 }
228
229 async fn receive(&self) -> Option<M> {
230 self.incoming_receiver.recv().await.ok()
231 }
232}
233
234struct P2PConnectionStateMachine<M> {
235 state: P2PConnectionSMState<M>,
236 common: P2PConnectionSMCommon<M>,
237}
238
239struct P2PConnectionSMCommon<M> {
240 incoming_sender: async_channel::Sender<M>,
241 outgoing_receiver: async_channel::Receiver<M>,
242 our_id: PeerId,
243 our_id_str: String,
244 peer_id: PeerId,
245 peer_id_str: String,
246 connector: DynP2PConnector<M>,
247 incoming_connections: Receiver<DynP2PConnection<M>>,
248 status_sender: watch::Sender<P2PConnectionState>,
249}
250
251enum P2PConnectionSMState<M> {
252 Disconnected {
253 backoff: FibonacciBackoff,
254 last_error: Option<String>,
255 },
256 Connected(DynP2PConnection<M>),
257}
258
259impl<M: Send + 'static> P2PConnectionStateMachine<M> {
260 async fn state_transition(mut self) -> Option<Self> {
261 match self.state {
262 P2PConnectionSMState::Disconnected {
263 backoff,
264 last_error,
265 } => {
266 self.common.status_sender.send_replace(P2PConnectionState {
267 connected: None,
268 last_error,
269 });
270
271 self.common.transition_disconnected(backoff).await
272 }
273 P2PConnectionSMState::Connected(connection) => {
274 let status_updates = connection.connection_status_updates();
277 let status = P2PConnectionStatus {
278 conn_type: connection
279 .connection_type()
280 .or_else(|| self.common.connector.connection_type(self.common.peer_id)),
281 rtt: connection.rtt(),
282 };
283
284 self.common.status_sender.send_replace(P2PConnectionState {
285 connected: Some(status),
286 last_error: None,
287 });
288
289 self.common
290 .transition_connected(connection, status_updates)
291 .await
292 }
293 }
294 .map(|state| P2PConnectionStateMachine {
295 common: self.common,
296 state,
297 })
298 }
299}
300
301impl<M: Send + 'static> P2PConnectionSMCommon<M> {
302 async fn transition_connected(
303 &mut self,
304 mut connection: DynP2PConnection<M>,
305 mut status_updates: Option<DynConnectionStatusUpdates>,
306 ) -> Option<P2PConnectionSMState<M>> {
307 tokio::select! {
308 Some(()) = async {
309 match status_updates.as_mut() {
310 Some(status_updates) => status_updates.next().await,
311 None => std::future::pending().await,
312 }
313 } => {
314 Some(P2PConnectionSMState::Connected(connection))
315 },
316 message = self.outgoing_receiver.recv() => {
317 Some(self.send_message(connection, message.ok()?).await)
318 },
319 connection = self.incoming_connections.recv() => {
320 info!(target: LOG_NET_PEER, "Connected to peer");
321
322 Some(P2PConnectionSMState::Connected(connection.ok()?))
323 },
324 message = connection.receive() => {
325 let mut message = match message {
326 Ok(message) => message,
327 Err(e) => return Some(self.disconnect(e)),
328 };
329
330 match message.read_to_end().await {
331 Ok(message) => {
332 PEER_MESSAGES_COUNT
333 .with_label_values(&[self.our_id_str.as_str(), self.peer_id_str.as_str(), "incoming"])
334 .inc();
335
336 if self.incoming_sender.try_send(message).is_err() {
337 debug!(target: LOG_NET_PEER, "Incoming message channel is full");
338 }
339 },
340 Err(e) => return Some(self.disconnect(e)),
341 }
342
343 Some(P2PConnectionSMState::Connected(connection))
344 },
345 }
346 }
347
348 fn disconnect(&self, error: anyhow::Error) -> P2PConnectionSMState<M> {
349 let last_error = error.fmt_compact_anyhow().to_string();
350
351 info!(
352 target: LOG_NET_PEER,
353 error = %last_error,
354 "Disconnected from peer"
355 );
356
357 PEER_DISCONNECT_COUNT
358 .with_label_values(&[&self.our_id_str, &self.peer_id_str])
359 .inc();
360
361 P2PConnectionSMState::Disconnected {
362 backoff: api_networking_backoff(),
363 last_error: Some(last_error),
364 }
365 }
366
367 async fn send_message(
368 &mut self,
369 mut connection: DynP2PConnection<M>,
370 peer_message: M,
371 ) -> P2PConnectionSMState<M> {
372 PEER_MESSAGES_COUNT
373 .with_label_values(&[
374 self.our_id_str.as_str(),
375 self.peer_id_str.as_str(),
376 "outgoing",
377 ])
378 .inc();
379
380 if let Err(e) = connection.send(peer_message).await {
381 return self.disconnect(e);
382 }
383
384 P2PConnectionSMState::Connected(connection)
385 }
386
387 async fn transition_disconnected(
388 &mut self,
389 mut backoff: FibonacciBackoff,
390 ) -> Option<P2PConnectionSMState<M>> {
391 tokio::select! {
392 connection = self.incoming_connections.recv() => {
393 PEER_CONNECT_COUNT
394 .with_label_values(&[self.our_id_str.as_str(), self.peer_id_str.as_str(), "incoming"])
395 .inc();
396
397 info!(target: LOG_NET_PEER, "Connected to peer");
398
399 Some(P2PConnectionSMState::Connected(connection.ok()?))
400 },
401 () = sleep(backoff.next().expect("Unlimited retries")), if self.our_id < self.peer_id => {
403 info!(target: LOG_NET_PEER, "Attempting to reconnect to peer");
404
405 match self.connector.connect(self.peer_id).await {
406 Ok(connection) => {
407 PEER_CONNECT_COUNT
408 .with_label_values(&[self.our_id_str.as_str(), self.peer_id_str.as_str(), "outgoing"])
409 .inc();
410
411 info!(target: LOG_NET_PEER, "Connected to peer");
412
413 Some(P2PConnectionSMState::Connected(connection))
414 }
415 Err(e) => {
416 let last_error = e.fmt_compact_anyhow().to_string();
417
418 warn!(
419 target: LOG_CONSENSUS,
420 error = %last_error,
421 "Failed to connect to peer"
422 );
423
424 Some(P2PConnectionSMState::Disconnected {
425 backoff,
426 last_error: Some(last_error),
427 })
428 }
429 }
430 },
431 }
432 }
433}