edit time feature

Author: erhant

compute/src/reqres/heartbeat.rs

@@ -15,7 +15,7 @@ pub struct HeartbeatRequester;
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct HeartbeatRequest {
     /// A unique ID for the heartbeat request.
-    pub(crate) uuid: Uuid,
+    pub(crate) heartbeat_id: Uuid,
     /// Deadline for the heartbeat request, in nanoseconds.
     pub(crate) deadline: chrono::DateTime<chrono::Utc>,
     /// Models available in the node.
@@ -31,7 +31,7 @@ pub struct HeartbeatRequest {
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct HeartbeatResponse {
     /// UUID as given in the request.
-    pub(crate) uuid: Uuid,
+    pub(crate) heartbeat_id: Uuid,
     /// Acknowledgement of the heartbeat.
     pub(crate) ack: bool,
 }
@@ -53,7 +53,7 @@ impl HeartbeatRequester {
         let deadline = chrono::Utc::now() + HEARTBEAT_DEADLINE_SECS;
 
         let heartbeat_request = HeartbeatRequest {
-            uuid,
+            heartbeat_id: uuid,
             deadline,
             models: node.config.workflows.models.clone(),
             pending_tasks: node.get_pending_task_count(),
@@ -79,7 +79,7 @@ impl HeartbeatRequester {
         node: &mut DriaComputeNode,
         res: HeartbeatResponse,
     ) -> Result<()> {
-        if let Some(deadline) = node.heartbeats.remove(&res.uuid) {
+        if let Some(deadline) = node.heartbeats.remove(&res.heartbeat_id) {
             if !res.ack {
                 Err(eyre!("Heartbeat was not acknowledged."))
             } else if chrono::Utc::now() > deadline {
@@ -93,7 +93,7 @@ impl HeartbeatRequester {
         } else {
             Err(eyre!(
                 "Received an unknown heartbeat response with UUID {}.",
-                res.uuid
+                res.heartbeat_id
             ))
         }
     }

compute/src/reqres/specs.rs

@@ -2,17 +2,18 @@ use crate::utils::Specs;
 
 use super::IsResponder;
 use serde::{Deserialize, Serialize};
+use uuid::Uuid;
 
 #[derive(Serialize, Deserialize)]
 pub struct SpecRequest {
-    /// UUID of the specs request, prevents replay attacks.
-    pub request_id: String,
+    /// UUID of the specs request, prevents replays.
+    pub request_id: Uuid,
 }
 
 #[derive(Serialize, Deserialize)]
 pub struct SpecResponse {
-    /// UUID of the specs request, prevents replay attacks.
-    pub request_id: String,
+    /// UUID of the specs request, prevents replays.
+    pub request_id: Uuid,
     /// Node specs, will be flattened during serialization.
     #[serde(flatten)]
     specs: Specs,

p2p/tests/request_test.rs

@@ -1,4 +1,6 @@
 use std::str::FromStr;
+use std::thread::sleep;
+use std::time::Duration;
 
 use dkn_p2p::DriaNetworkType::Community;
 use dkn_p2p::{DriaNodes, DriaP2PClient, DriaP2PProtocol};
@@ -23,15 +25,13 @@ async fn test_request_message() -> Result<()> {
         .is_test(true)
         .try_init();
 
-    let listen_addr = "/ip4/0.0.0.0/tcp/4001".parse()?;
-
     // prepare nodes
     let nodes = DriaNodes::new(Community).with_statics();
 
     // spawn P2P client in another task
     let (client, mut commander, mut req_rx) = DriaP2PClient::new(
         Keypair::generate_secp256k1(),
-        listen_addr,
+        "/ip4/127.0.0.1/tcp/0".parse().unwrap(),
         &nodes,
         DriaP2PProtocol::default(),
     )
@@ -41,7 +41,7 @@ async fn test_request_message() -> Result<()> {
     let task_handle = tokio::spawn(async move { client.run().await });
 
     log::info!("Waiting a bit until we have enough peers");
-    tokio::time::sleep(tokio::time::Duration::from_secs(10)).await;
+    sleep(Duration::from_secs(10));
 
     let peer_id =
         PeerId::from_str("16Uiu2HAmB5HGdwLNHX81u7ey1fvDx5Mr4ofa2PdSSVxFKrrcErAN").unwrap();