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