app/backend/fluksio/flow/executor.py

"""The execution service: what turns journaled work into node runs.

A consumer thread claims items from the work queue and hands each one to a
dispatch pool, which drives the wave it starts. Node bodies run on a second,
separate pool: if cascade drivers and node bodies shared one, a wave waiting
for its own nodes could occupy every thread and deadlock.

A reaper takes back items claimed by an engine that died before acknowledging
them, which is the mechanism that makes a crash mid-cascade recoverable rather
than lossy.
"""

from __future__ import annotations

import logging
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from typing import TYPE_CHECKING, Any

from fluksio.flow.queue import MAX_DELIVERIES, WorkItem, WorkQueue

if TYPE_CHECKING:
    from fluksio.flow.events import EventBus
    from fluksio.flow.pipeline import Pipeline

logger = logging.getLogger(__name__)

CLAIM_COUNT = 4
CLAIM_BLOCK_MS = 1000
# Long enough that a busy cascade is not mistaken for a dead one.
RECLAIM_IDLE_MS = 60_000
RECLAIM_INTERVAL_S = 30.0
# How often to tell the queue that what we hold is still being worked on. A
# node may run for as long as it likes, so what marks an item abandoned is this
# stopping — which is what an engine that died does.
TOUCH_INTERVAL_S = 20.0
DELAYED_INTERVAL_S = 1.0
MAX_CASCADES = 4
# How long a reload waits for claimed work to finish before rebuilding anyway.
DRAIN_TIMEOUT_S = 10.0


