distinct_agent/fanout.py

"""One worker, several servers, one queue.

A volunteer with a spare machine has no reason to offer it to exactly one
server. They might run one for themselves, join a friend's, and join a public
one, and none of those should require a second copy of the worker, a second
model download or a second slice of the machine.

**Why this is a transport and not a second loop.** The obvious design is one
:class:`~distinct_agent.worker.WorkerLoop` per server, and it is wrong. Each
loop would own its own queue, so a machine advertising "two free slots" to
three servers would be advertising six, and three servers would each fill what
they believed was an empty queue. Capacity is a property of the machine, so
there is one queue, one loop and one capacity figure; what fans out is the
conversation with each server.

So this implements the same :class:`~distinct_agent.worker.AgentTransport`
interface the loop already talks to, and multiplexes underneath. The loop is
unchanged and does not know how many servers exist.

**Routing.** A job arrives from one server and its result belongs to that
server and to no other. Every offer accepted is recorded against the transport
that supplied it, and status, results and acknowledgements follow that record.
A result must never be published to a server that did not ask for it: it would
be handing one server's user's answer to another server's operator.

**Failure.** One unreachable server must not stop the others. A transport that
raises is put aside for a back-off interval and retried; the worker keeps
serving everyone else meanwhile, and says so once rather than on every poll.
"""

from __future__ import annotations

import threading
import time
from collections.abc import Callable, Iterable, Sequence
from dataclasses import dataclass, field
from typing import Any

from distinct_protocol import AgentSnapshot, JobResult, JobSpec, JobStatus

#: How long a server that raised is left alone before it is tried again, and
#: the ceiling that back-off climbs to. Short enough that a server restarting
#: is picked up quickly, long enough that a server that is simply gone does
#: not dominate every poll cycle.
BACKOFF_START_SECONDS = 5.0
BACKOFF_CEILING_SECONDS = 120.0


@dataclass
class ServerLink:
    """One paired server, and what the worker knows about its health."""

    name: str
    url: str
    transport: Any
    failures: int = 0
    resume_at: float = 0.0
    last_error: str = ""
    #: Jobs currently owned by this server. Kept so a result can be routed
    #: back, and so a link's outstanding work can be described.
    job_ids: set[str] = field(default_factory=set)

    def healthy(self, now: float) -> bool:
        return now >= self.resume_at


