Major refactor in progress, unsafe at the moment.

1 files changed, 48 insertions, 12 deletions

diff --git a/kazoo/retry.py b/kazoo/retry.py
index 63d98f6..73eb42c 100644
--- a/kazoo/retry.py
+++ b/kazoo/retry.py
@@ -14,6 +14,50 @@ class ForceRetryError(Exception):
     """Raised when some recipe logic wants to force a retry"""
 
 
+class RetySleeper(object):
+    """A retry sleeper that will track its jitter, backoff and
+    sleep appropriately when asked"""
+    def __init__(self, max_tries=1, delay=0.1, backoff=2, max_jitter=0.8,
+                 ignore_expire=True, sleep_func=time.sleep):
+        """Create a :class:`KazooRetry` instance
+
+        :param max_tries: How many times to retry the command.
+        :param delay: Initial delay between retry attempts
+        :param backoff: Backoff multiplier between retry attempts. Defaults
+                        to 2 for exponential backoff.
+        :param max_jitter: Additional max jitter period to wait between retry
+                           attempts to avoid slamming the server.
+
+        """
+        self.sleep_func = sleep_func
+        self.max_tries = max_tries
+        self.delay = delay
+        self.backoff = backoff
+        self.max_jitter = int(max_jitter * 100)
+        self._attempts = 0
+        self._cur_delay = delay
+
+    def reset(self):
+        """Reset the attempt counter"""
+        self._attempts = 0
+        self._cur_delay = self.delay
+
+    def increment(self):
+        """Increment the failed count, and sleep appropriately before
+        continuing"""
+        if self._attempts == self.max_tries:
+            raise
+        self._attempts += 1
+        jitter = random.randint(0, self.max_jitter) / 100.0
+        self.sleep_func(self.delay + jitter)
+        self._cur_delay *= self.backoff
+
+    def copy(self):
+        """Return a clone of this retry sleeper"""
+        return RetySleeper(self.max_tries, self.delay, self.backoff,
+                           self.max_jitter, self.sleep_func)
+
+
 class KazooRetry(object):
     """Helper for retrying a method in the face of retry-able exceptions"""
     RETRY_EXCEPTIONS = (
@@ -46,11 +90,9 @@ class KazooRetry(object):
                               and treated as a retry-able command.
 
         """
+        self.retry = RetySleeper(max_tries, delay, backoff, max_jitter,
+                                 sleep_func)
         self.sleep_func = sleep_func
-        self.max_tries = max_tries
-        self.delay = delay
-        self.backoff = backoff
-        self.max_jitter = int(max_jitter * 100)
         self.retry_exceptions = self.RETRY_EXCEPTIONS
         if ignore_expire:
             self.retry_exceptions += self.EXPIRED_EXCEPTIONS
@@ -59,17 +101,11 @@ class KazooRetry(object):
         self(func, *args, **kwargs)
 
     def __call__(self, func, *args, **kwargs):
-        tries = 1
-        delay = self.delay
+        self.retry.reset()
 
         while True:
             try:
                 return func(*args, **kwargs)
 
             except self.retry_exceptions:
-                if tries == self.max_tries:
-                    raise
-                tries += 1
-                jitter = random.randint(0, self.max_jitter) / 100.0
-                self.sleep_func(delay + jitter)
-                delay *= self.backoff
+                self.retry.increment()