97 lines
3.0 KiB
Python
97 lines
3.0 KiB
Python
"""Small persistent-queue dispatcher for skill knowledge items."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import socket
|
|
from typing import Any
|
|
from uuid import uuid4
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SkillKnowledgeJobDispatcher:
|
|
def __init__(
|
|
self,
|
|
store: Any,
|
|
executor: Any,
|
|
*,
|
|
interval_seconds: float = 2.0,
|
|
concurrency: int = 3,
|
|
lease_seconds: int = 900,
|
|
) -> None:
|
|
self.store = store
|
|
self.executor = executor
|
|
self.interval_seconds = interval_seconds
|
|
self.lease_seconds = lease_seconds
|
|
self.owner = f"{socket.gethostname()}:{uuid4()}"
|
|
self.semaphore = asyncio.Semaphore(concurrency)
|
|
self._wake = asyncio.Event()
|
|
self._task: asyncio.Task | None = None
|
|
self._running: set[str] = set()
|
|
|
|
def start(self) -> None:
|
|
if self._task is None or self._task.done():
|
|
self._task = asyncio.create_task(self._loop())
|
|
|
|
def nudge(self) -> None:
|
|
self._wake.set()
|
|
|
|
async def close(self) -> None:
|
|
if self._task is not None:
|
|
self._task.cancel()
|
|
try:
|
|
await self._task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
async def _run(self, item: dict) -> None:
|
|
item_id = str(item["id"])
|
|
async with self.semaphore:
|
|
heartbeat = asyncio.create_task(self._heartbeat(item_id))
|
|
try:
|
|
await self.executor.run_item(item)
|
|
finally:
|
|
heartbeat.cancel()
|
|
try:
|
|
await heartbeat
|
|
except asyncio.CancelledError:
|
|
pass
|
|
self._running.discard(item_id)
|
|
self._wake.set()
|
|
|
|
async def _heartbeat(self, item_id: str) -> None:
|
|
while True:
|
|
await asyncio.sleep(min(30, max(5, self.lease_seconds // 3)))
|
|
if not await self.store.renew_lease(
|
|
item_id,
|
|
owner=self.owner,
|
|
lease_seconds=self.lease_seconds,
|
|
):
|
|
return
|
|
|
|
async def _loop(self) -> None:
|
|
while True:
|
|
try:
|
|
for item in await self.store.list_queued_items(limit=20):
|
|
item_id = str(item["id"])
|
|
if item_id in self._running:
|
|
continue
|
|
claimed = await self.store.claim_item(
|
|
item_id,
|
|
owner=self.owner,
|
|
lease_seconds=self.lease_seconds,
|
|
)
|
|
if claimed is None:
|
|
continue
|
|
self._running.add(item_id)
|
|
asyncio.create_task(self._run(claimed))
|
|
except Exception:
|
|
logger.exception("Skill knowledge dispatcher tick failed")
|
|
try:
|
|
await asyncio.wait_for(self._wake.wait(), timeout=self.interval_seconds)
|
|
except TimeoutError:
|
|
pass
|
|
self._wake.clear()
|