feat(p2p/webrtc): improve disconnection logic and remove mutex for active connection limiting

binier · binier · commit a21977baf65a · 2024-12-12T16:01:14.000+04:00
diff --git a/node/src/event_source/event_source_effects.rs b/node/src/event_source/event_source_effects.rs
@@ -198,6 +198,7 @@ pub fn event_source_effects<S: Service>(store: &mut Store<S>, action: EventSourc
                     },
                     P2pConnectionEvent::Closed(peer_id) => {
                         store.dispatch(P2pDisconnectionAction::PeerClosed { peer_id });
+                        store.dispatch(P2pDisconnectionAction::Finish { peer_id });
                     }
                 },
                 P2pEvent::Channel(e) => match e {
diff --git a/p2p/src/disconnection/p2p_disconnection_actions.rs b/p2p/src/disconnection/p2p_disconnection_actions.rs
@@ -42,14 +42,22 @@ impl redux::EnablingCondition<P2pState> for P2pDisconnectionAction {
             P2pDisconnectionAction::RandomTry => time
                 .checked_sub(state.last_random_disconnection_try)
                 .map_or(false, |dur| dur >= RANDOM_DISCONNECTION_TRY_FREQUENCY),
-            P2pDisconnectionAction::Init { peer_id, .. }
-            | P2pDisconnectionAction::PeerClosed { peer_id, .. }
-            | P2pDisconnectionAction::Finish { peer_id } => {
+            P2pDisconnectionAction::Init { peer_id, .. } => {
+                state.peers.get(peer_id).map_or(false, |peer| {
+                    !peer.status.is_disconnected_or_disconnecting() && !peer.status.is_error()
+                })
+            }
+            P2pDisconnectionAction::Finish { peer_id } => {
                 state.peers.get(peer_id).map_or(false, |peer| {
                     !matches!(peer.status, P2pPeerStatus::Disconnected { .. })
                         && !peer.status.is_error()
                 })
             }
+            P2pDisconnectionAction::PeerClosed { peer_id, .. } => {
+                state.peers.get(peer_id).map_or(false, |peer| {
+                    !peer.status.is_disconnected_or_disconnecting() && !peer.status.is_error()
+                })
+            }
             P2pDisconnectionAction::FailedCleanup { peer_id } => state
                 .peers
                 .get(peer_id)
diff --git a/p2p/src/disconnection/p2p_disconnection_reducer.rs b/p2p/src/disconnection/p2p_disconnection_reducer.rs
@@ -52,8 +52,14 @@ impl P2pDisconnectedState {
                 Ok(())
             }
             P2pDisconnectionAction::Init { peer_id, reason } => {
+                let Some(peer) = p2p_state.peers.get_mut(&peer_id) else {
+                    bug_condition!("Invalid state for: `P2pDisconnectionAction::Init`");
+                    return Ok(());
+                };
+                peer.status = P2pPeerStatus::Disconnecting { time: meta.time() };
+
                 #[cfg(feature = "p2p-libp2p")]
-                if p2p_state.is_libp2p_peer(&peer_id) {
+                if peer.is_libp2p() {
                     let connections = p2p_state
                         .network
                         .scheduler
@@ -63,12 +69,6 @@ impl P2pDisconnectedState {
                         .map(|(addr, _)| *addr)
                         .collect::<Vec<_>>();
 
-                    let Some(peer) = p2p_state.peers.get_mut(&peer_id) else {
-                        bug_condition!("Invalid state for: `P2pDisconnectionAction::Finish`");
-                        return Ok(());
-                    };
-                    peer.status = P2pPeerStatus::Disconnecting { time: meta.time() };
-
                     let dispatcher = state_context.into_dispatcher();
                     for addr in connections {
                         dispatcher.push(P2pNetworkSchedulerAction::Disconnect {
@@ -91,6 +91,12 @@ impl P2pDisconnectedState {
                 Ok(())
             }
             P2pDisconnectionAction::FailedCleanup { peer_id } => {
+                let Some(peer) = p2p_state.peers.get_mut(&peer_id) else {
+                    bug_condition!("Invalid state for: `P2pDisconnectionAction::FailedCleanup`");
+                    return Ok(());
+                };
+                peer.status = P2pPeerStatus::Disconnecting { time: meta.time() };
+
                 let dispatcher = state_context.into_dispatcher();
                 dispatcher.push(P2pDisconnectionEffectfulAction::Init { peer_id });
                 Ok(())
diff --git a/p2p/src/disconnection_effectful/p2p_disconnection_effectful_effects.rs b/p2p/src/disconnection_effectful/p2p_disconnection_effectful_effects.rs
@@ -1,7 +1,5 @@
 use redux::ActionMeta;
 
-use crate::disconnection::P2pDisconnectionAction;
-
 use super::{P2pDisconnectionEffectfulAction, P2pDisconnectionService};
 
 impl P2pDisconnectionEffectfulAction {
@@ -13,7 +11,6 @@ impl P2pDisconnectionEffectfulAction {
         match self {
             P2pDisconnectionEffectfulAction::Init { peer_id } => {
                 store.service().disconnect(peer_id);
-                store.dispatch(P2pDisconnectionAction::Finish { peer_id });
             }
         }
     }
diff --git a/p2p/src/p2p_state.rs b/p2p/src/p2p_state.rs
@@ -1,6 +1,6 @@
 use openmina_core::{
     block::{ArcBlockWithHash, BlockWithHash},
-    impl_substate_access,
+    bug_condition, impl_substate_access,
     requests::RpcId,
     snark::{Snark, SnarkInfo, SnarkJobCommitment},
     transaction::{TransactionInfo, TransactionWithHash},
@@ -395,10 +395,18 @@ impl P2pPeerState {
                 P2pPeerStatus::Connecting(P2pConnectionState::Outgoing(
                     P2pConnectionOutgoingState::Error { time, .. },
                 )) => is_time_passed(now, *time, timeouts.outgoing_error_reconnect_timeout),
-                P2pPeerStatus::Disconnected { time } | P2pPeerStatus::Disconnecting { time } => {
+                P2pPeerStatus::Disconnected { time } => {
                     *time == Timestamp::ZERO
                         || is_time_passed(now, *time, timeouts.reconnect_timeout)
                 }
+                P2pPeerStatus::Disconnecting { time } => {
+                    if !is_time_passed(now, *time, timeouts.reconnect_timeout) {
+                        false
+                    } else {
+                        bug_condition!("peer stuck in `P2pPeerStatus::Disconnecting` state?");
+                        true
+                    }
+                }
                 _ => false,
             }
     }
diff --git a/p2p/src/service_impl/webrtc/api.rs b/p2p/src/service_impl/webrtc/api.rs
@@ -0,0 +1,27 @@
+use std::{collections::BTreeMap, sync::Arc};
+
+use openmina_core::channels::broadcast;
+use tokio::sync::Mutex;
+
+use crate::PeerId;
+
+use super::{ApiInner, build_api};
+
+#[derive(Clone)]
+pub struct Api {
+    inner: ApiInner,
+    active_peers: Arc<Mutex<BTreeMap<PeerId, broadcast::Receiver<()>>>>
+}
+
+impl Api {
+    pub fn new() -> Self {
+        Self {
+            inner: build_api(),
+            active_peers: Mutex::new(BTreeMap::new()).into()
+        }
+    }
+
+    pub(super) fn inner(&self) -> &ApiInner {
+        &self.inner
+    }
+}
diff --git a/p2p/src/service_impl/webrtc/mod.rs b/p2p/src/service_impl/webrtc/mod.rs
@@ -12,7 +12,7 @@ use std::{collections::BTreeMap, time::Duration};
 
 use openmina_core::bug_condition;
 use serde::Serialize;
-use tokio::sync::Mutex;
+use tokio::sync::Semaphore;
 
 #[cfg(not(target_arch = "wasm32"))]
 use tokio::task::spawn_local;
@@ -226,7 +226,12 @@ async fn wait_for_ice_gathering_complete(pc: &mut RTCConnection) {
     }
 }
 
-async fn peer_start(api: Api, args: PeerAddArgs, abort: broadcast::Receiver<()>) {
+async fn peer_start(
+    api: Api,
+    args: PeerAddArgs,
+    abort: broadcast::Receiver<()>,
+    closed: broadcast::Sender<()>,
+) {
     let PeerAddArgs {
         peer_id,
         kind,
@@ -337,7 +342,6 @@ async fn peer_start(api: Api, args: PeerAddArgs, abort: broadcast::Receiver<()>)
         connected_tx.send(Ok(())).unwrap();
     } else {
         let mut connected_tx = Some(connected_tx);
-        let event_sender = event_sender.clone();
         pc.on_connection_state_change(Box::new(move |state| {
             match state {
                 RTCConnectionState::Connected => {
@@ -349,7 +353,7 @@ async fn peer_start(api: Api, args: PeerAddArgs, abort: broadcast::Receiver<()>)
                     if let Some(connected_tx) = connected_tx.take() {
                         let _ = connected_tx.send(Err("disconnected"));
                     } else {
-                        let _ = event_sender(P2pConnectionEvent::Closed(peer_id).into());
+                        let _ = closed.send(());
                     }
                 }
                 _ => {}
@@ -365,6 +369,7 @@ async fn peer_start(api: Api, args: PeerAddArgs, abort: broadcast::Receiver<()>)
         Ok(_) => {}
         Err(err) => {
             let _ = event_sender(P2pConnectionEvent::Finalized(peer_id, Err(err)).into());
+            return;
         }
     }
 
@@ -727,50 +732,45 @@ pub trait P2pServiceWebrtc: redux::Service {
     fn peers(&mut self) -> &mut BTreeMap<PeerId, PeerState>;
 
     fn init<S: TaskSpawner>(secret_key: SecretKey, spawner: S) -> P2pServiceCtx {
+        const MAX_PEERS: usize = 500;
         let (cmd_sender, mut cmd_receiver) = mpsc::unbounded_channel();
 
         let _ = secret_key;
 
         spawner.spawn_main("webrtc", async move {
             #[allow(clippy::all)]
             let api = build_api();
-            let active_peers = Arc::new(Mutex::new(BTreeMap::new()));
+            let conn_permits = Arc::new(Semaphore::const_new(MAX_PEERS));
             while let Some(cmd) = cmd_receiver.recv().await {
                 match cmd {
-                    Cmd::PeerAdd { args, abort } => {
+                    Cmd::PeerAdd { args, mut abort } => {
                         #[allow(clippy::all)]
                         let api = api.clone();
-                        let active_peers = active_peers.clone();
+                        let conn_permits = conn_permits.clone();
+                        let peer_id = args.peer_id;
+                        let event_sender = args.event_sender.clone();
                         spawn_local(async move {
-                            let peer_id = args.peer_id;
-                            let (abort_finished_tx, abort_finished) = broadcast::channel::<()>(1);
-                            let mut aborted = abort.resubscribe();
-
-                            if let Some(mut abort_finished) =
-                                active_peers.lock().await.insert(peer_id, abort_finished)
-                            {
-                                // wait for peer conn with same id to finish.
-                                let _ = abort_finished.recv().await;
-                            }
-                            if active_peers.lock().await.len() > 100 {
-                                sleep(Duration::from_millis(500)).await;
-                            }
-
+                            let Ok(_permit) = conn_permits.try_acquire() else {
+                                // state machine shouldn't allow this to happen.
+                                bug_condition!("P2P WebRTC Semaphore acquisition failed!");
+                                return;
+                            };
+                            let (closed_tx, mut closed) = broadcast::channel(1);
+                            let event_sender_clone = event_sender.clone();
+                            spawn_local(async move {
+                                // to avoid sending closed multiple times
+                                let _ = closed.recv().await;
+                                event_sender_clone(P2pConnectionEvent::Closed(peer_id).into());
+                            });
                             tokio::select! {
-                                _ = aborted.recv() => {}
-                                _ = peer_start(api, args, abort) => {}
+                                _ = peer_start(api, args, abort.resubscribe(), closed_tx.clone()) => {}
+                                _ = abort.recv() => {
+                                }
                             }
 
-                            // delay removing active_peer to give some time for cleanup.
-                            sleep(Duration::from_millis(500)).await;
-                            let rx1 = abort_finished_tx.subscribe();
-                            let mut active_peers = active_peers.lock().await;
-                            if let Some(rx) = active_peers
-                                .remove(&peer_id)
-                                .filter(|rx2| !rx2.same_channel(&rx1))
-                            {
-                                active_peers.insert(peer_id, rx);
-                            }
+                            // delay dropping permit to give some time for cleanup.
+                            sleep(Duration::from_millis(100)).await;
+                            let _ = closed_tx.send(());
                         });
                     }
                 }

Original file line number	Diff line number	Diff line change
`@@ -198,6 +198,7 @@ pub fn event_source_effects<S: Service>(store: &mut Store<S>, action: EventSourc`
`198`	`198`	`},`
`199`	`199`	`P2pConnectionEvent::Closed(peer_id) => {`
`200`	`200`	`store.dispatch(P2pDisconnectionAction::PeerClosed { peer_id });`
	`201`	`+ store.dispatch(P2pDisconnectionAction::Finish { peer_id });`
`201`	`202`	`}`
`202`	`203`	`},`
`203`	`204`	`P2pEvent::Channel(e) => match e {`
Original file line number	Diff line number	Diff line change
`@@ -1,7 +1,5 @@`
`1`	`1`	`use redux::ActionMeta;`
`2`	`2`
`3`		`-use crate::disconnection::P2pDisconnectionAction;`
`4`		`-`
`5`	`3`	`use super::{P2pDisconnectionEffectfulAction, P2pDisconnectionService};`
`6`	`4`
`7`	`5`	`impl P2pDisconnectionEffectfulAction {`
`@@ -13,7 +11,6 @@ impl P2pDisconnectionEffectfulAction {`
`13`	`11`	`match self {`
`14`	`12`	`P2pDisconnectionEffectfulAction::Init { peer_id } => {`
`15`	`13`	`store.service().disconnect(peer_id);`
`16`		`- store.dispatch(P2pDisconnectionAction::Finish { peer_id });`
`17`	`14`	`}`
`18`	`15`	`}`
`19`	`16`	`}`