offwork 0.4.0__py3-none-any.whl

offwork/worker/backends/http.py

@@ -0,0 +1,237 @@
+"""HTTP(S) backend for hosted broker deployments."""
+
+import json
+import time
+import base64
+import asyncio
+from typing import Any
+from urllib.error import HTTPError, URLError
+from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
+from urllib.request import Request, urlopen
+from collections.abc import AsyncIterator
+
+from offwork.worker.backends.base import Backend
+
+_DEFAULT_BROKER_PATH = "/api/v1/broker"
+_DEFAULT_LONG_POLL_SECONDS = 30.0
+
+
+class HttpBackend(Backend):
+    """HTTP(S)-based backend for hosted offwork brokers.
+
+    The base URL can point either at the broker root or at the service root;
+    when no path is provided, ``/api/v1/broker`` is assumed.
+
+    Authentication is intentionally simple for the proof-of-concept:
+    include ``?api_key=...`` in the URL and the backend will move it into the
+    ``X-Offwork-API-Key`` request header.
+    """
+
+    def __init__(self, url: str) -> None:
+        parsed = urlparse(url)
+        if parsed.scheme not in {"http", "https"}:
+            raise ValueError(f"Unsupported HTTP backend scheme: {parsed.scheme!r}")
+
+        query_items = parse_qsl(parsed.query, keep_blank_values=True)
+        api_key = ""
+        filtered_query: list[tuple[str, str]] = []
+        for key, value in query_items:
+            if key == "api_key" and not api_key:
+                api_key = value
+                continue
+            filtered_query.append((key, value))
+
+        path = parsed.path.rstrip("/") or _DEFAULT_BROKER_PATH
+        self._base_url = urlunparse(parsed._replace(path=path, query=urlencode(filtered_query)))
+        self._api_key = api_key or None
+
+    def _headers(self) -> dict[str, str]:
+        headers = {"Content-Type": "application/json"}
+        if self._api_key:
+            headers["X-Offwork-API-Key"] = self._api_key
+        return headers
+
+    def _url(self, suffix: str, query: dict[str, str | float | int] | None = None) -> str:
+        url = f"{self._base_url}{suffix}"
+        if not query:
+            return url
+        encoded = urlencode({key: str(value) for key, value in query.items()})
+        separator = "&" if "?" in url else "?"
+        return f"{url}{separator}{encoded}"
+
+    def _do_request(
+        self,
+        method: str,
+        suffix: str,
+        *,
+        payload: dict[str, Any] | None = None,
+        query: dict[str, str | float | int] | None = None,
+        timeout: float | None = None,
+        allow_not_found: bool = False,
+    ) -> tuple[int, Any | None]:
+        data = None if payload is None else json.dumps(payload).encode("utf-8")
+        request = Request(
+            self._url(suffix, query=query),
+            data=data,
+            method=method,
+            headers=self._headers(),
+        )
+        try:
+            with urlopen(request, timeout=timeout) as response:
+                raw = response.read()
+                if not raw:
+                    return response.status, None
+                return response.status, json.loads(raw.decode("utf-8"))
+        except HTTPError as exc:
+            if exc.code in {204, 404} and allow_not_found:
+                return exc.code, None
+            message = exc.read().decode("utf-8", errors="replace")
+            raise RuntimeError(
+                f"HTTP backend request failed: {method} {suffix} -> {exc.code} {message}"
+            ) from exc
+        except URLError as exc:
+            raise ConnectionError(f"HTTP backend connection failed: {exc.reason}") from exc
+
+    async def _request(
+        self,
+        method: str,
+        suffix: str,
+        *,
+        payload: dict[str, Any] | None = None,
+        query: dict[str, str | float | int] | None = None,
+        timeout: float | None = None,
+        allow_not_found: bool = False,
+    ) -> tuple[int, Any | None]:
+        return await asyncio.to_thread(
+            self._do_request,
+            method,
+            suffix,
+            payload=payload,
+            query=query,
+            timeout=timeout,
+            allow_not_found=allow_not_found,
+        )
+
+    async def submit(self, task_json: str) -> None:
+        await self._request("POST", "/tasks", payload={"task_json": task_json})
+
+    async def listen(self) -> AsyncIterator[str]:
+        while True:
+            _status, body = await self._request(
+                "POST",
+                "/tasks/claim",
+                payload={"wait_seconds": _DEFAULT_LONG_POLL_SECONDS},
+                timeout=_DEFAULT_LONG_POLL_SECONDS + 5.0,
+                allow_not_found=True,
+            )
+            if body is None:
+                continue
+            task_json = body.get("task_json")
+            if isinstance(task_json, str):
+                yield task_json
+
+    async def send_result(self, task_id: str, result_json: str) -> None:
+        await self._request(
+            "POST",
+            f"/tasks/{task_id}/result",
+            payload={"result_json": result_json},
+        )
+
+    async def get_result(self, task_id: str, timeout: float | None = None) -> str:
+        deadline = None if timeout is None else time.monotonic() + timeout
+        while True:
+            remaining = None if deadline is None else max(0.0, deadline - time.monotonic())
+            wait_seconds = _DEFAULT_LONG_POLL_SECONDS if remaining is None else min(_DEFAULT_LONG_POLL_SECONDS, remaining)
+            status, body = await self._request(
+                "GET",
+                f"/tasks/{task_id}/result",
+                query={"wait_seconds": wait_seconds},
+                timeout=(wait_seconds + 5.0) if wait_seconds else 5.0,
+                allow_not_found=True,
+            )
+            if body is not None:
+                result_json = body.get("result_json")
+                if isinstance(result_json, str):
+                    return result_json
+            if deadline is not None and (status == 204 or time.monotonic() >= deadline):
+                raise TimeoutError(f"Timed out waiting for result of task {task_id}")
+
+    async def try_get_result(self, task_id: str) -> str | None:
+        _status, body = await self._request(
+            "GET",
+            f"/tasks/{task_id}/result",
+            query={"wait_seconds": 0},
+            allow_not_found=True,
+        )
+        if body is None:
+            return None
+        result_json = body.get("result_json")
+        return result_json if isinstance(result_json, str) else None
+
+    async def send_heartbeat(self, task_id: str) -> None:
+        await self._request("POST", f"/tasks/{task_id}/heartbeat")
+
+    async def get_heartbeat(self, task_id: str) -> float | None:
+        _status, body = await self._request(
+            "GET", f"/tasks/{task_id}/heartbeat", allow_not_found=True,
+        )
+        if body is None:
+            return None
+        raw = body.get("heartbeat")
+        return float(raw) if isinstance(raw, (int, float)) else None
+
+    async def cancel_task(self, task_id: str) -> None:
+        await self._request("POST", f"/tasks/{task_id}/cancel")
+
+    async def is_cancelled(self, task_id: str) -> bool:
+        _status, body = await self._request(
+            "GET", f"/tasks/{task_id}/cancel", allow_not_found=True,
+        )
+        return bool(body and body.get("cancelled"))
+
+    async def send_progress(self, task_id: str, progress_json: str) -> None:
+        await self._request(
+            "POST",
+            f"/tasks/{task_id}/progress",
+            payload={"progress_json": progress_json},
+        )
+
+    async def get_progress(self, task_id: str) -> str | None:
+        _status, body = await self._request(
+            "GET", f"/tasks/{task_id}/progress", allow_not_found=True,
+        )
+        if body is None:
+            return None
+        progress_json = body.get("progress_json")
+        return progress_json if isinstance(progress_json, str) else None
+
+    async def cancel_schedule(self, schedule_id: str) -> None:
+        await self._request("POST", f"/schedules/{schedule_id}/cancel")
+
+    async def is_schedule_cancelled(self, schedule_id: str) -> bool:
+        _status, body = await self._request(
+            "GET", f"/schedules/{schedule_id}/cancel", allow_not_found=True,
+        )
+        return bool(body and body.get("cancelled"))
+
+    async def check_throttle(self, function_name: str) -> bool:
+        _status, body = await self._request(
+            "GET",
+            "/throttle/check",
+            query={"function_name": function_name},
+            allow_not_found=True,
+        )
+        return True if body is None else bool(body.get("allowed", True))
+
+    async def record_throttle(self, function_name: str, throttle_seconds: float) -> None:
+        await self._request(
+            "POST",
+            "/throttle/record",
+            payload={
+                "function_name": function_name,
+                "throttle_seconds": throttle_seconds,
+            },
+        )
+
+    async def close(self) -> None:
+        return None

