[PATCH 02/11] block-coroutine-wrapper: support BlockBackend first argument

[Vladimir Sementsov-Ogievskiy <vsementsov@virtuozzo.com>]

We'll need to wrap functions with first argument of BlockBackend *
type. For this let's generalize core function and struct to work with
pure AioContext.

Signed-off-by: Vladimir Sementsov-Ogievskiy <vsementsov@virtuozzo.com>
---
 block/block-gen.h                  | 12 ++++++------
 scripts/block-coroutine-wrapper.py | 23 ++++++++++++++++++-----
 2 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/block/block-gen.h b/block/block-gen.h
index f80cf4897d..c1fd3f40de 100644
--- a/block/block-gen.h
+++ b/block/block-gen.h
@@ -29,19 +29,19 @@
 #include "block/block_int.h"
 
 /* Base structure for argument packing structures */
-typedef struct BdrvPollCo {
-    BlockDriverState *bs;
+typedef struct AioPollCo {
+    AioContext *ctx;
     bool in_progress;
     int ret;
     Coroutine *co; /* Keep pointer here for debugging */
-} BdrvPollCo;
+} AioPollCo;
 
-static inline int bdrv_poll_co(BdrvPollCo *s)
+static inline int aio_poll_co(AioPollCo *s)
 {
     assert(!qemu_in_coroutine());
 
-    bdrv_coroutine_enter(s->bs, s->co);
-    BDRV_POLL_WHILE(s->bs, s->in_progress);
+    aio_co_enter(s->ctx, s->co);
+    AIO_WAIT_WHILE(s->ctx, s->in_progress);
 
     return s->ret;
 }
diff --git a/scripts/block-coroutine-wrapper.py b/scripts/block-coroutine-wrapper.py
index 85dbeb9ecf..114a54fcce 100644
--- a/scripts/block-coroutine-wrapper.py
+++ b/scripts/block-coroutine-wrapper.py
@@ -42,6 +42,8 @@ def gen_header():
 #include "qemu/osdep.h"
 #include "block/coroutines.h"
 #include "block/block-gen.h"
+#include "qemu-io.h"
+#include "sysemu/block-backend.h"
 #include "block/block_int.h"\
 """
 
@@ -100,12 +102,23 @@ def snake_to_camel(func_name: str) -> str:
 def gen_wrapper(func: FuncDecl) -> str:
     assert not '_co_' in func.name
     assert func.return_type == 'int'
-    assert func.args[0].type in ['BlockDriverState *', 'BdrvChild *']
+    assert func.args[0].type in ['BlockDriverState *', 'BdrvChild *',
+                                 'BlockBackend *']
 
     subsystem, subname = func.name.split('_', 1)
 
     name = f'{subsystem}_co_{subname}'
-    bs = 'bs' if func.args[0].type == 'BlockDriverState *' else 'child->bs'
+
+    first_arg_type = func.args[0].type
+    if first_arg_type == 'BlockDriverState *':
+        ctx = 'bdrv_get_aio_context(bs)'
+    elif first_arg_type == 'BdrvChild *':
+        ctx = '(child ? bdrv_get_aio_context(child->bs) : ' \
+            'qemu_get_aio_context())'
+    else:
+        assert first_arg_type == 'BlockBackend *'
+        ctx = '(blk ? blk_get_aio_context(blk) : qemu_get_aio_context())'
+
     struct_name = snake_to_camel(name)
 
     return f"""\
@@ -114,7 +127,7 @@ def gen_wrapper(func: FuncDecl) -> str:
  */
 
 typedef struct {struct_name} {{
-    BdrvPollCo poll_state;
+    AioPollCo poll_state;
 { func.gen_block('    {decl};') }
 }} {struct_name};
 
@@ -134,7 +147,7 @@ def gen_wrapper(func: FuncDecl) -> str:
         return {name}({ func.gen_list('{name}') });
     }} else {{
         {struct_name} s = {{
-            .poll_state.bs = {bs},
+            .poll_state.ctx = {ctx},
             .poll_state.in_progress = true,
 
 { func.gen_block('            .{name} = {name},') }
@@ -142,7 +155,7 @@ def gen_wrapper(func: FuncDecl) -> str:
 
         s.poll_state.co = qemu_coroutine_create({name}_entry, &s);
 
-        return bdrv_poll_co(&s.poll_state);
+        return aio_poll_co(&s.poll_state);
     }}
 }}"""
 
-- 
2.29.2