vmm: Add a final memory pass after capturing snapshot

During migration send one final set of changed memory after capturing
the snapshot/state. This captures any memory changed as a side effect of
capturing that state. In particular with vhost-user capturing the device
state can lead to inflight requests being drained/flushed which could
change memory. As this is related to the snapshot account for this
memory transfer in the snapshot metrics.

No equivalent change is needed for snapshot as the memory is written
after the state is snapshotted.

Signed-off-by: Rob Bradford <rbradford@meta.com>

vmm/src/lib.rs

@@ -1527,8 +1527,22 @@ impl Vmm {
         vm.release_disk_locks()
             .map_err(|e| MigratableError::UnlockError(anyhow!("{e}")))?;
 
-        // Capture snapshot and send it
-        let (vm_snapshot, snapshot_duration) = measure_ok(|| vm.snapshot())?;
+        let (vm_snapshot, snapshot_duration) = measure_ok(|| {
+            // Capture snapshot. This may have side effects, e.g. vhost-user backend inflight drain
+            let snapshot = vm.snapshot()?;
+
+            // One final memory iteration to handle side effects from snapshot.
+            if !send_data_migration.local {
+                let memory_ranges = vm.dirty_log()?;
+                migration_transport::send_memory_ranges(
+                    &vm.guest_memory(),
+                    &memory_ranges,
+                    &mut socket,
+                )?;
+            }
+            Ok(snapshot)
+        })?;
+
         let (_, send_snapshot_duration) =
             measure_ok(|| migration_transport::send_state(&mut socket, &vm_snapshot))?;