musicfetch/server/jobs.py

"""In-memory async job store. Personal-scale: jobs are lost on restart.
Generic — knows nothing about musicfetch; callers pass a no-arg `fn`.

A job can also park at `awaiting_choice`: the worker produced a list of
`choices` and waits for the client to pick one (`take_choice`). A parked job
expires after CHOICE_TTL seconds so the store never fills with dead prompts."""
import os
import time
import uuid
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from typing import Any, Callable, Optional

_EXECUTOR = ThreadPoolExecutor(max_workers=2)
JOBS: "dict[str, Job]" = {}
_MAX_JOBS = 200  # cap to bound memory
CHOICE_TTL = float(os.environ.get("MUSICFETCH_CHOICE_TTL", "900"))  # 15 minutes


class ChoiceError(Exception):
    """The client's pick cannot be used. The text is safe to show a user."""


@dataclass
class Job:
    id: str
    # queued | running | awaiting_choice | done | failed | expired
    status: str
    hit: Any
    message: str
    result: Optional[dict] = None
    error: Optional[str] = None
    choices: Optional[list] = None    # [{"n": 1, "label": "..."}] shown to the client
    pending: Optional[list] = None    # server-side data behind each choice
    expires_at: Optional[float] = None
    created_at: float = field(default_factory=time.time)
    updated_at: float = field(default_factory=time.time)


def _touch(job: "Job", **changes):
    for k, v in changes.items():
        setattr(job, k, v)
    job.updated_at = time.time()


def _evict_if_needed():
    # Post-condition: len(JOBS) <= _MAX_JOBS (evicts oldest overflow entries).
    if len(JOBS) <= _MAX_JOBS:
        return
    for jid in sorted(JOBS, key=lambda j: JOBS[j].created_at)[: len(JOBS) - _MAX_JOBS]:
        JOBS.pop(jid, None)


def create_job(hit: Any, message: str) -> "Job":
    job = Job(id=uuid.uuid4().hex[:8], status="queued", hit=hit, message=message)
    JOBS[job.id] = job
    _evict_if_needed()
    return job


def get_job(job_id: str) -> Optional["Job"]:
    job = JOBS.get(job_id)
    if job is not None and _is_stale(job):
        _touch(job, status="expired", choices=None, pending=None,
               message="That choice list has expired. Please search again.")
    return job


def _is_stale(job: "Job") -> bool:
    return (job.status == "awaiting_choice"
            and job.expires_at is not None
            and time.time() > job.expires_at)


def _fail_text(fail_message, exc: Exception) -> str:
    return fail_message(exc) if callable(fail_message) else fail_message


def run_job(job_id: str, fn: Callable[[], dict],
            done_message: "str | Callable[[dict], str]",
            fail_message: "str | Callable[[Exception], str]" = "Something went wrong while fetching.") -> None:
    def _task():
        job = JOBS.get(job_id)
        if job is None:
            return
        _touch(job, status="running")
        try:
            result = fn()
            msg = done_message(result) if callable(done_message) else done_message
            _touch(job, status="done", result=result, message=msg)
        except Exception as e:  # noqa: BLE001 — record any failure on the job
            _touch(job, status="failed", error=f"{type(e).__name__}: {e}",
                   message=_fail_text(fail_message, e))
    _EXECUTOR.submit(_task)


def await_choice(job_id: str, choices: list, pending: list, message: str) -> None:
    """Park a running job until the client picks one of `choices`."""
    job = JOBS.get(job_id)
    if job is None:
        return
    _touch(job, status="awaiting_choice", choices=choices, pending=pending,
           message=message, expires_at=time.time() + CHOICE_TTL)


def requeue(job: "Job", message: str) -> None:
    """Re-arm a parked job for a second piece of work (the client's pick)."""
    _touch(job, status="queued", message=message, error=None)


def take_choice(job_id: str, n: int) -> tuple["Job", Any]:
    """Return (job, the picked pending item) and clear the prompt.
    Raises ChoiceError with a user-safe reason when the pick cannot be used."""
    job = get_job(job_id)
    if job is None:
        raise KeyError(job_id)
    if job.status == "expired":
        raise ChoiceError("That choice list has expired. Please search again.")
    if job.status != "awaiting_choice":
        raise ChoiceError(f"That search is not waiting for a choice (it is {job.status}).")
    pending = job.pending or []
    if not isinstance(n, int) or n < 1 or n > len(pending):
        raise ChoiceError(f"Pick a number between 1 and {len(pending)}.")
    item = pending[n - 1]
    # Leave awaiting_choice in the same touch that clears the list, so a second
    # request cannot take the same choice.
    _touch(job, status="queued", choices=None, pending=None, expires_at=None)
    return job, item


def run_choice_job(job_id: str, fn: Callable[[], dict],
                   choice_message: Callable[[dict], str],
                   done_message: "str | Callable[[dict], str]",
                   fail_message: "str | Callable[[Exception], str]" = "Something went wrong while searching.") -> None:
    """Like run_job, but a result carrying a non-empty "choices" list parks the
    job at awaiting_choice instead of finishing it."""
    def _task():
        job = JOBS.get(job_id)
        if job is None:
            return
        _touch(job, status="running")
        try:
            result = fn()
            # `pending` holds indexer guids and `choices` is already top-level,
            # so neither belongs in the public `result` body.
            public = {k: v for k, v in result.items() if k not in ("pending", "choices")}
            choices = result.get("choices") or []
            if choices:
                await_choice(job_id, choices, result.get("pending") or [],
                             choice_message(result))
                _touch(job, result=public)
                return
            msg = done_message(result) if callable(done_message) else done_message
            _touch(job, status="done", result=public, message=msg)
        except Exception as e:  # noqa: BLE001 — record any failure on the job
            _touch(job, status="failed", error=f"{type(e).__name__}: {e}",
                   message=_fail_text(fail_message, e))
    _EXECUTOR.submit(_task)