Issue 10260

Adding the wait_for() method to threading.Condition

2 files changed, 72 insertions, 7 deletions

diff --git a/Lib/test/lock_tests.py b/Lib/test/lock_tests.py
index d956bb6..b6d818e 100644
--- a/Lib/test/lock_tests.py
+++ b/Lib/test/lock_tests.py
@@ -446,6 +446,46 @@ class ConditionTests(BaseTestCase):
             # In practice, this implementation has no spurious wakeups.
             self.assertFalse(result)
 
+    def test_waitfor(self):
+        cond = self.condtype()
+        state = 0
+        def f():
+            with cond:
+                result = cond.wait_for(lambda : state==4)
+                self.assertTrue(result)
+                self.assertEqual(state, 4)
+        b = Bunch(f, 1)
+        b.wait_for_started()
+        for i in range(5):
+            time.sleep(0.01)
+            with cond:
+                state += 1
+                cond.notify()
+        b.wait_for_finished()
+
+    def test_waitfor_timeout(self):
+        cond = self.condtype()
+        state = 0
+        success = []
+        def f():
+            with cond:
+                dt = time.time()
+                result = cond.wait_for(lambda : state==4, timeout=0.1)
+                dt = time.time() - dt
+                self.assertFalse(result)
+                self.assertTimeout(dt, 0.1)
+                success.append(None)
+        b = Bunch(f, 1)
+        b.wait_for_started()
+        # Only increment 3 times, so state == 4 is never reached.
+        for i in range(3):
+            time.sleep(0.01)
+            with cond:
+                state += 1
+                cond.notify()
+        b.wait_for_finished()
+        self.assertEqual(len(success), 1)
+
 
 class BaseSemaphoreTests(BaseTestCase):
     """
diff --git a/Lib/threading.py b/Lib/threading.py
index 41956ed..b6c1e5d 100644
--- a/Lib/threading.py
+++ b/Lib/threading.py
@@ -254,6 +254,32 @@ class _Condition(_Verbose):
         finally:
             self._acquire_restore(saved_state)
 
+    def wait_for(self, predicate, timeout=None):
+        endtime = None
+        waittime = timeout
+        result = predicate()
+        while not result:
+            if waittime is not None:
+                if endtime is None:
+                    endtime = _time() + waittime
+                else:
+                    waittime = endtime - _time()
+                    if waittime <= 0:
+                        if __debug__:
+                            self._note("%s.wait_for(%r, %r): Timed out.",
+                                       self, predicate, timeout)
+                        break
+            if __debug__:
+                self._note("%s.wait_for(%r, %r): Waiting with timeout=%s.",
+                           self, predicate, timeout, waittime)
+            self.wait(waittime)
+            result = predicate()
+        else:
+            if __debug__:
+                self._note("%s.wait_for(%r, %r): Success.",
+                           self, predicate, timeout)
+        return result
+
     def notify(self, n=1):
         if not self._is_owned():
             raise RuntimeError("cannot notify on un-acquired lock")
@@ -482,13 +508,12 @@ class Barrier(_Verbose):
     # Wait in the barrier until we are relased.  Raise an exception
     # if the barrier is reset or broken.
     def _wait(self, timeout):
-        while self._state == 0:
-            if self._cond.wait(timeout) is False:
-                #timed out.  Break the barrier
-                self._break()
-                raise BrokenBarrierError
-            if self._state < 0:
-                raise BrokenBarrierError
+        if not self._cond.wait_for(lambda : self._state != 0, timeout):
+            #timed out.  Break the barrier
+            self._break()
+            raise BrokenBarrierError
+        if self._state < 0:
+            raise BrokenBarrierError
         assert self._state == 1
 
     # If we are the last thread to exit the barrier, signal any threads