...

Author: mikeshardmind

_misc/_l.py

@@ -0,0 +1,22 @@
+import asyncio
+import random
+
+from async_utils._simple_lock import AsyncLock  # noqa: PLC2701
+from async_utils.bg_loop import threaded_loop
+
+
+async def check(lock: AsyncLock):
+    async with lock:
+        v = random.random()
+        print(await asyncio.sleep(v, v), flush=True)  # noqa: T201
+
+
+async def amain():
+    lock = AsyncLock()
+    with threaded_loop() as tl1, threaded_loop() as tl2:
+        tsks = {loop.run(check(lock)) for loop in (tl1, tl2) for _ in range(10)}
+        await asyncio.gather(*tsks)
+
+
+if __name__ == "__main__":
+    asyncio.run(amain())

src/async_utils/_simple_lock.py

@@ -16,10 +16,38 @@
 import concurrent.futures as cf
 import threading
 from collections import deque
+from collections.abc import Generator
 
 from . import _typings as t
 
 # TODO: pick what public namespace to re-export this from.
+# TODO: write a test for the odd behavior noticed with prior version in discord
+
+
+class _Waiter:
+    __slots__ = ("future",)
+
+    def __init__(self, future: cf.Future[None], /) -> None:
+        self.future: cf.Future[None] = future
+
+    def cancelled(self) -> bool:
+        return self.future.cancelled()
+
+    def done(self) -> bool:
+        return self.future.done()
+
+    def set_result(self, val: None) -> None:
+        self.future.set_result(val)
+
+    def __await__(self) -> Generator[t.Any, t.Any, None]:
+        f = asyncio.wrap_future(self.future)
+        return (yield from f.__await__())
+
+    __final__ = True
+
+    def __init_subclass__(cls) -> t.Never:
+        msg = "Don't subclass this"
+        raise RuntimeError(msg)
 
 
 class AsyncLock:
@@ -32,44 +60,65 @@ def __init_subclass__(cls) -> t.Never:
     __final__ = True
 
     def __init__(self) -> None:
-        self._waiters: deque[cf.Future[None]] = deque()
+        self._waiters: deque[_Waiter] | None = None
+        self._lockv: bool = False
         self._internal_lock: threading.RLock = threading.RLock()
-        self._locked: bool = False
 
-    async def __aenter__(self, /) -> None:
+    def __locked(self) -> bool:
         with self._internal_lock:
-            if not self._locked and (all(w.cancelled() for w in self._waiters)):
-                self._locked = True
-                return
+            return self._lockv or (any(not w.cancelled() for w in (self._waiters or ())))
+
+    async def __aenter__(self) -> None:
+        await self.__acquire()
+
+    async def __aexit__(self, *dont_care: object) -> t.Literal[False]:
+        self.__release()
+        return False
+
+    async def __acquire(self) -> bool:
+        with self._internal_lock:
+            if not self.__locked():
+                self._lockv = True
+                return True
+
+        with self._internal_lock:
+            if self._waiters is None:
+                self._waiters = deque()
 
         fut: cf.Future[None] = cf.Future()
 
+        waiter = _Waiter(fut)
+
         with self._internal_lock:
-            self._waiters.append(fut)
+            self._waiters.append(waiter)
 
         try:
-            await asyncio.wrap_future(fut)
-        except (asyncio.CancelledError, cf.CancelledError):
-            with self._internal_lock:
-                if self._locked:
-                    self._maybe_wake()
+            await waiter
+        except asyncio.CancelledError:
+            if fut.done() and not fut.cancelled():
+                self._lockv = False
+            raise
+
         finally:
-            self._waiters.remove(fut)
+            self._maybe_wake()
+        return True
 
-    async def __aexit__(self, *_dont_care: object) -> t.Literal[False]:
+    def _maybe_wake(self) -> None:
         with self._internal_lock:
-            if self._locked:
-                self._locked = False
-                self._maybe_wake()
+            while (not self._lockv) and self._waiters:
+                next_waiter = self._waiters.popleft()
 
-        return False
+                if not (next_waiter.done() or next_waiter.cancelled()):
+                    self._lockv = True
+                    next_waiter.set_result(None)
 
-    def _maybe_wake(self) -> None:
+            while self._waiters:
+                next_waiter = self._waiters.popleft()
+                if not (next_waiter.done() or next_waiter.cancelled()):
+                    self._waiters.appendleft(next_waiter)
+                    break
+
+    def __release(self) -> None:
         with self._internal_lock:
-            if self._waiters:
-                try:
-                    fut = next(iter(self._waiters))
-                except StopIteration:
-                    return
-                if not fut.done():
-                    fut.set_result(None)
+            self._lockv = False
+            self._maybe_wake()