some naming fixes, todo check tests

Author: erhant

compute/src/node/core.rs

@@ -1,8 +1,8 @@
-use eyre::Result;
+use eyre::{eyre, Result};
 use std::time::Duration;
 use tokio_util::sync::CancellationToken;
 
-use crate::{node::PingpongHandler, reqres::TaskResponder, utils::DriaMessage, DriaComputeNode};
+use crate::{node::PingpongHandler, utils::DriaMessage, DriaComputeNode};
 
 impl DriaComputeNode {
     /// Runs the main loop of the compute node.
@@ -27,72 +27,45 @@ impl DriaComputeNode {
 
         loop {
             tokio::select! {
-                // a Workflow message to be published is received from the channel
-                // this is expected to be sent by the workflow worker
-                publish_msg_opt = self.publish_rx.recv() => {
-                    if let Some(publish_msg) = publish_msg_opt {
-                        // remove the task from pending tasks based on its batchability
-                        let task_metadata = match publish_msg.batchable {
-                            true => {
-                                self.completed_tasks_batch += 1;
-                                self.pending_tasks_batch.remove(&publish_msg.task_id)
-                            },
-                            false => {
-                                self.completed_tasks_single += 1;
-                                self.pending_tasks_single.remove(&publish_msg.task_id)
-                            }
-                        };
-
-                        // respond to the request
-                        match task_metadata {
-                            Some(channel) => {
-                              TaskResponder::handle_respond(self, publish_msg, channel).await?;
-                            }
-                            None => log::error!("Channel not found for task id: {}", publish_msg.task_id),
-                        }
-                    } else {
-                        log::error!("Publish channel closed unexpectedly, we still have {} batch and {} single tasks.", self.pending_tasks_batch.len(), self.pending_tasks_single.len());
-                        break;
-                    };
+                // a task is completed by the worker & should be responded to the requesting peer
+                task_response_msg_opt = self.task_output_rx.recv() => {
+                    let task_response_msg = task_response_msg_opt.ok_or(
+                      eyre!("Publish channel closed unexpectedly, we still have {} batch and {} single tasks.", self.pending_tasks_batch.len(), self.pending_tasks_single.len())
+                    )?; {
+                        self.handle_task_response(task_response_msg).await?;
+                    }
                 },
 
-                // check peer count every now and then
-                _ = diagnostic_refresh_interval.tick() => self.handle_diagnostic_refresh().await,
-
-                // available nodes are refreshed every now and then
-                _ = available_node_refresh_interval.tick() => self.handle_available_nodes_refresh().await,
-
                 // a GossipSub message is received from the channel
                 // this is expected to be sent by the p2p client
-                gossipsub_msg_opt = self.message_rx.recv() => {
-                    if let Some((peer_id, message_id, message)) = gossipsub_msg_opt {
-                        // handle the message, returning a message acceptance for the received one
-                        let acceptance = self.handle_message((peer_id, &message_id, message)).await;
-
-                        // validate the message based on the acceptance
-                        // cant do anything but log if this gives an error as well
-                        if let Err(e) = self.p2p.validate_message(&message_id, &peer_id, acceptance).await {
-                            log::error!("Error validating message {}: {:?}", message_id, e);
-                        }
-                    } else {
-                        log::error!("message_rx channel closed unexpectedly.");
-                        break;
-                    };
+                gossipsub_msg_opt = self.gossip_message_rx.recv() => {
+                    let (peer_id, message_id, message) = gossipsub_msg_opt.ok_or(eyre!("message_rx channel closed unexpectedly."))?;
+
+                    // handle the message, returning a message acceptance for the received one
+                    let acceptance = self.handle_message((peer_id, &message_id, message)).await;
+
+                    // validate the message based on the acceptance
+                    // cant do anything but log if this gives an error as well
+                    if let Err(e) = self.p2p.validate_message(&message_id, &peer_id, acceptance).await {
+                        log::error!("Error validating message {}: {:?}", message_id, e);
+                    }
+
                 },
 
-                // a Response message is received from the channel
-                // this is expected to be sent by the p2p client
+                // a Request is received from the channel, sent by p2p client
                 request_msg_opt = self.request_rx.recv() => {
-                    if let Some((peer_id, data, channel)) = request_msg_opt {
-                        if let Err(e) = self.handle_request((peer_id, data, channel)).await {
-                            log::error!("Error handling request: {:?}", e);
-                        }
-                    } else {
-                        log::error!("request_rx channel closed unexpectedly.");
-                        break;
-                    };
+                  let request = request_msg_opt.ok_or(eyre!("request_rx channel closed unexpectedly."))?;
+                  if let Err(e) = self.handle_request(request).await {
+                      log::error!("Error handling request: {:?}", e);
+                  }
                 },
 
+                // check peer count every now and then
+                _ = diagnostic_refresh_interval.tick() => self.handle_diagnostic_refresh().await,
+
+                // available nodes are refreshed every now and then
+                _ = available_node_refresh_interval.tick() => self.handle_available_nodes_refresh().await,
+
                 // check if the cancellation token is cancelled
                 // this is expected to be cancelled by the main thread with signal handling
                 _ = cancellation.cancelled() => break,
@@ -119,15 +92,18 @@ impl DriaComputeNode {
     }
 
     /// Shutdown channels between p2p, worker and yourself.
+    ///
+    /// Can be inlined as it is called only once from very few places.
+    #[inline]
     pub async fn shutdown(&mut self) -> Result<()> {
         log::debug!("Sending shutdown command to p2p client.");
         self.p2p.shutdown().await?;
 
-        log::debug!("Closing message channel.");
-        self.message_rx.close();
+        log::debug!("Closing gossip message receipt channel.");
+        self.gossip_message_rx.close();
 
-        log::debug!("Closing publish channel.");
-        self.publish_rx.close();
+        log::debug!("Closing task response channel.");
+        self.task_output_rx.close();
 
         Ok(())
     }

compute/src/node/diagnostic.rs

@@ -48,11 +48,17 @@ impl DriaComputeNode {
 
         log::info!("{}", diagnostics.join("\n  "));
 
+        // check liveness of the node w.r.t last ping-pong time
         if self.last_pinged_at < Instant::now() - Duration::from_secs(PING_LIVENESS_SECS) {
             log::error!(
-            "Node has not received any pings for at least {} seconds & it may be unreachable!\nPlease restart your node!",
-            PING_LIVENESS_SECS
-        );
+                "Node has not received any pings for at least {} seconds & it may be unreachable!\nPlease restart your node!",
+                PING_LIVENESS_SECS
+            );
+        }
+
+        // added rpc nodes check, sometimes this happens when API is down / bugs for some reason
+        if self.dria_nodes.rpc_peerids.is_empty() {
+            log::error!("No RPC peerids were found to be available, please restart your node!",);
         }
     }
 

compute/src/node/gossipsub.rs

@@ -152,8 +152,10 @@ impl DriaComputeNode {
 mod tests {
     use super::*;
     use crate::DriaComputeNodeConfig;
+    use tokio::time::sleep;
     use tokio_util::sync::CancellationToken;
 
+    // FIXME: test is failing
     #[tokio::test]
     #[ignore = "run this manually"]
     async fn test_publish_message() -> eyre::Result<()> {
@@ -173,11 +175,16 @@ mod tests {
         let p2p_task = tokio::spawn(async move { p2p.run().await });
 
         // launch & wait for a while for connections
-        log::info!("Waiting a bit for peer setup.");
         let run_cancellation = cancellation.clone();
+        let wait_duration = tokio::time::Duration::from_secs(20);
+        log::info!(
+            "Waiting a bit ({}ms) for peer setup.",
+            wait_duration.as_millis()
+        );
+        sleep(wait_duration).await;
         tokio::select! {
             _ = node.run(run_cancellation) => (),
-            _ = tokio::time::sleep(tokio::time::Duration::from_secs(20)) => cancellation.cancel(),
+            _ = tokio::time::sleep(wait_duration) => cancellation.cancel(),
         }
         log::info!("Connected Peers:\n{:#?}", node.peers().await?);
 

compute/src/node/mod.rs

@@ -35,15 +35,17 @@ pub struct DriaComputeNode {
     /// If this is too much, we can say that the node is not reachable by RPC.
     pub last_pinged_at: Instant,
     /// Gossipsub message receiver, used by peer-to-peer client in a separate thread.
-    message_rx: mpsc::Receiver<(PeerId, MessageId, Message)>,
+    ///
+    /// It will publish messages sent to this channel to the network.
+    gossip_message_rx: mpsc::Receiver<(PeerId, MessageId, Message)>,
     /// Request-response request receiver.
     request_rx: mpsc::Receiver<(PeerId, Vec<u8>, ResponseChannel<Vec<u8>>)>,
-    /// Publish receiver to receive messages to be published.
-    publish_rx: mpsc::Receiver<TaskWorkerOutput>,
+    /// Task response receiver, will respond to the request-response channel with the given result.
+    task_output_rx: mpsc::Receiver<TaskWorkerOutput>,
     /// Task worker transmitter to send batchable tasks.
-    task_batch_tx: Option<mpsc::Sender<TaskWorkerInput>>,
+    task_request_batch_tx: Option<mpsc::Sender<TaskWorkerInput>>,
     /// Task worker transmitter to send single tasks.
-    task_single_tx: Option<mpsc::Sender<TaskWorkerInput>>,
+    task_request_single_tx: Option<mpsc::Sender<TaskWorkerInput>>,
     // Single tasks
     pending_tasks_single: HashMap<String, TaskWorkerMetadata>,
     // Batchable tasks
@@ -118,12 +120,12 @@ impl DriaComputeNode {
                 p2p: p2p_commander,
                 dria_nodes,
                 // receivers
-                publish_rx,
-                message_rx,
+                task_output_rx: publish_rx,
+                gossip_message_rx: message_rx,
                 request_rx,
                 // transmitters
-                task_batch_tx,
-                task_single_tx,
+                task_request_batch_tx: task_batch_tx,
+                task_request_single_tx: task_single_tx,
                 // task trackers
                 pending_tasks_single: HashMap::new(),
                 pending_tasks_batch: HashMap::new(),

compute/src/node/reqres.rs

@@ -1,7 +1,7 @@
 use dkn_p2p::libp2p::{request_response::ResponseChannel, PeerId};
 use eyre::{eyre, Result};
 
-use crate::reqres::*;
+use crate::{reqres::*, workers::task::TaskWorkerOutput};
 
 use super::DriaComputeNode;
 
@@ -38,6 +38,7 @@ impl DriaComputeNode {
         Ok(())
     }
 
+    /// Handles a Specifications request received from the network.
     async fn handle_spec_request(
         &mut self,
         peer_id: PeerId,
@@ -73,6 +74,11 @@ impl DriaComputeNode {
         Ok(())
     }
 
+    /// Handles a Task request received from the network.
+    ///
+    /// Based on the task type, the task is sent to the appropriate worker & metadata is stored in memory.
+    /// This metadata will be used during response as well, and we can count the number of tasks at hand by
+    /// looking at the number metadata stored.
     async fn handle_task_request(
         &mut self,
         peer_id: PeerId,
@@ -86,7 +92,7 @@ impl DriaComputeNode {
         if let Err(e) = match task_input.batchable {
             // this is a batchable task, send it to batch worker
             // and keep track of the task id in pending tasks
-            true => match self.task_batch_tx {
+            true => match self.task_request_batch_tx {
                 Some(ref mut tx) => {
                     self.pending_tasks_batch
                         .insert(task_input.task_id.clone(), task_metadata);
@@ -101,7 +107,7 @@ impl DriaComputeNode {
 
             // this is a single task, send it to single worker
             // and keep track of the task id in pending tasks
-            false => match self.task_single_tx {
+            false => match self.task_request_single_tx {
                 Some(ref mut tx) => {
                     self.pending_tasks_single
                         .insert(task_input.task_id.clone(), task_metadata);
@@ -117,4 +123,31 @@ impl DriaComputeNode {
 
         Ok(())
     }
+
+    pub(crate) async fn handle_task_response(
+        &mut self,
+        task_response: TaskWorkerOutput,
+    ) -> Result<()> {
+        // remove the task from pending tasks, and get its metadata
+        let task_metadata = match task_response.batchable {
+            true => {
+                self.completed_tasks_batch += 1; // TODO: this should be done in success
+                self.pending_tasks_batch.remove(&task_response.task_id)
+            }
+            false => {
+                self.completed_tasks_single += 1; // TODO: this should be done in success
+                self.pending_tasks_single.remove(&task_response.task_id)
+            }
+        };
+
+        // respond to the response channel with the result
+        match task_metadata {
+            Some(channel) => {
+                TaskResponder::handle_respond(self, task_response, channel).await?;
+            }
+            None => log::error!("Channel not found for task id: {}", task_response.task_id),
+        };
+
+        Ok(())
+    }
 }