Added wait parameter to disconnect.

If wait=True, then the disconnect call will block until
the send queue has emptied.

WARNING: Using wait=True when more stanzas are being added to the
queue than can be processed such that the queue is never empty
will cause the disconnect call to block indefinitely without actually
disconnecting.

1 files changed, 16 insertions, 3 deletions

diff --git a/sleekxmpp/xmlstream/xmlstream.py b/sleekxmpp/xmlstream/xmlstream.py
index c7d0d3a8..8b6c08b7 100644
--- a/sleekxmpp/xmlstream/xmlstream.py
+++ b/sleekxmpp/xmlstream/xmlstream.py
@@ -356,22 +356,34 @@ class XMLStream(object):
             self.reconnect_delay = delay
             return False
 
-    def disconnect(self, reconnect=False):
+    def disconnect(self, reconnect=False, wait=False):
         """
         Terminate processing and close the XML streams.
 
         Optionally, the connection may be reconnected and
         resume processing afterwards.
 
+        If the disconnect should take place after all items
+        in the send queue have been sent, use wait=True. However,
+        take note: If you are constantly adding items to the queue
+        such that it is never empty, then the disconnect will
+        not occur and the call will continue to block.
+
         Arguments:
             reconnect -- Flag indicating if the connection
                          and processing should be restarted.
                          Defaults to False.
+            wait      -- Flag indicating if the send queue should
+                         be emptied before disconnecting.
         """
         self.state.transition('connected', 'disconnected', wait=0.0,
-                              func=self._disconnect, args=(reconnect,))
+                              func=self._disconnect, args=(reconnect, wait))
+
+    def _disconnect(self, reconnect=False, wait=False):
+        # Wait for the send queue to empty.
+        if wait:
+            self.send_queue.join()
 
-    def _disconnect(self, reconnect=False):
         # Send the end of stream marker.
         self.send_raw(self.stream_footer, now=True)
         self.session_started_event.clear()
@@ -1036,6 +1048,7 @@ class XMLStream(object):
                 log.debug("SEND: %s" % data)
                 try:
                     self.socket.send(data.encode('utf-8'))
+                    self.send_queue.task_done()
                 except Socket.error as serr:
                     self.event('socket_error', serr)
                     log.warning("Failed to send %s" % data)