VM healthcheck async routine should be cancellable

guvenc · guvenc · commit 8e7a39dda3d6 · 2025-08-20T16:16:59.000+02:00
Signed-off-by: Guvenc Gulce &lt;guevenc.guelce@sap.com&gt;
diff --git a/feos/services/vm-service/src/dispatcher.rs b/feos/services/vm-service/src/dispatcher.rs
@@ -27,12 +27,14 @@ pub struct VmServiceDispatcher {
     status_channel_tx: broadcast::Sender<VmEventWrapper>,
     hypervisor: Arc<dyn Hypervisor>,
     repository: VmRepository,
+    healthcheck_cancel_bus: broadcast::Sender<Uuid>,
 }
 
 impl VmServiceDispatcher {
     pub async fn new(rx: mpsc::Receiver<Command>, db_url: &str) -> Result<Self> {
         let (event_bus_tx, event_bus_rx_for_dispatcher) = mpsc::channel(32);
         let (status_channel_tx, _) = broadcast::channel(32);
+        let (healthcheck_cancel_bus, _) = broadcast::channel::<Uuid>(32);
         let hypervisor = Arc::from(factory(VmmType::CloudHypervisor));
         info!("VM_DISPATCHER: Connecting to persistence layer at {db_url}...");
         let repository = VmRepository::connect(db_url).await?;
@@ -44,6 +46,7 @@ impl VmServiceDispatcher {
             status_channel_tx,
             hypervisor,
             repository,
+            healthcheck_cancel_bus,
         })
     }
 
