one-ring-loop 0.1.0__py3-none-any.whl

one_ring_loop/__init__.py

@@ -0,0 +1,12 @@
+"""One Ring Loop package."""
+
+__version__ = "0.1.0"
+
+from one_ring_loop.loop import run
+from one_ring_loop.task import Task, TaskGroup
+
+__all__ = [
+    "Task",
+    "TaskGroup",
+    "run",
+]

one_ring_loop/_utils.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import threading
+from collections import deque
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, cast
+
+from one_ring_core.results import IOResult
+
+if TYPE_CHECKING:
+    from one_ring_core.operations import IOOperation
+    from one_ring_loop.loop import Loop
+    from one_ring_loop.typedefs import Coro, TaskID
+
+
+def _get_new_operation_id() -> TaskID:
+    """Gets an unused ID to submit to the IO worker."""
+    ret = _local.free_operation_id
+    _local.free_operation_id += 1
+
+    return ret
+
+
+def _execute[T: IOResult](op: IOOperation[T]) -> Coro[T]:
+    """Unwrap an IO completion into the expected result type."""
+    expected = op.result_type
+    completion = yield cast("IOOperation[IOResult]", op)
+    if completion is not None and isinstance(result := completion.unwrap(), expected):
+        return result
+    elif completion is None:
+        raise RuntimeError("Low level coroutine was sent None")
+
+    msg = f"Expected {expected.__name__}, got {type(completion)}. Expected {expected}"
+    raise TypeError(msg)
+
+
+@dataclass(kw_only=True)
+class _Local(threading.local):
+    """Wrapper around threading.local for proper type annotations."""
+
+    loop: Loop | None = None
+    free_operation_id: int = 1
+
+    # TODO: Move the below two to be attributes on Loop.
+    cancel_queue: deque[TaskID] = field(default_factory=deque)
+    unpark_queue: deque[TaskID] = field(default_factory=deque)
+
+    def cleanup(self) -> None:
+        """Resets all attributes."""
+        self.loop = None
+        self.free_operation_id = 1
+
+        self.cancel_queue = deque()
+        self.unpark_queue = deque()
+
+
+_local = _Local()
+
+__all__ = [
+    "_execute",
+    "_get_new_operation_id",
+    "_local",
+]

one_ring_loop/cancellation.py

@@ -0,0 +1,49 @@
+from contextlib import contextmanager, suppress
+from typing import TYPE_CHECKING
+
+from one_ring_loop.exceptions import Cancelled
+from one_ring_loop.log import get_logger
+from one_ring_loop.task import CancelScope, _create_standalone_task
+from one_ring_loop.timerio import sleep
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+    from one_ring_loop.typedefs import Coro
+
+logger = get_logger(__name__)
+
+
+@contextmanager
+def fail_after(delay: float, *, shield: bool = False) -> Generator[CancelScope]:
+    """Cancels cancel scope and throws Cancelled after delay."""
+
+    def cancellation_task(cancel_scope: CancelScope) -> Coro[None]:
+        """Background task that sleeps for delay.
+
+        Cancels the cancel scope if not finished after sleep.
+        """
+        with suppress(Cancelled):
+            yield from sleep(delay)
+
+        if not finished:
+            cancel_scope.cancel()
+
+    finished = False
+    with CancelScope(shielded=shield) as scope:
+        task = _create_standalone_task(cancellation_task(scope), None, None)
+        yield scope
+        task.current_cancel_scope().cancel()
+    finished = True
+
+
+@contextmanager
+def move_on_after(delay: float, *, shield: bool = False) -> Generator[CancelScope]:
+    """Moves on after delay by catching Cancelled."""
+    with fail_after(delay, shield=shield) as cancel_scope:
+        try:
+            yield cancel_scope
+        except Cancelled:
+            if not cancel_scope.cancelled:
+                # Another scope cancelled from above. Re-raise.
+                raise

one_ring_loop/exceptions.py

@@ -0,0 +1,2 @@
+class Cancelled(BaseException):
+    """Thrown when coroutine is cancelled, like asyncio.CancelledError."""

