fix: suspend/resume on worker deactivate/activate (#2350)

fluencelabs · Aug 16, 2024 · ac72e69 · ac72e69
1 parent 9d66b16
commit ac72e69
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 4 deletions.
diff --git a/crates/vm-utils/src/lib.rs b/crates/vm-utils/src/lib.rs
@@ -2,9 +2,11 @@ mod vm_utils;
 
 pub use nonempty::NonEmpty;
 pub use vm_utils::create_domain;
+pub use vm_utils::pause_vm;
 pub use vm_utils::reboot_vm;
 pub use vm_utils::remove_domain;
 pub use vm_utils::reset_vm;
+pub use vm_utils::resume_vm;
 pub use vm_utils::start_vm;
 pub use vm_utils::status_vm;
 pub use vm_utils::stop_vm;

diff --git a/crates/vm-utils/src/vm_utils.rs b/crates/vm-utils/src/vm_utils.rs
@@ -88,6 +88,19 @@ pub enum VmError {
         err: virt::error::Error,
         name: String,
     },
+    #[error("Failed to suspend VM {name}: {err}")]
+    FailedToSuspendVM {
+        name: String,
+        #[source]
+        err: virt::error::Error,
+    },
+
+    #[error("Failed to resume VM {name}: {err}")]
+    FailedToResumeVM {
+        name: String,
+        #[source]
+        err: virt::error::Error,
+    },
 }
 
 // The list of states is taken from the libvirt documentation
@@ -258,6 +271,34 @@ pub fn status_vm(uri: &str, name: &str) -> Result<VmStatus, VmError> {
     Ok(VmStatus::from_u32(info.state))
 }
 
+pub fn pause_vm(uri: &str, name: &str) -> Result<(), VmError> {
+    tracing::info!(target: "vm-utils","Pausing VM with name {name}");
+    let conn = Connect::open(Some(uri)).map_err(|err| VmError::FailedToConnect { err })?;
+    let domain = Domain::lookup_by_name(&conn, name).map_err(|err| VmError::VmNotFound {
+        name: name.to_string(),
+        err,
+    })?;
+    domain.suspend().map_err(|err| VmError::FailedToSuspendVM {
+        name: name.to_string(),
+        err,
+    })?;
+    Ok(())
+}
+
+pub fn resume_vm(uri: &str, name: &str) -> Result<(), VmError> {
+    tracing::info!(target: "vm-utils","Resuming VM with name {name}");
+    let conn = Connect::open(Some(uri)).map_err(|err| VmError::FailedToConnect { err })?;
+    let domain = Domain::lookup_by_name(&conn, name).map_err(|err| VmError::VmNotFound {
+        name: name.to_string(),
+        err,
+    })?;
+    domain.resume().map_err(|err| VmError::FailedToResumeVM {
+        name: name.to_string(),
+        err,
+    })?;
+    Ok(())
+}
+
 fn generate_random_mac() -> MacAddress {
     let mut rng = rand::thread_rng();
     let mut result = [0u8; 6];

diff --git a/crates/workers/src/workers.rs b/crates/workers/src/workers.rs
@@ -432,6 +432,7 @@ impl Workers {
     /// - `Err(WorkersError)` if an error occurs during the activation process.
     ///
     pub async fn activate_worker(&self, worker_id: WorkerId) -> Result<(), WorkersError> {
+        self.resume_vm(worker_id)?;
         self.set_worker_status(worker_id, true).await?;
         Ok(())
     }
@@ -452,6 +453,7 @@ impl Workers {
     /// - `Err(WorkersError)` if an error occurs during the deactivation process.
     ///
     pub async fn deactivate_worker(&self, worker_id: WorkerId) -> Result<(), WorkersError> {
+        self.pause_vm(worker_id)?;
         self.set_worker_status(worker_id, false).await?;
         Ok(())
     }
@@ -696,10 +698,12 @@ impl Workers {
 
     fn remove_vm(&self, worker_id: WorkerId) -> Result<(), WorkersError> {
         if let Some(vm_config) = &self.config.vm {
-            vm_utils::remove_domain(
-                vm_config.libvirt_uri.as_str(),
-                worker_id.to_string().as_str(),
-            )?;
+            if self.has_vm(worker_id)? {
+                vm_utils::remove_domain(
+                    vm_config.libvirt_uri.as_str(),
+                    worker_id.to_string().as_str(),
+                )?;
+            }
         }
         Ok(())
     }
@@ -727,6 +731,24 @@ impl Workers {
         self.on_vm(worker_id, vm_utils::reset_vm)
     }
 
+    fn pause_vm(&self, worker_id: WorkerId) -> Result<(), WorkersError> {
+        if let Some(vm_config) = &self.config.vm {
+            if self.has_vm(worker_id)? {
+                vm_utils::pause_vm(vm_config.libvirt_uri.as_str(), &worker_id.to_string())?;
+            }
+        }
+        Ok(())
+    }
+
+    fn resume_vm(&self, worker_id: WorkerId) -> Result<(), WorkersError> {
+        if let Some(vm_config) = &self.config.vm {
+            if self.has_vm(worker_id)? {
+                vm_utils::resume_vm(vm_config.libvirt_uri.as_str(), &worker_id.to_string())?;
+            }
+        }
+        Ok(())
+    }
+
     pub fn status_vm(&self, worker_id: WorkerId) -> Result<VmStatus, WorkersError> {
         self.on_vm(worker_id, vm_utils::status_vm)
     }