use DriaMessage for heartbeat

Author: erhant

compute/src/reqres/heartbeat.rs

@@ -1,14 +1,12 @@
-use std::time::Duration;
-
 use dkn_p2p::libp2p::{request_response::OutboundRequestId, PeerId};
-use dkn_workflows::{Model, ModelProvider};
 use eyre::{eyre, Result};
+use serde::{Deserialize, Serialize};
+use std::time::Duration;
 use uuid::Uuid;
 
 use super::IsResponder;
-use serde::{Deserialize, Serialize};
 
-use crate::DriaComputeNode;
+use crate::{utils::DriaMessage, DriaComputeNode};
 
 pub struct HeartbeatRequester;
 
@@ -18,8 +16,8 @@ pub struct HeartbeatRequest {
     pub(crate) heartbeat_id: Uuid,
     /// Deadline for the heartbeat request, in nanoseconds.
     pub(crate) deadline: chrono::DateTime<chrono::Utc>,
-    /// Models available in the node.
-    pub(crate) models: Vec<(ModelProvider, Model)>,
+    /// Model names available in the node.
+    pub(crate) models: Vec<String>,
     /// Number of tasks in the channel currently, `single` and `batch`.
     pub(crate) pending_tasks: [usize; 2],
 }
@@ -39,13 +37,16 @@ pub struct HeartbeatResponse {
 }
 
 impl IsResponder for HeartbeatRequester {
-    type Request = HeartbeatRequest;
-    type Response = HeartbeatResponse;
+    type Request = DriaMessage; // TODO: HeartbeatRequest;
+    type Response = DriaMessage; // TODO: HeartbeatResponse;
 }
 
 /// Any acknowledged heartbeat that is older than this duration is considered dead.
 const HEARTBEAT_DEADLINE_SECS: Duration = Duration::from_secs(20);
 
+/// Topic for the [`DriaMessage`].
+const HEARTBEAT_TOPIC: &str = "heartbeat";
+
 impl HeartbeatRequester {
     pub(crate) async fn send_heartbeat(
         node: &mut DriaComputeNode,
@@ -57,16 +58,23 @@ impl HeartbeatRequester {
         let heartbeat_request = HeartbeatRequest {
             heartbeat_id: uuid,
             deadline,
-            models: node.config.workflows.models.clone(),
+            models: node
+                .config
+                .workflows
+                .models
+                .iter()
+                .map(|m| m.1.to_string())
+                .collect(),
             pending_tasks: node.get_pending_task_count(),
         };
 
+        let heartbeat_message = node.new_message(
+            serde_json::to_vec(&heartbeat_request).expect("should be serializable"),
+            HEARTBEAT_TOPIC,
+        );
         let request_id = node
             .p2p
-            .request(
-                peer_id,
-                serde_json::to_vec(&heartbeat_request).expect("should be serializable"),
-            )
+            .request(peer_id, heartbeat_message.to_bytes()?)
             .await?;
 
         // add it to local heartbeats set
@@ -78,8 +86,10 @@ impl HeartbeatRequester {
     /// Handles the heartbeat request received from the network.
     pub(crate) async fn handle_ack(
         node: &mut DriaComputeNode,
-        res: HeartbeatResponse,
+        ack_message: DriaMessage,
     ) -> Result<()> {
+        let res = ack_message.parse_payload::<HeartbeatResponse>()?;
+
         if let Some(deadline) = node.heartbeats.remove(&res.heartbeat_id) {
             if let Some(err) = res.error {
                 Err(eyre!("Heartbeat was not acknowledged: {}", err))