Start adding N:M concurrency building blocks

Author: mikeshardmind

src/async_utils/_nm_building_blocks.py

@@ -0,0 +1,157 @@
+#   Copyright 2020-present Michael Hall
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+# This used to include CPython code, some minor performance losses have been
+# taken to not tightly include upstream code
+"""Building blocks used to create thread-safe multi-event loop async things."""
+
+from __future__ import annotations
+
+import asyncio
+import heapq
+import threading
+from collections.abc import Callable, Generator
+
+from . import _typings as t
+
+
+def _wake_nm_waiter(fut: asyncio.Future[None]) -> None:
+    if not fut.done():
+        fut.set_result(None)
+
+
+class _NMWaiter:
+    __slots__ = ("_future", "_loop", "_thread_id", "_val")
+
+    def __init__(self, val: int = 1, /) -> None:
+        self._loop = asyncio.get_running_loop()
+        self._future: asyncio.Future[None] = self._loop.create_future()
+        self._thread_id: int = threading.get_ident()
+        self._val = val
+
+    def __lt__(self, other: object) -> bool:
+        if isinstance(other, _NMWaiter):
+            return self._val < other._val
+        return NotImplemented
+
+    def __init_subclass__(cls) -> t.Never:
+        msg = "Don't subclass this"
+        raise RuntimeError(msg)
+
+    __final__ = True
+
+    def __await__(self) -> Generator[t.Any, None, None]:
+        if threading.get_ident() == self._thread_id:
+            return self._future.__await__()
+
+        msg = "Attempted to await future from the wrong thread"
+        raise RuntimeError(msg)
+
+    def wake(self) -> None:
+        self._loop.call_soon_threadsafe(_wake_nm_waiter, self._future)
+
+    def cancelled(self):
+        return self._future.cancelled()
+
+    def done(self):
+        return self._future.done()
+
+    def add_done_callback(self, cb: Callable[[asyncio.Future[None]], None]) -> None:
+        self._loop.call_soon_threadsafe(self._future.add_done_callback, cb)
+
+
+class _WrappedRLock:
+    __slots__ = ("_lock",)
+
+    def __init__(self, lock: threading.RLock) -> None:
+        self._lock = lock
+
+    async def __aenter__(self) -> None:
+        acquired = False
+        while not acquired:
+            acquired = self._lock.acquire(blocking=False)
+            await asyncio.sleep(0)
+
+    async def __aexit__(self, *dont_care: object) -> None:
+        self._lock.release()
+
+
+class NMSemaphoreBase:
+    def __init__(self, value: int = 1, /) -> None:
+        self._waiters: list[_NMWaiter] = []
+        self._lock = _WrappedRLock(threading.RLock())
+        self._value: int = value
+
+    def _get_priority(self) -> int:
+        # This and value are the only differences between a semaphore,
+        # lock, and priority semaphore, and priority lock
+        # TODO:
+        #   - split this before it is in a public importable namespace.
+        #   - simplify aquire for locks
+        return 0 - self._value
+
+    def __repr__(self) -> str:
+        res = super().__repr__()
+        extra = "locked" if self.__locked() else f"unlocked, value:{self._value}"
+        if self._waiters:
+            extra = f"{extra}, waiters:{len(self._waiters)}"
+        return f"<{res[1:-1]} [{extra}]>"
+
+    def __locked(self) -> bool:
+        return self._value == 0 or (any(not w.cancelled() for w in (self._waiters or ())))
+
+    async def __aexit__(self, *dont_care: object) -> None:
+        async with self._lock:
+            self._value += 1
+        await self._maybe_wake()
+
+    async def __aenter__(self) -> None:
+        async with self._lock:
+            if not self.__locked():
+                self._value -= 1
+                return
+
+            waiter = _NMWaiter(self._get_priority())
+            heapq.heappush(self._waiters, waiter)
+
+        try:
+            await waiter
+            # we don't remove ourselves
+            # we need to maintain the heap invariants
+        except asyncio.CancelledError:
+            if waiter.done() and not waiter.cancelled():
+                self._value += 1
+            raise
+
+        finally:
+            await self._maybe_wake()
+
+    async def _maybe_wake(self) -> None:
+        async with self._lock:
+            while self._value > 0 and self._waiters:
+                next_waiter = heapq.heappop(self._waiters)
+
+                if not (next_waiter.done() or next_waiter.cancelled()):
+                    self._value -= 1
+                    next_waiter.wake()
+
+            while self._waiters:
+                # cleanup maintaining heap invariant
+                # This will only fully empty the heap when
+                # all things remaining in the heap after waking tasks in
+                # above section are all done.
+                next_waiter = heapq.heappop(self._waiters)
+                if not (next_waiter.done() or next_waiter.cancelled()):
+                    heapq.heappush(self._waiters, next_waiter)
+                    break