4 files changed, 95 insertions, 1 deletions

diff --git a/aio-posix.c b/aio-posix.c
index 44247224e2..65b26073f0 100644
--- a/aio-posix.c
+++ b/aio-posix.c
@@ -56,6 +56,8 @@ void aio_set_fd_handler(AioContext *ctx,
     /* Are we deleting the fd handler? */
     if (!io_read && !io_write) {
         if (node) {
+            g_source_remove_poll(&ctx->source, &node->pfd);
+
             /* If the lock is held, just mark the node as deleted */
             if (ctx->walking_handlers) {
                 node->deleted = 1;
@@ -75,6 +77,8 @@ void aio_set_fd_handler(AioContext *ctx,
             node = g_malloc0(sizeof(AioHandler));
             node->pfd.fd = fd;
             QLIST_INSERT_HEAD(&ctx->aio_handlers, node, node);
+
+            g_source_add_poll(&ctx->source, &node->pfd);
         }
         /* Update handler with latest information */
         node->io_read = io_read;
diff --git a/aio-win32.c b/aio-win32.c
index 9881fdbca7..e460bd848a 100644
--- a/aio-win32.c
+++ b/aio-win32.c
@@ -45,6 +45,8 @@ void aio_set_event_notifier(AioContext *ctx,
     /* Are we deleting the fd handler? */
     if (!io_notify) {
         if (node) {
+            g_source_remove_poll(&ctx->source, &node->pfd);
+
             /* If the lock is held, just mark the node as deleted */
             if (ctx->walking_handlers) {
                 node->deleted = 1;
@@ -66,6 +68,8 @@ void aio_set_event_notifier(AioContext *ctx,
             node->pfd.fd = (uintptr_t)event_notifier_get_handle(e);
             node->pfd.events = G_IO_IN;
             QLIST_INSERT_HEAD(&ctx->aio_handlers, node, node);
+
+            g_source_add_poll(&ctx->source, &node->pfd);
         }
         /* Update handler with latest information */
         node->io_notify = io_notify;
diff --git a/async.c b/async.c
index 513bdd7aa2..4ffdd986f1 100644
--- a/async.c
+++ b/async.c
@@ -136,10 +136,73 @@ void aio_bh_update_timeout(AioContext *ctx, uint32_t *timeout)
     }
 }
 
+static gboolean
+aio_ctx_prepare(GSource *source, gint    *timeout)
+{
+    AioContext *ctx = (AioContext *) source;
+    uint32_t wait = -1;
+    aio_bh_update_timeout(ctx, &wait);
+
+    if (wait != -1) {
+        *timeout = MIN(*timeout, wait);
+        return wait == 0;
+    }
+
+    return false;
+}
+
+static gboolean
+aio_ctx_check(GSource *source)
+{
+    AioContext *ctx = (AioContext *) source;
+    QEMUBH *bh;
+
+    for (bh = ctx->first_bh; bh; bh = bh->next) {
+        if (!bh->deleted && bh->scheduled) {
+            return true;
+	}
+    }
+    return aio_pending(ctx);
+}
+
+static gboolean
+aio_ctx_dispatch(GSource     *source,
+                 GSourceFunc  callback,
+                 gpointer     user_data)
+{
+    AioContext *ctx = (AioContext *) source;
+
+    assert(callback == NULL);
+    aio_poll(ctx, false);
+    return true;
+}
+
+static GSourceFuncs aio_source_funcs = {
+    aio_ctx_prepare,
+    aio_ctx_check,
+    aio_ctx_dispatch,
+    NULL
+};
+
+GSource *aio_get_g_source(AioContext *ctx)
+{
+    g_source_ref(&ctx->source);
+    return &ctx->source;
+}
 
 AioContext *aio_context_new(void)
 {
-    return g_new0(AioContext, 1);
+    return (AioContext *) g_source_new(&aio_source_funcs, sizeof(AioContext));
+}
+
+void aio_context_ref(AioContext *ctx)
+{
+    g_source_ref(&ctx->source);
+}
+
+void aio_context_unref(AioContext *ctx)
+{
+    g_source_unref(&ctx->source);
 }
 
 void aio_flush(AioContext *ctx)
diff --git a/qemu-aio.h b/qemu-aio.h
index ac248962fe..aedf66cfa1 100644
--- a/qemu-aio.h
+++ b/qemu-aio.h
@@ -44,6 +44,8 @@ typedef void QEMUBHFunc(void *opaque);
 typedef void IOHandler(void *opaque);
 
 typedef struct AioContext {
+    GSource source;
+
     /* The list of registered AIO handlers */
     QLIST_HEAD(, AioHandler) aio_handlers;
 
@@ -75,6 +77,22 @@ typedef int (AioFlushEventNotifierHandler)(EventNotifier *e);
 AioContext *aio_context_new(void);
 
 /**
+ * aio_context_ref:
+ * @ctx: The AioContext to operate on.
+ *
+ * Add a reference to an AioContext.
+ */
+void aio_context_ref(AioContext *ctx);
+
+/**
+ * aio_context_unref:
+ * @ctx: The AioContext to operate on.
+ *
+ * Drop a reference to an AioContext.
+ */
+void aio_context_unref(AioContext *ctx);
+
+/**
  * aio_bh_new: Allocate a new bottom half structure.
  *
  * Bottom halves are lightweight callbacks whose invocation is guaranteed
@@ -188,6 +206,11 @@ void aio_set_event_notifier(AioContext *ctx,
                             EventNotifierHandler *io_read,
                             AioFlushEventNotifierHandler *io_flush);
 
+/* Return a GSource that lets the main loop poll the file descriptors attached
+ * to this AioContext.
+ */
+GSource *aio_get_g_source(AioContext *ctx);
+
 /* Functions to operate on the main QEMU AioContext.  */
 
 void qemu_aio_flush(void);