[PATCH 15/15] block: only get out of coroutine context for polling

[Paolo Bonzini <pbonzini@redhat.com>]

The drained_begin callbacks are not polling, all they do is schedule
any operations that must complete before bdrv_drained_end().  As such,
they can be called before bdrv_co_yield_to_drain().  Thus, the only
remaining task left for bdrv_co_drain_bh_cb() is the BDRV_POLL_WHILE()
loop.  This patch extracts it into two new functions bdrv_drain_wait()
and bdrv_drain_all_wait(), so that it is easily accessible from
bdrv_co_drain_bh_cb().

Signed-off-by: Paolo Bonzini <pbonzini@redhat.com>
---
 block/io.c | 77 +++++++++++++++++++++++++++---------------------------
 1 file changed, 38 insertions(+), 39 deletions(-)

diff --git a/block/io.c b/block/io.c
index a75f42ee13cb..fc4bc38ee9d5 100644
--- a/block/io.c
+++ b/block/io.c
@@ -245,13 +245,8 @@ bool bdrv_drain_poll(BlockDriverState *bs, bool ignore_bds_parents)
     return false;
 }
 
-static bool bdrv_drain_poll_top_level(BlockDriverState *bs)
-{
-    return bdrv_drain_poll(bs, false);
-}
-
-static void bdrv_do_drained_begin(BlockDriverState *bs, bool poll);
-
+static void bdrv_drain_wait(BlockDriverState *bs);
+static void bdrv_drain_all_wait(void);
 static void bdrv_co_drain_bh_cb(void *opaque)
 {
     BdrvCoDrainData *data = opaque;
@@ -262,18 +257,17 @@ static void bdrv_co_drain_bh_cb(void *opaque)
         AioContext *ctx = bdrv_get_aio_context(bs);
         aio_context_acquire(ctx);
         bdrv_dec_in_flight(bs);
-        bdrv_do_drained_begin(bs, data->poll);
+        bdrv_drain_wait(bs);
         aio_context_release(ctx);
     } else {
-        bdrv_drain_all_begin();
+        bdrv_drain_all_wait();
     }
 
     data->done = true;
     aio_co_wake(co);
 }
 
-static void coroutine_fn bdrv_co_yield_to_drain(BlockDriverState *bs,
-                                                bool poll)
+static void coroutine_fn bdrv_co_yield_to_drain(BlockDriverState *bs)
 {
     BdrvCoDrainData data;
     Coroutine *self = qemu_coroutine_self();
@@ -288,7 +282,6 @@ static void coroutine_fn bdrv_co_yield_to_drain(BlockDriverState *bs,
         .co = self,
         .bs = bs,
         .done = false,
-        .poll = poll,
     };
 
     if (bs) {
@@ -319,15 +312,10 @@ static void coroutine_fn bdrv_co_yield_to_drain(BlockDriverState *bs,
     }
 }
 
-static void bdrv_do_drained_begin(BlockDriverState *bs, bool poll)
+void bdrv_do_drained_begin_quiesce(BlockDriverState *bs)
 {
     IO_OR_GS_CODE();
 
-    if (qemu_in_coroutine()) {
-        bdrv_co_yield_to_drain(bs, poll);
-        return;
-    }
-
     /* Stop things in parent-to-child order */
     if (qatomic_fetch_inc(&bs->quiesce_counter) == 0) {
         aio_disable_external(bdrv_get_aio_context(bs));
@@ -336,6 +324,23 @@ static void bdrv_do_drained_begin(BlockDriverState *bs, bool poll)
             bs->drv->bdrv_drain_begin(bs);
         }
     }
+}
+
+void bdrv_drained_begin(BlockDriverState *bs)
+{
+    IO_OR_GS_CODE();
+    bdrv_do_drained_begin_quiesce(bs);
+    bdrv_drain_wait(bs);
+}
+
+static void bdrv_drain_wait(BlockDriverState *bs)
+{
+    IO_OR_GS_CODE();
+
+    if (qemu_in_coroutine()) {
+        bdrv_co_yield_to_drain(bs);
+        return;
+    }
 
     /*
      * Wait for drained requests to finish.
@@ -346,20 +351,7 @@ static void bdrv_do_drained_begin(BlockDriverState *bs, bool poll)
      * includes other nodes in the same AioContext and therefore all child
      * nodes.
      */
-    if (poll) {
-        BDRV_POLL_WHILE(bs, bdrv_drain_poll_top_level(bs));
-    }
-}
-
-void bdrv_do_drained_begin_quiesce(BlockDriverState *bs)
-{
-    bdrv_do_drained_begin(bs, false);
-}
-
-void bdrv_drained_begin(BlockDriverState *bs)
-{
-    IO_OR_GS_CODE();
-    bdrv_do_drained_begin(bs, true);
+    BDRV_POLL_WHILE(bs, bdrv_drain_poll(bs, false));
 }
 
 /**
@@ -438,11 +430,6 @@ void bdrv_drain_all_begin(void)
     BlockDriverState *bs = NULL;
     GLOBAL_STATE_CODE();
 
-    if (qemu_in_coroutine()) {
-        bdrv_co_yield_to_drain(NULL, true);
-        return;
-    }
-
     /*
      * bdrv queue is managed by record/replay,
      * waiting for finishing the I/O requests may
@@ -464,18 +451,30 @@ void bdrv_drain_all_begin(void)
         AioContext *aio_context = bdrv_get_aio_context(bs);
 
         aio_context_acquire(aio_context);
-        bdrv_do_drained_begin(bs, false);
+        bdrv_do_drained_begin_quiesce(bs);
         aio_context_release(aio_context);
     }
 
     /* Now poll the in-flight requests */
-    AIO_WAIT_WHILE(NULL, bdrv_drain_all_poll());
+    bdrv_drain_all_wait();
 
     while ((bs = bdrv_next_all_states(bs))) {
         bdrv_drain_assert_idle(bs);
     }
 }
 
+static void bdrv_drain_all_wait(void)
+{
+    GLOBAL_STATE_CODE();
+
+    if (qemu_in_coroutine()) {
+        bdrv_co_yield_to_drain(NULL);
+        return;
+    }
+
+    AIO_WAIT_WHILE(NULL, bdrv_drain_all_poll());
+}
+
 void bdrv_drain_all_end_quiesce(BlockDriverState *bs)
 {
     GLOBAL_STATE_CODE();
-- 
2.38.1