From 9f43d31bf511e0917f11b7edcac93921100383b7 Mon Sep 17 00:00:00 2001
From: Lance Stout <lancestout@gmail.com>
Date: Mon, 12 Mar 2012 16:19:18 -0700
Subject: Add setting for maximum number of reconnection attempts.

Setting self.reconnect_max_attempts to a non-None value will limit
the number of times a connection attempt will be made before quiting
and raising a 'connection_failed' event.
---
 sleekxmpp/xmlstream/xmlstream.py | 66 ++++++++++++++++++++++++++++------------
 1 file changed, 47 insertions(+), 19 deletions(-)

(limited to 'sleekxmpp/xmlstream')

diff --git a/sleekxmpp/xmlstream/xmlstream.py b/sleekxmpp/xmlstream/xmlstream.py
index 81ba2305..f3e3faa8 100644
--- a/sleekxmpp/xmlstream/xmlstream.py
+++ b/sleekxmpp/xmlstream/xmlstream.py
@@ -81,6 +81,12 @@ SSL_RETRY_MAX = 10
 #: Maximum time to delay between connection attempts is one hour.
 RECONNECT_MAX_DELAY = 600
 
+#: Maximum number of attempts to connect to the server before quitting
+#: and raising a 'connect_failed' event. Setting this to ``None`` will
+#: allow infinite reconnection attempts, and using ``0`` will disable
+#: reconnections. Defaults to ``None``.
+RECONNECT_MAX_ATTEMPTS = None
+
 
 log = logging.getLogger(__name__)
 
@@ -157,6 +163,12 @@ class XMLStream(object):
         #: Maximum time to delay between connection attempts is one hour.
         self.reconnect_max_delay = RECONNECT_MAX_DELAY
 
+        #: Maximum number of attempts to connect to the server before 
+        #: quitting and raising a 'connect_failed' event. Setting to
+        #: ``None`` allows infinite reattempts, while setting it to ``0``
+        #: will disable reconnection attempts. Defaults to ``None``.
+        self.reconnect_max_attempts = RECONNECT_MAX_ATTEMPTS
+
         #: The time in seconds to delay between attempts to resend data
         #: after an SSL error.
         self.ssl_retry_max = SSL_RETRY_MAX
@@ -383,13 +395,21 @@ class XMLStream(object):
         if use_tls is not None:
             self.use_tls = use_tls
 
+
         # Repeatedly attempt to connect until a successful connection
         # is established.
+        attempts = self.reconnect_max_attempts
         connected = self.state.transition('disconnected', 'connected',
                                           func=self._connect)
         while reattempt and not connected and not self.stop.is_set():
             connected = self.state.transition('disconnected', 'connected',
                                               func=self._connect)
+            if not connected:
+                if attempts is not None:
+                    attempts -= 1
+                    if attempts <= 0:
+                        self.event('connection_failed', direct=True)
+                        return False
         return connected
 
     def _connect(self):
@@ -399,6 +419,24 @@ class XMLStream(object):
             self.address = self.pick_dns_answer(self.default_domain,
                                                 self.address[1])
         
+        if self.reconnect_delay is None:
+            delay = 1.0
+        else:
+            delay = min(self.reconnect_delay * 2, self.reconnect_max_delay)
+            delay = random.normalvariate(delay, delay * 0.1)
+            log.debug('Waiting %s seconds before connecting.', delay)
+            elapsed = 0
+            try:
+                while elapsed < delay and not self.stop.is_set():
+                    time.sleep(0.1)
+                    elapsed += 0.1
+            except KeyboardInterrupt:
+                self.stop.set()
+                return False
+            except SystemExit:
+                self.stop.set()
+                return False
+
         try:
             # Look for IPv6 addresses, in addition to IPv4
             for res in Socket.getaddrinfo(self.address[0],
@@ -415,29 +453,11 @@ class XMLStream(object):
         except Socket.gaierror:
             log.warning("Socket could not be opened: no connectivity" + \
                         " or wrong IP versions.")
-            self.stop.set()
+            self.reconnect_delay = delay
             return False
 
         self.configure_socket()
 
-        if self.reconnect_delay is None:
-            delay = 1.0
-        else:
-            delay = min(self.reconnect_delay * 2, self.reconnect_max_delay)
-            delay = random.normalvariate(delay, delay * 0.1)
-            log.debug('Waiting %s seconds before connecting.', delay)
-            elapsed = 0
-            try:
-                while elapsed < delay and not self.stop.is_set():
-                    time.sleep(0.1)
-                    elapsed += 0.1
-            except KeyboardInterrupt:
-                self.stop.set()
-                return False
-            except SystemExit:
-                self.stop.set()
-                return False
-
         if self.use_proxy:
             connected = self._connect_proxy()
             if not connected:
@@ -618,6 +638,8 @@ class XMLStream(object):
             self.state.transition('connected', 'disconnected', wait=2.0,
                                   func=self._disconnect, args=(True,))
 
+        attempts = self.reconnect_max_attempts
+
         log.debug("connecting...")
         connected = self.state.transition('disconnected', 'connected',
                                           wait=2.0, func=self._connect)
@@ -625,6 +647,12 @@ class XMLStream(object):
             connected = self.state.transition('disconnected', 'connected',
                                               wait=2.0, func=self._connect)
             connected = connected or self.state.ensure('connected')
+            if not connected:
+                if attempts is not None:
+                    attempts -= 1
+                    if attempts <= 0:
+                        self.event('connection_failed', direct=True)
+                        return False
         return connected
 
     def set_socket(self, socket, ignore=False):
-- 
cgit v1.2.3