offwork/worker/backends/local.py

@@ -0,0 +1,452 @@
+"""Async-native TCP backend for same-machine IPC.
+
+A lightweight broker server built on :mod:`asyncio` TCP streams handles
+task dispatch, result routing, and heartbeats -- no threads, no
+:mod:`multiprocessing`, no external services.
+
+URL scheme: ``local://host:port``  (default ``local://127.0.0.1:9748``)
+"""
+
+import sys
+import json
+import time
+import atexit
+import socket
+import struct
+import asyncio
+import logging
+import contextlib
+import subprocess
+from typing import Any
+from urllib.parse import urlparse
+from collections.abc import AsyncIterator
+
+from offwork.worker.backends.base import Backend
+
+logger = logging.getLogger(__name__)
+
+_HEADER = struct.Struct("!I")  # 4-byte big-endian length prefix
+_DEFAULT_HOST = "127.0.0.1"
+_DEFAULT_PORT = 9748
+
+
+# ---------------------------------------------------------------------------
+# Wire protocol
+# ---------------------------------------------------------------------------
+
+
+async def _send_msg(writer: asyncio.StreamWriter, obj: dict[str, Any]) -> None:
+    """Send a length-prefixed JSON message."""
+    payload = json.dumps(obj, separators=(",", ":")).encode()
+    writer.write(_HEADER.pack(len(payload)) + payload)
+    await writer.drain()
+
+
+async def _recv_msg(reader: asyncio.StreamReader) -> dict[str, Any]:
+    """Receive a length-prefixed JSON message.
+
+    Raises :class:`asyncio.IncompleteReadError` on EOF.
+    """
+    raw = await reader.readexactly(_HEADER.size)
+    (length,) = _HEADER.unpack(raw)
+    data = await reader.readexactly(length)
+    result: dict[str, Any] = json.loads(data)
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Broker server (pure asyncio)
+# ---------------------------------------------------------------------------
+
+
+class _Broker:
+    """Task broker backed entirely by asyncio primitives."""
+
+    def __init__(self) -> None:
+        self._tasks: asyncio.Queue[str] = asyncio.Queue()
+        self._results: dict[str, asyncio.Queue[str]] = {}
+        self._heartbeats: dict[str, float] = {}
+        self._cancelled: set[str] = set()
+        self._cancelled_schedules: set[str] = set()
+        self._throttles: dict[str, float] = {}
+        self._progress: dict[str, str] = {}
+        self._result_subs: list[asyncio.Queue[str]] = []
+
+    def _result_slot(self, task_id: str) -> asyncio.Queue[str]:
+        if task_id not in self._results:
+            self._results[task_id] = asyncio.Queue(maxsize=1)
+        return self._results[task_id]
+
+    # -- connection handler ----------------------------------------------------
+
+    async def handle(
+        self,
+        reader: asyncio.StreamReader,
+        writer: asyncio.StreamWriter,
+    ) -> None:
+        try:
+            msg = await _recv_msg(reader)
+            op = msg.get("op", "")
+            if op == "listen":
+                await self._stream_tasks(writer)
+            elif op == "subscribe":
+                await self._stream_results(writer)
+            else:
+                await self._rpc_loop(msg, reader, writer)
+        except (asyncio.IncompleteReadError, ConnectionError, OSError):
+            pass
+        finally:
+            writer.close()
+            with contextlib.suppress(ConnectionError, OSError):
+                await writer.wait_closed()
+
+    async def _rpc_loop(
+        self,
+        first: dict[str, Any],
+        reader: asyncio.StreamReader,
+        writer: asyncio.StreamWriter,
+    ) -> None:
+        await _send_msg(writer, self._dispatch(first))
+        while True:
+            msg = await _recv_msg(reader)
+            await _send_msg(writer, self._dispatch(msg))
+
+    # -- dispatch (sync -- no awaits, safe for single-threaded asyncio) --------
+
+    def _dispatch(self, msg: dict[str, Any]) -> dict[str, Any]:
+        op = msg["op"]
+        if op == "submit":
+            self._tasks.put_nowait(msg["data"])
+            return {"ok": True}
+        if op == "result_put":
+            try:
+                self._result_slot(msg["task_id"]).put_nowait(msg["data"])
+            except asyncio.QueueFull:
+                pass  # first result wins (e.g., cancel before worker result)
+            for sub in self._result_subs:
+                sub.put_nowait(msg["task_id"])
+            return {"ok": True}
+        if op == "result_try":
+            try:
+                return {"ok": True, "data": self._result_slot(msg["task_id"]).get_nowait()}
+            except asyncio.QueueEmpty:
+                return {"ok": True, "data": None}
+        if op == "hb_put":
+            self._heartbeats[msg["task_id"]] = msg["ts"]
+            return {"ok": True}
+        if op == "hb_get":
+            return {"ok": True, "data": self._heartbeats.get(msg["task_id"])}
+        if op == "hb_batch":
+            return {
+                "ok": True,
+                "data": {tid: self._heartbeats.get(tid) for tid in msg["task_ids"]},
+            }
+        if op == "cancel":
+            self._cancelled.add(msg["task_id"])
+            return {"ok": True}
+        if op == "is_cancelled":
+            return {"ok": True, "data": msg["task_id"] in self._cancelled}
+        if op == "progress_put":
+            self._progress[msg["task_id"]] = msg["data"]
+            return {"ok": True}
+        if op == "progress_get":
+            return {"ok": True, "data": self._progress.get(msg["task_id"])}
+        if op == "schedule_cancel":
+            self._cancelled_schedules.add(msg["schedule_id"])
+            return {"ok": True}
+        if op == "schedule_check":
+            return {"ok": True, "data": msg["schedule_id"] in self._cancelled_schedules}
+        if op == "throttle_check":
+            fn = msg["function_name"]
+            expiry = self._throttles.get(fn)
+            if expiry is not None and time.time() < expiry:
+                return {"ok": True, "data": False}
+            self._throttles.pop(fn, None)
+            return {"ok": True, "data": True}
+        if op == "throttle_record":
+            fn = msg["function_name"]
+            self._throttles[fn] = time.time() + msg["seconds"]
+            return {"ok": True}
+        return {"ok": False, "error": f"unknown op: {op}"}
+
+    # -- streaming handlers ----------------------------------------------------
+
+    async def _stream_tasks(self, writer: asyncio.StreamWriter) -> None:
+        while True:
+            try:
+                task = await asyncio.wait_for(self._tasks.get(), timeout=1.0)
+            except asyncio.TimeoutError:
+                continue
+            try:
+                await _send_msg(writer, {"data": task})
+            except (ConnectionError, OSError):
+                # Client gone -- put the task back so another listener gets it.
+                self._tasks.put_nowait(task)
+                return
+
+    async def _stream_results(self, writer: asyncio.StreamWriter) -> None:
+        q: asyncio.Queue[str] = asyncio.Queue()
+        self._result_subs.append(q)
+        try:
+            while True:
+                try:
+                    task_id = await asyncio.wait_for(q.get(), timeout=1.0)
+                except asyncio.TimeoutError:
+                    continue
+                try:
+                    await _send_msg(writer, {"data": task_id})
+                except (ConnectionError, OSError):
+                    return
+        finally:
+            self._result_subs.remove(q)
+
+
+async def run_broker(host: str, port: int) -> None:
+    """Start the broker TCP server (runs forever)."""
+    broker = _Broker()
+    server = await asyncio.start_server(broker.handle, host, port)
+    logger.info("Local broker listening on %s:%d", host, port)
+    async with server:
+        await server.serve_forever()
+
+
+def _broker_main(host: str, port: int) -> None:
+    """Entry point for the broker subprocess."""
+    asyncio.run(run_broker(host, port))
+
+
+# ---------------------------------------------------------------------------
+# URL parsing
+# ---------------------------------------------------------------------------
+
+
+def _parse_local_url(url: str) -> tuple[str, int]:
+    parsed = urlparse(url)
+    host = parsed.hostname or _DEFAULT_HOST
+    port = parsed.port or _DEFAULT_PORT
+    return host, port
+
+
+# ---------------------------------------------------------------------------
+# LocalBackend
+# ---------------------------------------------------------------------------
+
+
+class LocalBackend(Backend):
+    """Async-native TCP backend for same-machine IPC.
+
+    A lightweight broker process handles task dispatch, result routing,
+    and heartbeats over TCP on localhost.  Every I/O operation is a
+    native :mod:`asyncio` coroutine -- no threads anywhere.
+
+    Parameters
+    ----------
+    url
+        ``local://host:port``  (default ``local://127.0.0.1:9748``).
+    server
+        ``True`` to start the broker, ``False`` to connect to an
+        existing one, ``None`` (default) to auto-detect.
+    """
+
+    def __init__(
+        self,
+        url: str = "local://localhost",
+        *,
+        server: bool | None = None,
+    ) -> None:
+        self._host, self._port = _parse_local_url(url)
+        self._reader: asyncio.StreamReader | None = None
+        self._writer: asyncio.StreamWriter | None = None
+        self._lock = asyncio.Lock()
+        self._server_proc: subprocess.Popen[bytes] | None = None
+
+        self._ensure_broker(server)
+        logger.info(
+            "LocalBackend ready (server=%s, %s:%d)",
+            self._server_proc is not None, self._host, self._port,
+        )
+
+    # -- broker lifecycle ------------------------------------------------------
+
+    def _ensure_broker(self, server: bool | None) -> None:
+        if server is False:
+            return
+        if self._probe():
+            return
+        self._start_broker()
+
+    def _probe(self) -> bool:
+        """Check whether a broker is already accepting connections."""
+        try:
+            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            s.settimeout(0.5)
+            s.connect((self._host, self._port))
+            s.close()
+            return True
+        except (ConnectionRefusedError, OSError):
+            return False
+
+    def _start_broker(self) -> None:
+        self._server_proc = subprocess.Popen(
+            [
+                sys.executable, "-c",
+                "from offwork.worker.backends.local import _broker_main; "
+                f"_broker_main({self._host!r}, {self._port})",
+            ],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+        atexit.register(self._kill_broker)
+        for _ in range(50):  # up to 5 s
+            if self._probe():
+                return
+            time.sleep(0.1)
+        raise ConnectionError(
+            f"Broker failed to start on {self._host}:{self._port}"
+        )
+
+    def _kill_broker(self) -> None:
+        p = self._server_proc
+        if p is not None and p.poll() is None:
+            p.terminate()
+            try:
+                p.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                p.kill()
+        self._server_proc = None
+
+    # -- TCP connection --------------------------------------------------------
+
+    async def _connect(self) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
+        if self._reader is None or self._writer is None:
+            self._reader, self._writer = await asyncio.open_connection(
+                self._host, self._port,
+            )
+        return self._reader, self._writer
+
+    async def _request(self, msg: dict[str, Any]) -> dict[str, Any]:
+        async with self._lock:
+            try:
+                reader, writer = await self._connect()
+                await _send_msg(writer, msg)
+                return await _recv_msg(reader)
+            except (asyncio.IncompleteReadError, ConnectionError, OSError):
+                if self._writer is not None:
+                    self._writer.close()
+                self._reader = self._writer = None
+                raise
+
+    # -- Backend interface -----------------------------------------------------
+
+    async def submit(self, task_json: str) -> None:
+        await self._request({"op": "submit", "data": task_json})
+
+    async def listen(self) -> AsyncIterator[str]:
+        reader, writer = await asyncio.open_connection(self._host, self._port)
+        try:
+            await _send_msg(writer, {"op": "listen"})
+            while True:
+                msg = await _recv_msg(reader)
+                yield msg["data"]
+        except (asyncio.IncompleteReadError, ConnectionError, OSError):
+            return
+        finally:
+            writer.close()
+            with contextlib.suppress(ConnectionError, OSError):
+                await writer.wait_closed()
+
+    async def send_result(self, task_id: str, result_json: str) -> None:
+        await self._request({
+            "op": "result_put", "task_id": task_id, "data": result_json,
+        })
+
+    async def get_result(self, task_id: str, timeout: float | None = None) -> str:
+        deadline = None if timeout is None else time.monotonic() + timeout
+        while True:
+            raw = await self.try_get_result(task_id)
+            if raw is not None:
+                return raw
+            if deadline is not None and time.monotonic() >= deadline:
+                raise TimeoutError(
+                    f"Timed out waiting for result of task {task_id}"
+                )
+            await asyncio.sleep(0.05)
+
+    async def try_get_result(self, task_id: str) -> str | None:
+        resp = await self._request({"op": "result_try", "task_id": task_id})
+        return resp.get("data")
+
+    async def send_heartbeat(self, task_id: str) -> None:
+        await self._request({
+            "op": "hb_put", "task_id": task_id, "ts": time.time(),
+        })
+
+    async def get_heartbeat(self, task_id: str) -> float | None:
+        resp = await self._request({"op": "hb_get", "task_id": task_id})
+        return resp.get("data")
+
+    async def get_heartbeats(self, task_ids: list[str]) -> dict[str, float | None]:
+        resp = await self._request({"op": "hb_batch", "task_ids": task_ids})
+        result: dict[str, float | None] = resp.get("data", {})
+        return result
+
+    async def cancel_task(self, task_id: str) -> None:
+        await self._request({"op": "cancel", "task_id": task_id})
+
+    async def is_cancelled(self, task_id: str) -> bool:
+        resp = await self._request({"op": "is_cancelled", "task_id": task_id})
+        return bool(resp.get("data", False))
+
+    async def send_progress(self, task_id: str, progress_json: str) -> None:
+        await self._request({
+            "op": "progress_put", "task_id": task_id, "data": progress_json,
+        })
+
+    async def get_progress(self, task_id: str) -> str | None:
+        resp = await self._request({"op": "progress_get", "task_id": task_id})
+        return resp.get("data")
+
+    async def cancel_schedule(self, schedule_id: str) -> None:
+        await self._request({"op": "schedule_cancel", "schedule_id": schedule_id})
+
+    async def is_schedule_cancelled(self, schedule_id: str) -> bool:
+        resp = await self._request({"op": "schedule_check", "schedule_id": schedule_id})
+        return bool(resp.get("data", False))
+
+    async def check_throttle(self, function_name: str) -> bool:
+        resp = await self._request({"op": "throttle_check", "function_name": function_name})
+        return bool(resp.get("data", True))
+
+    async def record_throttle(
+        self, function_name: str, throttle_seconds: float,
+    ) -> None:
+        await self._request({
+            "op": "throttle_record",
+            "function_name": function_name,
+            "seconds": throttle_seconds,
+        })
+
+    async def notify_result(self, task_id: str) -> None:
+        pass  # broker dispatches notifications inside result_put
+
+    async def subscribe_results(self) -> AsyncIterator[str]:
+        reader, writer = await asyncio.open_connection(self._host, self._port)
+        try:
+            await _send_msg(writer, {"op": "subscribe"})
+            while True:
+                msg = await _recv_msg(reader)
+                yield msg["data"]
+        except (asyncio.IncompleteReadError, ConnectionError, OSError):
+            return
+        finally:
+            writer.close()
+            with contextlib.suppress(ConnectionError, OSError):
+                await writer.wait_closed()
+
+    async def close(self) -> None:
+        if self._writer is not None:
+            self._writer.close()
+            with contextlib.suppress(ConnectionError, OSError):
+                await self._writer.wait_closed()
+            self._reader = self._writer = None
+        self._kill_broker()