2003-04-09 Havoc Pennington <hp@redhat.com>

	* test/test-utils.c: use dispatch status function to fix this up

	* bus/connection.c (connection_watch_callback): don't dispatch
	from here
	(connection_timeout_callback): don't dispatch from here
	(bus_connections_setup_connection): set the dispatch status function
	(bus_connection_disconnected): unset it

	* dbus/dbus-mainloop.c (_dbus_loop_queue_dispatch): new function
	used to add a connection to be dispatched
	(_dbus_loop_iterate): do the dispatching at the end of each
	iteration

	* dbus/dbus-connection.c
	(dbus_connection_set_dispatch_status_function): new function
	allowing us to fix up main loop usage
	(_dbus_connection_last_unref): free all the various function
	user data
	(dbus_connection_dispatch): call the DispatchStatusFunction
	whenever this function returns
	(dbus_connection_handle_watch): call DispatchStatusFunction
	(dbus_connection_send_with_reply_and_block): call DispatchStatusFunction
	(reply_handler_timeout): call DispatchStatusFunction
	(dbus_connection_flush): call DispatchStatusFunction

7 files changed, 351 insertions, 118 deletions

diff --git a/ChangeLog b/ChangeLog
index b520f1fd..0ff4b11f 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,5 +1,32 @@
 2003-04-09  Havoc Pennington  <hp@redhat.com>
 
+	* test/test-utils.c: use dispatch status function to fix this up
+
+	* bus/connection.c (connection_watch_callback): don't dispatch
+	from here
+	(connection_timeout_callback): don't dispatch from here
+	(bus_connections_setup_connection): set the dispatch status function
+	(bus_connection_disconnected): unset it
+
+	* dbus/dbus-mainloop.c (_dbus_loop_queue_dispatch): new function
+	used to add a connection to be dispatched
+	(_dbus_loop_iterate): do the dispatching at the end of each
+	iteration
+
+	* dbus/dbus-connection.c
+	(dbus_connection_set_dispatch_status_function): new function
+	allowing us to fix up main loop usage
+	(_dbus_connection_last_unref): free all the various function
+	user data
+	(dbus_connection_dispatch): call the DispatchStatusFunction 
+	whenever this function returns
+	(dbus_connection_handle_watch): call DispatchStatusFunction
+	(dbus_connection_send_with_reply_and_block): call DispatchStatusFunction
+	(reply_handler_timeout): call DispatchStatusFunction
+	(dbus_connection_flush): call DispatchStatusFunction
+
+2003-04-09  Havoc Pennington  <hp@redhat.com>
+
 	* dbus/dbus-bus.c (dbus_bus_register): fix up error handling and 
 	a memory leak
 
diff --git a/bus/connection.c b/bus/connection.c
index f5a0ac16..aeb4f6e0 100644
--- a/bus/connection.c
+++ b/bus/connection.c
@@ -170,6 +170,9 @@ bus_connection_disconnected (DBusConnection *connection)
   
   dbus_connection_set_unix_user_function (connection,
                                           NULL, NULL, NULL);
+
+  dbus_connection_set_dispatch_status_function (connection,
+                                                NULL, NULL, NULL);
   
   bus_connection_remove_transactions (connection);
 
