app/jobqueue.py

"""A minimal job transport on top of redis-py.

Hand-written instead of rq/arq/celery: the only added dependency is `redis`
and the whole protocol fits in ~100 readable lines.

    API    : LPUSH pc:jobs:segment '{"job_id": ..., "payload": ...}'
             then poll the status key -- segment/analyze jobs run for
             minutes, so nobody waits on a reply list
    worker : BRPOP pc:jobs:segment | pc:jobs:analyze -> handler(payload)
             LPUSH pc:result:{job_id} + EXPIRE (the reply is garbage
             collected if the API already walked away) and status updates
             along the way

Mode comes from the env var `REDIS_URL`, mirroring `DATABASE_URL` in
app/db.py: unset -> "inprocess", the CV layer has no way to run and the
routes say so; set -> "queue". There is deliberately NO silent in-process
fallback when the queue dies -- hiding a dead worker is worse than a clear
error.

There is no priority queue yet. Room was left for one by naming the list
after the JOB TYPE (`pc:jobs:segment`) instead of using a single shared
list: a new job type is a new list, not a new protocol.

Every Redis command here BLOCKS (sync redis-py), so routes must call into
this module through `run_in_executor`. The worker-side functions
(`serve_one`/`beat`/`clear_heartbeat`) accept an optional `client` so tests
can inject a SECOND fakeredis connection -- two connections behaving like
two real processes, while the real worker uses the module client created by
`setup()`.
"""

from __future__ import annotations

import json
import os
import time
import uuid

import redis

RESULT_KEY = "pc:result:{job_id}"
CV_HEARTBEAT_KEY = "pc:worker:heartbeat:cv"

# Analyze one clip: a LONG job (minutes), so it does not use a request-reply
# round trip. The API fires `enqueue` and forgets, then polls the status key;
# the worker updates that key while it runs.
ANALYZE_JOBS_KEY = "pc:jobs:analyze"
ANALYZE_STATUS_KEY = "pc:analyze:{analyze_id}"
ANALYZE_STATUS_TTL_S = 3600   # results live 1h, then collect themselves

# Segment a multi-shot video: same fire-and-forget protocol and the same CV
# worker, with progress per VIDEO under VIDEO_STATUS_KEY. The DURABLE result
# (shot list plus per-shot JSON) is on disk as soon as it exists -- Redis is
# only a progress channel, because job results were genuinely lost once when
# they only lived in Redis under a 1 hour TTL.
SEGMENT_JOBS_KEY = "pc:jobs:segment"
VIDEO_STATUS_KEY = "pc:video:{video_id}"

RESULT_TTL_S = 60        # an orphaned reply (API timed out) evaporates
HEARTBEAT_TTL_S = 15     # worker SETs every ~5s -> health sees death in <=15s
DEFAULT_TIMEOUT_S = 30.0

# redis-py 8.x sets a DEFAULT socket_timeout of 5s on EVERY command (it was
# None in the 7.x line). Measured against a real Redis on 04/08/2026: a
# BLPOP/BRPOP that blocks server-side for >= 5s makes the client raise
# TimeoutError exactly at 5.0s, before the server can return nil (fakeredis
# has no socket layer, so tests never see it). The defence is to wait in
# SLICES shorter than SOCKET_TIMEOUT_S through `_bwait` -- never call
# blpop/brpop directly with a long timeout. socket_timeout is pinned
# explicitly in setup() so a library default change cannot quietly alter the
# behaviour here.
SOCKET_TIMEOUT_S = 5.0
WAIT_SLICE_S = 2.0


class QueueDown(Exception):
    """Cannot talk to Redis (connection lost, or never enabled)."""


class QueueTimeout(Exception):
    """The job was pushed but produced no result in time -- the worker is
    dead or stuck."""


_client: redis.Redis | None = None
_mode: str = "inprocess"


