From 5efb170e1d7d210304d5bc31f343743770d82e13 Mon Sep 17 00:00:00 2001 From: Lance Stout Date: Mon, 4 Jul 2011 18:47:57 -0700 Subject: [PATCH] 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. --- sleekxmpp/xmlstream/xmlstream.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/sleekxmpp/xmlstream/xmlstream.py b/sleekxmpp/xmlstream/xmlstream.py index c7d0d3a..8b6c08b 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)