@@ -190,15 +193,9 @@ connection_watch_callback (DBusWatch     *watch,
 {
   DBusConnection *connection = data;
   dbus_bool_t retval;
-
-  dbus_connection_ref (connection);
   
   retval = dbus_connection_handle_watch (connection, watch, condition);
 
-  bus_connection_dispatch_all_messages (connection);
-  
-  dbus_connection_unref (connection);
-
   return retval;
 }
 
@@ -229,14 +226,8 @@ connection_timeout_callback (DBusTimeout   *timeout,
 {
   DBusConnection *connection = data;
 
-  dbus_connection_ref (connection);
-
   /* can return FALSE on OOM but we just let it fire again later */
   dbus_timeout_handle (timeout);
-
-  bus_connection_dispatch_all_messages (connection);
-  
-  dbus_connection_unref (connection);
 }
 
 static dbus_bool_t
@@ -259,6 +250,20 @@ remove_connection_timeout (DBusTimeout    *timeout,
                              timeout, connection_timeout_callback, connection);
 }
 
+static void
+dispatch_status_function (DBusConnection    *connection,
+                          DBusDispatchStatus new_status,
+                          void              *data)
+{
+  DBusLoop *loop = data;
+  
+  if (new_status != DBUS_DISPATCH_COMPLETE)
+    {
+      while (!_dbus_loop_queue_dispatch (loop, connection))
+        _dbus_wait_for_memory ();
+    }
+}
+
 static dbus_bool_t
 allow_user_function (DBusConnection *connection,
                      unsigned long   uid,
@@ -405,6 +410,11 @@ bus_connections_setup_connection (BusConnections *connections,
   dbus_connection_set_unix_user_function (connection,
                                           allow_user_function,
                                           NULL, NULL);
+
+  dbus_connection_set_dispatch_status_function (connection,
+                                                dispatch_status_function,
+                                                bus_context_get_loop (connections->context),
+                                                NULL);
   
   /* Setup the connection with the dispatcher */
   if (!bus_dispatch_add_connection (connection))
@@ -415,6 +425,15 @@ bus_connections_setup_connection (BusConnections *connections,
       bus_dispatch_remove_connection (connection);
       goto out;
     }
+
+  if (dbus_connection_get_dispatch_status (connection) != DBUS_DISPATCH_COMPLETE)
+    {
+      if (!_dbus_loop_queue_dispatch (bus_context_get_loop (connections->context), connection))
+        {
+          bus_dispatch_remove_connection (connection);
+          goto out;
+        }
+    }
   
   dbus_connection_ref (connection);
   retval = TRUE;
@@ -437,6 +456,9 @@ bus_connections_setup_connection (BusConnections *connections,
       dbus_connection_set_unix_user_function (connection,
                                               NULL, NULL, NULL);
 
+      dbus_connection_set_dispatch_status_function (connection,
+                                                    NULL, NULL, NULL);
+      
       if (!dbus_connection_set_data (connection,
                                      connection_data_slot,
                                      NULL, NULL))
diff --git a/dbus/dbus-connection.c b/dbus/dbus-connection.c
index 60c075fe..3be4ee45 100644
--- a/dbus/dbus-connection.c
+++ b/dbus/dbus-connection.c
@@ -109,6 +109,12 @@ struct DBusConnection
   DBusWakeupMainFunction wakeup_main_function; /**< Function to wake up the mainloop  */
   void *wakeup_main_data; /**< Application data for wakeup_main_function */
   DBusFreeFunction free_wakeup_main_data; /**< free wakeup_main_data */
+
+  DBusDispatchStatusFunction dispatch_status_function; /**< Function on dispatch status changes  */
+  void *dispatch_status_data; /**< Application data for dispatch_status_function */
+  DBusFreeFunction free_dispatch_status_data; /**< free dispatch_status_data */
+
+  DBusDispatchStatus last_dispatch_status;
 };
 
 typedef struct
@@ -126,9 +132,12 @@ typedef struct
 
 static void reply_handler_data_free (ReplyHandlerData *data);
 
-static void               _dbus_connection_remove_timeout_locked        (DBusConnection *connection,
-                                                                         DBusTimeout    *timeout);
-static DBusDispatchStatus _dbus_connection_get_dispatch_status_unlocked (DBusConnection *connection);
+static void               _dbus_connection_remove_timeout_locked         (DBusConnection     *connection,
+                                                                          DBusTimeout        *timeout);
+static DBusDispatchStatus _dbus_connection_get_dispatch_status_unlocked  (DBusConnection     *connection);
+static void               _dbus_connection_update_dispatch_status_locked (DBusConnection     *connection,
+                                                                          DBusDispatchStatus  new_status);
+
 
 /**
  * Acquires the connection lock.
@@ -231,7 +240,7 @@ _dbus_connection_queue_received_message_link (DBusConnection  *connection,
   connection->n_incoming += 1;
 
   _dbus_connection_wakeup_mainloop (connection);
-
+  
   _dbus_assert (dbus_message_get_name (message) != NULL);
   _dbus_verbose ("Message %p (%s) added to incoming queue %p, %d incoming\n",
                  message, dbus_message_get_name (message),
@@ -652,7 +661,8 @@ _dbus_connection_new_for_transport (DBusTransport *transport)
   connection->handler_table = handler_table;
   connection->pending_replies = pending_replies;
   connection->filter_list = NULL;
-
+  connection->last_dispatch_status = DBUS_DISPATCH_COMPLETE; /* so we're notified first time there's data */
+  
   _dbus_data_slot_list_init (&connection->slot_list);
 
   connection->client_serial = 1;
@@ -887,6 +897,11 @@ _dbus_connection_last_unref (DBusConnection *connection)
       _dbus_counter_unref (connection->connection_counter);
       connection->connection_counter = NULL;
     }
+
+  /* ---- We're going to call various application callbacks here, hope it doesn't break anything... */
+  dbus_connection_set_dispatch_status_function (connection, NULL, NULL, NULL);
+  dbus_connection_set_wakeup_main_function (connection, NULL, NULL, NULL);
+  dbus_connection_set_unix_user_function (connection, NULL, NULL, NULL);
   
   _dbus_watch_list_free (connection->watches);
   connection->watches = NULL;
@@ -894,8 +909,8 @@ _dbus_connection_last_unref (DBusConnection *connection)
   _dbus_timeout_list_free (connection->timeouts);
   connection->timeouts = NULL;
 
-  /* calls out to application code... */
   _dbus_data_slot_list_free (&connection->slot_list);
+  /* ---- Done with stuff that invokes application callbacks */
   
   _dbus_hash_iter_init (connection->handler_table, &iter);
   while (_dbus_hash_iter_next (&iter))
@@ -1184,6 +1199,7 @@ reply_handler_timeout (void *data)
 {
   DBusConnection *connection;
   ReplyHandlerData *reply_handler_data = data;
+  DBusDispatchStatus status;
 
   connection = reply_handler_data->connection;
   
@@ -1198,9 +1214,13 @@ reply_handler_timeout (void *data)
   _dbus_connection_remove_timeout (connection,
 				   reply_handler_data->timeout);
   reply_handler_data->timeout_added = FALSE;
+
+  status = _dbus_connection_get_dispatch_status_unlocked (connection);
   
   dbus_mutex_unlock (connection->mutex);
 
+  _dbus_connection_update_dispatch_status_locked (connection, status);
+  
   return TRUE;
 }
 
@@ -1488,11 +1508,15 @@ dbus_connection_send_with_reply_and_block (DBusConnection     *connection,
       
       reply = check_for_reply_unlocked (connection, client_serial);
       if (reply != NULL)
-        {
+        {          
+          status = _dbus_connection_get_dispatch_status_unlocked (connection);
+          
           dbus_mutex_unlock (connection->mutex);
 
           _dbus_verbose ("dbus_connection_send_with_reply_and_block(): got reply %s\n",
                          dbus_message_get_name (reply));
+
+          _dbus_connection_update_dispatch_status_locked (connection, status);
           
           return reply;
         }
@@ -1549,6 +1573,8 @@ dbus_connection_send_with_reply_and_block (DBusConnection     *connection,
   
   dbus_mutex_unlock (connection->mutex);
 
+  _dbus_connection_update_dispatch_status_locked (connection, status);
+
   return NULL;
 }
 
@@ -1560,11 +1586,12 @@ dbus_connection_send_with_reply_and_block (DBusConnection     *connection,
 void
 dbus_connection_flush (DBusConnection *connection)
 {
-  /* We have to specify DBUS_ITERATION_DO_READING here
-   * because otherwise we could have two apps deadlock
-   * if they are both doing a flush(), and the kernel
-   * buffers fill up.
+  /* We have to specify DBUS_ITERATION_DO_READING here because
+   * otherwise we could have two apps deadlock if they are both doing
+   * a flush(), and the kernel buffers fill up. This could change the
+   * dispatch status.
    */
+  DBusDispatchStatus status;
   
   dbus_mutex_lock (connection->mutex);
   while (connection->n_outgoing > 0)
@@ -1573,7 +1600,12 @@ dbus_connection_flush (DBusConnection *connection)
                                    DBUS_ITERATION_DO_WRITING |
                                    DBUS_ITERATION_BLOCK,
                                    -1);
+
+  status = _dbus_connection_get_dispatch_status_unlocked (connection);
+  
   dbus_mutex_unlock (connection->mutex);
+
+  _dbus_connection_update_dispatch_status_locked (connection, status);
 }
 
 /* Call with mutex held. Will drop it while waiting and re-acquire
@@ -1818,6 +1850,40 @@ _dbus_connection_get_dispatch_status_unlocked (DBusConnection *connection)
     }
 }
 
+static void
+_dbus_connection_update_dispatch_status_locked (DBusConnection    *connection,
+                                                DBusDispatchStatus new_status)
+{
+  dbus_bool_t changed;
+  DBusDispatchStatusFunction function;
+  void *data;
+  
+  dbus_mutex_lock (connection->mutex);
+  _dbus_connection_ref_unlocked (connection);
+
+  changed = new_status != connection->last_dispatch_status;
+
+  connection->last_dispatch_status = new_status;
+
+  function = connection->dispatch_status_function;
+  data = connection->dispatch_status_data;
+  
+  dbus_mutex_unlock (connection->mutex);
+  
+  if (changed && function)
+    {
+      _dbus_verbose ("Notifying of change to dispatch status of %p now %d (%s)\n",
+                     connection, new_status,
+                     new_status == DBUS_DISPATCH_COMPLETE ? "complete" :
+                     new_status == DBUS_DISPATCH_DATA_REMAINS ? "data remains" :
+                     new_status == DBUS_DISPATCH_NEED_MEMORY ? "need memory" :
+                     "???");
+      (* function) (connection, new_status, data);      
+    }
+  
+  dbus_connection_unref (connection);
+}
+
 /**
  * Gets the current state (what we would currently return
  * from dbus_connection_dispatch()) but doesn't actually
@@ -1864,7 +1930,10 @@ dbus_connection_dispatch (DBusConnection *connection)
   
   status = dbus_connection_get_dispatch_status (connection);
   if (status != DBUS_DISPATCH_DATA_REMAINS)
-    return status;
+    {
+      _dbus_connection_update_dispatch_status_locked (connection, status);
+      return status;
+    }
 
   dbus_mutex_lock (connection->mutex);
   
@@ -1891,8 +1960,10 @@ dbus_connection_dispatch (DBusConnection *connection)
 
       status = dbus_connection_get_dispatch_status (connection);
 
+      _dbus_connection_update_dispatch_status_locked (connection, status);
+      
       dbus_connection_unref (connection);
-
+      
       return status;
     }
 
@@ -1909,7 +1980,11 @@ dbus_connection_dispatch (DBusConnection *connection)
       _dbus_connection_release_dispatch (connection);
       dbus_mutex_unlock (connection->mutex);
       _dbus_connection_failed_pop (connection, message_link);
+
+      _dbus_connection_update_dispatch_status_locked (connection, DBUS_DISPATCH_NEED_MEMORY);
+
       dbus_connection_unref (connection);
+      
       return DBUS_DISPATCH_NEED_MEMORY;
     }
   
@@ -2012,6 +2087,8 @@ dbus_connection_dispatch (DBusConnection *connection)
                                  */
   
   status = dbus_connection_get_dispatch_status (connection);
+
+  _dbus_connection_update_dispatch_status_locked (connection, status);
   
   dbus_connection_unref (connection);
   
@@ -2201,6 +2278,46 @@ dbus_connection_set_wakeup_main_function (DBusConnection            *connection,
 }
 
 /**
+ * Set a function to be invoked when the dispatch status changes.
+ * If the dispatch status is #DBUS_DISPATCH_DATA_REMAINS, then
+ * dbus_connection_dispatch() needs to be called to process incoming
+ * messages. However, dbus_connection_dispatch() MUST NOT BE CALLED
+ * from inside the DBusDispatchStatusFunction. Indeed, almost
+ * any reentrancy in this function is a bad idea. Instead,
+ * the DBusDispatchStatusFunction should simply save an indication
+ * that messages should be dispatched later, when the main loop
+ * is re-entered.
+ *
+ * @param connection the connection
+ * @param function function to call on dispatch status changes
+ * @param data data for function
+ * @param free_data_function free the function data
+ */
+void
+dbus_connection_set_dispatch_status_function (DBusConnection             *connection,
+                                              DBusDispatchStatusFunction  function,
+                                              void                       *data,
+                                              DBusFreeFunction            free_data_function)
+{
+  void *old_data;
+  DBusFreeFunction old_free_data;
+  
+  dbus_mutex_lock (connection->mutex);
+  old_data = connection->dispatch_status_data;
+  old_free_data = connection->free_dispatch_status_data;
+
+  connection->dispatch_status_function = function;
+  connection->dispatch_status_data = data;
+  connection->free_dispatch_status_data = free_data_function;
+  
+  dbus_mutex_unlock (connection->mutex);
+
+  /* Callback outside the lock */
+  if (old_free_data)
+    (*old_free_data) (old_data);
+}
+
+/**
  * Called to notify the connection when a previously-added watch
  * is ready for reading or writing, or has an exception such
  * as a hangup.
@@ -2223,14 +2340,20 @@ dbus_connection_handle_watch (DBusConnection              *connection,
                               unsigned int                 condition)
 {
   dbus_bool_t retval;
+  DBusDispatchStatus status;
   
   dbus_mutex_lock (connection->mutex);
   _dbus_connection_acquire_io_path (connection, -1);
   retval = _dbus_transport_handle_watch (connection->transport,
                                          watch, condition);
   _dbus_connection_release_io_path (connection);
+
+  status = _dbus_connection_get_dispatch_status_unlocked (connection);
+  
   dbus_mutex_unlock (connection->mutex);
 
+  _dbus_connection_update_dispatch_status_locked (connection, status);
+  
   return retval;
 }
 
diff --git a/dbus/dbus-connection.h b/dbus/dbus-connection.h
index 1acbb1ad..e945c856 100644
--- a/dbus/dbus-connection.h
+++ b/dbus/dbus-connection.h
@@ -75,63 +75,71 @@ typedef void        (* DBusTimeoutToggledFunction) (DBusTimeout    *timeout,
                                                     void           *data);
 typedef void        (* DBusRemoveTimeoutFunction)  (DBusTimeout    *timeout,
                                                     void           *data);
+typedef void        (* DBusDispatchStatusFunction) (DBusConnection *connection,
+                                                    DBusDispatchStatus new_status,
+                                                    void           *data);
 typedef void        (* DBusWakeupMainFunction)     (void           *data);
 typedef dbus_bool_t (* DBusAllowUnixUserFunction)  (DBusConnection *connection,
                                                     unsigned long   uid,
                                                     void           *data);
 
-DBusConnection*    dbus_connection_open                      (const char                 *address,
-                                                              DBusError                  *error);
-void               dbus_connection_ref                       (DBusConnection             *connection);
-void               dbus_connection_unref                     (DBusConnection             *connection);
-void               dbus_connection_disconnect                (DBusConnection             *connection);
-dbus_bool_t        dbus_connection_get_is_connected          (DBusConnection             *connection);
-dbus_bool_t        dbus_connection_get_is_authenticated      (DBusConnection             *connection);
-void               dbus_connection_flush                     (DBusConnection             *connection);
-DBusMessage*       dbus_connection_borrow_message            (DBusConnection             *connection);
-void               dbus_connection_return_message            (DBusConnection             *connection,
-                                                              DBusMessage                *message);
-void               dbus_connection_steal_borrowed_message    (DBusConnection             *connection,
-                                                              DBusMessage                *message);
-DBusMessage*       dbus_connection_pop_message               (DBusConnection             *connection);
-DBusDispatchStatus dbus_connection_get_dispatch_status       (DBusConnection             *connection);
-DBusDispatchStatus dbus_connection_dispatch                  (DBusConnection             *connection);
-dbus_bool_t        dbus_connection_send                      (DBusConnection             *connection,
-                                                              DBusMessage                *message,
-                                                              dbus_int32_t               *client_serial);
-dbus_bool_t        dbus_connection_send_with_reply           (DBusConnection             *connection,
-                                                              DBusMessage                *message,
-                                                              DBusMessageHandler         *reply_handler,
-                                                              int                         timeout_milliseconds);
-DBusMessage *      dbus_connection_send_with_reply_and_block (DBusConnection             *connection,
-                                                              DBusMessage                *message,
-                                                              int                         timeout_milliseconds,
-                                                              DBusError                  *error);
-dbus_bool_t        dbus_connection_set_watch_functions       (DBusConnection             *connection,
-                                                              DBusAddWatchFunction        add_function,
-                                                              DBusRemoveWatchFunction     remove_function,
-                                                              DBusWatchToggledFunction    toggled_function,
-                                                              void                       *data,
-                                                              DBusFreeFunction            free_data_function);
-dbus_bool_t        dbus_connection_set_timeout_functions     (DBusConnection             *connection,
-                                                              DBusAddTimeoutFunction      add_function,
-                                                              DBusRemoveTimeoutFunction   remove_function,
-                                                              DBusTimeoutToggledFunction  toggled_function,
-                                                              void                       *data,
-                                                              DBusFreeFunction            free_data_function);
-void               dbus_connection_set_wakeup_main_function  (DBusConnection             *connection,
-                                                              DBusWakeupMainFunction      wakeup_main_function,
-                                                              void                       *data,
-                                                              DBusFreeFunction            free_data_function);
-dbus_bool_t        dbus_connection_handle_watch              (DBusConnection             *connection,
-                                                              DBusWatch                  *watch,
-                                                              unsigned int                condition);
-dbus_bool_t        dbus_connection_get_unix_user             (DBusConnection             *connection,
-                                                              unsigned long              *uid);
-void               dbus_connection_set_unix_user_function    (DBusConnection             *connection,
-                                                              DBusAllowUnixUserFunction   function,
-                                                              void                       *data,
-                                                              DBusFreeFunction            free_data_function);
+DBusConnection*    dbus_connection_open                         (const char                 *address,
+                                                                 DBusError                  *error);
+void               dbus_connection_ref                          (DBusConnection             *connection);
+void               dbus_connection_unref                        (DBusConnection             *connection);
+void               dbus_connection_disconnect                   (DBusConnection             *connection);
+dbus_bool_t        dbus_connection_get_is_connected             (DBusConnection             *connection);
+dbus_bool_t        dbus_connection_get_is_authenticated         (DBusConnection             *connection);
+void               dbus_connection_flush                        (DBusConnection             *connection);
+DBusMessage*       dbus_connection_borrow_message               (DBusConnection             *connection);
+void               dbus_connection_return_message               (DBusConnection             *connection,
+                                                                 DBusMessage                *message);
+void               dbus_connection_steal_borrowed_message       (DBusConnection             *connection,
+                                                                 DBusMessage                *message);
+DBusMessage*       dbus_connection_pop_message                  (DBusConnection             *connection);
+DBusDispatchStatus dbus_connection_get_dispatch_status          (DBusConnection             *connection);
+DBusDispatchStatus dbus_connection_dispatch                     (DBusConnection             *connection);
+dbus_bool_t        dbus_connection_send                         (DBusConnection             *connection,
+                                                                 DBusMessage                *message,
+                                                                 dbus_int32_t               *client_serial);
+dbus_bool_t        dbus_connection_send_with_reply              (DBusConnection             *connection,
+                                                                 DBusMessage                *message,
+                                                                 DBusMessageHandler         *reply_handler,
+                                                                 int                         timeout_milliseconds);
+DBusMessage *      dbus_connection_send_with_reply_and_block    (DBusConnection             *connection,
+                                                                 DBusMessage                *message,
+                                                                 int                         timeout_milliseconds,
+                                                                 DBusError                  *error);
+dbus_bool_t        dbus_connection_set_watch_functions          (DBusConnection             *connection,
+                                                                 DBusAddWatchFunction        add_function,
+                                                                 DBusRemoveWatchFunction     remove_function,
+                                                                 DBusWatchToggledFunction    toggled_function,
+                                                                 void                       *data,
+                                                                 DBusFreeFunction            free_data_function);
+dbus_bool_t        dbus_connection_set_timeout_functions        (DBusConnection             *connection,
+                                                                 DBusAddTimeoutFunction      add_function,
+                                                                 DBusRemoveTimeoutFunction   remove_function,
+                                                                 DBusTimeoutToggledFunction  toggled_function,
+                                                                 void                       *data,
+                                                                 DBusFreeFunction            free_data_function);
+void               dbus_connection_set_wakeup_main_function     (DBusConnection             *connection,
+                                                                 DBusWakeupMainFunction      wakeup_main_function,
+                                                                 void                       *data,
+                                                                 DBusFreeFunction            free_data_function);
+void               dbus_connection_set_dispatch_status_function (DBusConnection             *connection,
+                                                                 DBusDispatchStatusFunction  function,
+                                                                 void                       *data,
+                                                                 DBusFreeFunction            free_data_function);
+dbus_bool_t        dbus_connection_handle_watch                 (DBusConnection             *connection,
+                                                                 DBusWatch                  *watch,
+                                                                 unsigned int                condition);
+dbus_bool_t        dbus_connection_get_unix_user                (DBusConnection             *connection,
+                                                                 unsigned long              *uid);
+void               dbus_connection_set_unix_user_function       (DBusConnection             *connection,
+                                                                 DBusAllowUnixUserFunction   function,
+                                                                 void                       *data,
+                                                                 DBusFreeFunction            free_data_function);
+
 
 int          dbus_watch_get_fd      (DBusWatch        *watch);
 unsigned int dbus_watch_get_flags   (DBusWatch        *watch);
diff --git a/dbus/dbus-mainloop.c b/dbus/dbus-mainloop.c
index b1706fe7..0ad2618a 100644
--- a/dbus/dbus-mainloop.c
+++ b/dbus/dbus-mainloop.c
@@ -34,6 +34,7 @@ struct DBusLoop
   int watch_count;
   int timeout_count;
   int depth; /**< number of recursive runs */
+  DBusList *need_dispatch;
 };
 
 typedef enum
@@ -199,6 +200,12 @@ _dbus_loop_unref (DBusLoop *loop)
   loop->refcount -= 1;
   if (loop->refcount == 0)
     {
+      while (loop->need_dispatch)
+        {
+          DBusConnection *connection = _dbus_list_pop_first (&loop->need_dispatch);
+
+          dbus_connection_unref (connection);
+        }
       
       dbus_free (loop);
     }
@@ -399,13 +406,55 @@ check_timeout (unsigned long    tv_sec,
   return msec == 0;
 }
 
+static void
+_dbus_loop_dispatch (DBusLoop *loop)
+{
+ next:
+  while (loop->need_dispatch != NULL)
+    {
+      DBusConnection *connection = _dbus_list_pop_first (&loop->need_dispatch);
+
+      while (TRUE)
+        {
+          DBusDispatchStatus status;
+          
+          status = dbus_connection_dispatch (connection);
+
+          if (status == DBUS_DISPATCH_COMPLETE)
+            {
+              dbus_connection_unref (connection);
+              goto next;
+            }
+          else
+            {
+              if (status == DBUS_DISPATCH_NEED_MEMORY)
+                _dbus_wait_for_memory ();
+            }
+        }
+    }
+}
+
+dbus_bool_t
+_dbus_loop_queue_dispatch (DBusLoop       *loop,
+                           DBusConnection *connection)
+{
+  
+  if (_dbus_list_append (&loop->need_dispatch, connection))
+    {
+      dbus_connection_ref (connection);
+      return TRUE;
+    }
+  else
+    return FALSE;
+}
+
 /* Returns TRUE if we have any timeouts or ready file descriptors,
  * which is just used in test code as a debug hack
  */
 
 dbus_bool_t
 _dbus_loop_iterate (DBusLoop     *loop,
-                    dbus_bool_t  block)
+                    dbus_bool_t   block)
 {
   dbus_bool_t retval;
   DBusPollFD *fds;
@@ -553,7 +602,8 @@ _dbus_loop_iterate (DBusLoop     *loop,
         }
     }
 
-  if (!block)
+  /* Never block if we have stuff to dispatch */
+  if (!block || loop->need_dispatch != NULL)
     {
       timeout = 0;
 #if 0
@@ -673,6 +723,12 @@ _dbus_loop_iterate (DBusLoop     *loop,
   dbus_free (fds);
   dbus_free (watches_for_fds);
 
+  if (loop->need_dispatch != NULL)
+    {
+      retval = TRUE;
+      _dbus_loop_dispatch (loop);
+    }
+  
   return retval;
 }
 
@@ -716,3 +772,4 @@ _dbus_wait_for_memory (void)
 {
   _dbus_sleep_milliseconds (_dbus_get_oom_wait ());
 }
+
diff --git a/dbus/dbus-mainloop.h b/dbus/dbus-mainloop.h
index bca4317d..7aaf9de0 100644
--- a/dbus/dbus-mainloop.h
+++ b/dbus/dbus-mainloop.h
@@ -55,6 +55,10 @@ void        _dbus_loop_remove_timeout (DBusLoop            *loop,
                                        DBusTimeout         *timeout,
                                        DBusTimeoutFunction  function,
                                        void                *data);
+
+dbus_bool_t _dbus_loop_queue_dispatch (DBusLoop            *loop,
+                                       DBusConnection      *connection);
+
 void        _dbus_loop_run            (DBusLoop            *loop);
 void        _dbus_loop_quit           (DBusLoop            *loop);
 dbus_bool_t _dbus_loop_iterate        (DBusLoop            *loop,
diff --git a/test/test-utils.c b/test/test-utils.c
index c7d4aa59..6372314d 100644
--- a/test/test-utils.c
+++ b/test/test-utils.c
@@ -7,24 +7,6 @@ typedef struct
 
 } CData;
 
-dbus_bool_t
-test_connection_dispatch_one_message  (DBusConnection *connection)
-{
-  DBusDispatchStatus status;
-
-  while ((status = dbus_connection_dispatch (connection)) == DBUS_DISPATCH_NEED_MEMORY)
-    _dbus_wait_for_memory ();
-  
-  return status == DBUS_DISPATCH_DATA_REMAINS;
-}
-
-void
-test_connection_dispatch_all_messages (DBusConnection *connection)
-{
-  while (test_connection_dispatch_one_message (connection))
-    ;
-}
-
 static dbus_bool_t
 connection_watch_callback (DBusWatch     *watch,
                            unsigned int   condition,
@@ -32,21 +14,9 @@ connection_watch_callback (DBusWatch     *watch,
 {
   CData *cd = data;
   dbus_bool_t retval;
-
-  dbus_connection_ref (cd->connection);
-
-  _dbus_verbose (" Handling watch\n");
   
   retval = dbus_connection_handle_watch (cd->connection, watch, condition);
 
-  _dbus_verbose (" Watch handled\n");
-  
-  test_connection_dispatch_all_messages (cd->connection);
-
-  _dbus_verbose (" Dispatched all\n");
-  
-  dbus_connection_unref (cd->connection);
-
   return retval;
 }
 
@@ -78,14 +48,8 @@ connection_timeout_callback (DBusTimeout   *timeout,
 {
   CData *cd = data;
 
-  dbus_connection_ref (cd->connection);
-
   /* Can return FALSE on OOM but we just let it fire again later */
   dbus_timeout_handle (timeout);
-  
-  test_connection_dispatch_all_messages (cd->connection);
-  
-  dbus_connection_unref (cd->connection);
 }
 
 static dbus_bool_t
@@ -109,6 +73,20 @@ remove_timeout (DBusTimeout *timeout,
 }
 
 static void
+dispatch_status_function (DBusConnection    *connection,
+                          DBusDispatchStatus new_status,
+                          void              *data)
+{
+  DBusLoop *loop = data;
+  
+  if (new_status != DBUS_DISPATCH_COMPLETE)
+    {
+      while (!_dbus_loop_queue_dispatch (loop, connection))
+        _dbus_wait_for_memory ();
+    }
+}
+
+static void
 cdata_free (void *data)
 {
   CData *cd = data;
@@ -144,6 +122,11 @@ test_connection_setup (DBusLoop       *loop,
 {
   CData *cd;
 
+  cd = NULL;
+  
+  dbus_connection_set_dispatch_status_function (connection, dispatch_status_function,
+                                                loop, NULL);
+  
   cd = cdata_new (loop, connection);
   if (cd == NULL)
     goto nomem;
@@ -170,16 +153,24 @@ test_connection_setup (DBusLoop       *loop,
                                               remove_timeout,
                                               NULL,
                                               cd, cdata_free))
+    goto nomem;
+
+  if (dbus_connection_get_dispatch_status (connection) != DBUS_DISPATCH_COMPLETE)
     {
-      dbus_connection_set_watch_functions (connection, NULL, NULL, NULL, NULL, NULL);
-      goto nomem;
+      if (!_dbus_loop_queue_dispatch (loop, connection))
+        goto nomem;
     }
-
+  
   return TRUE;
   
  nomem:
   if (cd)
     cdata_free (cd);
+  
+  dbus_connection_set_dispatch_status_function (connection, NULL, NULL, NULL);
+  dbus_connection_set_watch_functions (connection, NULL, NULL, NULL, NULL, NULL);
+  dbus_connection_set_timeout_functions (connection, NULL, NULL, NULL, NULL, NULL);
+  
   return FALSE;
 }
 
@@ -201,4 +192,5 @@ test_connection_shutdown (DBusLoop       *loop,
                                               NULL, NULL))
     _dbus_assert_not_reached ("setting timeout functions to NULL failed");
 
+  dbus_connection_set_dispatch_status_function (connection, NULL, NULL, NULL);
 }