def setup(url: str | None = None, client: redis.Redis | None = None) -> str:
    """Initialise the transport from `url` / env `REDIS_URL` (production) or
    from an injected `client` (tests use fakeredis). Returns the mode.
    Idempotent, like db.setup().

    A third mode, **"local"**, exists for the single-container Hugging Face
    Space: no `REDIS_URL` but `POOLCOACH_LOCAL_CV` enabled means the
    transport runs on fakeredis IN PROCESS and the app hosts a CV worker
    thread of its own (app/localcv.py). fakeredis rather than a hand-written
    in-memory queue because it is EXACTLY the backend every transport test
    already runs through -- a second implementation would be a second place
    to drift. If fakeredis is missing we fall back to "inprocess" instead of
    failing to boot.
    """
    global _client, _mode
    teardown()
    if client is not None:
        _client, _mode = client, "queue"
        return _mode
    url = url or os.environ.get("REDIS_URL")
    if url:
        _client = redis.Redis.from_url(url, decode_responses=True,
                                       socket_timeout=SOCKET_TIMEOUT_S)
        _mode = "queue"
        return _mode
    if os.environ.get("POOLCOACH_LOCAL_CV", "").strip().lower() not in (
            "", "0", "false", "no"):
        try:
            import fakeredis
        except ImportError:
            print("[poolcoach] POOLCOACH_LOCAL_CV is set but the fakeredis "
                  "package is missing -- falling back to inprocess mode",
                  flush=True)
            return _mode
        _client = fakeredis.FakeStrictRedis(decode_responses=True)
        _mode = "local"
    return _mode


def teardown() -> None:
    """Back to inprocess mode (closing the connection if there is one)."""
    global _client, _mode
    if _client is not None:
        try:
            _client.close()
        except Exception:  # noqa: BLE001 -- closing a dead Redis must still finish
            pass
    _client = None
    _mode = "inprocess"


def mode() -> str:
    return _mode


def cv_enabled() -> bool:
    """Does the CV layer have a way to run? Routes ask this instead of
    `mode() == "queue"`: both "queue" (a worker process) and "local" (an
    embedded thread) have a worker serving jobs; only "inprocess" does
    not."""
    return _mode in ("queue", "local")


def timeout_s() -> float:
    """How long to wait for a result -- read from the env EVERY time so
    tests can override it with monkeypatch."""
    return float(os.environ.get("POOLCOACH_QUEUE_TIMEOUT_S",
                                DEFAULT_TIMEOUT_S))


def _require(client: redis.Redis | None = None) -> redis.Redis:
    c = client or _client
    if c is None:
        raise RuntimeError("transport not set up -- inprocess mode has no "
                           "Redis client (check mode() first)")
    return c


def _bwait(op, key: str, timeout: float):
    """BLPOP/BRPOP waiting up to `timeout` seconds in short SLICES.

    Each slice is shorter than SOCKET_TIMEOUT_S so the server always gets to
    return nil before the client drops the socket (see the SOCKET_TIMEOUT_S
    note above). Returns the item, or None when the deadline passes; Redis
    errors propagate for the caller to classify.
    """
    deadline = time.monotonic() + timeout
    while True:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            return None
        # clamp to >= 0.05: a BLPOP timeout of 0 means "wait forever", so 0
        # must never reach Redis (overshooting the deadline by at most 50ms
        # is harmless)
        item = op(key, timeout=max(min(WAIT_SLICE_S, remaining), 0.05))
        if item is not None:
            return item


# --------------------------------------------------------- the API side

def enqueue(payload: dict, jobs_key: str,
            job_id: str | None = None) -> str:
    """Push ONE job without waiting for a result -- the shape every job here
    has, since they all run for minutes.

    The wire format is `{"job_id", "payload"}`, which is what `serve_one`
    reads on the worker side; progress and results travel through the status
    keys (`set_analyze_status` / `set_video_status`). `job_id` can be passed
    in so the route can reuse it as the public id the client polls with.
    Raises QueueDown when Redis cannot be reached.
    """
    c = _require()
    job_id = job_id or uuid.uuid4().hex
    try:
        c.lpush(jobs_key, json.dumps({"job_id": job_id, "payload": payload}))
    except redis.RedisError as e:
        raise QueueDown(f"{type(e).__name__}: {e}") from e
    return job_id


def set_analyze_status(analyze_id: str, data: dict,
                       client: redis.Redis | None = None) -> None:
    """Write the status of an analyze job (SET + TTL, so the key evaporates
    after an hour).

    `data` is a plain JSON dict {"status": queued|running|done|error, ...} --
    its shape is agreed between route and worker, the transport does not
    judge it. Raises QueueDown when Redis dies; the caller decides whether to
    swallow that (the worker does, so a long job does not die over one
    dropped beat; the route returns 503).
    """
    try:
        _require(client).set(ANALYZE_STATUS_KEY.format(analyze_id=analyze_id),
                             json.dumps(data), ex=ANALYZE_STATUS_TTL_S)
    except redis.RedisError as e:
        raise QueueDown(f"{type(e).__name__}: {e}") from e


