Use a mini-watch helper to ensure timeout is decreased

After the initial lock is acquired the timeout can now
be different (due to time used getting the lock) so add
a little helper class to ensure that the timeout is
changed due to the elapsed time.

Author: Joshua Harlow

kazoo/recipe/lock.py

@@ -16,6 +16,10 @@
 """
 
 import sys
+try:
+    from time import monotonic as now
+except ImportError:
+    from time import time as now
 import uuid
 
 import six
@@ -32,6 +36,22 @@
 from kazoo.protocol.states import KazooState
 
 
+class _Watch(object):
+    def __init__(self, duration=None):
+        self.duration = duration
+        self.started_at = None
+
+    def start(self):
+        self.started_at = now()
+
+    def leftover(self):
+        if self.duration is None:
+            return None
+        else:
+            elapsed = now() - self.started_at
+            return max(0, self.duration - elapsed)
+
+
 class Lock(object):
     """Kazoo Lock
 
@@ -448,8 +468,10 @@ def _inner_acquire(self, blocking, timeout=None):
         if self.client.exists(self.create_path):
             return True
 
+        w = _Watch(duration=timeout)
+        w.start()
         lock = self.client.Lock(self.lock_path, self.data)
-        gotten = lock.acquire(blocking=blocking, timeout=timeout)
+        gotten = lock.acquire(blocking=blocking, timeout=w.leftover())
         if not gotten:
             return False
         try:
@@ -463,7 +485,7 @@ def _inner_acquire(self, blocking, timeout=None):
                 if blocking:
                     # If blocking, wait until self._watch_lease_change() is
                     # called before returning
-                    self.wake_event.wait(timeout)
+                    self.wake_event.wait(w.leftover())
                     if not self.wake_event.isSet():
                         raise LockTimeout(
                             "Failed to acquire semaphore on %s "