FabledCurator/backend/app/services/worker_control.py

"""Read and change a lane's live pool, over the broker.

Milestone 422 step 2. The half of the milestone that does something.

## No docker socket is involved, and that is the point

Milestone 365 put "acting on the state" out of scope because restarting a
dead worker needs a docker socket the web container deliberately does not
have. That is true of RESTARTING a container. It is not true of changing how
much work a RUNNING worker does: celery's remote control sends a message over
the broker and the worker resizes its own pool. Same Redis the app already
uses, no new privilege, no new surface.

    pool_grow / pool_shrink   how many slots a lane runs
    add_consumer / cancel_consumer   whether it consumes its queues at all

The operator ruled the socket out independently (2026-09-22: *"this feature
is a very invasive idea in my mind and I'd like to avoid it"*), and nothing
here raises the question.

## The setting is PER PROCESS, not per lane total

`pool_grow(n, destination=[...])` adds n slots to EACH destination it names.
While the stack still runs several containers per lane — the operator's
production `worker` is `replicas: 2` — a single delta applied to a lane's
total would be wrong for every replica.

So `slots` means what `CELERY_CONCURRENCY` means: the pool size of one
process. The reconcile below drives EACH replica to that number
independently, computing its own delta from that replica's current pool, so
replicas that have drifted apart (one restarted, one was grown) converge
rather than being moved in lockstep from a shared baseline.

After step 5 there is one process per lane and the distinction disappears.
It matters now, and getting it wrong now would be invisible — the totals
would simply be double what the UI claimed.

## Why reserved() is read alongside the queue depth

Celery PREFETCHES: a worker pulls more messages than it can run and holds
them in memory. Those have already left the Redis list, so `LLEN` — which is
what `/api/system/activity/queues` reports — can read 0 while thirty tasks
are waiting inside a worker. Any judgement about backlog that uses only LLEN
under-reports, which matters for the UI and is disqualifying for step 7's
autoscaler.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

from .worker_lanes import LANES, LANES_BY_QUEUE_KEY, Lane

log = logging.getLogger(__name__)

# celery control is a broker round trip on a request path, so it gets a
# deadline (rule 156) — the same reasoning and the same budget as
# service_roster's inspect. A broker that stopped answering must make this
# report "not present", which is true, rather than hang the page.
CONTROL_TIMEOUT_SECONDS = 2.0


@dataclass
class LaneLiveState:
    """What `celery inspect` says about one lane right now.

    `present=False` is NOT "zero slots" — it is "nothing answered". A lane
    whose worker is restarting, or whose broker is unreachable, must read as
    unknown rather than as stopped: an unswept absence is not a verdict
    (snippet #3969). The reconcile in step 3 skips an absent lane rather than
    correcting it, which is only safe because this distinction is kept.
    """

    present: bool = False
    replicas: int = 0
    # Per-process pool size. Equal across replicas unless one has drifted.
    pool: int | None = None
    active: int = 0
    reserved: int = 0
    hostnames: list[str] = field(default_factory=list)


def _lane_for_queues(queues: tuple[str, ...]) -> Lane | None:
    return LANES_BY_QUEUE_KEY.get(tuple(sorted(queues)))


def inspect_lanes_sync() -> dict[str, LaneLiveState]:
    """Live state per lane name. Sync — callers wrap in asyncio.to_thread.

    Never raises. Every lane is present in the result; ones nothing answered
    for carry `present=False`, so a caller cannot accidentally read a missing
    lane as an empty one by iterating only what came back.
    """
    out = {lane.name: LaneLiveState() for lane in LANES}
    try:
        from ..celery_app import celery as celery_app

        insp = celery_app.control.inspect(timeout=CONTROL_TIMEOUT_SECONDS)
        active_queues = insp.active_queues() or {}
        stats = insp.stats() or {}
        active = insp.active() or {}
        reserved = insp.reserved() or {}
    except Exception:
        log.warning("worker_control: celery inspect failed", exc_info=True)
        return out

    for hostname, queues in active_queues.items():
        lane = _lane_for_queues(tuple(q["name"] for q in queues))
        if lane is None:
            # A deployment slicing CELERY_QUEUES differently. Reported by the
            # roster under its raw queue list; it simply has no lane row to
            # control, which is honest rather than an error.
            continue
        state = out[lane.name]
        state.present = True
        state.replicas += 1
        state.hostnames.append(hostname)
        state.active += len(active.get(hostname, []))
        state.reserved += len(reserved.get(hostname, []))

        # `pool.max-concurrency` is the number pool_grow/pool_shrink move and
        # the number the UI shows. Absent on a worker whose stats did not
        # answer, which leaves pool=None — unknown, not zero.
        pool = (stats.get(hostname) or {}).get("pool", {}).get("max-concurrency")
        if isinstance(pool, int):
            state.pool = pool if state.pool is None else max(state.pool, pool)

    for state in out.values():
        state.hostnames.sort()
    return out


def set_lane_slots_sync(lane: Lane, target: int) -> tuple[bool, str | None]:
    """Drive every replica of `lane` to `target` slots. Returns (applied, err).

    Per-replica deltas rather than one shared delta: see the module docstring.
    A replica already at the target is issued nothing at all, which is what
    makes step 3's periodic reconcile converge instead of re-sending a grow of
    zero forever (lesson #4183 — an enforcer without a reachable fixed point
    re-does its own work every tick).

    `applied=False` is not a failure of the SETTING. The caller has already
    stored the value; this says only that the live push did not land, and the
    reconcile will carry it when the lane answers again.
    """
    try:
        from ..celery_app import celery as celery_app

        live = inspect_lanes_sync()[lane.name]
        if not live.present:
            return False, "lane is not running"
        if live.pool is None:
            return False, "worker did not report its pool size"

        control = celery_app.control
        for hostname in live.hostnames:
            delta = target - live.pool
            if delta > 0:
                control.pool_grow(delta, destination=[hostname])
            elif delta < 0:
                control.pool_shrink(-delta, destination=[hostname])
        return True, None
    except Exception as exc:  # noqa: BLE001 — reported, never raised at a caller
        log.warning("worker_control: could not resize %s", lane.name, exc_info=True)
        return False, str(exc)


def set_lane_enabled_sync(lane: Lane, enabled: bool) -> tuple[bool, str | None]:
    """Start or stop `lane` consuming its queues, without killing the process.

    `cancel_consumer` rather than a shutdown: a stopped consumer keeps its
    worker alive and answering `inspect`, so a disabled lane stays visible and
    can be turned back on. A killed worker would read as absent, which is the
    same signal as a crash — and the whole point of the roster (#365) is that
    those two must not look alike.
    """
    try:
        from ..celery_app import celery as celery_app

        live = inspect_lanes_sync()[lane.name]
        if not live.present:
            return False, "lane is not running"
        control = celery_app.control
        for queue in lane.queues:
            if enabled:
                control.add_consumer(queue, destination=live.hostnames)
            else:
                control.cancel_consumer(queue, destination=live.hostnames)
        return True, None
    except Exception as exc:  # noqa: BLE001
        log.warning(
            "worker_control: could not %s %s",
            "enable" if enabled else "disable", lane.name, exc_info=True,
        )
        return False, str(exc)