[Qemu-devel] [PATCH v7 6/9] qemu-coroutine-sleep: introduce qemu_co_sleep_wake

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

Introduce a function to gracefully wake a coroutine sleeping in
qemu_co_sleep_ns().

Signed-off-by: Vladimir Sementsov-Ogievskiy <vsementsov@virtuozzo.com>
---
 include/qemu/coroutine.h    | 17 ++++++++++++--
 block/null.c                |  2 +-
 block/sheepdog.c            |  2 +-
 tests/test-bdrv-drain.c     |  6 ++---
 tests/test-block-iothread.c |  2 +-
 util/qemu-coroutine-sleep.c | 47 +++++++++++++++++++++++++++----------
 6 files changed, 55 insertions(+), 21 deletions(-)

diff --git a/include/qemu/coroutine.h b/include/qemu/coroutine.h
index 9801e7f5a4..96780a4902 100644
--- a/include/qemu/coroutine.h
+++ b/include/qemu/coroutine.h
@@ -274,9 +274,22 @@ void qemu_co_rwlock_wrlock(CoRwlock *lock);
 void qemu_co_rwlock_unlock(CoRwlock *lock);
 
 /**
- * Yield the coroutine for a given duration
+ * Yield the coroutine for a given duration. During this yield @sleep_state (if
+ * not NULL) is set to opaque pointer, which may be used for
+ * qemu_co_sleep_wake(). Be careful, the pointer is set back to zero when timer
+ * shoots. Don't save obtained value to other variables and don't call
+ * qemu_co_sleep_wake from another aio context.
  */
-void coroutine_fn qemu_co_sleep_ns(QEMUClockType type, int64_t ns);
+void coroutine_fn qemu_co_sleep_ns(QEMUClockType type, int64_t ns,
+                                   void **sleep_state);
+
+/**
+ * Wake a coroutine if it is sleeping by qemu_co_sleep_ns. Timer will be
+ * deleted. @sleep_state must be the variable which address was given to
+ * qemu_co_sleep_ns() and should be checked to be non-NULL before calling
+ * qemu_co_sleep_wake().
+ */
+void qemu_co_sleep_wake(void *sleep_state);
 
 /**
  * Yield until a file descriptor becomes readable
diff --git a/block/null.c b/block/null.c
index 699aa295cb..1e3f26b07e 100644
--- a/block/null.c
+++ b/block/null.c
@@ -109,7 +109,7 @@ static coroutine_fn int null_co_common(BlockDriverState *bs)
     BDRVNullState *s = bs->opaque;
 
     if (s->latency_ns) {
-        qemu_co_sleep_ns(QEMU_CLOCK_REALTIME, s->latency_ns);
+        qemu_co_sleep_ns(QEMU_CLOCK_REALTIME, s->latency_ns, NULL);
     }
     return 0;
 }
diff --git a/block/sheepdog.c b/block/sheepdog.c
index 6f402e5d4d..e7cab41c95 100644
--- a/block/sheepdog.c
+++ b/block/sheepdog.c
@@ -742,7 +742,7 @@ static coroutine_fn void reconnect_to_sdog(void *opaque)
         if (s->fd < 0) {
             trace_sheepdog_reconnect_to_sdog();
             error_report_err(local_err);
-            qemu_co_sleep_ns(QEMU_CLOCK_REALTIME, 1000000000ULL);
+            qemu_co_sleep_ns(QEMU_CLOCK_REALTIME, 1000000000ULL, NULL);
         }
     };
 
diff --git a/tests/test-bdrv-drain.c b/tests/test-bdrv-drain.c
index 12e2ecf517..aabe9e4264 100644
--- a/tests/test-bdrv-drain.c
+++ b/tests/test-bdrv-drain.c
@@ -42,7 +42,7 @@ static void coroutine_fn bdrv_test_co_drain_begin(BlockDriverState *bs)
     BDRVTestState *s = bs->opaque;
     s->drain_count++;
     if (s->sleep_in_drain_begin) {
-        qemu_co_sleep_ns(QEMU_CLOCK_REALTIME, 100000);
+        qemu_co_sleep_ns(QEMU_CLOCK_REALTIME, 100000, NULL);
     }
 }
 
@@ -73,7 +73,7 @@ static int coroutine_fn bdrv_test_co_preadv(BlockDriverState *bs,
      * it to complete. We need to sleep a while as bdrv_drain_invoke() comes
      * first and polls its result, too, but it shouldn't accidentally complete
      * this request yet. */
