refactor available node checks, todo add network modes

Author: erhant

Cargo.lock

@@ -7,7 +7,7 @@ default-members = ["compute"]
 
 [workspace.package]
 edition = "2021"
-version = "0.2.22"
+version = "0.2.23"
 license = "Apache-2.0"
 readme = "README.md"
 

Cargo.toml

@@ -1,6 +1,5 @@
-use dkn_p2p::{libp2p::gossipsub, DriaP2PClient};
+use dkn_p2p::{libp2p::gossipsub, DriaP2PClient, DriaP2PProtocol};
 use eyre::{eyre, Result};
-use std::time::Duration;
 use tokio_util::sync::CancellationToken;
 
 use crate::{
@@ -9,9 +8,6 @@ use crate::{
     utils::{crypto::secret_to_keypair, AvailableNodes, DKNMessage},
 };
 
-/// Number of seconds between refreshing the Admin RPC PeerIDs from Dria server.
-const RPC_PEER_ID_REFRESH_INTERVAL_SECS: u64 = 30;
-
 /// **Dria Compute Node**
 ///
 /// Internally, the node will create a new P2P client with the given secret key.
@@ -29,7 +25,6 @@ pub struct DriaComputeNode {
     pub config: DriaComputeNodeConfig,
     pub p2p: DriaP2PClient,
     pub available_nodes: AvailableNodes,
-    pub available_nodes_last_refreshed: tokio::time::Instant,
     pub cancellation: CancellationToken,
 }
 
@@ -41,31 +36,21 @@ impl DriaComputeNode {
         let keypair = secret_to_keypair(&config.secret_key);
 
         // get available nodes (bootstrap, relay, rpc) for p2p
-        let available_nodes = AvailableNodes::default()
-            .join(AvailableNodes::new_from_statics())
-            .join(AvailableNodes::new_from_env())
-            .join(
-                AvailableNodes::get_available_nodes()
-                    .await
-                    .unwrap_or_default(),
-            )
-            .sort_dedup();
+        let mut available_nodes =
+            AvailableNodes::new_from_statics().join(AvailableNodes::new_from_env());
+        available_nodes.refresh().await;
 
         // we are using the major.minor version as the P2P version
         // so that patch versions do not interfere with the protocol
-        const P2P_VERSION: &str = concat!(
-            env!("CARGO_PKG_VERSION_MAJOR"),
-            ".",
-            env!("CARGO_PKG_VERSION_MINOR")
-        );
+        let protocol = DriaP2PProtocol::new_major_minor("dria");
 
         // create p2p client
         let mut p2p = DriaP2PClient::new(
             keypair,
             config.p2p_listen_addr.clone(),
             &available_nodes.bootstrap_nodes,
             &available_nodes.relay_nodes,
-            P2P_VERSION,
+            protocol,
         )?;
 
         // dial rpc nodes
@@ -83,7 +68,6 @@ impl DriaComputeNode {
             config,
             cancellation,
             available_nodes,
-            available_nodes_last_refreshed: tokio::time::Instant::now(),
         })
     }
 
