Skip to content

Commit 2321561

Browse files
gh-97545: Make Semaphore run faster. (GH-97549)
(cherry picked from commit 68c46ae) Co-authored-by: Cyker Way <[email protected]>
1 parent 72a7815 commit 2321561

File tree

3 files changed

+19
-23
lines changed

3 files changed

+19
-23
lines changed

Lib/asyncio/locks.py

Lines changed: 17 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -360,8 +360,9 @@ def __repr__(self):
360360
return f'<{res[1:-1]} [{extra}]>'
361361

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

366367
async def acquire(self):
367368
"""Acquire a semaphore.
@@ -372,8 +373,7 @@ async def acquire(self):
372373
called release() to make it larger than 0, and then return
373374
True.
374375
"""
375-
if (not self.locked() and (self._waiters is None or
376-
all(w.cancelled() for w in self._waiters))):
376+
if not self.locked():
377377
self._value -= 1
378378
return True
379379

@@ -391,13 +391,13 @@ async def acquire(self):
391391
finally:
392392
self._waiters.remove(fut)
393393
except exceptions.CancelledError:
394-
if not self.locked():
395-
self._wake_up_first()
394+
if not fut.cancelled():
395+
self._value += 1
396+
self._wake_up_next()
396397
raise
397398

398-
self._value -= 1
399-
if not self.locked():
400-
self._wake_up_first()
399+
if self._value > 0:
400+
self._wake_up_next()
401401
return True
402402

403403
def release(self):
@@ -407,22 +407,18 @@ def release(self):
407407
become larger than zero again, wake up that coroutine.
408408
"""
409409
self._value += 1
410-
self._wake_up_first()
410+
self._wake_up_next()
411411

412-
def _wake_up_first(self):
413-
"""Wake up the first waiter if it isn't done."""
412+
def _wake_up_next(self):
413+
"""Wake up the first waiter that isn't done."""
414414
if not self._waiters:
415415
return
416-
try:
417-
fut = next(iter(self._waiters))
418-
except StopIteration:
419-
return
420416

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

427423

428424
class BoundedSemaphore(Semaphore):

Lib/test/test_asyncio/test_locks.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -857,9 +857,8 @@ async def c4(result):
857857

858858
sem.release()
859859
sem.release()
860-
self.assertEqual(2, sem._value)
860+
self.assertEqual(0, sem._value)
861861

862-
await asyncio.sleep(0)
863862
await asyncio.sleep(0)
864863
self.assertEqual(0, sem._value)
865864
self.assertEqual(3, len(result))
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Make Semaphore run faster.

0 commit comments

Comments
 (0)