class MultiServerTransport:
    """Poll several servers in turn; route each job's traffic back to its own.

    Thread safety matters here and is not incidental: the poll thread calls
    :meth:`poll_offer` while worker threads call :meth:`publish_result` for
    jobs that arrived several cycles ago. The routing table is guarded, and
    the underlying transports are assumed to be individually safe in the same
    way the single-server case already assumes.
    """

    def __init__(
        self,
        links: Sequence[ServerLink],
        *,
        monotonic: Callable[[], float] = time.monotonic,
        report: Callable[[str], None] | None = None,
    ) -> None:
        if not links:
            raise ValueError("a worker needs at least one server")
        self._links = list(links)
        self._monotonic = monotonic
        self._report = report or (lambda message: None)
        self._lock = threading.RLock()
        self._owner: dict[str, ServerLink] = {}
        self._cursor = 0

    # Introspection ------------------------------------------------------

    @property
    def links(self) -> tuple[ServerLink, ...]:
        with self._lock:
            return tuple(self._links)

    def describe(self) -> str:
        """One line per server, for the operator's start-up banner."""

        now = self._monotonic()
        lines = []
        for link in self.links:
            if link.healthy(now):
                state = f"{len(link.job_ids)} running" if link.job_ids else "idle"
            else:
                state = f"unreachable, retrying ({link.last_error[:60]})"
            lines.append(f"  {link.name}: {link.url} [{state}]")
        return "\n".join(lines)

    # Health -------------------------------------------------------------

    def _fail(self, link: ServerLink, exc: BaseException) -> None:
        with self._lock:
            first = link.failures == 0
            link.failures += 1
            delay = min(
                BACKOFF_CEILING_SECONDS, BACKOFF_START_SECONDS * (2 ** (link.failures - 1))
            )
            link.resume_at = self._monotonic() + delay
            link.last_error = f"{type(exc).__name__}: {exc}"
        if first:
            # Once, not on every cycle. A server that is down for an hour
            # should not produce an hour of identical lines.
            self._report(
                f"{link.name} is not responding ({link.last_error[:120]}). "
                f"Still serving the others; retrying in {delay:.0f}s."
            )

    def _recover(self, link: ServerLink) -> None:
        with self._lock:
            if not link.failures:
                return
            link.failures = 0
            link.resume_at = 0.0
            link.last_error = ""
        self._report(f"{link.name} is responding again.")

    # AgentTransport -----------------------------------------------------

    def poll_offer(self, snapshot: AgentSnapshot) -> JobSpec | None:
        """Ask each healthy server in turn, starting where the last poll left off.

        Round-robin rather than always-first-server, so a busy server cannot
        starve the others of this worker's attention. The first offer wins and
        the cursor advances past it, which is the fairness this needs: no
        server is asked twice before every other has been asked once.
        """

        now = self._monotonic()
        with self._lock:
            order = self._links[self._cursor :] + self._links[: self._cursor]
        for index, link in enumerate(order):
            if not link.healthy(now):
                continue
            try:
                offer = link.transport.poll_offer(snapshot)
            except Exception as exc:  # noqa: BLE001 - one server must not stop the rest
                self._fail(link, exc)
                continue
            self._recover(link)
            if offer is None:
                continue
            with self._lock:
                self._owner[offer.id] = link
                link.job_ids.add(offer.id)
                self._cursor = (self._cursor + index + 1) % len(self._links)
            return offer
        with self._lock:
            if self._links:
                self._cursor = (self._cursor + 1) % len(self._links)
        return None

    def poll_cancellations(self) -> Iterable[str]:
        """Every server's cancellations, together.

        A cancellation names a job id, and job ids are unique across servers
        because the server that issued one is the only one that knows it. So
        the union is safe: an id from server A cannot collide with an id from
        server B, and if it somehow did, the queue would cancel a job that
        server A had every right to cancel anyway.
        """

        now = self._monotonic()
        cancelled: list[str] = []
        for link in self.links:
            if not link.healthy(now):
                continue
            try:
                cancelled.extend(str(job_id) for job_id in link.transport.poll_cancellations())
            except Exception as exc:  # noqa: BLE001
                self._fail(link, exc)
        return cancelled

    def _link_for(self, job_id: str) -> ServerLink | None:
        with self._lock:
            return self._owner.get(job_id)

    def acknowledge_offer(self, decision: Any) -> None:
        link = self._link_for(getattr(decision, "job_id", ""))
        if link is None:
            return
        try:
            link.transport.acknowledge_offer(decision)
        except Exception as exc:  # noqa: BLE001
            self._fail(link, exc)
        if not getattr(decision, "accepted", True):
            self._release(getattr(decision, "job_id", ""))

    def publish_status(
        self, job_id: str, status: JobStatus, reason: str | None = None
    ) -> None:
        link = self._link_for(job_id)
        if link is None:
            return
        try:
            link.transport.publish_status(job_id, status, reason)
        except Exception as exc:  # noqa: BLE001
            self._fail(link, exc)

    def publish_result(self, result: JobResult) -> None:
        """To the server that asked, and to no other.

        A missing owner is dropped rather than broadcast. Publishing an answer
        to every server would hand one server's user's text to operators who
        never asked for it, which is a worse outcome than losing a result the
        worker cannot attribute.
        """

        link = self._link_for(result.job_id)
        if link is None:
            self._report(
                f"Result for {result.job_id} has no server to return to; discarding it "
                "rather than sending it to a server that did not ask."
            )
            return
        try:
            link.transport.publish_result(result)
        except Exception as exc:  # noqa: BLE001
            self._fail(link, exc)
        finally:
            self._release(result.job_id)

    def _release(self, job_id: str) -> None:
        with self._lock:
            link = self._owner.pop(job_id, None)
            if link is not None:
                link.job_ids.discard(job_id)

    def close(self) -> None:
        for link in self.links:
            close = getattr(link.transport, "close", None)
            if callable(close):
                try:
                    close()
                except Exception:  # noqa: BLE001 - closing must not raise
                    pass


__all__ = [
    "BACKOFF_CEILING_SECONDS",
    "BACKOFF_START_SECONDS",
    "MultiServerTransport",
    "ServerLink",
]