throttle: add throttle_detach/attach_aio_context()

Block I/O throttling uses timers and currently always adds them to the
main loop.  Throttling will break if bdrv_set_aio_context() is used to
move a BlockDriverState to a different AioContext.

This patch adds throttle_detach/attach_aio_context() interfaces so the
throttling timers and uses them to move timers to the new AioContext.
Note that bdrv_set_aio_context() already drains all requests so we're
sure no throttled requests are pending.

The test cases need to be updated since the throttle_init() interface
has changed.

Signed-off-by: Stefan Hajnoczi <stefanha@redhat.com>
Reviewed-by: Benoit Canet <benoit@irqsave.net>

4 files changed, 60 insertions, 9 deletions

diff --git a/block.c b/block.c
index 1fd38159e2..17f763db79 100644
--- a/block.c
+++ b/block.c
@@ -179,6 +179,7 @@ void bdrv_io_limits_enable(BlockDriverState *bs)
 {
     assert(!bs->io_limits_enabled);
     throttle_init(&bs->throttle_state,
+                  bdrv_get_aio_context(bs),
                   QEMU_CLOCK_VIRTUAL,
                   bdrv_throttle_read_timer_cb,
                   bdrv_throttle_write_timer_cb,
@@ -5671,6 +5672,9 @@ void bdrv_detach_aio_context(BlockDriverState *bs)
         return;
     }
 
+    if (bs->io_limits_enabled) {
+        throttle_detach_aio_context(&bs->throttle_state);
+    }
     if (bs->drv->bdrv_detach_aio_context) {
         bs->drv->bdrv_detach_aio_context(bs);
     }
@@ -5702,6 +5706,9 @@ void bdrv_attach_aio_context(BlockDriverState *bs,
     if (bs->drv->bdrv_attach_aio_context) {
         bs->drv->bdrv_attach_aio_context(bs, new_context);
     }
+    if (bs->io_limits_enabled) {
+        throttle_attach_aio_context(&bs->throttle_state, new_context);
+    }
 }
 
 void bdrv_set_aio_context(BlockDriverState *bs, AioContext *new_context)
diff --git a/include/qemu/throttle.h b/include/qemu/throttle.h
index ab29b0b918..b890613a9c 100644
--- a/include/qemu/throttle.h
+++ b/include/qemu/throttle.h
@@ -67,6 +67,11 @@ typedef struct ThrottleState {
     int64_t previous_leak;    /* timestamp of the last leak done */
     QEMUTimer * timers[2];    /* timers used to do the throttling */
     QEMUClockType clock_type; /* the clock used */
+
+    /* Callbacks */
+    QEMUTimerCB *read_timer_cb;
+    QEMUTimerCB *write_timer_cb;
+    void *timer_opaque;
 } ThrottleState;
 
 /* operations on single leaky buckets */
@@ -82,6 +87,7 @@ bool throttle_compute_timer(ThrottleState *ts,
 
 /* init/destroy cycle */
 void throttle_init(ThrottleState *ts,
+                   AioContext *aio_context,
                    QEMUClockType clock_type,
                    void (read_timer)(void *),
                    void (write_timer)(void *),
@@ -89,6 +95,10 @@ void throttle_init(ThrottleState *ts,
 
 void throttle_destroy(ThrottleState *ts);
 
+void throttle_detach_aio_context(ThrottleState *ts);
+
+void throttle_attach_aio_context(ThrottleState *ts, AioContext *new_context);
+
 bool throttle_have_timer(ThrottleState *ts);
 
 /* configuration */
diff --git a/tests/test-throttle.c b/tests/test-throttle.c
index 1d4ffd3603..5fa5000124 100644
--- a/tests/test-throttle.c
+++ b/tests/test-throttle.c
@@ -12,8 +12,10 @@
 
 #include <glib.h>
 #include <math.h>
+#include "block/aio.h"
 #include "qemu/throttle.h"
 
+AioContext     *ctx;
 LeakyBucket    bkt;
 ThrottleConfig cfg;
 ThrottleState  ts;
@@ -104,7 +106,8 @@ static void test_init(void)
     memset(&ts, 1, sizeof(ts));
 
     /* init the structure */
-    throttle_init(&ts, QEMU_CLOCK_VIRTUAL, read_timer_cb, write_timer_cb, &ts);
+    throttle_init(&ts, ctx, QEMU_CLOCK_VIRTUAL,
+                  read_timer_cb, write_timer_cb, &ts);
 
     /* check initialized fields */
     g_assert(ts.clock_type == QEMU_CLOCK_VIRTUAL);
@@ -126,7 +129,8 @@ static void test_init(void)
 static void test_destroy(void)
 {
     int i;
-    throttle_init(&ts, QEMU_CLOCK_VIRTUAL, read_timer_cb, write_timer_cb, &ts);
+    throttle_init(&ts, ctx, QEMU_CLOCK_VIRTUAL,
+                  read_timer_cb, write_timer_cb, &ts);
     throttle_destroy(&ts);
     for (i = 0; i < 2; i++) {
         g_assert(!ts.timers[i]);
@@ -165,7 +169,8 @@ static void test_config_functions(void)
 
     orig_cfg.op_size = 1;
 
-    throttle_init(&ts, QEMU_CLOCK_VIRTUAL, read_timer_cb, write_timer_cb, &ts);
+    throttle_init(&ts, ctx, QEMU_CLOCK_VIRTUAL,
+                  read_timer_cb, write_timer_cb, &ts);
     /* structure reset by throttle_init previous_leak should be null */
     g_assert(!ts.previous_leak);
     throttle_config(&ts, &orig_cfg);
@@ -324,7 +329,8 @@ static void test_have_timer(void)
     g_assert(!throttle_have_timer(&ts));
 
     /* init the structure */
-    throttle_init(&ts, QEMU_CLOCK_VIRTUAL, read_timer_cb, write_timer_cb, &ts);
+    throttle_init(&ts, ctx, QEMU_CLOCK_VIRTUAL,
+                  read_timer_cb, write_timer_cb, &ts);
 
     /* timer set by init should return true */
     g_assert(throttle_have_timer(&ts));
@@ -357,7 +363,8 @@ static bool do_test_accounting(bool is_ops, /* are we testing bps or ops */
 
     cfg.op_size = op_size;
 
-    throttle_init(&ts, QEMU_CLOCK_VIRTUAL, read_timer_cb, write_timer_cb, &ts);
+    throttle_init(&ts, ctx, QEMU_CLOCK_VIRTUAL,
+                  read_timer_cb, write_timer_cb, &ts);
     throttle_config(&ts, &cfg);
 
     /* account a read */
@@ -461,7 +468,15 @@ static void test_accounting(void)
 
 int main(int argc, char **argv)
 {
+    GSource *src;
+
     init_clocks();
+
+    ctx = aio_context_new();
+    src = aio_get_g_source(ctx);
+    g_source_attach(src, NULL);
+    g_source_unref(src);
+
     do {} while (g_main_context_iteration(NULL, false));
 
     /* tests in the same order as the header function declarations */
diff --git a/util/throttle.c b/util/throttle.c
index 02e6f15587..f976ac7de5 100644
--- a/util/throttle.c
+++ b/util/throttle.c
@@ -22,6 +22,7 @@
 
 #include "qemu/throttle.h"
 #include "qemu/timer.h"
+#include "block/aio.h"
 
 /* This function make a bucket leak
  *
@@ -157,8 +158,18 @@ bool throttle_compute_timer(ThrottleState *ts,
     return false;
 }
 
+/* Add timers to event loop */
+void throttle_attach_aio_context(ThrottleState *ts, AioContext *new_context)
+{
+    ts->timers[0] = aio_timer_new(new_context, ts->clock_type, SCALE_NS,
+                                  ts->read_timer_cb, ts->timer_opaque);
+    ts->timers[1] = aio_timer_new(new_context, ts->clock_type, SCALE_NS,
+                                  ts->write_timer_cb, ts->timer_opaque);
+}
+
 /* To be called first on the ThrottleState */
 void throttle_init(ThrottleState *ts,
+                   AioContext *aio_context,
                    QEMUClockType clock_type,
                    QEMUTimerCB *read_timer_cb,
                    QEMUTimerCB *write_timer_cb,
@@ -167,8 +178,10 @@ void throttle_init(ThrottleState *ts,
     memset(ts, 0, sizeof(ThrottleState));
 
     ts->clock_type = clock_type;
-    ts->timers[0] = timer_new_ns(clock_type, read_timer_cb, timer_opaque);
-    ts->timers[1] = timer_new_ns(clock_type, write_timer_cb, timer_opaque);
+    ts->read_timer_cb = read_timer_cb;
+    ts->write_timer_cb = write_timer_cb;
+    ts->timer_opaque = timer_opaque;
+    throttle_attach_aio_context(ts, aio_context);
 }
 
 /* destroy a timer */
@@ -181,8 +194,8 @@ static void throttle_timer_destroy(QEMUTimer **timer)
     *timer = NULL;
 }
 
-/* To be called last on the ThrottleState */
-void throttle_destroy(ThrottleState *ts)
+/* Remove timers from event loop */
+void throttle_detach_aio_context(ThrottleState *ts)
 {
     int i;
 
@@ -191,6 +204,12 @@ void throttle_destroy(ThrottleState *ts)
     }
 }
 
+/* To be called last on the ThrottleState */
+void throttle_destroy(ThrottleState *ts)
+{
+    throttle_detach_aio_context(ts);
+}
+
 /* is any throttling timer configured */
 bool throttle_have_timer(ThrottleState *ts)
 {