Merge "Improve multilock class and its associated unit test"

2 files changed, 20 insertions, 7 deletions

diff --git a/taskflow/utils/lock_utils.py b/taskflow/utils/lock_utils.py
index ea2023e..1ee924b 100644
--- a/taskflow/utils/lock_utils.py
+++ b/taskflow/utils/lock_utils.py
@@ -363,21 +363,19 @@ class MultiLock(object):
         # Cleans off one level of the stack (this is done so that if there
         # are multiple __enter__() and __exit__() pairs active that this will
         # only remove one level (the last one), and not all levels...
-        leftover = self._lock_stacks[-1]
-        while leftover:
-            lock = self._locks[leftover - 1]
+        for left in misc.countdown_iter(self._lock_stacks[-1]):
+            lock_idx = left - 1
+            lock = self._locks[lock_idx]
             try:
                 lock.release()
             except (threading.ThreadError, RuntimeError) as e:
                 # Ensure that we adjust the lock stack under failure so that
                 # if release is attempted again that we do not try to release
                 # the locks we already released...
-                self._lock_stacks[-1] = leftover
+                self._lock_stacks[-1] = left
                 raise threading.ThreadError(
                     "Unable to release lock %s/%s due to '%s'"
-                    % (leftover, len(self._locks), e))
-            else:
-                leftover -= 1
+                    % (left, len(self._locks), e))
         # At the end only clear it off, so that under partial failure we don't
         # lose any locks...
         self._lock_stacks.pop()
diff --git a/taskflow/utils/misc.py b/taskflow/utils/misc.py
index 299082a..b2e967d 100644
--- a/taskflow/utils/misc.py
+++ b/taskflow/utils/misc.py
@@ -74,6 +74,21 @@ def find_monotonic(allow_time_time=False):
         return None
 
 
+def countdown_iter(start_at, decr=1):
+    """Generator that decrements after each generation until <= zero.
+
+    NOTE(harlowja): we can likely remove this when we can use an
+    ``itertools.count`` that takes a step (on py2.6 which we still support
+    that step parameter does **not** exist and therefore can't be used).
+    """
+    if decr <= 0:
+        raise ValueError("Decrement value must be greater"
+                         " than zero and not %s" % decr)
+    while start_at > 0:
+        yield start_at
+        start_at -= decr
+
+
 def merge_uri(uri, conf):
     """Merges a parsed uri into the given configuration dictionary.