Merge pull request #502 from NillionNetwork/feat/robust-heartbeats

feat: add a few robustness-es features around htx submission

Author: mfontanini

cvm-agent/src/heartbeat.rs

@@ -1,4 +1,4 @@
-use crate::heartbeat::NilAVRouter::NilAVRouterInstance;
+use crate::{heartbeat::NilAVRouter::NilAVRouterInstance, monitors::caddy::CaddyStatus};
 use alloy::{primitives::Address, providers::ProviderBuilder, signers::local::PrivateKeySigner, sol_types::sol};
 use alloy_provider::{Provider, WsConnect};
 use anyhow::Context;
@@ -32,6 +32,7 @@ pub(crate) struct HeartbeatEmitterArgs {
     pub(crate) measurement_hash_url: String,
     pub(crate) cpu_count: u64,
     pub(crate) gpu_count: u64,
+    pub(crate) caddy_status: CaddyStatus,
 }
 
 pub(crate) struct HeartbeatEmitter {
@@ -62,6 +63,7 @@ impl HeartbeatEmitter {
             measurement_hash_url,
             cpu_count,
             gpu_count,
+            caddy_status,
         } = args;
         let contract_address: Address = contract_address.parse().context("Invalid contract address")?;
         let attestation_url = format!("https://{workload_domain}{ATTESTATION_PATH}");
@@ -81,11 +83,15 @@ impl HeartbeatEmitter {
             cpu_count,
             gpu_count,
         };
-        tokio::spawn(async move { submitter.run().await });
+        tokio::spawn(async move { submitter.run(caddy_status).await });
         Ok(())
     }
 
-    async fn run(self) {
+    async fn run(self, caddy_status: CaddyStatus) {
+        info!("Waiting for caddy to generate a TLS certificate before emitting heartbeats");
+        caddy_status.wait_tls_certificate().await;
+        info!("Starting heartbeat generation");
+
         let provider = loop {
             match self.connect().await {
                 Ok(out) => break out,
@@ -111,8 +117,12 @@ impl HeartbeatEmitter {
 
     async fn connect(&self) -> anyhow::Result<impl Provider> {
         let ws = WsConnect::new(&self.rpc_endpoint).with_max_retries(u32::MAX);
-        let provider =
-            ProviderBuilder::new().wallet(self.wallet.clone()).with_simple_nonce_management().connect_ws(ws).await?;
+        let provider = ProviderBuilder::new()
+            .wallet(self.wallet.clone())
+            .with_simple_nonce_management()
+            .with_gas_estimation()
+            .connect_ws(ws)
+            .await?;
         Ok(provider)
     }
 
@@ -133,7 +143,8 @@ impl HeartbeatEmitter {
         let pending_tx = call.send().await?;
         let receipt = pending_tx.get_receipt().await?;
         let tx_hash = receipt.transaction_hash;
-        info!("HTX submitted in transaction {tx_hash}");
+        let status = if receipt.status() { "success" } else { "failure" };
+        info!("HTX submitted in transaction {tx_hash} with status {status}");
         Ok(())
     }
 }

cvm-agent/src/monitors/caddy.rs

@@ -8,7 +8,7 @@ use cvm_agent_models::health::EventKind;
 use futures::{Stream, StreamExt};
 use serde::Deserialize;
 use std::{borrow::Cow, mem, sync::Arc, time::Duration};
-use tokio::sync::Mutex;
+use tokio::sync::{Mutex, oneshot};
 use tokio::time::sleep;
 use tracing::{debug, error, info, warn};
 
@@ -23,15 +23,18 @@ pub struct CaddyMonitor {
 }
 
 impl CaddyMonitor {
-    pub fn spawn(docker: Docker, system_state: Arc<Mutex<SystemState>>, event_holder: EventHolder) {
+    pub fn spawn(docker: Docker, system_state: Arc<Mutex<SystemState>>, event_holder: EventHolder) -> CaddyStatus {
         let monitor = Self { docker, system_state, event_holder };
+        let (sender, receiver) = oneshot::channel();
         info!("Spawning caddy monitor");
         tokio::spawn(async move {
-            monitor.run().await;
+            monitor.run(sender).await;
         });
+        CaddyStatus(receiver)
     }
 
-    async fn run(self) {
+    async fn run(self, sender: oneshot::Sender<()>) {
+        let mut sender = Some(sender);
         let mut threshold_timestamp = 0.0;
         loop {
             let builder = LogsOptionsBuilder::new().tail("10").stderr(true);
@@ -45,6 +48,10 @@ impl CaddyMonitor {
                         SystemState::WaitingBootstrap => error!("System is still waiting for bootstrap"),
                         SystemState::Starting | SystemState::Ready => {
                             info!("Caddy fetched TLS certificate and is running successfully");
+                            // Notify the listener that we're ready
+                            if let Some(sender) = sender.take() {
+                                let _ = sender.send(());
+                            }
                             *system_state = SystemState::Ready
                         }
                     }
@@ -99,6 +106,16 @@ impl CaddyMonitor {
     }
 }
 
+pub(crate) struct CaddyStatus(oneshot::Receiver<()>);
+
+impl CaddyStatus {
+    pub(crate) async fn wait_tls_certificate(self) {
+        if self.0.await.is_err() {
+            error!("Caddy status receiver dropped");
+        }
+    }
+}
+
 #[derive(Deserialize)]
 struct LogLine<'a> {
     ts: f64,

cvm-agent/src/routes/system/bootstrap.rs

@@ -17,6 +17,8 @@ pub(crate) async fn handler(state: SharedState, request: Json<BootstrapRequest>)
     let event_holder = ctx.event_holder.clone();
     *system_state = SystemState::Starting;
 
+    let caddy_status = CaddyMonitor::spawn(state.docker.clone(), state.system_state.clone(), event_holder);
+
     match (request.workload_id, request.heartbeat) {
         (Some(workload_id), Some(heartbeat)) => {
             let args = HeartbeatEmitterArgs {
@@ -31,6 +33,7 @@ pub(crate) async fn handler(state: SharedState, request: Json<BootstrapRequest>)
                 measurement_hash_url: heartbeat.measurement_hash_url,
                 cpu_count: ctx.cpus,
                 gpu_count: ctx.gpus,
+                caddy_status,
             };
             if let Err(e) = HeartbeatEmitter::spawn(args).await {
                 error!("Failed setting up heartbeat emitter: {e}");
@@ -41,6 +44,5 @@ pub(crate) async fn handler(state: SharedState, request: Json<BootstrapRequest>)
     };
 
     ComposeMonitor::spawn(ctx, request.acme, request.docker, request.domain);
-    CaddyMonitor::spawn(state.docker.clone(), state.system_state.clone(), event_holder);
     StatusCode::OK
 }