diff mbox

[1/3] throttle: add throttle_detach/attach_aio_context()

Message ID 1400077367-23409-2-git-send-email-stefanha@redhat.com
State New
Headers show

Commit Message

Stefan Hajnoczi May 14, 2014, 2:22 p.m. UTC
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.

Cc: Benoît Canet <benoit.canet@irqsave.net>
Signed-off-by: Stefan Hajnoczi <stefanha@redhat.com>
---
 block.c                 |  7 +++++++
 include/qemu/throttle.h | 10 ++++++++++
 tests/test-throttle.c   | 25 ++++++++++++++++++++-----
 util/throttle.c         | 27 +++++++++++++++++++++++----
 4 files changed, 60 insertions(+), 9 deletions(-)

Comments

Benoît Canet May 14, 2014, 3:05 p.m. UTC | #1
The Wednesday 14 May 2014 à 16:22:45 (+0200), Stefan Hajnoczi wrote :
> 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
s/so the/to 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.
> 
> Cc: Benoît Canet <benoit.canet@irqsave.net>
> Signed-off-by: Stefan Hajnoczi <stefanha@redhat.com>
> ---
>  block.c                 |  7 +++++++
>  include/qemu/throttle.h | 10 ++++++++++
>  tests/test-throttle.c   | 25 ++++++++++++++++++++-----
>  util/throttle.c         | 27 +++++++++++++++++++++++----
>  4 files changed, 60 insertions(+), 9 deletions(-)
> 
> diff --git a/block.c b/block.c
> index 189fc0d..b30bcd5 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,
> @@ -5532,6 +5533,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);
>      }
> @@ -5563,6 +5567,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 ab29b0b..b890613 100644
> --- a/include/qemu/throttle.h
> +++ b/include/qemu/throttle.h
> @@ -67,6 +67,11 @@ typedef struct ThrottleState {

Can we make sure this header knows about AioContext in case
there is another throttle user not block related ?

>      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 1d4ffd3..5fa5000 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"

And remove this one eventually.

>  #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 02e6f15..f976ac7 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)
>  {
> -- 
> 1.9.0
>
Stefan Hajnoczi May 15, 2014, 7:57 a.m. UTC | #2
On Wed, May 14, 2014 at 05:05:58PM +0200, Benoît Canet wrote:
> The Wednesday 14 May 2014 à 16:22:45 (+0200), Stefan Hajnoczi wrote :
> > 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
> s/so the/to the/ ? 

Thanks, let me know how you feel about my replies below.  If there is no
other need to respin then let's let this commit description typo slide.

> > diff --git a/include/qemu/throttle.h b/include/qemu/throttle.h
> > index ab29b0b..b890613 100644
> > --- a/include/qemu/throttle.h
> > +++ b/include/qemu/throttle.h
> > @@ -67,6 +67,11 @@ typedef struct ThrottleState {
> 
> Can we make sure this header knows about AioContext in case
> there is another throttle user not block related ?

It already does because qemu-common.h includes qemu-typedefs.h.  We get
an opaque AioContext.

The point of the typedefs.h is to avoid pulling in all the headers when
callers just pass around pointers and don't need the actual definition.

In this case, I think the split makes sense because callers might pass
AioContext without actually using the AioContext API themselves.

> > diff --git a/tests/test-throttle.c b/tests/test-throttle.c
> > index 1d4ffd3..5fa5000 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"
> 
> And remove this one eventually.

This include pulls in the AioContext API because the test case itself
manipulates the AioContext.
Benoît Canet May 15, 2014, 11:25 a.m. UTC | #3
The Wednesday 14 May 2014 à 16:22:45 (+0200), Stefan Hajnoczi wrote :
> 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.
> 
> Cc: Benoît Canet <benoit.canet@irqsave.net>
> Signed-off-by: Stefan Hajnoczi <stefanha@redhat.com>
> ---
>  block.c                 |  7 +++++++
>  include/qemu/throttle.h | 10 ++++++++++
>  tests/test-throttle.c   | 25 ++++++++++++++++++++-----
>  util/throttle.c         | 27 +++++++++++++++++++++++----
>  4 files changed, 60 insertions(+), 9 deletions(-)
> 
> diff --git a/block.c b/block.c
> index 189fc0d..b30bcd5 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,
> @@ -5532,6 +5533,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);
>      }
> @@ -5563,6 +5567,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 ab29b0b..b890613 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 1d4ffd3..5fa5000 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 02e6f15..f976ac7 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)
>  {
> -- 
> 1.9.0
> 
Reviewed-by: Benoit Canet <benoit@irqsave.net>
diff mbox

Patch

diff --git a/block.c b/block.c
index 189fc0d..b30bcd5 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,
@@ -5532,6 +5533,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);
     }
@@ -5563,6 +5567,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 ab29b0b..b890613 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 1d4ffd3..5fa5000 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 02e6f15..f976ac7 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)
 {