one_ring_loop/fileio/__init__.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from one_ring_core.operations import Close, FileOpen, FileRead, FileWrite
+from one_ring_loop._utils import _execute
+
+if TYPE_CHECKING:
+    from one_ring_loop.typedefs import Coro
+
+
+@dataclass(slots=True, kw_only=True)
+class File:
+    """Utility wrapper for file operations. Represents a single regular file."""
+
+    fd: int
+
+    def read(self) -> Coro[str]:
+        """Read file low-level coroutine."""
+        result = yield from _execute(FileRead(fd=self.fd))
+        return result.content.decode()
+
+    def write(self, data: bytes | str) -> Coro[int]:
+        """Write file low-level coroutine."""
+        _data = data.encode() if isinstance(data, str) else data
+        result = yield from _execute(FileWrite(fd=self.fd, data=_data))
+        return result.size
+
+    def close(self) -> Coro[None]:
+        """Close file low-level coroutine."""
+        yield from _execute(Close(fd=self.fd))
+        return None
+
+
+def open_file(path: str | Path, mode: str = "r") -> Coro[File]:
+    """Open file coroutine."""
+    _path = str(path) if isinstance(path, Path) else path
+
+    result = yield from _execute(FileOpen(path=_path.encode(), mode=mode))
+    return File(fd=result.fd)

one_ring_loop/log.py

@@ -0,0 +1,56 @@
+"""Structured logging configuration using structlog."""
+
+from __future__ import annotations
+
+import os
+import sys
+
+import structlog
+
+
+def setup_logging() -> None:
+    """Configure structlog with sensible defaults.
+
+    Uses JSON output when ``LOG_FORMAT=json`` (e.g. production),
+    otherwise uses colored console output for development.
+    """
+    shared_processors: list[structlog.types.Processor] = [
+        structlog.contextvars.merge_contextvars,
+        structlog.stdlib.add_log_level,
+        # structlog.stdlib.add_logger_name,
+        structlog.processors.TimeStamper(fmt="iso"),
+        structlog.processors.StackInfoRenderer(),
+        structlog.processors.UnicodeDecoder(),
+    ]
+
+    if os.environ.get("LOG_FORMAT") == "json":
+        renderer: structlog.types.Processor = structlog.processors.JSONRenderer()
+    else:
+        renderer = structlog.dev.ConsoleRenderer()
+
+    structlog.configure(
+        processors=[
+            *shared_processors,
+            renderer,
+        ],
+        wrapper_class=structlog.make_filtering_bound_logger(0),
+        context_class=dict,
+        logger_factory=structlog.PrintLoggerFactory(file=sys.stderr),
+        cache_logger_on_first_use=True,
+    )
+
+
+_configured = False
+
+
+def get_logger(*args: object, **kwargs: object) -> structlog.stdlib.BoundLogger:
+    """Return a structlog logger, configuring on first call.
+
+    This avoids running ``setup_logging()`` at import time, which would
+    interfere with tests and multi-process setups.
+    """
+    global _configured  # noqa: PLW0603
+    if not _configured:
+        setup_logging()
+        _configured = True
+    return structlog.get_logger(*args, **kwargs)

one_ring_loop/loop.py

