From c194f63cf45d0ab99257102a665ab6a8adc7c3f2 Mon Sep 17 00:00:00 2001
From: Rob Bradford <rbradford@meta.com>
Date: Wed, 10 Jun 2026 12:27:06 +0100
Subject: [PATCH] virtio-devices: vhost_user: Save the dirty log before
 shutdown

Query and save the dirty log before shutting down the vhost-user device.
This allows any final dirty memory ranges to be recorded before it
becomes impossible to do that as the vhost-user handle has been closed.
This is required to ensure that all memory writes have been correctly
recorded that may be triggered by inflight I/O drains from vhost-user
device state capture.

One small implementation wrinkle: with local migrations there is no
dirty logging (since we just pass the memory FD over the socket) so
calling dirty_log() would generate an error. As there is no clean way to
query if dirty logging has been started add a boolean to track if its
active.

Signed-off-by: Rob Bradford <rbradford@meta.com>
---
 virtio-devices/src/vhost_user/mod.rs | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/virtio-devices/src/vhost_user/mod.rs b/virtio-devices/src/vhost_user/mod.rs
index 957ecc7c9..1974ca80d 100644
--- a/virtio-devices/src/vhost_user/mod.rs
+++ b/virtio-devices/src/vhost_user/mod.rs
@@ -468,6 +468,8 @@ pub struct VhostUserCommon {
     pub vring_bases: Option<Vec<u64>>,
     /// Indicates that the backend is no longer reachable. Shared with EPollHandler.
     pub disconnected: Arc<AtomicBool>,
+    saved_dirty_log: Option<MemoryRangeTable>,
+    dirty_logging: bool,
 }
 
 impl VhostUserCommon {
@@ -776,6 +778,10 @@ impl VhostUserCommon {
         let snapshot = Snapshot::new_from_state(state)?;
 
         if self.migration_started {
+            // Local migration does not enable dirty logging.
+            if self.dirty_logging {
+                self.saved_dirty_log = Some(self.dirty_log()?);
+            }
             self.shutdown();
         }
 
@@ -793,7 +799,9 @@ impl VhostUserCommon {
                         MigratableError::StartDirtyLog(anyhow!(
                             "Error starting migration for vhost-user backend: {e:?}"
                         ))
-                    })
+                    })?;
+                self.dirty_logging = true;
+                Ok(())
             } else {
                 Err(MigratableError::StartDirtyLog(anyhow!(
                     "Missing guest memory"
@@ -810,10 +818,11 @@ impl VhostUserCommon {
                 MigratableError::StopDirtyLog(anyhow!(
                     "Error stopping migration for vhost-user backend: {e:?}"
                 ))
-            })
-        } else {
-            Ok(())
+            })?;
         }
+
+        self.dirty_logging = false;
+        Ok(())
     }
 
     pub fn dirty_log(&mut self) -> std::result::Result<MemoryRangeTable, MigratableError> {
@@ -829,7 +838,7 @@ impl VhostUserCommon {
                 Err(MigratableError::DirtyLog(anyhow!("Missing guest memory")))
             }
         } else {
-            Ok(MemoryRangeTable::default())
+            Ok(self.saved_dirty_log.take().unwrap_or_default())
         }
     }
 
@@ -840,6 +849,7 @@ impl VhostUserCommon {
 
     pub fn complete_migration(&mut self) -> std::result::Result<(), MigratableError> {
         self.migration_started = false;
+        self.dirty_logging = false;
 
         // Make sure the device thread is killed in order to prevent from
         // reconnections to the socket.