Correct handling of acks for XEP-0198 under heavy load.

2 files changed, 41 insertions, 33 deletions

diff --git a/sleekxmpp/plugins/xep_0198/stream_management.py b/sleekxmpp/plugins/xep_0198/stream_management.py
index 6ed1ea26..04aa0001 100644
--- a/sleekxmpp/plugins/xep_0198/stream_management.py
+++ b/sleekxmpp/plugins/xep_0198/stream_management.py
@@ -60,6 +60,8 @@ class XEP_0198(BasePlugin):
         #: the server. Setting this to ``1`` will send an ack request after
         #: every sent stanza. Defaults to ``5``.
         self.window = self.config.get('window', 5)
+        self.window_counter = self.window
+        self.window_counter_lock = threading.Lock()
 
         #: Control whether or not the ability to resume the stream will be
         #: requested when enabling stream management. Defaults to ``True``.
@@ -132,12 +134,12 @@ class XEP_0198(BasePlugin):
         ack = stanza.Ack(self.xmpp)
         with self.handled_lock:
             ack['h'] = self.handled
-        ack.send()
+        self.xmpp.send_raw(str(ack), now=True)
 
     def request_ack(self, e=None):
         """Request an ack from the server."""
         req = stanza.RequestAck(self.xmpp)
-        req.send()
+        self.xmpp.send_queue.put(str(req))
 
     def _handle_sm_feature(self, features):
         """
@@ -158,7 +160,7 @@ class XEP_0198(BasePlugin):
                 self.enabled.set()
                 enable = stanza.Enable(self.xmpp)
                 enable['resume'] = self.allow_resume
-                enable.send()
+                enable.send(now=True)
                 self.handled = 0
         elif self.sm_id and self.allow_resume:
             self.enabled.set()
@@ -261,6 +263,9 @@ class XEP_0198(BasePlugin):
                 self.seq = (self.seq + 1) % MAX_SEQ
                 seq = self.seq
             self.unacked_queue.append((seq, stanza))
-            if len(self.unacked_queue) > self.window:
-                self.xmpp.event('need_ack')
+            with self.window_counter_lock:
+                self.window_counter -= 1
+                if self.window_counter == 0:
+                    self.window_counter = self.window
+                    self.request_ack()
         return stanza
diff --git a/sleekxmpp/xmlstream/xmlstream.py b/sleekxmpp/xmlstream/xmlstream.py
index 8dfae7f4..125d4993 100644
--- a/sleekxmpp/xmlstream/xmlstream.py
+++ b/sleekxmpp/xmlstream/xmlstream.py
@@ -268,6 +268,7 @@ class XMLStream(object):
         #: A queue of string data to be sent over the stream.
         self.send_queue = queue.Queue()
         self.send_queue_lock = threading.Lock()
+        self.send_lock = threading.RLock()
 
         #: A :class:`~sleekxmpp.xmlstream.scheduler.Scheduler` instance for
         #: executing callbacks in the future based on time delays.
@@ -1180,21 +1181,22 @@ class XMLStream(object):
                 sent = 0
                 count = 0
                 tries = 0
-                while sent < total and not self.stop.is_set():
-                    try:
-                        sent += self.socket.send(data[sent:])
-                        count += 1
-                    except ssl.SSLError as serr:
-                        if tries >= self.ssl_retry_max:
-                            log.debug('SSL error - max retries reached')
-                            self.exception(serr)
-                            log.warning("Failed to send %s", data)
-                            if reconnect is None:
-                                reconnect = self.auto_reconnect
-                            self.disconnect(reconnect)
-                        log.warning('SSL write error - reattempting')
-                        time.sleep(self.ssl_retry_delay)
-                        tries += 1
+                with self.send_lock:
+                    while sent < total and not self.stop.is_set():
+                        try:
+                            sent += self.socket.send(data[sent:])
+                            count += 1
+                        except ssl.SSLError as serr:
+                            if tries >= self.ssl_retry_max:
+                                log.debug('SSL error - max retries reached')
+                                self.exception(serr)
+                                log.warning("Failed to send %s", data)
+                                if reconnect is None:
+                                    reconnect = self.auto_reconnect
+                                self.disconnect(reconnect)
+                            log.warning('SSL write error - reattempting')
+                            time.sleep(self.ssl_retry_delay)
+                            tries += 1
                 if count > 1:
                     log.debug('SENT: %d chunks', count)
             except Socket.error as serr:
@@ -1531,19 +1533,20 @@ class XMLStream(object):
                 count = 0
                 tries = 0
                 try:
-                    while sent < total and not self.stop.is_set():
-                        try:
-                            sent += self.socket.send(enc_data[sent:])
-                            count += 1
-                        except ssl.SSLError as serr:
-                            if tries >= self.ssl_retry_max:
-                                log.debug('SSL error - max retries reached')
-                                self.exception(serr)
-                                log.warning("Failed to send %s", data)
-                                self.disconnect(self.auto_reconnect)
-                            log.warning('SSL write error - reattempting')
-                            time.sleep(self.ssl_retry_delay)
-                            tries += 1
+                    with self.send_lock:
+                        while sent < total and not self.stop.is_set():
+                            try:
+                                sent += self.socket.send(enc_data[sent:])
+                                count += 1
+                            except ssl.SSLError as serr:
+                                if tries >= self.ssl_retry_max:
+                                    log.debug('SSL error - max retries reached')
+                                    self.exception(serr)
+                                    log.warning("Failed to send %s", data)
+                                    self.disconnect(self.auto_reconnect)
+                                log.warning('SSL write error - reattempting')
+                                time.sleep(self.ssl_retry_delay)
+                                tries += 1
                     if count > 1:
                         log.debug('SENT: %d chunks', count)
                     self.send_queue.task_done()