Ported remaining callbacks in p2p layer and snark layer

Author: 0xMimir

node/src/rpc/rpc_reducer.rs

@@ -1,11 +1,12 @@
 use openmina_core::{
     block::AppliedBlock,
     bug_condition,
-    requests::{RequestId, RpcIdType},
+    requests::{RequestId, RpcId, RpcIdType},
 };
 use p2p::{
     connection::{incoming::P2pConnectionIncomingAction, outgoing::P2pConnectionOutgoingAction},
     webrtc::P2pConnectionResponse,
+    PeerId,
 };
 use redux::ActionWithMeta;
 
@@ -79,11 +80,21 @@ impl RpcState {
                 state.requests.insert(*rpc_id, rpc_state);
 
                 let dispatcher = state_context.into_dispatcher();
+
+                let callback = redux::callback!(
+                    on_p2p_connection_outgoing_rpc_connection_success((peer_id: PeerId, rpc_id: Option<RpcId>)) -> crate::Action {
+                        let Some(rpc_id) = rpc_id else {
+                            panic!("RPC ID not provided");
+                        };
+
+                        RpcAction::P2pConnectionOutgoingPending{ rpc_id }
+                    }
+                );
                 dispatcher.push(P2pConnectionOutgoingAction::Init {
                     opts: opts.clone(),
                     rpc_id: Some(*rpc_id),
+                    on_success: Some(callback),
                 });
-                dispatcher.push(RpcAction::P2pConnectionOutgoingPending { rpc_id: *rpc_id });
             }
             RpcAction::P2pConnectionOutgoingPending { rpc_id } => {
                 let Some(rpc) = state.requests.get_mut(rpc_id) else {

node/src/snark/snark_effects.rs

@@ -1,6 +1,4 @@
-use snark::work_verify::SnarkWorkVerifyAction;
-
-use crate::{snark_pool::candidate::SnarkPoolCandidateAction, Service, Store};
+use crate::{Service, Store};
 
 use super::{SnarkAction, SnarkActionWithMeta};
 
@@ -12,35 +10,7 @@ pub fn snark_effects<S: Service>(store: &mut Store<S>, action: SnarkActionWithMe
         SnarkAction::BlockVerifyEffect(a) => {
             a.effects(&meta, store);
         }
-        SnarkAction::WorkVerify(a) => match a {
-            // TODO(tizoc): handle this logic with the on_error callback passed on the Init action
-            SnarkWorkVerifyAction::Error { req_id, .. } => {
-                let req = store.state().snark.work_verify.jobs.get(req_id);
-                let Some(req) = req else { return };
-                let sender = req.sender().parse().unwrap();
-
-                store.dispatch(SnarkPoolCandidateAction::WorkVerifyError {
-                    peer_id: sender,
-                    verify_id: req_id,
-                });
-            }
-            // TODO(tizoc): handle this logic with the on_success callback passed on the Init action
-            SnarkWorkVerifyAction::Success { req_id } => {
-                let req = store.state().snark.work_verify.jobs.get(req_id);
-                let Some(req) = req else { return };
-                let sender = req.sender().parse().unwrap();
-                let batch = req.batch().to_vec();
-
-                store.dispatch(SnarkPoolCandidateAction::WorkVerifySuccess {
-                    peer_id: sender,
-                    verify_id: req_id,
-                    batch,
-                });
-            }
-            SnarkWorkVerifyAction::Init { .. } => {}
-            SnarkWorkVerifyAction::Pending { .. } => {}
-            SnarkWorkVerifyAction::Finish { .. } => {}
-        },
+        SnarkAction::WorkVerify(_) => {}
         SnarkAction::WorkVerifyEffect(a) => {
             a.effects(&meta, store);
         }

p2p/src/channels/signaling/discovery/p2p_channels_signaling_discovery_reducer.rs

@@ -269,6 +269,7 @@ impl P2pChannelsSignalingDiscoveryState {
                         },
                     },
                     rpc_id: None,
+                    on_success: None,
                 };
                 let accepted = redux::EnablingCondition::is_enabled(&action, state, meta.time());
                 if accepted {

p2p/src/connection/incoming/p2p_connection_incoming_reducer.rs

@@ -60,6 +60,7 @@ impl P2pConnectionIncomingState {
                         }),
                         status: P2pPeerStatus::Connecting(P2pConnectionState::incoming_init(&opts)),
                         identify: None,
+                        on_connect_success: None,
                     });
 
                 state.status =
