feat: initial event loop

Author: brian14708

.github/workflows/ci.yaml

@@ -12,12 +12,16 @@ permissions:
 jobs:
   python:
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
+
     steps:
       - uses: actions/checkout@v5
       - name: "Set up Python"
         uses: actions/setup-python@v6
         with:
-          python-version-file: "pyproject.toml"
+          python-version: ${{ matrix.python-version }}
       - name: Install uv
         uses: astral-sh/setup-uv@v6
         with:

flake.nix

@@ -30,12 +30,11 @@
             default = pkgs.mkShell {
               buildInputs = with pkgs; [
                 gnumake
-                (python313.withPackages (
-                  p: with p; [
-                    uv
-                  ]
-                ))
+                uv
               ];
+              env = {
+                UV_PYTHON = pkgs.python310.interpreter;
+              };
             };
           };
 

main.py

@@ -6,6 +6,7 @@ readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
     "pydantic>=2.11.7",
+    "typing-extensions>=4.15.0",
 ]
 
 [build-system]
@@ -53,3 +54,4 @@ venv = ".venv"
 venvPath = "."
 reportAny = false
 reportExplicitAny = false
+reportUnreachable = false

pyproject.toml

@@ -0,0 +1,339 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import contextvars
+import heapq
+import logging
+import threading
+from asyncio import AbstractEventLoop, Handle, Task, TimerHandle, events
+from collections import deque
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, cast, overload
+
+from typing_extensions import (
+    Any,
+    TypeVar,
+    TypeVarTuple,
+    Unpack,
+    override,
+)
+
+from duron.utils import mix_id
+
+if TYPE_CHECKING:
+    import sys
+    from asyncio.futures import Future
+    from collections.abc import Callable, Coroutine, Generator
+    from contextvars import Context
+
+    _T = TypeVar("_T")
+    _Ts = TypeVarTuple("_Ts")
+
+    if sys.version_info >= (3, 12):
+        _TaskCompatibleCoro = Coroutine[Any, Any, _T]
+    else:
+        _TaskCompatibleCoro = Generator[Any, None, _T] | Coroutine[Any, Any, _T]
+
+
+logger = logging.getLogger(__name__)
+
+
+_task_id_ctx: contextvars.ContextVar[TaskCtx] = contextvars.ContextVar("task_id")
+
+
+class OpFuture(asyncio.Future[object]):
+    id: str
+    params: object
+
+    def __init__(self, id: str, params: object, loop: EventLoop) -> None:
+        super().__init__(loop=loop)
+        self.id = id
+        self.params = params
+
+
+@dataclass(eq=True, frozen=True, slots=True)
+class WaitSet:
+    ops: frozenset[OpFuture]
+    timer: float | None
+    event: asyncio.Event
+
+    async def wait(self, now_ns: int) -> None:
+        if self.timer is None:
+            _ = await self.event.wait()
+            return
+        with contextlib.suppress(asyncio.TimeoutError):
+            t = self.timer - (now_ns / 1e9)
+            if t <= 0:
+                return
+            _ = await asyncio.wait_for(self.event.wait(), timeout=t)
+
+
+@dataclass(slots=True)
+class TaskCtx:
+    parent_id: str
+    value: int
+
+
+class EventLoop(AbstractEventLoop):
+    def __init__(self) -> None:
+        self._ready: deque[Handle] = deque()
+        self._debug: bool = False
+        self._exc_handler: (
+            Callable[[AbstractEventLoop, dict[str, object]], object] | None
+        ) = None
+        self._blocked: dict[str, OpFuture] = {}
+        self._root_seq: int = 0
+        self._now_ns: int = 0
+        self._closed: bool = False
+        self._event: asyncio.Event = asyncio.Event()
+
+        self._lock: threading.Lock = threading.Lock()
+        self._timers: list[TimerHandle] = []
+
+    def _generate_id(self) -> str:
+        ctx = _task_id_ctx.get(None)
+        if not ctx:
+            parent_id = ""
+            seq = self._root_seq
+            self._root_seq += 1
+        else:
+            parent_id = ctx.parent_id
+            seq = ctx.value
+            ctx.value += 1
+
+        return mix_id(parent_id.encode(), int(seq).to_bytes(8, "big"))
+
+    @override
+    def call_soon(
+        self,
+        callback: Callable[[Unpack[_Ts]], object],
+        *args: Unpack[_Ts],
+        context: Context | None = None,
+    ) -> Handle:
+        h = Handle(
+            callback,
+            args,
+            self,
+            context=context,
+        )
+        self._ready.append(h)
+        return h
+
+    @override
+    def call_soon_threadsafe(
+        self,
+        callback: Callable[[Unpack[_Ts]], object],
+        *args: Unpack[_Ts],
+        context: Context | None = None,
+        task_id: str | None = None,
+    ) -> Handle:
+        self._event.set()
+        h = TimerHandle(
+            0,
+            callback,
+            args,
+            self,
+            context=self._context_with_task_id(context, task_id=task_id),
+        )
+        with self._lock:
+            heapq.heappush(self._timers, h)
+        return h
+
+    @override
+    def call_at(
+        self,
+        when: float,
+        callback: Callable[[Unpack[_Ts]], object],
+        *args: Unpack[_Ts],
+        context: Context | None = None,
+    ) -> TimerHandle:
+        th = TimerHandle(
+            when,
+            callback,
+            args,
+            loop=self,
+            context=self._context_with_task_id(context),
+        )
+        with self._lock:
+            heapq.heappush(self._timers, th)
+        return th
+
+    @override
+    def call_later(
+        self,
+        delay: float,
+        callback: Callable[[Unpack[_Ts]], object],
+        *args: Unpack[_Ts],
+        context: Context | None = None,
+    ) -> TimerHandle:
+        return self.call_at(self.time() + delay, callback, *args, context=context)
+
+    @override
+    def time(self) -> float:
+        return self._now_ns / 1e9
+
+    def tick(self, time: int) -> None:
+        self._now_ns = time
+
+    @override
+    def create_future(self) -> asyncio.Future[object]:
+        return asyncio.Future(loop=self)
+
+    @override
+    def create_task(
+        self,
+        coro: _TaskCompatibleCoro[_T],
+        *,
+        name: str | None = None,
+        context: Context | None = None,
+        **kwargs: object,
+    ) -> Task[_T]:
+        ctx = self._context_with_task_id(context, task_id=self._generate_id())
+        return ctx.run(
+            cast("type[Task[_T]]", Task), coro, name=name, loop=self, **kwargs
+        )
+
+    def _run_once(self) -> bool:
+        found = False
+        now = self.time()
+        with self._lock:
+            while self._timers and (
+                self._timers[0].when() <= now or self._timers[0].cancelled()
+            ):
+                found = True
+                ht = heapq.heappop(self._timers)
+                if not ht.cancelled():
+                    self._ready.append(ht)
+
+        while self._ready:
+            h = self._ready.popleft()
+            found = True
+            if h.cancelled():
+                continue
+            try:
+                h._run()
+            except BaseException as exc:
+                self.call_exception_handler(
+                    {"message": "exception in callback", "exception": exc, "handle": h}
+                )
+        return found
+
+    def poll_completion(self, task: Future[_T]) -> None | WaitSet:
+        old = events._get_running_loop()
+        events._set_running_loop(self)
+        try:
+            self._event.clear()
+            while self._run_once():
+                if task.done():
+                    return None
+
+            with self._lock:
+                return WaitSet(
+                    ops=frozenset(self._blocked.values()),
+                    timer=self._timers[0].when() if self._timers else None,
+                    event=self._event,
+                )
+        finally:
+            events._set_running_loop(old)
+
+    def create_op(self, params: object) -> OpFuture:
+        id = self._generate_id()
+        s = OpFuture(id, params, self)
+        self._blocked[id] = s
+        return s
+
+    @overload
+    def post_completion_threadsafe(
+        self,
+        id: str,
+        *,
+        result: object,
+    ) -> None: ...
+    @overload
+    def post_completion_threadsafe(
+        self,
+        id: str,
+        *,
+        exception: BaseException,
+    ) -> None: ...
+    def post_completion_threadsafe(
+        self,
+        id: str,
+        *,
+        result: object = None,
+        exception: BaseException | None = None,
+    ) -> None:
+        sc = self._blocked.pop(id, None)
+        if sc is None:
+            return
+
+        tid = mix_id(sc.id.encode(), b"end")
+        if exception is not None:
+            _ = self.call_soon_threadsafe(sc.set_exception, exception, task_id=tid)
+        else:
+            _ = self.call_soon_threadsafe(sc.set_result, result, task_id=tid)
+
+    @override
+    def is_closed(self) -> bool:
+        return self._closed
+
+    @override
+    def close(self) -> None:
+        self._ready.clear()
+        self._timers.clear()
+        self._closed = True
+
+    @override
+    def get_debug(self) -> bool:
+        return self._debug
+
+    @override
+    def set_debug(self, enabled: bool) -> None:
+        self._debug = enabled
+
+    @override
+    def default_exception_handler(self, context: dict[str, object]) -> None:
+        msg = context.get("message", "Unhandled exception")
+        exc = context.get("exception")
+        if exc:
+            logger.error("%s: %r", msg, exc)
+        else:
+            logger.error("%s", msg)
+
+    @override
+    def set_exception_handler(
+        self, handler: Callable[[AbstractEventLoop, dict[str, object]], object] | None
+    ) -> None:
+        self._exc_handler = handler
+
+    @override
+    def call_exception_handler(self, context: dict[str, object]) -> None:
+        if self._exc_handler is None:
+            self.default_exception_handler(context)
+        else:
+            _ = self._exc_handler(self, context)
+
+    @override
+    async def shutdown_asyncgens(self):
+        pass
+
+    @override
+    async def shutdown_default_executor(self):
+        pass
+
+    def _timer_handle_cancelled(self, _th: TimerHandle) -> None:
+        pass
+
+    def _context_with_task_id(
+        self, context: Context | None, task_id: str | None = None
+    ) -> Context:
+        base = context or contextvars.copy_context()
+        if task_id is None:
+            task_id = self._generate_id()
+            _ = base.run(_task_id_ctx.set, TaskCtx(parent_id=task_id, value=0))
+        return base
+
+
+def create_loop() -> EventLoop:
+    return EventLoop()  # type: ignore[abstract]