@@ -0,0 +1,246 @@
+from __future__ import annotations
+
+import errno
+from collections import defaultdict
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from one_ring_core.log import get_logger
+from one_ring_core.operations import Cancel, IOOperation
+from one_ring_core.worker import IOWorker
+from one_ring_loop._utils import _get_new_operation_id, _local
+from one_ring_loop.exceptions import Cancelled
+from one_ring_loop.operations import Park, WaitsOn
+from one_ring_loop.task.state import (
+    Ready,
+    Submitted,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+    from one_ring_core.results import IOCompletion, IOResult
+    from one_ring_loop.task import Task
+    from one_ring_loop.typedefs import Coro, TaskID
+
+logger = get_logger(__name__)
+
+
+@dataclass(slots=True, kw_only=True)
+class Loop:
+    """The one-ring-loop. Bask in it's glory."""
+
+    # TODO: See if there's a nice way to consolidate the below three attributes.
+    """The tasks currently running"""
+    tasks: dict[TaskID, Task] = field(default_factory=dict, init=False)
+
+    """Says which other tasks depends on a given task."""
+    task_dependencies: defaultdict[TaskID, set[TaskID]] = field(
+        default_factory=lambda: defaultdict(set), init=False
+    )
+
+    """The task which is currently executing synchronously"""
+    _current_task: Task | None = None
+
+    """Maps operation id to task id for in-flight operations"""
+    operation_to_task: dict[int, TaskID] = field(default_factory=dict, init=False)
+
+    def run_until_complete(self) -> None:
+        """Runs the event loop until all tasks are complete."""
+        with IOWorker() as worker:
+            while self.tasks:
+                self._start_tasks()  # Start new tasks
+                self._cancel_ready_tasks()  # Cancel tasks ready to be submitted
+                self._register_io_cancellations(worker)  # Register I/O cancellations
+                self._register_ready_tasks(worker)  # Submit new I/O
+                self._drive_unparked_tasks()  # Drive wakeups
+                self._drive_completed_tasks(worker)  # Drive kernel completions
+                self._drive_checkpointed_tasks()  # Drive checkpoints
+                self._remove_done_tasks()  # Clean up and wake dependent tasks
+
+    def _start_tasks(self) -> None:
+        """Starts new tasks."""
+        unstarted_tasks = [task for task in self.tasks.values() if not task.is_started]
+        for task in unstarted_tasks:
+            with self.set_current_task(task):
+                task.start()
+
+    def _register_io_cancellations(self, worker: IOWorker) -> None:
+        """Drains cancellation queue and registers and submits IO cancellations ops."""
+        should_submit = False
+        while _local.cancel_queue:
+            task_id = _local.cancel_queue.popleft()
+            if task_id not in self.tasks:
+                continue
+            task = self.tasks[task_id]
+            if task.is_waiting_on:
+                continue
+            if task.should_cancel():
+                if task.is_parked:
+                    # No kernel op, throw directly
+                    with self.set_current_task(task):
+                        task.throw(Cancelled())
+                elif (op_id := task.pending_cancel_op_id()) is not None:
+                    should_submit = True
+                    cancel_op = Cancel(target_identifier=op_id)
+                    worker.register(cancel_op, _get_new_operation_id())
+
+        if should_submit:
+            worker.submit()
+
+    def _get_ready_tasks(self) -> list[Task]:
+        """Gets the tasks ready to register with the IO worker."""
+        return [task for task in self.tasks.values() if task.is_ready]
+
+    def _cancel_ready_tasks(self) -> None:
+        """Cancel tasks with cancelled scopes, otherwise ready for I/O submition."""
+        for task in self._get_ready_tasks():
+            # Check for cancelled, non-shielded cancel scopes
+            if task.should_cancel():
+                with self.set_current_task(task):
+                    task.throw(Cancelled(f"Task {task.task_id} was cancelled"))
+
+    def _register_ready_tasks(self, worker: IOWorker) -> None:
+        """Register ready tasks with the I/O worker."""
+        should_submit = False
+        for task in self._get_ready_tasks():
+            # If a .throw call finished the task, don't register it.
+            if task.is_done:
+                continue
+
+            match task.state:
+                case Ready(operation=IOOperation() as op):
+                    should_submit = True
+                    op_id = _get_new_operation_id()
+                    worker.register(op, op_id)
+                    self.operation_to_task[op_id] = task.task_id
+                    task.state = Submitted(operation=op, op_id=op_id)
+                case Ready(operation=WaitsOn(task_ids=ids)):
+                    for task_id in ids:
+                        self.task_dependencies[task_id].add(task.task_id)
+                    task.state = Submitted(operation=WaitsOn(task_ids=ids))
+                case Ready(operation=Park()):
+                    task.state = Submitted(operation=Park())
+                case _:
+                    continue
+
+        if should_submit:
+            worker.submit()
+
+    def _collect_completions(self, worker: IOWorker) -> set[IOCompletion[IOResult]]:
+        """Waits for completions if all tasks are waiting, otherwise peeks."""
+        completions: set[IOCompletion] = set()
+
+        if all(task.is_submitted for task in self.tasks.values()):
+            if not any(task.has_pending_io for task in self.tasks.values()):
+                raise RuntimeError("Deadlock: all tasks blocked, no pending I/O")
+            completion = worker.wait()
+            completions.add(completion)
+        else:
+            # Peek until we get None.
+            while (completion := worker.peek()) is not None:
+                completions.add(completion)
+
+        return completions
+
+    def _drive_completed_tasks(self, worker: IOWorker) -> None:
+        completions = self._collect_completions(worker)
+        for completion in completions:
+            op_id = completion.user_data
+            task_id = self.operation_to_task.pop(op_id, None)
+
+            if task_id is None:
+                continue
+            task = self.tasks.get(task_id)
+            if task is None:
+                continue
+
+            with self.set_current_task(task):
+                if (
+                    isinstance(oserror := completion.result, OSError)
+                    and oserror.errno is not None
+                    and oserror.errno == errno.ECANCELED
+                ):
+                    task.throw(Cancelled())
+                else:
+                    task.drive(completion)
+
+    def _drive_unparked_tasks(self) -> None:
+        """Drives tasks that have been unparked.
+
+        Needs to run before _drive_completed_tasks to avoid deadlocks.
+        """
+        while _local.unpark_queue:
+            unparked_task_id = _local.unpark_queue.popleft()
+            unparked_task = self.tasks[unparked_task_id]
+
+            with self.set_current_task(unparked_task):
+                unparked_task.drive(None)
+
+    def _drive_checkpointed_tasks(self) -> None:
+        """Drives tasks that have been checkpointed."""
+        checkpointed_tasks = [
+            task for task in self.tasks.values() if task.is_checkpointed
+        ]
+
+        for task in checkpointed_tasks:
+            with self.set_current_task(task):
+                task.drive(None)
+
+    def _remove_done_tasks(self) -> None:
+        done_tasks = [task for task in self.tasks.values() if task.is_done]
+        for done_task in done_tasks:
+            # Now drive tasks that were dependant on the done tasks.
+            while self.task_dependencies[done_task.task_id]:
+                waiting_task_id = self.task_dependencies[done_task.task_id].pop()
+                waiting_task = self.tasks[waiting_task_id]
+                if waiting_task.is_done:
+                    # The loop has already driven this forward, for example if all tasks
+                    # waiting_task depended on finished in the same loop iteration.
+                    continue
+                with self.set_current_task(waiting_task):
+                    waiting_task.drive(None)
+
+            self.task_dependencies.pop(done_task.task_id)
+
+        self.tasks = {
+            task_id: task for task_id, task in self.tasks.items() if not task.is_done
+        }
+
+    @property
+    def current_task(self) -> Task:
+        """Gets currently executing task."""
+        if self._current_task is None:
+            raise RuntimeError("No task currently executing")
+
+        return self._current_task
+
+    @contextmanager
+    def set_current_task(self, task: Task) -> Generator[None]:
+        """Utility wrapper for setting and removing currently executing task."""
+        self._current_task = task
+        try:
+            yield
+        finally:
+            self._current_task = None
+
+    def add_task(self, task: Task) -> None:
+        """Adds a task to be run by the event loop."""
+        self.tasks[task.task_id] = task
+
+
+def run(gen: Coro) -> None:
+    """Entry point for running the event loop.
+
+    Creates a Task from the generator on the event loop, and runs the loop.
+
+    Args:
+        gen: the entry coroutine
+    """
+    from one_ring_loop.task import _create_standalone_task  # noqa: PLC0415
+
+    _local.loop = Loop()
+    _create_standalone_task(gen, None, None)
+    _local.loop.run_until_complete()
+    _local.cleanup()

one_ring_loop/lowlevel.py

@@ -0,0 +1,36 @@
+from typing import TYPE_CHECKING
+
+from one_ring_loop._utils import _local
+from one_ring_loop.operations import Checkpoint
+
+if TYPE_CHECKING:
+    from one_ring_loop.loop import Loop
+    from one_ring_loop.task import Task
+    from one_ring_loop.typedefs import Coro, TaskID
+
+
+def get_running_loop() -> Loop:
+    """Gets the currently running event loop. I don't want to expose this easily."""
+    if _local.loop is None:
+        raise RuntimeError("No event loop running")
+
+    return _local.loop
+
+
+def get_current_task() -> Task:
+    """Gets the currently executing task from the loop."""
+    return get_running_loop().current_task
+
+
+def unpark(task_id: TaskID) -> None:
+    """Unparks a parked task.
+
+    Args:
+        task_id: The id of the task to unpark
+    """
+    _local.unpark_queue.append(task_id)
+
+
+def checkpoint() -> Coro[None]:
+    """Nop that yields control back to event loop."""
+    yield Checkpoint()

one_ring_loop/operations.py

@@ -0,0 +1,29 @@
+"""This namespace includes extra operations for the event loop.
+
+This extends the operations provided by one_ring_core.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from one_ring_loop.typedefs import TaskID
+
+
+@dataclass(slots=True, kw_only=True)
+class WaitsOn:
+    """For dependency relationships between tasks."""
+
+    task_ids: tuple[TaskID, ...]
+
+
+@dataclass(slots=True, kw_only=True)
+class Park:
+    """Parks the yielding task until resumed by another task."""
+
+
+@dataclass(slots=True, kw_only=True)
+class Checkpoint:
+    """Sentinel that yields control back to event loop."""

one_ring_loop/socketio/__init__.py

@@ -0,0 +1,108 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from one_ring_core.operations import (
+    Close,
+    SocketAccept,
+    SocketBind,
+    SocketConnect,
+    SocketCreate,
+    SocketListen,
+    SocketRecv,
+    SocketSend,
+    SocketSetOpt,
+)
+from one_ring_loop._utils import _execute
+
+if TYPE_CHECKING:
+    from one_ring_loop.typedefs import Coro
+
+# TODO: Clean up
+# Better names
+# Centralize "close" together with fileio
+
+
+def _create() -> Coro[int]:
+    result = yield from _execute(SocketCreate())
+    return result.fd
+
+
+def _set_options(fd: int) -> Coro[None]:
+    yield from _execute(SocketSetOpt(fd=fd))
+    return None
+
+
+def _bind(fd: int, host: bytes, port: int) -> Coro[None]:
+    yield from _execute(SocketBind(fd=fd, ip=host, port=port))
+    return None
+
+
+def _listen(fd: int) -> Coro[None]:
+    yield from _execute(SocketListen(fd=fd))
+    return None
+
+
+def _connect(fd: int, host: bytes, port: int) -> Coro[None]:
+    yield from _execute(SocketConnect(fd=fd, ip=host, port=port))
+    return None
+
+
+def create_server(host: bytes, port: int) -> Coro[Server]:
+    """Creates a socket (server), sets options, binds it and wraps in SocketListener."""
+    fd = yield from _create()
+    yield from _set_options(fd)
+    yield from _bind(fd, host, port)
+    yield from _listen(fd)
+    return Server(fd=fd)
+
+
+def connect(host: bytes, port: int) -> Coro[Connection]:
+    """Connects to a listening socket."""
+    fd = yield from _create()
+    yield from _connect(fd, host, port)
+
+    return Connection(fd=fd)
+
+
+@dataclass(slots=True, kw_only=True)
+class Server:
+    """Used to accept new connections."""
+
+    """The socket's file descriptor"""
+    fd: int
+
+    def accept(self) -> Coro[Connection]:
+        """Waits until there's a connection to accept.
+
+        Returns:
+            client file descriptor
+        """
+        result = yield from _execute(SocketAccept(fd=self.fd))
+        return Connection(fd=result.fd)
+
+    def close(self) -> Coro[None]:
+        """Close socket."""
+        yield from _execute(Close(fd=self.fd))
+        return None
+
+
+@dataclass(slots=True, kw_only=True)
+class Connection:
+    """Corresponds a socket connection. Either from server, or client."""
+
+    """Either server file descriptor, or client file descriptor."""
+    fd: int
+
+    def receive(self, max_bytes: int = 65536) -> Coro[bytes]:
+        """Reads data from socket."""
+        result = yield from _execute(SocketRecv(fd=self.fd, size=max_bytes))
+        return result.content
+
+    def send(self, data: bytes, /) -> Coro[None]:
+        """Sends data to socket."""
+        yield from _execute(SocketSend(fd=self.fd, data=data))
+
+    def close(self) -> Coro[None]:
+        """Close socket."""
+        yield from _execute(Close(fd=self.fd))
+        return None