@@ -427,6 +428,7 @@ impl P2pConnectionIncomingState {
                             )),
                             status: P2pPeerStatus::Disconnected { time: meta.time() },
                             identify: None,
+                            on_connect_success: None,
                         });
 
                     Self::reduce_finalize_libp2p_pending(state, addr, time, my_id, peer_id);

p2p/src/connection/outgoing/p2p_connection_outgoing_actions.rs

@@ -1,4 +1,5 @@
 use openmina_core::ActionEvent;
+use redux::Callback;
 use serde::{Deserialize, Serialize};
 
 use openmina_core::requests::RpcId;
@@ -18,6 +19,7 @@ pub enum P2pConnectionOutgoingAction {
     Init {
         opts: P2pConnectionOutgoingInitOpts,
         rpc_id: Option<RpcId>,
+        on_success: Option<Callback<(PeerId, Option<RpcId>)>>,
     },
     /// Reconnect to an existing peer.
     // TODO: rename `Init` and `Reconnect` to `New` and `Connect` or something

p2p/src/connection/outgoing/p2p_connection_outgoing_reducer.rs

@@ -40,7 +40,11 @@ impl P2pConnectionOutgoingState {
                 dispatcher.push(P2pConnectionOutgoingEffectfulAction::RandomInit);
                 Ok(())
             }
-            P2pConnectionOutgoingAction::Init { opts, rpc_id } => {
+            P2pConnectionOutgoingAction::Init {
+                opts,
+                rpc_id,
+                on_success,
+            } => {
                 let peer_state =
                     p2p_state
                         .peers
@@ -52,8 +56,13 @@ impl P2pConnectionOutgoingState {
                                 &opts,
                             )),
                             identify: None,
+                            on_connect_success: None,
                         });
 
