utils as package, code tidyups and refactors

Author: erhant

Cargo.lock

@@ -1,6 +1,6 @@
 [workspace]
 resolver = "2"
-members = ["compute", "p2p", "workflows"]
+members = ["compute", "p2p", "workflows", "utils"]
 # compute node is the default member, until Oracle comes in
 # then, a Launcher will be the default member
 default-members = ["compute"]

Cargo.toml

@@ -44,6 +44,7 @@ fastbloom-rs = "0.5.9"
 
 # dria subcrates
 dkn-p2p = { path = "../p2p" }
+dkn-utils = { path = "../utils" }
 dkn-workflows = { path = "../workflows" }
 
 # vendor OpenSSL so that its easier to build cross-platform packages

compute/Cargo.toml

@@ -2,40 +2,13 @@ use crate::utils::{
     address_in_use,
     crypto::{secret_to_keypair, to_address},
 };
-use dkn_p2p::libp2p::Multiaddr;
+use dkn_p2p::{libp2p::Multiaddr, DriaNetworkType};
 use dkn_workflows::DriaWorkflowsConfig;
 use eyre::{eyre, Result};
 use libsecp256k1::{PublicKey, SecretKey};
 
 use std::{env, str::FromStr};
 
-/// Network type.
-#[derive(Default, Debug, Clone, Copy)]
-pub enum DriaNetworkType {
-    #[default]
-    Community,
-    Pro,
-}
-
-impl From<&str> for DriaNetworkType {
-    fn from(s: &str) -> Self {
-        match s {
-            "community" => DriaNetworkType::Community,
-            "pro" => DriaNetworkType::Pro,
-            _ => Default::default(),
-        }
-    }
-}
-
-impl DriaNetworkType {
-    pub fn protocol_name(&self) -> &str {
-        match self {
-            DriaNetworkType::Community => "dria",
-            DriaNetworkType::Pro => "dria-sdk",
-        }
-    }
-}
-
 #[derive(Debug, Clone)]
 pub struct DriaComputeNodeConfig {
     /// Wallet secret/private key.

compute/src/config.rs

@@ -1,8 +1,6 @@
-use crate::{
-    utils::{get_current_time_nanos, DKNMessage},
-    DriaComputeNode,
-};
+use crate::{utils::DriaMessage, DriaComputeNode};
 use dkn_p2p::libp2p::gossipsub::MessageAcceptance;
+use dkn_utils::get_current_time_nanos;
 use dkn_workflows::{Model, ModelProvider};
 use eyre::{Context, Result};
 use serde::{Deserialize, Serialize};
@@ -42,7 +40,7 @@ impl PingpongHandler {
     /// 7. Returns `MessageAcceptance::Accept` so that ping is propagated to others as well.
     pub(crate) async fn handle_ping(
         node: &mut DriaComputeNode,
-        ping_message: &DKNMessage,
+        ping_message: &DriaMessage,
     ) -> Result<MessageAcceptance> {
         let pingpong = ping_message
             .parse_payload::<PingpongPayload>(true)
@@ -70,7 +68,7 @@ impl PingpongHandler {
         };
 
         // publish message
-        let message = DKNMessage::new_signed(
+        let message = DriaMessage::new_signed(
             serde_json::json!(response_body).to_string(),
             Self::RESPONSE_TOPIC,
             &node.config.secret_key,

compute/src/handlers/pingpong.rs

@@ -1,12 +1,13 @@
 use dkn_p2p::libp2p::gossipsub::MessageAcceptance;
+use dkn_utils::get_current_time_nanos;
 use dkn_workflows::{Entry, Executor, ModelProvider, Workflow};
 use eyre::{Context, Result};
 use libsecp256k1::PublicKey;
 use serde::Deserialize;
 use tokio_util::either::Either;
 
 use crate::payloads::{TaskErrorPayload, TaskRequestPayload, TaskResponsePayload, TaskStats};
-use crate::utils::{get_current_time_nanos, DKNMessage};
+use crate::utils::DriaMessage;
 use crate::workers::workflow::*;
 use crate::DriaComputeNode;
 
@@ -31,7 +32,7 @@ impl WorkflowHandler {
 
     pub(crate) async fn handle_compute(
         node: &mut DriaComputeNode,
-        compute_message: &DKNMessage,
+        compute_message: &DriaMessage,
     ) -> Result<Either<MessageAcceptance, WorkflowsWorkerInput>> {
         let stats = TaskStats::new().record_received_at();
         let task = compute_message
@@ -133,7 +134,7 @@ impl WorkflowHandler {
                     task.task_id,
                     payload_str
                 );
-                DKNMessage::new(payload_str, Self::RESPONSE_TOPIC)
+                DriaMessage::new(payload_str, Self::RESPONSE_TOPIC)
             }
             Err(err) => {
                 // use pretty display string for error logging with causes
@@ -150,7 +151,7 @@ impl WorkflowHandler {
                 let error_payload_str = serde_json::json!(error_payload).to_string();
 
                 // prepare signed message
-                DKNMessage::new_signed(
+                DriaMessage::new_signed(
                     error_payload_str,
                     Self::RESPONSE_TOPIC,
                     &node.config.secret_key,
@@ -167,7 +168,7 @@ impl WorkflowHandler {
                 "taskId": task.task_id,
                 "error": err_msg,
             });
-            let message = DKNMessage::new_signed(
+            let message = DriaMessage::new_signed(
                 payload.to_string(),
                 Self::RESPONSE_TOPIC,
                 &node.config.secret_key,

compute/src/handlers/workflow.rs

@@ -9,5 +9,5 @@ pub(crate) mod workers;
 /// This value is attached within the published messages.
 pub const DRIA_COMPUTE_NODE_VERSION: &str = env!("CARGO_PKG_VERSION");
 
-pub use config::{DriaComputeNodeConfig, DriaNetworkType};
+pub use config::DriaComputeNodeConfig;
 pub use node::DriaComputeNode;

compute/src/lib.rs

@@ -3,7 +3,7 @@ use dkn_p2p::{
         gossipsub::{Message, MessageAcceptance, MessageId},
         PeerId,
     },
-    DriaP2PClient, DriaP2PCommander, DriaP2PProtocol,
+    DriaNodes, DriaP2PClient, DriaP2PCommander, DriaP2PProtocol,
 };
 use eyre::Result;
 use std::collections::HashSet;
@@ -13,7 +13,7 @@ use tokio_util::{either::Either, sync::CancellationToken};
 use crate::{
     config::*,
     handlers::*,
-    utils::{crypto::secret_to_keypair, AvailableNodes, DKNMessage},
+    utils::{crypto::secret_to_keypair, refresh_dria_nodes, DriaMessage},
     workers::workflow::{WorkflowsWorker, WorkflowsWorkerInput, WorkflowsWorkerOutput},
 };
 
@@ -26,7 +26,8 @@ const PUBLISH_CHANNEL_BUFSIZE: usize = 1024;
 
 pub struct DriaComputeNode {
     pub config: DriaComputeNodeConfig,
-    pub available_nodes: AvailableNodes,
+    /// Pre-defined nodes that belong to Dria, e.g. bootstraps, relays and RPCs.
+    pub dria_nodes: DriaNodes,
     /// Peer-to-peer client commander to interact with the network.
     pub p2p: DriaP2PCommander,
     /// Gossipsub message receiver, used by peer-to-peer client in a separate thread.
@@ -63,10 +64,10 @@ impl DriaComputeNode {
         let keypair = secret_to_keypair(&config.secret_key);
 
         // get available nodes (bootstrap, relay, rpc) for p2p
-        let mut available_nodes = AvailableNodes::new(config.network_type);
-        available_nodes.populate_with_statics();
-        available_nodes.populate_with_env();
-        if let Err(e) = available_nodes.populate_with_api().await {
+        let mut available_nodes = DriaNodes::new(config.network_type)
+            .with_statics()
+            .with_envs();
+        if let Err(e) = refresh_dria_nodes(&mut available_nodes).await {
             log::error!("Error populating available nodes: {:?}", e);
         };
 
@@ -81,7 +82,7 @@ impl DriaComputeNode {
             config.p2p_listen_addr.clone(),
             available_nodes.bootstrap_nodes.clone().into_iter(),
             available_nodes.relay_nodes.clone().into_iter(),
-            available_nodes.rpc_addrs.clone().into_iter(),
+            available_nodes.rpc_nodes.clone().into_iter(),
             protocol,
         )?;
 
@@ -110,7 +111,7 @@ impl DriaComputeNode {
             DriaComputeNode {
                 config,
                 p2p: p2p_commander,
-                available_nodes,
+                dria_nodes: available_nodes,
                 message_rx,
                 publish_rx,
                 workflow_batch_tx,
@@ -160,7 +161,7 @@ impl DriaComputeNode {
     ///
     /// Internally, identity is attached to the the message which is then JSON serialized to bytes
     /// and then published to the network as is.
-    pub async fn publish(&mut self, mut message: DKNMessage) -> Result<()> {
+    pub async fn publish(&mut self, mut message: DriaMessage) -> Result<()> {
         // attach protocol name to the message
         message = message.with_identity(self.p2p.protocol().name.clone());
 
@@ -203,17 +204,17 @@ impl DriaComputeNode {
                 );
 
                 // ensure that message is from the known RPCs
-                if !self.available_nodes.rpc_nodes.contains(&source_peer_id) {
+                if !self.dria_nodes.rpc_peerids.contains(&source_peer_id) {
                     log::warn!(
                         "Received message from unauthorized source: {}",
                         source_peer_id
                     );
-                    log::debug!("Allowed sources: {:#?}", self.available_nodes.rpc_nodes);
+                    log::debug!("Allowed sources: {:#?}", self.dria_nodes.rpc_peerids);
                     return MessageAcceptance::Ignore;
                 }
 
                 // parse the raw gossipsub message to a prepared DKN message
-                let message = match DKNMessage::try_from_gossipsub_message(
+                let message = match DriaMessage::try_from_gossipsub_message(
                     &message,
                     &self.config.admin_public_key,
                 ) {
@@ -414,20 +415,22 @@ impl DriaComputeNode {
     /// Updates the local list of available nodes by refreshing it.
     /// Dials the RPC nodes again for better connectivity.
     async fn handle_available_nodes_refresh(&mut self) {
-        log::info!("Refreshing available nodes.");
+        log::info!("Refreshing available Dria nodes.");
 
         // refresh available nodes
-        if let Err(e) = self.available_nodes.populate_with_api().await {
+        if let Err(e) = refresh_dria_nodes(&mut self.dria_nodes).await {
             log::error!("Error refreshing available nodes: {:?}", e);
         };
 
         // dial all rpc nodes
-        for rpc_addr in self.available_nodes.rpc_addrs.iter() {
-            log::debug!("Dialling RPC node: {}", rpc_addr);
+        for rpc_addr in self.dria_nodes.rpc_nodes.iter() {
+            log::info!("Dialling RPC node: {}", rpc_addr);
             if let Err(e) = self.p2p.dial(rpc_addr.clone()).await {
                 log::warn!("Error dialling RPC node: {:?}", e);
             };
         }
+
+        log::info!("Finished refreshing!");
     }
 }
 
@@ -462,7 +465,7 @@ mod tests {
 
         // publish a dummy message
         let topic = "foo";
-        let message = DKNMessage::new("hello from the other side", topic);
+        let message = DriaMessage::new("hello from the other side", topic);
         node.subscribe(topic).await.expect("should subscribe");
         node.publish(message).await.expect("should publish");
         node.unsubscribe(topic).await.expect("should unsubscribe");

compute/src/node.rs

@@ -1,8 +1,7 @@
+use dkn_utils::get_current_time_nanos;
 use serde::{Deserialize, Serialize};
 use std::time::Instant;
 
-use crate::utils::get_current_time_nanos;
-
 /// Task stats for diagnostics.
 /// Returning this as the payload helps to debug the errors received at client side, and latencies.
 #[derive(Default, Debug, Clone, Serialize, Deserialize)]