1 files changed, 28 insertions, 5 deletions

diff --git a/util/vhost-user-server.c b/util/vhost-user-server.c
index 783d847a6d..f66fbba710 100644
--- a/util/vhost-user-server.c
+++ b/util/vhost-user-server.c
@@ -74,6 +74,20 @@ static void panic_cb(VuDev *vu_dev, const char *buf)
     error_report("vu_panic: %s", buf);
 }
 
+void vhost_user_server_ref(VuServer *server)
+{
+    assert(!server->wait_idle);
+    server->refcount++;
+}
+
+void vhost_user_server_unref(VuServer *server)
+{
+    server->refcount--;
+    if (server->wait_idle && !server->refcount) {
+        aio_co_wake(server->co_trip);
+    }
+}
+
 static bool coroutine_fn
 vu_message_read(VuDev *vu_dev, int conn_fd, VhostUserMsg *vmsg)
 {
@@ -177,6 +191,14 @@ static coroutine_fn void vu_client_trip(void *opaque)
         /* Keep running */
     }
 
+    if (server->refcount) {
+        /* Wait for requests to complete before we can unmap the memory */
+        server->wait_idle = true;
+        qemu_coroutine_yield();
+        server->wait_idle = false;
+    }
+    assert(server->refcount == 0);
+
     vu_deinit(vu_dev);
 
     /* vu_deinit() should have called remove_watch() */
@@ -250,7 +272,7 @@ set_watch(VuDev *vu_dev, int fd, int vu_evt,
         vu_fd_watch->cb = cb;
         qemu_set_nonblock(fd);
         aio_set_fd_handler(server->ioc->ctx, fd, true, kick_handler,
-                           NULL, NULL, vu_fd_watch);
+                           NULL, NULL, NULL, vu_fd_watch);
         vu_fd_watch->vu_dev = vu_dev;
         vu_fd_watch->pvt = pvt;
     }
@@ -270,7 +292,8 @@ static void remove_watch(VuDev *vu_dev, int fd)
     if (!vu_fd_watch) {
         return;
     }
-    aio_set_fd_handler(server->ioc->ctx, fd, true, NULL, NULL, NULL, NULL);
+    aio_set_fd_handler(server->ioc->ctx, fd, true,
+                       NULL, NULL, NULL, NULL, NULL);
 
     QTAILQ_REMOVE(&server->vu_fd_watches, vu_fd_watch, next);
     g_free(vu_fd_watch);
@@ -334,7 +357,7 @@ void vhost_user_server_stop(VuServer *server)
 
         QTAILQ_FOREACH(vu_fd_watch, &server->vu_fd_watches, next) {
             aio_set_fd_handler(server->ctx, vu_fd_watch->fd, true,
-                               NULL, NULL, NULL, vu_fd_watch);
+                               NULL, NULL, NULL, NULL, vu_fd_watch);
         }
 
         qio_channel_shutdown(server->ioc, QIO_CHANNEL_SHUTDOWN_BOTH, NULL);
@@ -377,7 +400,7 @@ void vhost_user_server_attach_aio_context(VuServer *server, AioContext *ctx)
 
     QTAILQ_FOREACH(vu_fd_watch, &server->vu_fd_watches, next) {
         aio_set_fd_handler(ctx, vu_fd_watch->fd, true, kick_handler, NULL,
-                           NULL, vu_fd_watch);
+                           NULL, NULL, vu_fd_watch);
     }
 
     aio_co_schedule(ctx, server->co_trip);
@@ -391,7 +414,7 @@ void vhost_user_server_detach_aio_context(VuServer *server)
 
         QTAILQ_FOREACH(vu_fd_watch, &server->vu_fd_watches, next) {
             aio_set_fd_handler(server->ctx, vu_fd_watch->fd, true,
-                               NULL, NULL, NULL, vu_fd_watch);
+                               NULL, NULL, NULL, NULL, vu_fd_watch);
         }
 
         qio_channel_detach_aio_context(server->ioc);