Add a function to move the current coroutine to the AioContext of a
given BlockDriverState.

Signed-off-by: Kevin Wolf <kwolf@redhat.com>
---
 include/block/block.h |  6 ++++++
 block.c               | 10 ++++++++++
 2 files changed, 16 insertions(+)

diff --git a/include/block/block.h b/include/block/block.h
index 25e299605e..fb77062a1e 100644
--- a/include/block/block.h
+++ b/include/block/block.h
@@ -631,6 +631,12 @@ bool bdrv_debug_is_suspended(BlockDriverState *bs, const char *tag);
  */
 AioContext *bdrv_get_aio_context(BlockDriverState *bs);
 
+/**
+ * Move the current coroutine to the AioContext of @bs and return the old
+ * AioContext of the coroutine.
+ */
+AioContext *coroutine_fn bdrv_co_move_to_aio_context(BlockDriverState *bs);
+
 /**
  * Transfer control to @co in the aio context of @bs
  */
diff --git a/block.c b/block.c
index 8416376c9b..cd576b991c 100644
--- a/block.c
+++ b/block.c
@@ -6215,6 +6215,16 @@ AioContext *bdrv_get_aio_context(BlockDriverState *bs)
     return bs ? bs->aio_context : qemu_get_aio_context();
 }
 
+AioContext *coroutine_fn bdrv_co_move_to_aio_context(BlockDriverState *bs)
+{
+    Coroutine *self = qemu_coroutine_self();
+    AioContext *old_ctx = qemu_coroutine_get_aio_context(self);
+    AioContext *new_ctx = bdrv_get_aio_context(bs);
+
+    aio_co_reschedule_self(new_ctx);
+    return old_ctx;
+}
+
 void bdrv_coroutine_enter(BlockDriverState *bs, Coroutine *co)
 {
     aio_co_enter(bdrv_get_aio_context(bs), co);
-- 
2.25.4