fix(vsock): save state after sending a notification

src/vmm/src/device_manager/persist.rs

@@ -365,11 +365,6 @@ impl<'a> Persist<'a> for MMIODeviceManager {
                         .downcast_mut::<Vsock<VsockUnixBackend>>()
                         .unwrap();
 
-                    let vsock_state = VsockState {
-                        backend: vsock.backend().save(),
-                        frontend: vsock.save(),
-                    };
-
                     // Send Transport event to reset connections if device
                     // is activated.
                     if vsock.is_activated() {
@@ -378,6 +373,13 @@ impl<'a> Persist<'a> for MMIODeviceManager {
                         });
                     }
 
+                    // Save state after potential notification to the guest. This
+                    // way we save changes to the queue the notification can cause.
+                    let vsock_state = VsockState {
+                        backend: vsock.backend().save(),
+                        frontend: vsock.save(),
+                    };
+
                     states.vsock_device = Some(ConnectedVsockState {
                         device_id: devid.clone(),
                         device_state: vsock_state,

tests/integration_tests/functional/test_vsock.py

@@ -225,54 +225,61 @@ def test_vsock_transport_reset_g2h(uvm_nano, microvm_factory):
     test_vm.api.vsock.put(vsock_id="vsock0", guest_cid=3, uds_path=f"/{VSOCK_UDS_PATH}")
     test_vm.start()
 
-    host_socket_path = os.path.join(
-        test_vm.path, f"{VSOCK_UDS_PATH}_{ECHO_SERVER_PORT}"
-    )
-    host_socat_commmand = [
-        "socat",
-        "-dddd",
-        f"UNIX-LISTEN:{host_socket_path},fork",
-        "STDOUT",
-    ]
-    host_socat = subprocess.Popen(
-        host_socat_commmand, stdout=subprocess.PIPE, stderr=subprocess.PIPE
-    )
-
-    # Give some time for host socat to create socket
-    time.sleep(0.5)
-    assert Path(host_socket_path).exists()
-    test_vm.create_jailed_resource(host_socket_path)
-
-    # Create a socat process in the guest which will connect to the host socat
-    guest_socat_commmand = f"tmux new -d 'socat - vsock-connect:2:{ECHO_SERVER_PORT}'"
-    test_vm.ssh.run(guest_socat_commmand)
-
-    # socat should be running in the guest now
-    code, _, _ = test_vm.ssh.run("pidof socat")
-    assert code == 0
-
-    # Create snapshot.
+    # Create snapshot and terminate a VM.
     snapshot = test_vm.snapshot_full()
-    test_vm.resume()
-
-    # After `create_snapshot` + 'restore' calls, connection should be dropped
-    code, _, _ = test_vm.ssh.run("pidof socat")
-    assert code == 1
-
-    # Kill host socat as it is not useful anymore
-    host_socat.kill()
-    host_socat.communicate()
-
-    # Terminate VM.
     test_vm.kill()
 
-    # Load snapshot.
-    vm2 = microvm_factory.build()
-    vm2.spawn()
-    vm2.restore_from_snapshot(snapshot, resume=True)
-
-    # After snap restore all vsock connections should be
-    # dropped. This means guest socat should exit same way
-    # as it did after snapshot was taken.
-    code, _, _ = vm2.ssh.run("pidof socat")
-    assert code == 1
+    for _ in range(5):
+        # Load snapshot.
+        new_vm = microvm_factory.build()
+        new_vm.spawn()
+        new_vm.restore_from_snapshot(snapshot, resume=True)
+
+        # After snap restore all vsock connections should be
+        # dropped. This means guest socat should exit same way
+        # as it did after snapshot was taken.
+        code, _, _ = new_vm.ssh.run("pidof socat")
+        assert code == 1
+
+        host_socket_path = os.path.join(
+            new_vm.path, f"{VSOCK_UDS_PATH}_{ECHO_SERVER_PORT}"
+        )
+        host_socat_commmand = [
+            "socat",
+            "-dddd",
+            f"UNIX-LISTEN:{host_socket_path},fork",
+            "STDOUT",
+        ]
+        host_socat = subprocess.Popen(
+            host_socat_commmand, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        )
+
+        # Give some time for host socat to create socket
+        time.sleep(0.5)
+        assert Path(host_socket_path).exists()
+        new_vm.create_jailed_resource(host_socket_path)
+
+        # Create a socat process in the guest which will connect to the host socat
+        guest_socat_commmand = (
+            f"tmux new -d 'socat - vsock-connect:2:{ECHO_SERVER_PORT}'"
+        )
+        new_vm.ssh.run(guest_socat_commmand)
+
+        # socat should be running in the guest now
+        code, _, _ = new_vm.ssh.run("pidof socat")
+        assert code == 0
+
+        # Create snapshot.
+        snapshot = new_vm.snapshot_full()
+        new_vm.resume()
+
+        # After `create_snapshot` + 'restore' calls, connection should be dropped
+        code, _, _ = new_vm.ssh.run("pidof socat")
+        assert code == 1
+
+        # Kill host socat as it is not useful anymore
+        host_socat.kill()
+        host_socat.communicate()
+
+        # Terminate VM.
+        new_vm.kill()