@@ -111,8 +95,10 @@ impl DriaComputeNode {
 
     /// Publishes a given message to the network w.r.t the topic of it.
     ///
-    /// Internally, the message is JSON serialized to bytes and then published to the network as is.
-    pub fn publish(&mut self, message: DKNMessage) -> Result<()> {
+    /// Internally, identity is attached to the the message which is then JSON serialized to bytes
+    /// and then published to the network as is.
+    pub fn publish(&mut self, mut message: DKNMessage) -> Result<()> {
+        message = message.with_identity(self.p2p.protocol().identity());
         let message_bytes = serde_json::to_vec(&message)?;
         let message_id = self.p2p.publish(&message.topic, message_bytes)?;
         log::info!("Published message ({}) to {}", message_id, message.topic);
@@ -145,11 +131,19 @@ impl DriaComputeNode {
             tokio::select! {
                 event = self.p2p.process_events() => {
                     // refresh admin rpc peer ids
-                    if self.available_nodes_last_refreshed.elapsed() > Duration::from_secs(RPC_PEER_ID_REFRESH_INTERVAL_SECS) {
+                    if self.available_nodes.can_refresh() {
                         log::info!("Refreshing available nodes.");
 
-                        self.available_nodes = AvailableNodes::get_available_nodes().await.unwrap_or_default().join(self.available_nodes.clone()).sort_dedup();
-                        self.available_nodes_last_refreshed = tokio::time::Instant::now();
+                        self.available_nodes.refresh().await;
+
+                        // dial all rpc nodes for better connectivity
+                        for rpc_addr in self.available_nodes.rpc_addrs.iter() {
+                            log::debug!("Dialling RPC node: {}", rpc_addr);
+                            // TODO: does this cause resource issues?
+                            if let Err(e) = self.p2p.dial(rpc_addr.clone()) {
+                                log::warn!("Error dialling RPC node: {:?}", e);
+                            };
+                        }
 
                         // also print network info
                         log::debug!("{:?}", self.p2p.network_info().connection_counters());
@@ -173,12 +167,7 @@ impl DriaComputeNode {
                             }
                         };
 
-                        // log::info!(
-                        //     "Received {} message ({})\nFrom:   {}\nSource: {}",
-                        //     topic_str,
-                        //     message_id,
-                        //     peer_id,
-                        // );
+                        // log the received message
                         log::info!(
                             "Received {} message ({}) from {}",
                             topic_str,

compute/src/node.rs

@@ -25,6 +25,9 @@ const STATIC_RPC_PEER_IDS: [&str; 0] = [];
 /// API URL for refreshing the Admin RPC PeerIDs from Dria server.
 const RPC_PEER_ID_REFRESH_API_URL: &str = "https://dkn.dria.co/available-nodes";
 
+/// Number of seconds between refreshing the Admin RPC PeerIDs from Dria server.
+const RPC_PEER_ID_REFRESH_INTERVAL_SECS: u64 = 30;
+
 /// Available nodes within the hybrid P2P network.
 ///
 /// - Bootstrap: used for Kademlia DHT bootstrap.
@@ -33,12 +36,14 @@ const RPC_PEER_ID_REFRESH_API_URL: &str = "https://dkn.dria.co/available-nodes";
 ///
 /// Note that while bootstrap & relay nodes are `Multiaddr`, RPC nodes are `PeerId` because we communicate
 /// with them via GossipSub only.
-#[derive(Debug, Default, Clone)]
+#[derive(Debug, Clone)]
 pub struct AvailableNodes {
     pub bootstrap_nodes: Vec<Multiaddr>,
     pub relay_nodes: Vec<Multiaddr>,
     pub rpc_nodes: Vec<PeerId>,
     pub rpc_addrs: Vec<Multiaddr>,
+    // refreshing
+    pub last_refreshed: tokio::time::Instant,
 }
 
 impl AvailableNodes {
@@ -69,6 +74,7 @@ impl AvailableNodes {
             relay_nodes: parse_vec(relay_nodes),
             rpc_nodes: vec![],
             rpc_addrs: vec![],
+            last_refreshed: tokio::time::Instant::now(),
         }
     }
 
@@ -79,6 +85,7 @@ impl AvailableNodes {
             relay_nodes: parse_vec(STATIC_RELAY_NODES.to_vec()),
             rpc_nodes: parse_vec(STATIC_RPC_PEER_IDS.to_vec()),
             rpc_addrs: vec![],
+            last_refreshed: tokio::time::Instant::now(),
         }
     }
 
@@ -108,6 +115,10 @@ impl AvailableNodes {
         self
     }
 
+    pub fn can_refresh(&self) -> bool {
+        self.last_refreshed.elapsed().as_secs() > RPC_PEER_ID_REFRESH_INTERVAL_SECS
+    }
+
     /// Refreshes the available nodes for Bootstrap, Relay and RPC nodes.
     pub async fn get_available_nodes() -> Result<Self> {
         #[derive(serde::Deserialize, Debug)]
@@ -127,8 +138,14 @@ impl AvailableNodes {
             relay_nodes: parse_vec(response_body.relays),
             rpc_nodes: parse_vec(response_body.rpcs),
             rpc_addrs: parse_vec(response_body.rpc_addrs),
+            last_refreshed: tokio::time::Instant::now(),
         })
     }
+
+    /// Refresh available nodes using the API.
+    pub async fn refresh(&mut self) {
+        todo!("TODO: refresh the available nodes")
+    }
 }
 
 /// Like `parse` of `str` but for vectors.

compute/src/utils/available_nodes.rs

@@ -49,9 +49,7 @@ impl DKNMessage {
             payload: BASE64_STANDARD.encode(data),
             topic: topic.to_string(),
             version: DRIA_COMPUTE_NODE_VERSION.to_string(),
-            identity: dkn_p2p::P2P_IDENTITY_PREFIX
-                .trim_end_matches('/')
-                .to_string(),
+            identity: String::default(),
             timestamp: get_current_time_nanos(),
         }
     }
@@ -70,6 +68,12 @@ impl DKNMessage {
         Self::new(signed_data, topic)
     }
 
+    /// Sets the identity of the message.
+    pub fn with_identity(mut self, identity: String) -> Self {
+        self.identity = identity;
+        self
+    }
+
     /// Decodes the base64 payload into bytes.
     #[inline(always)]
     pub fn decode_payload(&self) -> Result<Vec<u8>, base64::DecodeError> {
@@ -102,11 +106,11 @@ impl DKNMessage {
         let (signature_hex_bytes, body) =
             (&data[..SIGNATURE_SIZE_HEX - 2], &data[SIGNATURE_SIZE_HEX..]);
         let signature_bytes =
-            hex::decode(signature_hex_bytes).wrap_err("Could not decode signature hex")?;
+            hex::decode(signature_hex_bytes).wrap_err("could not decode signature hex")?;
 
         // now obtain the signature itself
         let signature = Signature::parse_standard_slice(&signature_bytes)
-            .wrap_err("Could not parse signature bytes")?;
+            .wrap_err("could not parse signature bytes")?;
 
         // verify signature w.r.t the body and the given public key
         let digest = Message::parse(&sha256hash(body));