block: acquire AioContext in bdrv_drain_all()

Modify bdrv_drain_all() to take into account that BlockDriverState
instances may be running in different AioContexts.

This patch changes the implementation of bdrv_drain_all() while
preserving the semantics.  Previously kicking throttled requests and
checking for pending requests were done across all BlockDriverState
instances in sequence.  Now we process each BlockDriverState in turn,
making sure to acquire and release its AioContext.

This prevents race conditions between the thread executing
bdrv_drain_all() and the thread running the AioContext.

Signed-off-by: Stefan Hajnoczi <stefanha@redhat.com>

block.c

@@ -1885,17 +1885,6 @@ static bool bdrv_requests_pending(BlockDriverState *bs)
     return false;
 }
 
-static bool bdrv_requests_pending_all(void)
-{
-    BlockDriverState *bs;
-    QTAILQ_FOREACH(bs, &bdrv_states, device_list) {
-        if (bdrv_requests_pending(bs)) {
-            return true;
-        }
-    }
-    return false;
-}
-
 /*
  * Wait for pending requests to complete across all BlockDriverStates
  *
@@ -1915,12 +1904,20 @@ void bdrv_drain_all(void)
     BlockDriverState *bs;
 
     while (busy) {
+        busy = false;
+
         QTAILQ_FOREACH(bs, &bdrv_states, device_list) {
+            AioContext *aio_context = bdrv_get_aio_context(bs);
+            bool bs_busy;
+
+            aio_context_acquire(aio_context);
             bdrv_start_throttled_reqs(bs);
-        }
+            bs_busy = bdrv_requests_pending(bs);
+            bs_busy |= aio_poll(aio_context, bs_busy);
+            aio_context_release(aio_context);
 
-        busy = bdrv_requests_pending_all();
-        busy |= aio_poll(qemu_get_aio_context(), busy);
+            busy |= bs_busy;
+        }
     }
 }