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 16:57:52 -07:00 committed by GitHub
parent 72a78152f3
commit 232156144c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 19 additions and 23 deletions

View file

@ -360,8 +360,9 @@ class Semaphore(_ContextManagerMixin, mixins._LoopBoundMixin):
return f'<{res[1:-1]} [{extra}]>' return f'<{res[1:-1]} [{extra}]>'
def locked(self): def locked(self):
"""Returns True if semaphore counter is zero.""" """Returns True if semaphore cannot be acquired immediately."""
return self._value == 0 return self._value == 0 or (
any(not w.cancelled() for w in (self._waiters or ())))
async def acquire(self): async def acquire(self):
"""Acquire a semaphore. """Acquire a semaphore.
@ -372,8 +373,7 @@ class Semaphore(_ContextManagerMixin, mixins._LoopBoundMixin):
called release() to make it larger than 0, and then return called release() to make it larger than 0, and then return
True. True.
""" """
if (not self.locked() and (self._waiters is None or if not self.locked():
all(w.cancelled() for w in self._waiters))):
self._value -= 1 self._value -= 1
return True return True
@ -391,13 +391,13 @@ class Semaphore(_ContextManagerMixin, mixins._LoopBoundMixin):
finally: finally:
self._waiters.remove(fut) self._waiters.remove(fut)
except exceptions.CancelledError: except exceptions.CancelledError:
if not self.locked(): if not fut.cancelled():
self._wake_up_first() self._value += 1
self._wake_up_next()
raise raise
self._value -= 1 if self._value > 0:
if not self.locked(): self._wake_up_next()
self._wake_up_first()
return True return True
def release(self): def release(self):
@ -407,22 +407,18 @@ class Semaphore(_ContextManagerMixin, mixins._LoopBoundMixin):
become larger than zero again, wake up that coroutine. become larger than zero again, wake up that coroutine.
""" """
self._value += 1 self._value += 1
self._wake_up_first() self._wake_up_next()
def _wake_up_first(self): def _wake_up_next(self):
"""Wake up the first waiter if it isn't done.""" """Wake up the first waiter that isn't done."""
if not self._waiters: if not self._waiters:
return return
try:
fut = next(iter(self._waiters))
except StopIteration:
return
# .done() necessarily means that a waiter will wake up later on and for fut in self._waiters:
# either take the lock, or, if it was cancelled and lock wasn't if not fut.done():
# taken already, will hit this again and wake up a new waiter. self._value -= 1
if not fut.done(): fut.set_result(True)
fut.set_result(True) return
class BoundedSemaphore(Semaphore): class BoundedSemaphore(Semaphore):

View file

@ -857,9 +857,8 @@ class SemaphoreTests(unittest.IsolatedAsyncioTestCase):
sem.release() sem.release()
sem.release() sem.release()
self.assertEqual(2, sem._value) self.assertEqual(0, sem._value)
await asyncio.sleep(0)
await asyncio.sleep(0) await asyncio.sleep(0)
self.assertEqual(0, sem._value) self.assertEqual(0, sem._value)
self.assertEqual(3, len(result)) self.assertEqual(3, len(result))

View file

@ -0,0 +1 @@
Make Semaphore run faster.