feat: initial event loop

Author: brian14708

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]

pyproject.toml

@@ -0,0 +1,339 @@
+from __future__ import annotations
+
+import asyncio
+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
+
+from typing_extensions import (
+    Any,
+    TypeVar,
+    TypeVarTuple,
+    Unpack,
+    override,
+)
+
+from duron.utils import mix_id
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Coroutine, Generator
+    from contextvars import Context
+
+    T = TypeVar("T")
+    Ts = TypeVarTuple("Ts")
+
+logger = logging.getLogger(__name__)
+
+
+_task_id_ctx: contextvars.ContextVar[TaskCtx] = contextvars.ContextVar("task_id")
+
+
+@dataclass(slots=True)
+class SyscallFuture:
+    id: str
+    params: object
+    _future: asyncio.Future[Any]
+
+    def __await__(self) -> Generator[Any, None, Any]:
+        return self._future.__await__()
+
+    @override
+    def __hash__(self) -> int:
+        return hash(self.id)
+
+    @override
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, SyscallFuture):
+            return NotImplemented
+        return self.id == other.id
+
+
+@dataclass(eq=True, frozen=True, slots=True)
+class WaitSet:
+    syscalls: frozenset[SyscallFuture]
+    timer: float | None
+
+
+@dataclass(slots=True)
+class TaskCtx:
+    parent_id: str
+    value: int
+
+
+class EventLoop(AbstractEventLoop):
+    def __init__(self, event: asyncio.Event) -> None:
+        self._ready: deque[Handle] = deque()
+        self._timers: list[TimerHandle] = []
+        self._cv: threading.Condition = threading.Condition()
+        self._stopping: bool = False
+        self._closed: bool = False
+        self._debug: bool = False
+        self._exc_handler: (
+            Callable[[AbstractEventLoop, dict[str, Any]], object] | None
+        ) = None
+        self._blocked: dict[str, SyscallFuture] = {}
+        self._root_seq: int = 0
+        self._now: int = 0
+        self._event: asyncio.Event = event
+
+    def syscall(self, params: object) -> SyscallFuture:
+        id = self._generate_id()
+        fut = self.create_future()
+        s = SyscallFuture(_future=fut, id=id, params=params)
+        self._blocked[id] = s
+        return s
+
+    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[[*Ts], Any],
+        *args: Unpack[Ts],
+        context: Context | None = None,
+    ) -> Handle:
+        h = Handle(
+            callback,
+            args,
+            self,
+            context=context,
+        )
+        with self._cv:
+            self._ready.append(h)
+            self._cv.notify()
+        return h
+
+    @override
+    def call_soon_threadsafe(
+        self,
+        callback: Callable[[*Ts], Any],
+        *args: Unpack[Ts],
+        context: Context | None = None,
+        task_id: str | None = None,
+    ) -> Handle:
+        self._event.set()
+        h = Handle(
+            callback,
+            args,
+            self,
+            context=self._context_with_task_id(context, task_id=task_id),
+        )
+        with self._cv:
+            self._ready.append(h)
+            self._cv.notify()
+        return h
+
+    @override
+    def call_at(
+        self,
+        when: float,
+        callback: Callable[[*Ts], Any],
+        *args: Unpack[Ts],
+        context: Context | None = None,
+    ) -> TimerHandle:
+        th = TimerHandle(
+            when,
+            callback,
+            args,
+            loop=self,
+            context=self._context_with_task_id(context),
+        )
+        with self._cv:
+            heapq.heappush(self._timers, th)
+            self._cv.notify()
+        return th
+
+    @override
+    def call_later(
+        self,
+        delay: float,
+        callback: Callable[[*Ts], Any],
+        *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 / 1000.0
+
+    def tick(self, time: int) -> None:
+        self._now = time
+
+    @override
+    def create_future(self) -> asyncio.Future[Any]:
+        return asyncio.Future(loop=self)
+
+    @override
+    def create_task(
+        self,
+        coro: Generator[Any, None, T] | Coroutine[Any, Any, T],
+        *,
+        name: str | None = None,
+        context: Context | None = None,
+        **kwargs: Any,
+    ) -> 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) -> None:
+        now = self.time()
+        # promote due timers
+        with self._cv:
+            while self._timers and self._timers[0].when() <= now:
+                ht = heapq.heappop(self._timers)
+                if not ht.cancelled():
+                    self._ready.append(ht)
+
+        # drain ready queue in batches for better performance
+        batch_size = min(len(self._ready), 100)  # Process up to 100 at a time
+        for _ in range(batch_size):
+            if not self._ready:
+                break
+            h = self._ready.popleft()
+            if h.cancelled():
+                continue
+            try:
+                h._run()
+            except BaseException as exc:
+                self.call_exception_handler(
+                    {"message": "exception in callback", "exception": exc, "handle": h}
+                )
+
+    @override
+    def stop(self) -> None:
+        with self._cv:
+            self._stopping = True
+            self._cv.notify_all()
+
+    def poll_completion(self, task: Task[T]) -> None | WaitSet:
+        old = events._get_running_loop()
+        events._set_running_loop(self)
+        try:
+            while not task.done():
+                self._run_once()
+
+                with self._cv:
+                    if self._ready:
+                        continue
+                    if self._timers and self._timers[0].when() <= self.time():
+                        # timer is due
+                        continue
+
+                break
+
+            if task.done():
+                return None
+
+            return WaitSet(
+                syscalls=frozenset(self._blocked.values()),
+                timer=self._timers[0].when() if self._timers else None,
+            )
+        finally:
+            events._set_running_loop(old)
+
+    def complete_syscall(
+        self,
+        syscall_id: str,
+        *,
+        result: Any = None,
+        exception: BaseException | None = None,
+    ) -> None:
+        sc = self._blocked.pop(syscall_id, None)
+        if sc is None:
+            return
+
+        fut = sc._future  # pyright: ignore[reportPrivateUsage]
+        tid = mix_id(sc.id.encode(), b"end")
+        if exception is not None:
+            _ = self.call_soon_threadsafe(fut.set_exception, exception, task_id=tid)
+        else:
+            _ = self.call_soon_threadsafe(fut.set_result, result, task_id=tid)
+
+    @override
+    def is_running(self) -> bool:
+        return not self._stopping
+
+    @override
+    def is_closed(self) -> bool:
+        return self._closed
+
+    @override
+    def close(self) -> None:
+        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, Any]) -> 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, Any]], object] | None
+    ) -> None:
+        self._exc_handler = handler
+
+    @override
+    def call_exception_handler(self, context: dict[str, Any]) -> 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:
+        try:
+            self._timers.remove(th)
+            heapq.heapify(self._timers)
+        except ValueError:
+            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(event: asyncio.Event) -> EventLoop:
+    return EventLoop(event)  # type: ignore[abstract]

src/duron/event_loop.py

@@ -0,0 +1,10 @@
+import base64
+import hashlib
+
+_SEP = b"\x00"
+
+
+def mix_id(*parts: bytes) -> str:
+    buf = _SEP.join(parts)
+    h = hashlib.blake2b(buf, digest_size=12)
+    return base64.b64encode(h.digest()).decode("ascii")