+                if let Some(on_connect_success) = on_success {
+                    peer_state.on_connect_success = Some(on_connect_success);
+                }
+
                 peer_state.status =
                     P2pPeerStatus::Connecting(P2pConnectionState::Outgoing(Self::Init {
                         time,
@@ -509,19 +518,30 @@ impl P2pConnectionOutgoingState {
                     return Ok(());
                 }
 
+                let Some(peer_state) = p2p_state.peers.get_mut(&peer_id) else {
+                    bug_condition!("Outgoing peer state not found for: {}", peer_id);
+                    return Ok(());
+                };
+                let callback = peer_state.on_connect_success.take();
                 let (dispatcher, state) = state_context.into_dispatcher_and_state();
                 let p2p_state: &P2pState = state.substate()?;
+
                 dispatcher.push(P2pPeerAction::Ready {
                     peer_id,
                     incoming: false,
                 });
 
-                if let Some(rpc_id) = p2p_state.peer_connection_rpc_id(&peer_id) {
+                let rpc_id = p2p_state.peer_connection_rpc_id(&peer_id);
+                if let Some(rpc_id) = rpc_id {
                     if let Some(callback) = &p2p_state.callbacks.on_p2p_connection_outgoing_success
                     {
                         dispatcher.push_callback(callback.clone(), rpc_id);
                     }
                 }
+
+                if let Some(callback) = callback {
+                    dispatcher.push_callback(callback, (peer_id, rpc_id));
+                }
                 Ok(())
             }
         }

p2p/src/network/kad/request/p2p_network_kad_request_actions.rs

@@ -1,10 +1,20 @@
 use std::net::SocketAddr;
 
 use openmina_core::ActionEvent;
-use redux::EnablingCondition;
+use redux::{Callback, EnablingCondition};
 use serde::{Deserialize, Serialize};
 
-use crate::{ConnectionAddr, P2pAction, P2pNetworkKadEntry, P2pState, PeerId, StreamId};
+use crate::{
+    ConnectionAddr, P2pAction, P2pNetworkKadEntry, P2pNetworkKademliaRpcRequest, P2pState, PeerId,
+    StreamId,
+};
+
+type StreamReadyCallback = Callback<(
+    ConnectionAddr,
+    PeerId,
+    StreamId,
+    P2pNetworkKademliaRpcRequest,
+)>;
 
 #[derive(Clone, Debug, Serialize, Deserialize, ActionEvent)]
 #[action_event(fields(display(peer_id), display(addr), display(key), stream_id, error))]
@@ -29,6 +39,7 @@ pub enum P2pNetworkKadRequestAction {
         peer_id: PeerId,
         stream_id: StreamId,
         addr: ConnectionAddr,
+        callback: StreamReadyCallback,
     },
     RequestSent {
         peer_id: PeerId,

p2p/src/network/kad/request/p2p_network_kad_request_reducer.rs

@@ -1,11 +1,11 @@
-use openmina_core::{bug_condition, Substate, SubstateAccess};
+use openmina_core::{bug_condition, requests::RpcId, Substate, SubstateAccess};
 use redux::{ActionWithMeta, Dispatcher};
 
 use crate::{
     connection::outgoing::P2pConnectionOutgoingAction, ConnectionAddr,
     P2pNetworkConnectionMuxState, P2pNetworkKadBootstrapAction, P2pNetworkKadEffectfulAction,
     P2pNetworkKadState, P2pNetworkKademliaRpcRequest, P2pNetworkKademliaStreamAction,
-    P2pNetworkYamuxAction, P2pPeerState, P2pState,
+    P2pNetworkYamuxAction, P2pPeerState, P2pState, PeerId,
 };
 
 use super::{P2pNetworkKadRequestAction, P2pNetworkKadRequestState, P2pNetworkKadRequestStatus};
@@ -47,14 +47,19 @@ impl P2pNetworkKadRequestState {
                 let peer_state = p2p_state.peers.get(&peer_id);
 
                 let on_initialize_connection = |dispatcher: &mut Dispatcher<Action, State>| {
-                    // initialize connection to the peer.
-                    // when connection is establised and yamux layer is ready, we will continue with TODO
-                    // TODO: add callbacks
                     let opts = crate::connection::outgoing::P2pConnectionOutgoingInitOpts::LibP2P(
                         (peer_id, addr).into(),
                     );
-                    dispatcher.push(P2pConnectionOutgoingAction::Init { opts, rpc_id: None });
-                    dispatcher.push(P2pNetworkKadRequestAction::PeerIsConnecting { peer_id });
+                    let callback = redux::callback!(
+                        on_p2p_connection_outgoing_kad_connection_success((peer_id: PeerId, _rpc_id: Option<RpcId>)) -> crate::P2pAction {
+                            P2pNetworkKadRequestAction::PeerIsConnecting { peer_id }
+                        }
+                    );
+                    dispatcher.push(P2pConnectionOutgoingAction::Init {
+                        opts,
+                        rpc_id: None,
+                        on_success: Some(callback),
+                    });
                     Ok(())
                 };
 
@@ -172,6 +177,7 @@ impl P2pNetworkKadRequestState {
                 peer_id,
                 stream_id,
                 addr,
+                callback,
             } => {
                 let find_node = match P2pNetworkKademliaRpcRequest::find_node(request_state.key) {
                     Ok(find_node) => find_node,
@@ -193,20 +199,8 @@ impl P2pNetworkKadRequestState {
                     super::P2pNetworkKadRequestStatus::Request,
                 );
 
-                let key = request_state.key;
-
                 let dispatcher = state_context.into_dispatcher();
-                let data =
-                    P2pNetworkKademliaRpcRequest::find_node(key).map_err(|e| e.to_string())?;
-
-                // TODO: move action bellow to callback
-                dispatcher.push(P2pNetworkKademliaStreamAction::SendRequest {
-                    addr,
-                    peer_id,
-                    stream_id,
-                    data,
-                });
-                dispatcher.push(P2pNetworkKadRequestAction::RequestSent { peer_id });
+                dispatcher.push_callback(callback, (addr, peer_id, stream_id, find_node));
                 Ok(())
             }
             P2pNetworkKadRequestAction::RequestSent { .. } => {

p2p/src/network/kad/stream/mod.rs

@@ -1,8 +1,64 @@
 mod p2p_network_kad_stream_state;
+use redux::Callback;
+use serde::{Deserialize, Serialize};
+
+use crate::ConnectionAddr;
+use crate::P2pNetworkKademliaAction;
+use crate::PeerId;
+use crate::StreamId;
+
 pub use self::p2p_network_kad_stream_state::*;
 
 mod p2p_network_kad_stream_actions;
 pub use self::p2p_network_kad_stream_actions::*;
 
+use super::P2pNetworkKadEntry;
+use super::CID;
+
 #[cfg(feature = "p2p-libp2p")]
 mod p2p_network_kad_stream_reducer;
+
+#[derive(Deserialize, Serialize, Debug, Clone)]
+pub enum P2pNetworkKademliaStreamWaitOutgoingCallback {
+    AnswerFindNodeRequest {
+        callback: Callback<(ConnectionAddr, PeerId, StreamId, CID)>,
+        args: CID,
+    },
+    UpdateFindNodeRequest {
+        callback: Callback<(ConnectionAddr, PeerId, StreamId, Vec<P2pNetworkKadEntry>)>,
+        args: Vec<P2pNetworkKadEntry>,
+    },
+}
+
+impl P2pNetworkKademliaStreamWaitOutgoingCallback {
+    pub fn answer_find_node_request(cid: CID) -> Self {
+        Self::AnswerFindNodeRequest {
+            callback: redux::callback!(
+                on_p2p_network_stream_wait_outgoing_answer_find_node_request((
+                    addr: ConnectionAddr,
+                    peer_id: PeerId,
+                    stream_id: StreamId,
+                    cid: CID
+                )) -> crate::P2pAction{
+                    P2pNetworkKademliaAction::AnswerFindNodeRequest { addr, peer_id, stream_id, key: cid }
+                }
+            ),
+            args: cid,
+        }
+    }
+    pub fn update_find_node_request(peers: Vec<P2pNetworkKadEntry>) -> Self {
+        Self::UpdateFindNodeRequest {
+            callback: redux::callback!(
+                on_p2p_network_stream_wait_outgoing_answer_find_node_request((
+                    addr: ConnectionAddr,
+                    peer_id: PeerId,
+                    stream_id: StreamId,
+                    closest_peers: Vec<P2pNetworkKadEntry>
+                )) -> crate::P2pAction{
+                    P2pNetworkKademliaAction::UpdateFindNodeRequest { addr, peer_id, stream_id, closest_peers }
+                }
+            ),
+            args: peers,
+        }
+    }
+}

p2p/src/network/kad/stream/p2p_network_kad_stream_actions.rs

@@ -7,6 +7,8 @@ use crate::{
     P2pState, PeerId, StreamId,
 };
 
+use super::P2pNetworkKademliaStreamWaitOutgoingCallback;
+
 /// Kademlia stream related actions.
 #[derive(Debug, Clone, Serialize, Deserialize, ActionEvent)]
 #[action_event(fields(display(addr), display(peer_id), stream_id, incoming, debug(data)))]
@@ -44,6 +46,7 @@ pub enum P2pNetworkKademliaStreamAction {
         addr: ConnectionAddr,
         peer_id: PeerId,
         stream_id: StreamId,
+        callback: P2pNetworkKademliaStreamWaitOutgoingCallback,
     },
 
     /// Sends request to the stream.