@@ -170,13 +173,15 @@ impl VmServiceDispatcher {
                     let hypervisor = self.hypervisor.clone();
                     let event_bus_tx = self.event_bus_tx.clone();
                     let status_channel_tx = self.status_channel_tx.clone();
+                    let healthcheck_cancel_bus_tx = self.healthcheck_cancel_bus.clone();
 
                     match cmd {
                         Command::CreateVm(req, responder) => {
                             self.handle_create_vm_command(req, responder, hypervisor, event_bus_tx).await;
                         }
                         Command::StartVm(req, responder) => {
-                            tokio::spawn(worker::handle_start_vm(req, responder, hypervisor, event_bus_tx));
+                            let cancel_bus = healthcheck_cancel_bus_tx.subscribe();
+                            tokio::spawn(worker::handle_start_vm(req, responder, hypervisor, event_bus_tx, cancel_bus));
                         }
                         Command::GetVm(req, responder) => {
                             self.handle_get_vm_command(req, responder).await;
@@ -536,6 +541,12 @@ impl VmServiceDispatcher {
                 }
                 info!("VM_DISPATCHER: Deleted record for VM {vm_id} from database.");
 
+                if let Err(e) = self.healthcheck_cancel_bus.send(vm_id) {
+                    warn!(
+                        "VM_DISPATCHER: Failed to send healthcheck cancellation for {vm_id}: {e}"
+                    );
+                }
+
                 tokio::spawn(worker::handle_delete_vm(
                     req,
                     image_uuid_to_delete,
@@ -548,6 +559,13 @@ impl VmServiceDispatcher {
             Ok(None) => {
                 let msg = format!("VM with ID {vm_id} not found in database for deletion");
                 warn!("VM_DISPATCHER: {msg}. Still attempting hypervisor cleanup.");
+
+                if let Err(e) = self.healthcheck_cancel_bus.send(vm_id) {
+                    warn!(
+                        "VM_DISPATCHER: Failed to send healthcheck cancellation for {vm_id}: {e}"
+                    );
+                }
+
                 tokio::spawn(worker::handle_delete_vm(
                     req,
                     String::new(),
diff --git a/feos/services/vm-service/src/vmm/ch_adapter.rs b/feos/services/vm-service/src/vmm/ch_adapter.rs
@@ -15,15 +15,16 @@ use feos_proto::vm_service::{
 };
 use hyper_util::client::legacy::Client;
 use hyperlocal::{UnixClientExt, UnixConnector, Uri as HyperlocalUri};
-use log::{info, warn};
+use log::{error, info, warn};
 use nix::sys::signal::{kill, Signal};
 use nix::unistd::{self, Pid};
 use std::io;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use tokio::process::Command as TokioCommand;
-use tokio::sync::mpsc;
+use tokio::sync::{broadcast, mpsc};
 use tokio::time::{self, timeout, Duration};
+use uuid::Uuid;
 
 pub struct CloudHypervisorAdapter {
     ch_binary_path: PathBuf,
@@ -257,36 +258,56 @@ impl Hypervisor for CloudHypervisorAdapter {
         Ok(StartVmResponse {})
     }
 
-    async fn healthcheck_vm(&self, vm_id: String, broadcast_tx: mpsc::Sender<VmEventWrapper>) {
+    async fn healthcheck_vm(
+        &self,
+        vm_id: String,
+        broadcast_tx: mpsc::Sender<VmEventWrapper>,
+        mut cancel_bus: broadcast::Receiver<Uuid>,
+    ) {
         info!("CH_ADAPTER ({vm_id}): Starting healthcheck monitoring.");
         let mut interval = time::interval(Duration::from_secs(10));
+        let vm_id_uuid = match Uuid::parse_str(&vm_id) {
+            Ok(id) => id,
+            Err(e) => {
+                error!("CH_ADAPTER ({vm_id}): Invalid UUID format, cannot start healthcheck: {e}");
+                return;
+            }
+        };
 
         loop {
-            interval.tick().await;
-            log::debug!("CH_ADAPTER ({vm_id}): Performing healthcheck ping.");
-            let req = PingVmRequest {
-                vm_id: vm_id.clone(),
-            };
-
-            match self.ping_vm(req).await {
-                Ok(_) => {
-                    log::debug!("CH_ADAPTER ({vm_id}): Healthcheck ping successful.");
+            tokio::select! {
+                _ = interval.tick() => {
+                    log::debug!("CH_ADAPTER ({vm_id}): Performing healthcheck ping.");
+                    let req = PingVmRequest {
+                        vm_id: vm_id.clone(),
+                    };
+
+                    if let Err(e) = self.ping_vm(req).await {
+                        warn!("CH_ADAPTER ({vm_id}): Healthcheck failed: {e}. VM is considered unhealthy.");
+                        super::broadcast_state_change_event(
+                            &broadcast_tx,
+                            &vm_id,
+                            "vm-health-monitor",
+                            feos_proto::vm_service::VmStateChangedEvent {
+                                new_state: VmState::Crashed as i32,
+                                reason: format!("Healthcheck failed: {e}"),
+                            },
+                            None,
+                        )
+                        .await;
+                        break;
+                    } else {
+                        log::debug!("CH_ADAPTER ({vm_id}): Healthcheck ping successful.");
+                    }
+                }
+                Ok(cancelled_vm_id) = cancel_bus.recv() => {
+                    if cancelled_vm_id == vm_id_uuid {
+                        info!("CH_ADAPTER ({vm_id}): Received cancellation signal. Stopping healthcheck.");
+                        break;
+                    }
                 }
-                Err(e) => {
-                    warn!(
-                        "CH_ADAPTER ({vm_id}): Healthcheck failed: {e}. VM is considered unhealthy."
-                    );
-                    super::broadcast_state_change_event(
-                        &broadcast_tx,
-                        &vm_id,
-                        "vm-health-monitor",
-                        feos_proto::vm_service::VmStateChangedEvent {
-                            new_state: VmState::Crashed as i32,
-                            reason: format!("Healthcheck failed: {e}"),
-                        },
-                        None,
-                    )
-                    .await;
+                else => {
+                    info!("CH_ADAPTER ({vm_id}): Healthcheck cancellation channel closed. Stopping healthcheck.");
                     break;
                 }
             }
diff --git a/feos/services/vm-service/src/vmm/mod.rs b/feos/services/vm-service/src/vmm/mod.rs
@@ -8,7 +8,7 @@ use feos_proto::vm_service::{
 use prost::Message;
 use prost_types::Any;
 use std::path::{Path, PathBuf};
-use tokio::sync::mpsc;
+use tokio::sync::{broadcast, mpsc};
 use tonic::Status;
 use uuid::Uuid;
 
@@ -64,7 +64,12 @@ pub trait Hypervisor: Send + Sync {
 
     async fn start_vm(&self, req: StartVmRequest) -> Result<StartVmResponse, VmmError>;
 
-    async fn healthcheck_vm(&self, vm_id: String, broadcast_tx: mpsc::Sender<VmEventWrapper>);
+    async fn healthcheck_vm(
+        &self,
+        vm_id: String,
+        broadcast_tx: mpsc::Sender<VmEventWrapper>,
+        cancel_bus: broadcast::Receiver<Uuid>,
+    );
 
     async fn get_vm(&self, req: GetVmRequest) -> Result<VmInfo, VmmError>;
 
diff --git a/feos/services/vm-service/src/worker.rs b/feos/services/vm-service/src/worker.rs
@@ -15,6 +15,7 @@ use std::sync::Arc;
 use tokio::sync::{broadcast, mpsc, oneshot};
 use tokio_stream::StreamExt;
 use tonic::{Status, Streaming};
+use uuid::Uuid;
 
 async fn wait_for_image_ready(image_uuid: &str, image_ref: &str) -> Result<(), VmmError> {
     let mut client = vmservice_helper::get_image_service_client()
@@ -154,6 +155,7 @@ pub async fn handle_start_vm(
     responder: oneshot::Sender<Result<StartVmResponse, Status>>,
     hypervisor: Arc<dyn Hypervisor>,
     broadcast_tx: mpsc::Sender<VmEventWrapper>,
+    cancel_bus: broadcast::Receiver<Uuid>,
 ) {
     let vm_id = req.vm_id.clone();
     let result = hypervisor.start_vm(req).await;
@@ -175,7 +177,7 @@ pub async fn handle_start_vm(
         let health_broadcast_tx = broadcast_tx.clone();
         tokio::spawn(async move {
             health_hypervisor
-                .healthcheck_vm(vm_id, health_broadcast_tx)
+                .healthcheck_vm(vm_id, health_broadcast_tx, cancel_bus)
                 .await;
         });
     }
