1 files changed, 34 insertions, 2 deletions

diff --git a/block/nbd-client.c b/block/nbd-client.c
index f0ad54ce21..bfbaf7ebe9 100644
--- a/block/nbd-client.c
+++ b/block/nbd-client.c
@@ -76,8 +76,18 @@ static coroutine_fn void nbd_connection_entry(void *opaque)
     Error *local_err = NULL;
 
     while (!s->quit) {
+        /*
+         * The NBD client can only really be considered idle when it has
+         * yielded from qio_channel_readv_all_eof(), waiting for data. This is
+         * the point where the additional scheduled coroutine entry happens
+         * after nbd_client_attach_aio_context().
+         *
+         * Therefore we keep an additional in_flight reference all the time and
+         * only drop it temporarily here.
+         */
         assert(s->reply.handle == 0);
-        ret = nbd_receive_reply(s->ioc, &s->reply, &local_err);
+        ret = nbd_receive_reply(s->bs, s->ioc, &s->reply, &local_err);
+
         if (local_err) {
             trace_nbd_read_reply_entry_fail(ret, error_get_pretty(local_err));
             error_free(local_err);
@@ -116,6 +126,8 @@ static coroutine_fn void nbd_connection_entry(void *opaque)
 
     s->quit = true;
     nbd_recv_coroutines_wake_all(s);
+    bdrv_dec_in_flight(s->bs);
+
     s->connection_co = NULL;
     aio_wait_kick();
 }
@@ -965,12 +977,30 @@ void nbd_client_detach_aio_context(BlockDriverState *bs)
     qio_channel_detach_aio_context(QIO_CHANNEL(client->ioc));
 }
 
+static void nbd_client_attach_aio_context_bh(void *opaque)
+{
+    BlockDriverState *bs = opaque;
+    NBDClientSession *client = nbd_get_client_session(bs);
+
+    /* The node is still drained, so we know the coroutine has yielded in
+     * nbd_read_eof(), the only place where bs->in_flight can reach 0, or it is
+     * entered for the first time. Both places are safe for entering the
+     * coroutine.*/
+    qemu_aio_coroutine_enter(bs->aio_context, client->connection_co);
+    bdrv_dec_in_flight(bs);
+}
+
 void nbd_client_attach_aio_context(BlockDriverState *bs,
                                    AioContext *new_context)
 {
     NBDClientSession *client = nbd_get_client_session(bs);
     qio_channel_attach_aio_context(QIO_CHANNEL(client->ioc), new_context);
-    aio_co_schedule(new_context, client->connection_co);
+
+    bdrv_inc_in_flight(bs);
+
+    /* Need to wait here for the BH to run because the BH must run while the
+     * node is still drained. */
+    aio_wait_bh_oneshot(new_context, nbd_client_attach_aio_context_bh, bs);
 }
 
 void nbd_client_close(BlockDriverState *bs)
@@ -1076,6 +1106,7 @@ static int nbd_client_connect(BlockDriverState *bs,
      * kick the reply mechanism.  */
     qio_channel_set_blocking(QIO_CHANNEL(sioc), false, NULL);
     client->connection_co = qemu_coroutine_create(nbd_connection_entry, client);
+    bdrv_inc_in_flight(bs);
     nbd_client_attach_aio_context(bs, bdrv_get_aio_context(bs));
 
     logout("Established connection with NBD server\n");
@@ -1108,6 +1139,7 @@ int nbd_client_init(BlockDriverState *bs,
 {
     NBDClientSession *client = nbd_get_client_session(bs);
 
+    client->bs = bs;
     qemu_co_mutex_init(&client->send_mutex);
     qemu_co_queue_init(&client->free_sema);