nbd: Increase bs->in_flight during AioContext switch

bdrv_drain() must not leave connection_co scheduled, so bs->in_flight
needs to be increased while the coroutine is waiting to be scheduled
in the new AioContext after nbd_client_attach_aio_context().

Signed-off-by: Kevin Wolf <kwolf@redhat.com>

1 files changed, 18 insertions, 2 deletions

diff --git a/block/nbd-client.c b/block/nbd-client.c
index 60f38f0320..bfbaf7ebe9 100644
--- a/block/nbd-client.c
+++ b/block/nbd-client.c
@@ -977,14 +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);
 
-    /* FIXME Really need a bdrv_inc_in_flight() here */
-    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)