-    qemu_co_sleep_ns(QEMU_CLOCK_REALTIME, 100000);
+    qemu_co_sleep_ns(QEMU_CLOCK_REALTIME, 100000, NULL);
 
     if (s->bh_indirection_ctx) {
         aio_bh_schedule_oneshot(s->bh_indirection_ctx, co_reenter_bh,
@@ -818,7 +818,7 @@ static int coroutine_fn test_job_run(Job *job, Error **errp)
         /* Avoid job_sleep_ns() because it marks the job as !busy. We want to
          * emulate some actual activity (probably some I/O) here so that drain
          * has to wait for this activity to stop. */
-        qemu_co_sleep_ns(QEMU_CLOCK_REALTIME, 1000000);
+        qemu_co_sleep_ns(QEMU_CLOCK_REALTIME, 1000000, NULL);
 
         job_pause_point(&s->common.job);
     }
diff --git a/tests/test-block-iothread.c b/tests/test-block-iothread.c
index 79d9cf8a57..e0e871d53a 100644
--- a/tests/test-block-iothread.c
+++ b/tests/test-block-iothread.c
@@ -380,7 +380,7 @@ static int coroutine_fn test_job_run(Job *job, Error **errp)
          * emulate some actual activity (probably some I/O) here so that the
          * drain involved in AioContext switches has to wait for this activity
          * to stop. */
-        qemu_co_sleep_ns(QEMU_CLOCK_REALTIME, 1000000);
+        qemu_co_sleep_ns(QEMU_CLOCK_REALTIME, 1000000, NULL);
 
         job_pause_point(&s->common.job);
     }
diff --git a/util/qemu-coroutine-sleep.c b/util/qemu-coroutine-sleep.c
index 4bfdd30cbf..48a64bb8d8 100644
--- a/util/qemu-coroutine-sleep.c
+++ b/util/qemu-coroutine-sleep.c
@@ -17,31 +17,52 @@
 #include "qemu/timer.h"
 #include "block/aio.h"
 
-static void co_sleep_cb(void *opaque)
-{
-    Coroutine *co = opaque;
+const char *qemu_co_sleep_ns__scheduled = "qemu_co_sleep_ns";
+
+typedef struct QemuCoSleepState {
+    Coroutine *co;
+    QEMUTimer *ts;
+    void **user_state_pointer;
+} QemuCoSleepState;
 
+void qemu_co_sleep_wake(void *sleep_state)
+{
+    QemuCoSleepState *s = (QemuCoSleepState *)sleep_state;
     /* Write of schedule protected by barrier write in aio_co_schedule */
-    atomic_set(&co->scheduled, NULL);
-    aio_co_wake(co);
+    const char *scheduled = atomic_cmpxchg(&s->co->scheduled,
+                                           qemu_co_sleep_ns__scheduled, NULL);
+
+    assert(scheduled == qemu_co_sleep_ns__scheduled);
+    if (s->user_state_pointer) {
+        *s->user_state_pointer = NULL;
+    }
+    timer_del(s->ts);
+    aio_co_wake(s->co);
 }
 
-void coroutine_fn qemu_co_sleep_ns(QEMUClockType type, int64_t ns)
+void coroutine_fn qemu_co_sleep_ns(QEMUClockType type, int64_t ns,
+                                   void **sleep_state)
 {
     AioContext *ctx = qemu_get_current_aio_context();
-    QEMUTimer *ts;
-    Coroutine *co = qemu_coroutine_self();
+    QemuCoSleepState state = {
+        .co = qemu_coroutine_self(),
+        .ts = aio_timer_new(ctx, type, SCALE_NS, qemu_co_sleep_wake, &state),
+        .user_state_pointer = sleep_state,
+    };
 
-    const char *scheduled = atomic_cmpxchg(&co->scheduled, NULL, __func__);
+    const char *scheduled = atomic_cmpxchg(&state.co->scheduled, NULL,
+                                           qemu_co_sleep_ns__scheduled);
     if (scheduled) {
         fprintf(stderr,
                 "%s: Co-routine was already scheduled in '%s'\n",
                 __func__, scheduled);
         abort();
     }
-    ts = aio_timer_new(ctx, type, SCALE_NS, co_sleep_cb, co);
-    timer_mod(ts, qemu_clock_get_ns(type) + ns);
+
+    if (sleep_state) {
+        *sleep_state = &state;
+    }
+    timer_mod(state.ts, qemu_clock_get_ns(type) + ns);
     qemu_coroutine_yield();
-    timer_del(ts);
-    timer_free(ts);
+    timer_free(state.ts);
 }
-- 
2.18.0