def get_analyze_status(analyze_id: str) -> dict | None:
    """Read the status of an analyze job -- None for an unknown id or an
    expired key. Raises QueueDown when Redis dies, so the route can answer
    503 instead of blaming the id with a 404."""
    try:
        raw = _require().get(ANALYZE_STATUS_KEY.format(analyze_id=analyze_id))
    except redis.RedisError as e:
        raise QueueDown(f"{type(e).__name__}: {e}") from e
    return json.loads(raw) if raw else None


def set_video_status(video_id: str, data: dict,
                     client: redis.Redis | None = None) -> None:
    """Status of a segment job, per VIDEO -- same deal as
    set_analyze_status: SET + TTL, shape agreed between route and worker,
    QueueDown when Redis dies (worker swallows, route answers 503)."""
    try:
        _require(client).set(VIDEO_STATUS_KEY.format(video_id=video_id),
                             json.dumps(data), ex=ANALYZE_STATUS_TTL_S)
    except redis.RedisError as e:
        raise QueueDown(f"{type(e).__name__}: {e}") from e


def get_video_status(video_id: str) -> dict | None:
    """Read the status of a segment job -- None for an unknown id or an
    expired key. The DURABLE shot list is on disk and the route keeps
    reading it; Redis is only the progress channel."""
    try:
        raw = _require().get(VIDEO_STATUS_KEY.format(video_id=video_id))
    except redis.RedisError as e:
        raise QueueDown(f"{type(e).__name__}: {e}") from e
    return json.loads(raw) if raw else None


def ping() -> bool:
    """Is Redis alive? For /api/health -- errors are swallowed, because
    False is an answer, not a failure."""
    try:
        return bool(_require().ping())
    except redis.RedisError:
        return False


def worker_alive(key: str = CV_HEARTBEAT_KEY) -> bool:
    """Is the worker heartbeat still valid? For /api/health, swallowing
    errors the same way ping() does."""
    try:
        return _require().exists(key) == 1
    except redis.RedisError:
        return False


def worker_alive_or_raise(key: str = CV_HEARTBEAT_KEY) -> bool:
    """Heartbeat check BEFORE enqueueing (fast-fail). DIFFERENT from
    `worker_alive`: a dead Redis must RAISE QueueDown so the route reports
    "Redis" rather than blaming the worker. It is a blocking Redis command
    like every other one here, so routes call it through an executor."""
    try:
        return _require().exists(key) == 1
    except redis.RedisError as e:
        raise QueueDown(f"{type(e).__name__}: {e}") from e


# ------------------------------------------------------ the worker side

def serve_one(handler, jobs_key: str, timeout_s: float = 5.0,
              client: redis.Redis | None = None) -> bool:
    """ONE round of BRPOP -> ``handler(payload)`` -> reply + TTL.

    Returns False when `timeout_s` passes with no job, which is the outer
    worker loop's chance to beat its heartbeat and notice Ctrl+C. A job left
    over from a crashed worker is still served when the worker comes back
    (an orphaned reply expires on its own after RESULT_TTL_S) -- accepting a
    few wasted seconds rather than adding deadlines to the protocol.
    """
    c = _require(client)
    item = _bwait(c.brpop, jobs_key, timeout_s)
    if item is None:
        return False
    job = json.loads(item[1])
    reply = handler(job["payload"])
    key = RESULT_KEY.format(job_id=job["job_id"])
    pipe = c.pipeline()
    pipe.lpush(key, json.dumps(reply))
    pipe.expire(key, RESULT_TTL_S)
    pipe.execute()
    return True


def beat(client: redis.Redis | None = None,
         key: str = CV_HEARTBEAT_KEY) -> None:
    """The worker reporting that it is alive -- call it ONLY after model
    loading and the JIT are done: health must not report alive while the
    worker is still starting up."""
    _require(client).set(key, "alive", ex=HEARTBEAT_TTL_S)


def clear_heartbeat(client: redis.Redis | None = None,
                    key: str = CV_HEARTBEAT_KEY) -> None:
    """Clean worker exit (Ctrl+C) -- drop the heartbeat so health sees the
    death IMMEDIATELY instead of waiting out the TTL. If Redis is already
    gone, never mind: the TTL will clean up."""
    try:
        _require(client).delete(key)
    except redis.RedisError:
        pass