Modify lock internals

- There's a lot of internal Cpython behavior attached to why
  the specific yields here are needed.

  I need to go through it in depth later so that I can be more confident
  in the solution prior to a public export

Author: mikeshardmind

_misc/_l.py

@@ -0,0 +1,26 @@
+import asyncio
+import random
+import time
+
+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()
+        s = time.monotonic()
+        await asyncio.sleep(v)
+        e = time.monotonic()
+        print(s, v, e, 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

@@ -33,13 +33,19 @@ def __init_subclass__(cls) -> t.Never:
 
     def __init__(self) -> None:
         self._waiters: deque[cf.Future[None]] = deque()
+        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 self._lockv or (any(not w.cancelled() for w in (self._waiters)))
+
+    async def __aenter__(self) -> None:
+        await asyncio.sleep(0)  # This yield is non-optional.
+
+        with self._internal_lock:
+            if not self.__locked():
+                self._lockv = True
                 return
 
         fut: cf.Future[None] = cf.Future()
@@ -49,27 +55,33 @@ async def __aenter__(self, /) -> None:
 
         try:
             await asyncio.wrap_future(fut)
-        except (asyncio.CancelledError, cf.CancelledError):
-            with self._internal_lock:
-                if self._locked:
-                    self._maybe_wake()
+        except asyncio.CancelledError:
+            if fut.done() and not fut.cancelled():
+                self._lockv = False
+            raise
+
         finally:
-            self._waiters.remove(fut)
+            self._maybe_wake()
+        return
 
-    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
+
+    async def __aexit__(self, *dont_care: object) -> t.Literal[False]:
+        await asyncio.sleep(0)  # this yield is not optional
         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()
+        return False