Skip to content

Commit 68c46ae

Browse files
authored
gh-97545: Make Semaphore run faster. (#97549)
1 parent d68c37c commit 68c46ae

File tree

3 files changed

+19
-23
lines changed

3 files changed

+19
-23
lines changed

Lib/asyncio/locks.py

+17-21
Original file line numberDiff line numberDiff line change
@@ -356,8 +356,9 @@ def __repr__(self):
356356
return f'<{res[1:-1]} [{extra}]>'
357357

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

362363
async def acquire(self):
363364
"""Acquire a semaphore.
@@ -368,8 +369,7 @@ async def acquire(self):
368369
called release() to make it larger than 0, and then return
369370
True.
370371
"""
371-
if (not self.locked() and (self._waiters is None or
372-
all(w.cancelled() for w in self._waiters))):
372+
if not self.locked():
373373
self._value -= 1
374374
return True
375375

@@ -387,13 +387,13 @@ async def acquire(self):
387387
finally:
388388
self._waiters.remove(fut)
389389
except exceptions.CancelledError:
390-
if not self.locked():
391-
self._wake_up_first()
390+
if not fut.cancelled():
391+
self._value += 1
392+
self._wake_up_next()
392393
raise
393394

394-
self._value -= 1
395-
if not self.locked():
396-
self._wake_up_first()
395+
if self._value > 0:
396+
self._wake_up_next()
397397
return True
398398

399399
def release(self):
@@ -403,22 +403,18 @@ def release(self):
403403
become larger than zero again, wake up that coroutine.
404404
"""
405405
self._value += 1
406-
self._wake_up_first()
406+
self._wake_up_next()
407407

408-
def _wake_up_first(self):
409-
"""Wake up the first waiter if it isn't done."""
408+
def _wake_up_next(self):
409+
"""Wake up the first waiter that isn't done."""
410410
if not self._waiters:
411411
return
412-
try:
413-
fut = next(iter(self._waiters))
414-
except StopIteration:
415-
return
416412

417-
# .done() necessarily means that a waiter will wake up later on and
418-
# either take the lock, or, if it was cancelled and lock wasn't
419-
# taken already, will hit this again and wake up a new waiter.
420-
if not fut.done():
421-
fut.set_result(True)
413+
for fut in self._waiters:
414+
if not fut.done():
415+
self._value -= 1
416+
fut.set_result(True)
417+
return
422418

423419

424420
class BoundedSemaphore(Semaphore):

Lib/test/test_asyncio/test_locks.py

+1-2
Original file line numberDiff line numberDiff line change
@@ -844,9 +844,8 @@ async def c4(result):
844844

845845
sem.release()
846846
sem.release()
847-
self.assertEqual(2, sem._value)
847+
self.assertEqual(0, sem._value)
848848

849-
await asyncio.sleep(0)
850849
await asyncio.sleep(0)
851850
self.assertEqual(0, sem._value)
852851
self.assertEqual(3, len(result))
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Make Semaphore run faster.

0 commit comments

Comments
 (0)