class ExecutionService:
    """Owns the engine's worker threads and the queue they read from."""

    def __init__(
        self,
        queue: WorkQueue,
        max_workers: int | None = None,
        events: EventBus | None = None,
    ) -> None:
        self.queue = queue
        self._events = events
        self._pipeline: Pipeline | None = None
        self._stop = threading.Event()
        self._intake = threading.Event()
        self._intake.set()
        self._inflight = 0
        self._inflight_lock = threading.Condition()
        # Entry ids claimed and still running, under _inflight_lock.
        self._active: set[str] = set()
        self.node_pool = ThreadPoolExecutor(
            max_workers=max_workers or 4, thread_name_prefix="node"
        )
        self._cascade_pool = ThreadPoolExecutor(
            max_workers=MAX_CASCADES, thread_name_prefix="cascade"
        )
        self._consumer: threading.Thread | None = None
        self._timers: threading.Thread | None = None

    # -------------------------------------------------------------------------
    # Lifecycle
    # -------------------------------------------------------------------------

    def start(self) -> None:
        if self._consumer is not None:
            return
        self._consumer = threading.Thread(
            target=self._consume, name="queue-consumer", daemon=True
        )
        self._consumer.start()
        self._timers = threading.Thread(
            target=self._tick, name="queue-timers", daemon=True
        )
        self._timers.start()

    def stop(self) -> None:
        self._stop.set()
        for thread in (self._consumer, self._timers):
            if thread is not None:
                thread.join(timeout=5)
        self._consumer = None
        self._timers = None
        self._cascade_pool.shutdown(wait=False)
        self.node_pool.shutdown(wait=False)
        self.queue.close()

    def bind(self, pipeline: Pipeline) -> None:
        """Point the service at the pipeline it should execute against."""
        self._pipeline = pipeline

    def pause_intake(self) -> None:
        """Stop claiming, and wait for what is already claimed to finish.

        Called around a rebuild: items claimed against the old pipeline should
        finish there rather than half-run against the new one.
        """
        self._intake.clear()
        deadline = time.monotonic() + DRAIN_TIMEOUT_S
        with self._inflight_lock:
            while self._inflight:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    logger.warning(
                        "Rebuild did not wait out %d cascades", self._inflight
                    )
                    return
                self._inflight_lock.wait(remaining)

    def resume_intake(self) -> None:
        self._intake.set()

    def alive(self) -> bool:
        return self._consumer is not None and self._consumer.is_alive()

    # -------------------------------------------------------------------------
    # Threads
    # -------------------------------------------------------------------------

    def _consume(self) -> None:
        failures = 0
        while not self._stop.is_set():
            if not self._intake.is_set():
                self._intake.wait(timeout=0.5)
                continue
            try:
                items = self.queue.claim(CLAIM_COUNT, CLAIM_BLOCK_MS)
                failures = 0
            except Exception as exc:
                failures += 1
                logger.error("Could not claim work: %s", exc)
                self._publish_unavailable(exc)
                # Backing off hard: a queue that is down stays down for a while.
                self._stop.wait(min(30.0, 2.0**failures))
                continue

            for item in items:
                self._dispatch(item)

    def _tick(self) -> None:
        """Promote delayed items, and take back what a dead engine dropped."""
        last_reclaim = 0.0
        last_touch = 0.0
        while not self._stop.is_set():
            self._stop.wait(DELAYED_INTERVAL_S)
            if self._stop.is_set():
                break
            try:
                self.queue.move_due(time.time())
            except Exception as exc:
                logger.error("Could not promote delayed work: %s", exc)

            now = time.monotonic()
            if now - last_touch >= TOUCH_INTERVAL_S:
                last_touch = now
                with self._inflight_lock:
                    running = list(self._active)
                try:
                    self.queue.touch(running)
                except Exception as exc:
                    logger.error("Could not touch claimed work: %s", exc)


            if now - last_reclaim < RECLAIM_INTERVAL_S:
                continue
            last_reclaim = now
            try:
                for item in self.queue.reclaim_stale(RECLAIM_IDLE_MS):
                    logger.info(
                        "Reclaimed work for '%s' (delivery %d)",
                        item.node,
                        item.deliveries,
                    )
                    self._dispatch(item)
            except Exception as exc:
                logger.error("Could not reclaim stale work: %s", exc)

    def _dispatch(self, item: WorkItem) -> None:
        with self._inflight_lock:
            self._inflight += 1
            if item.entry_id:
                self._active.add(item.entry_id)
        try:
            self._cascade_pool.submit(self._handle, item)
        except RuntimeError:
            # Pool already shutting down.
            self._done(item)

    def _done(self, item: WorkItem) -> None:
        with self._inflight_lock:
            self._inflight -= 1
            self._active.discard(item.entry_id)
            self._inflight_lock.notify_all()

    # -------------------------------------------------------------------------
    # Handling one item
    # -------------------------------------------------------------------------

    def step(self, flow: str) -> str | None:
        """Run one item a pause is holding, and hold everything else still.

        Blocking, so the caller sees the wave finish. Returns the node the item
        came from, or None when nothing is parked for this flow.
        """
        pipeline = self._pipeline
        if pipeline is None:
            return None
        item = self.queue.unpark_one(flow)
        if item is None:
            return None
        with pipeline.stepping(flow):
            try:
                self._run_item(item)
            except Exception:
                logger.exception("Step of '%s' failed", item.node)
        return item.node

    def _handle(self, item: WorkItem) -> None:
        handled = True
        try:
            handled = self._run_item(item)
        except Exception:
            logger.exception("Work item for '%s' failed", item.node)
        finally:
            # Leaving it unacknowledged is how it comes back: the reaper hands
            # it to whoever can actually run it.
            if handled:
                try:
                    self.queue.ack(item)
                except Exception as exc:
                    logger.error(
                        "Could not acknowledge work for '%s': %s", item.node, exc
                    )
            self._done(item)

    def _run_item(self, item: WorkItem) -> bool:
        """Run one item. False means it was not handled and must come back."""
        pipeline = self._pipeline
        if pipeline is None:
            logger.warning("No pipeline bound; leaving work for '%s'", item.node)
            return False

        if item.deliveries > MAX_DELIVERIES:
            self.queue.dead_letter(item, f"{item.deliveries} deliveries")
            self._publish(
                {
                    "type": "cascade_dropped",
                    "flow": item.flow,
                    "node": item.node,
                    "deliveries": item.deliveries,
                    "ts": time.time(),
                }
            )
            return True

        node = pipeline.get_node_by_id(item.node)
        if node is None:
            # The flow was edited while this was queued; its values are already
            # in state, so there is nothing to salvage.
            logger.debug("Work item for unknown node '%s', dropped", item.node)
            return True

        if pipeline.is_disabled(node.flow):
            return True

        if pipeline.is_paused(node.flow) and not pipeline.is_stepping(node.flow):
            self.queue.park(node.flow, item)
            return True

        if item.kind == "flush":
            # A rate-limit window ended; nothing to replay, only to let out.
            # ponytail: no run record for a flush — it is the tail of the run
            # that scheduled it, not a run of its own.
            pipeline.flush(node)
            return True

        if item.guard_key and str(node.recall(item.guard_key, "")) != item.guard_value:
            # The node moved on while this waited — a restarted timer, say.
            logger.debug("Guard no longer holds for '%s', dropped", item.node)
            return True

        # Only here is the item certain to run, which is what a run record is.
        now = time.time()
        self._publish(
            {
                "type": "cascade_started",
                "run": item.entry_id,
                "flow": item.flow,
                "node": item.node,
                "cause": item.cause,
                "deliveries": item.deliveries,
                "ts": now,
            }
        )
        if item.deliveries == 1:
            # A redelivery waited for the reaper, not for the engine.
            self._publish(
                {
                    "type": "work_latency",
                    "flow": item.flow,
                    "node": item.node,
                    "lag_ms": max(
                        0.0,
                        (now - max(item.enqueued_at, item.not_before)) * 1000,
                    ),
                    "ts": now,
                }
            )

        try:
            pipeline.apply_outputs(node, item.outputs or None)
            pipeline.run_downstream(
                node, entry_id=item.entry_id, replay=item.deliveries > 1
            )
        finally:
            # Paired, or a cascade that raised — state backend gone, say — is a
            # run left open until the abandoned sweep ten minutes later.
            self._publish(
                {
                    "type": "cascade_finished",
                    "run": item.entry_id,
                    "flow": item.flow,
                    "ts": time.time(),
                }
            )
        return True

    # -------------------------------------------------------------------------
    # Reporting
    # -------------------------------------------------------------------------

    def stats(self) -> dict[str, Any]:
        try:
            stats = self.queue.stats()
        except Exception as exc:
            return {"error": str(exc), "consumer_alive": self.alive()}
        stats["consumer_alive"] = self.alive()
        stats["cascades_busy"] = self._inflight
        return stats

    def _publish_unavailable(self, exc: Exception) -> None:
        self._publish(
            {
                "type": "queue_unavailable",
                "error": f"{type(exc).__name__}: {exc}",
                "ts": time.time(),
            }
        )

    def _publish(self, event: dict[str, Any]) -> None:
        if self._events is not None:
            self._events.publish(event)