Skip to content

Commit

Permalink
gh-97545: Make Semaphore run faster. (#97549)
Browse files Browse the repository at this point in the history
  • Loading branch information
cykerway authored Sep 26, 2022
1 parent d68c37c commit 68c46ae
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 23 deletions.
38 changes: 17 additions & 21 deletions Lib/asyncio/locks.py
Original file line number Diff line number Diff line change
Expand Up @@ -356,8 +356,9 @@ def __repr__(self):
return f'<{res[1:-1]} [{extra}]>'

def locked(self):
"""Returns True if semaphore counter is zero."""
return self._value == 0
"""Returns True if semaphore cannot be acquired immediately."""
return self._value == 0 or (
any(not w.cancelled() for w in (self._waiters or ())))

async def acquire(self):
"""Acquire a semaphore.
Expand All @@ -368,8 +369,7 @@ async def acquire(self):
called release() to make it larger than 0, and then return
True.
"""
if (not self.locked() and (self._waiters is None or
all(w.cancelled() for w in self._waiters))):
if not self.locked():
self._value -= 1
return True

Expand All @@ -387,13 +387,13 @@ async def acquire(self):
finally:
self._waiters.remove(fut)
except exceptions.CancelledError:
if not self.locked():
self._wake_up_first()
if not fut.cancelled():
self._value += 1
self._wake_up_next()
raise

self._value -= 1
if not self.locked():
self._wake_up_first()
if self._value > 0:
self._wake_up_next()
return True

def release(self):
Expand All @@ -403,22 +403,18 @@ def release(self):
become larger than zero again, wake up that coroutine.
"""
self._value += 1
self._wake_up_first()
self._wake_up_next()

def _wake_up_first(self):
"""Wake up the first waiter if it isn't done."""
def _wake_up_next(self):
"""Wake up the first waiter that isn't done."""
if not self._waiters:
return
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)
for fut in self._waiters:
if not fut.done():
self._value -= 1
fut.set_result(True)
return


class BoundedSemaphore(Semaphore):
Expand Down
3 changes: 1 addition & 2 deletions Lib/test/test_asyncio/test_locks.py
Original file line number Diff line number Diff line change
Expand Up @@ -844,9 +844,8 @@ async def c4(result):

sem.release()
sem.release()
self.assertEqual(2, sem._value)
self.assertEqual(0, sem._value)

await asyncio.sleep(0)
await asyncio.sleep(0)
self.assertEqual(0, sem._value)
self.assertEqual(3, len(result))
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Make Semaphore run faster.

0 comments on commit 68c46ae

Please sign in to comment.