Add notification post-livemigration

openhcl/underhill_core/src/worker.rs

@@ -2562,6 +2562,9 @@ async fn new_underhill_vm(
     let debug_interrupt_callback = move |vtl: u8| p.assert_debug_interrupt(vtl);
     get_client.set_debug_interrupt_callback(Box::new(debug_interrupt_callback));
 
+    // Set do-nothing callback.
+    get_client.set_post_live_migration_callback(Box::new(|| {}));
+
     let mut input_distributor = InputDistributor::new(remote_console_cfg.input);
     resolver.add_async_resolver::<KeyboardInputHandleKind, _, MultiplexedInputHandle, _>(
         input_distributor.client().clone(),

vm/devices/get/get_protocol/src/lib.rs

@@ -95,6 +95,7 @@ open_enum! {
         // --- GE ---
         BATTERY_STATUS = 7,
         INJECT_DEBUG_INTERRUPT = 8,
+        NOTIFY_POST_LIVE_MIGRATION = 9,
     }
 }
 
@@ -1772,9 +1773,15 @@ pub struct InjectDebugInterruptNotification {
     pub vtl: u8,
     pub _pad: u8,
 }
-
 const_assert_eq!(6, size_of::<InjectDebugInterruptNotification>());
 
+#[repr(C)]
+#[derive(Copy, Clone, Debug, IntoBytes, FromBytes, Immutable, KnownLayout)]
+pub struct PostLiveMigrationNotification {
+    pub message_header: HeaderGuestNotification,
+}
+const_assert_eq!(4, size_of::<PostLiveMigrationNotification>());
+
 #[bitfield(u64)]
 #[derive(IntoBytes, FromBytes, Immutable, KnownLayout)]
 pub struct CreateRamGpaRangeFlags {

vm/devices/get/guest_emulation_transport/src/client.rs

@@ -412,6 +412,12 @@ impl GuestEmulationTransportClient {
             .notify(msg::Msg::SetDebugInterruptCallback(callback.into()));
     }
 
+    /// Set the the callback to handle PostLiveMigrationNotification.
+    pub fn set_post_live_migration_callback(&mut self, callback: Box<dyn Fn() + Send + Sync>) {
+        self.control
+            .notify(msg::Msg::SetPostLiveMigrationCallback(callback.into()));
+    }
+
     /// Send the attestation request to the IGVM agent on the host.
     pub async fn igvm_attest(
         &self,

vm/devices/get/guest_emulation_transport/src/process_loop.rs

@@ -10,6 +10,7 @@ use crate::client::ModifyVtl2SettingsRequest;
 use crate::error::IgvmAttestError;
 use crate::error::TryIntoProtocolBool;
 use chipset_resources::battery::HostBatteryUpdate;
+use cvm_tracing::CVM_ALLOWED;
 use futures::FutureExt;
 use futures::TryFutureExt;
 use futures_concurrency::future::Race;
@@ -219,6 +220,8 @@ pub(crate) mod msg {
         /// Store the callback to trigger the debug interrupt.
         // TODO: Consider a strategy that avoids LocalOnly here.
         SetDebugInterruptCallback(LocalOnly<Box<dyn Fn(u8) + Send + Sync>>),
+        /// Notify that livemigration has finished.
+        SetPostLiveMigrationCallback(LocalOnly<Box<dyn Fn() + Send + Sync>>),
 
         // Late bound receivers for Guest Notifications
         /// Take the late-bound GuestRequest receiver for Generation Id updates.
@@ -536,6 +539,8 @@ pub(crate) struct ProcessLoop<T: RingMem> {
     gpa_allocator: Option<Arc<dyn DmaClient>>,
     #[inspect(skip)]
     set_debug_interrupt: Option<Box<dyn Fn(u8) + Send + Sync>>,
+    #[inspect(skip)]
+    post_live_migration: Option<Box<dyn Fn() + Send + Sync>>,
     stats: Stats,
 
     guest_notification_listeners: GuestNotificationListeners,
@@ -744,6 +749,7 @@ impl<T: RingMem> ProcessLoop<T> {
             },
             gpa_allocator: None,
             set_debug_interrupt: None,
+            post_live_migration: None,
         }
     }
 
@@ -1076,6 +1082,9 @@ impl<T: RingMem> ProcessLoop<T> {
             Msg::SetDebugInterruptCallback(callback) => {
                 self.set_debug_interrupt = Some(callback.0);
             }
+            Msg::SetPostLiveMigrationCallback(callback) => {
+                self.post_live_migration = Some(callback.0);
+            }
 
             // Late bound receivers for Guest Notifications
             Msg::TakeVtl2SettingsReceiver(req) => req.handle_sync(|()| {
@@ -1349,6 +1358,9 @@ impl<T: RingMem> ProcessLoop<T> {
             GuestNotifications::INJECT_DEBUG_INTERRUPT => {
                 self.handle_debug_interrupt_notification(read_guest_notification(id, buf)?)?;
             }
+            GuestNotifications::NOTIFY_POST_LIVE_MIGRATION => {
+                self.handle_post_live_migration_notification(read_guest_notification(id, buf)?);
+            }
             invalid_notification => {
                 tracing::error!(
                     ?invalid_notification,
@@ -1574,6 +1586,16 @@ impl<T: RingMem> ProcessLoop<T> {
         Ok(())
     }
 
+    fn handle_post_live_migration_notification(
+        &mut self,
+        _notification: get_protocol::PostLiveMigrationNotification,
+    ) {
+        tracing::info!(CVM_ALLOWED, "notify_post_live_migration");
+        if let Some(callback) = self.post_live_migration.as_ref() {
+            callback()
+        }
+    }
+
     fn complete_modify_vtl2_settings(
         &mut self,
         result: Result<(), RpcError<Vec<Vtl2SettingsErrorInfo>>>,