gh-97545: Make Semaphore run faster. (GH-97549)

(cherry picked from commit 68c46ae68b)

Co-authored-by: Cyker Way <cykerway@gmail.com>
This commit is contained in:
Miss Islington (bot) 2022-09-26 17:00:53 -07:00 committed by GitHub
parent 82932b3ec9
commit 9a9bf88898
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 19 additions and 23 deletions

View file

@ -357,8 +357,9 @@ class Semaphore(_ContextManagerMixin, mixins._LoopBoundMixin):
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.
@ -369,8 +370,7 @@ class Semaphore(_ContextManagerMixin, mixins._LoopBoundMixin):
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
@ -388,13 +388,13 @@ class Semaphore(_ContextManagerMixin, mixins._LoopBoundMixin):
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):
@ -404,22 +404,18 @@ class Semaphore(_ContextManagerMixin, mixins._LoopBoundMixin):
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):