diff options
Diffstat (limited to 'Lib')
-rw-r--r-- | Lib/asyncio/locks.py | 32 | ||||
-rw-r--r-- | Lib/test/test_asyncio/test_locks.py | 50 |
2 files changed, 72 insertions, 10 deletions
diff --git a/Lib/asyncio/locks.py b/Lib/asyncio/locks.py index 6193837..508a214 100644 --- a/Lib/asyncio/locks.py +++ b/Lib/asyncio/locks.py @@ -183,16 +183,22 @@ class Lock(_ContextManagerMixin): fut = self._loop.create_future() self._waiters.append(fut) + + # Finally block should be called before the CancelledError + # handling as we don't want CancelledError to call + # _wake_up_first() and attempt to wake up itself. try: - await fut - self._locked = True - return True + try: + await fut + finally: + self._waiters.remove(fut) except futures.CancelledError: if not self._locked: self._wake_up_first() raise - finally: - self._waiters.remove(fut) + + self._locked = True + return True def release(self): """Release a lock. @@ -212,11 +218,17 @@ class Lock(_ContextManagerMixin): raise RuntimeError('Lock is not acquired.') def _wake_up_first(self): - """Wake up the first waiter who isn't cancelled.""" - for fut in self._waiters: - if not fut.done(): - fut.set_result(True) - break + """Wake up the first waiter if it isn't done.""" + try: + fut = next(iter(self._waiters)) + except StopIteration: + return + + # .done() necessarily means that a waiter will wake up later on and + # either take the lock, or, if it was cancelled and lock wasn't + # taken already, will hit this again and wake up a new waiter. + if not fut.done(): + fut.set_result(True) class Event: diff --git a/Lib/test/test_asyncio/test_locks.py b/Lib/test/test_asyncio/test_locks.py index 3c50697..3e3dd79 100644 --- a/Lib/test/test_asyncio/test_locks.py +++ b/Lib/test/test_asyncio/test_locks.py @@ -200,6 +200,56 @@ class LockTests(test_utils.TestCase): self.assertTrue(tb.cancelled()) self.assertTrue(tc.done()) + def test_cancel_release_race(self): + # Issue 32734 + # Acquire 4 locks, cancel second, release first + # and 2 locks are taken at once. + lock = asyncio.Lock(loop=self.loop) + lock_count = 0 + call_count = 0 + + async def lockit(): + nonlocal lock_count + nonlocal call_count + call_count += 1 + await lock.acquire() + lock_count += 1 + + async def lockandtrigger(): + await lock.acquire() + self.loop.call_soon(trigger) + + def trigger(): + t1.cancel() + lock.release() + + t0 = self.loop.create_task(lockandtrigger()) + t1 = self.loop.create_task(lockit()) + t2 = self.loop.create_task(lockit()) + t3 = self.loop.create_task(lockit()) + + # First loop acquires all + test_utils.run_briefly(self.loop) + self.assertTrue(t0.done()) + + # Second loop calls trigger + test_utils.run_briefly(self.loop) + # Third loop calls cancellation + test_utils.run_briefly(self.loop) + + # Make sure only one lock was taken + self.assertEqual(lock_count, 1) + # While 3 calls were made to lockit() + self.assertEqual(call_count, 3) + self.assertTrue(t1.cancelled() and t2.done()) + + # Cleanup the task that is stuck on acquire. + t3.cancel() + test_utils.run_briefly(self.loop) + self.assertTrue(t3.cancelled()) + + + def test_finished_waiter_cancelled(self): lock = asyncio.Lock(loop=self.loop) |