[PATCH v14 09/14] vhost_user: Make source quit after reporting migration state

[David Gibson <david@gibson.dropbear.id.au>]

From: Stefano Brivio <sbrivio@redhat.com>

This will close all the sockets we currently have open in repair mode,
and completes our migration tasks as source. If the hypervisor wants
to have us back at this point, somebody needs to restart us.

Signed-off-by: Stefano Brivio <sbrivio@redhat.com>
Message-ID: <20250209222005.1640077-5-sbrivio@redhat.com>
Signed-off-by: David Gibson <david@gibson.dropbear.id.au>
---
 vhost_user.c | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/vhost_user.c b/vhost_user.c
index 256c8ab6..9115fb50 100644
--- a/vhost_user.c
+++ b/vhost_user.c
@@ -998,6 +998,9 @@ static bool vu_send_rarp_exec(struct vu_dev *vdev,
 	return false;
 }
 
+/* If set, quit when we get a VHOST_USER_CHECK_DEVICE_STATE, after replying */
+static bool quit_on_device_state = false;
+
 /**
  * vu_set_device_state_fd_exec() - Set the device state migration channel
  * @vdev:	vhost-user device
@@ -1025,6 +1028,9 @@ static bool vu_set_device_state_fd_exec(struct vu_dev *vdev,
 	migrate_request(vdev->context, msg->fds[0],
 			direction == VHOST_USER_TRANSFER_STATE_DIRECTION_LOAD);
 
+	if (direction == VHOST_USER_TRANSFER_STATE_DIRECTION_SAVE)
+		quit_on_device_state = true;
+
 	/* We don't provide a new fd for the data transfer */
 	vmsg_set_reply_u64(msg, VHOST_USER_VRING_NOFD_MASK);
 
@@ -1203,4 +1209,10 @@ void vu_control_handler(struct vu_dev *vdev, int fd, uint32_t events)
 
 	if (reply_requested)
 		vu_send_reply(fd, &msg);
+
+	if (quit_on_device_state &&
+	    msg.hdr.request == VHOST_USER_CHECK_DEVICE_STATE) {
+		info("Migration complete, exiting");
+		_exit(EXIT_SUCCESS);
+	}
 }
-- 
@@ -998,6 +998,9 @@ static bool vu_send_rarp_exec(struct vu_dev *vdev,
 	return false;
 }
 
+/* If set, quit when we get a VHOST_USER_CHECK_DEVICE_STATE, after replying */
+static bool quit_on_device_state = false;
+
 /**
  * vu_set_device_state_fd_exec() - Set the device state migration channel
  * @vdev:	vhost-user device
@@ -1025,6 +1028,9 @@ static bool vu_set_device_state_fd_exec(struct vu_dev *vdev,
 	migrate_request(vdev->context, msg->fds[0],
 			direction == VHOST_USER_TRANSFER_STATE_DIRECTION_LOAD);
 
+	if (direction == VHOST_USER_TRANSFER_STATE_DIRECTION_SAVE)
+		quit_on_device_state = true;
+
 	/* We don't provide a new fd for the data transfer */
 	vmsg_set_reply_u64(msg, VHOST_USER_VRING_NOFD_MASK);
 
@@ -1203,4 +1209,10 @@ void vu_control_handler(struct vu_dev *vdev, int fd, uint32_t events)
 
 	if (reply_requested)
 		vu_send_reply(fd, &msg);
+
+	if (quit_on_device_state &&
+	    msg.hdr.request == VHOST_USER_CHECK_DEVICE_STATE) {
+		info("Migration complete, exiting");
+		_exit(EXIT_SUCCESS);
+	}
 }
-- 
2.48.1