FabledSteward/steward/core/holding.py

"""Collected readings held in memory until the database takes them (#5995).

Steward's Postgres sits on the NFS storage it watches. When that storage
stalls, every write waits, and a collector that stores what it found before
collecting again stops collecting for as long as the stall lasts: on
2026-10-17 the NAS readings went blind for the fifteen minutes they mattered.

A collector built on a Holding collects on schedule whatever the database is
doing, and hands each round to the holding. One writer at a time stores what
is held, oldest first, each item at the time it was collected, so a stall
delays the readings instead of losing them, and alert rules count them in the
order they happened.

    held = holding(app, "SNMP poll", store=_store, limit=120)
    held.hold(round)                     # one collected round
    await held.flush(wait_s=20)          # store it, or leave it to the writer

The holding is bounded: past `limit` rounds the oldest is dropped, and that
is logged with how many have gone. A store has `deadline_s`: generous, past
any stall seen so far, so a connection that will never answer is given up
and the rounds stay held for a fresh one rather than waiting forever.
"""
from __future__ import annotations

import asyncio
import logging
from collections import deque
from typing import Awaitable, Callable

logger = logging.getLogger(__name__)


class Holding:
    def __init__(self, name: str, store: Callable[[list], Awaitable[None]], limit: int,
                 deadline_s: float = 1800.0) -> None:
        self.name = name
        self.store = store
        self.limit = limit
        self.deadline_s = deadline_s
        self._held: deque[tuple[int, object]] = deque()
        self._seq = 0
        self._writer: asyncio.Task | None = None
        self.dropped = 0

    def __len__(self) -> int:
        return len(self._held)

    def hold(self, item) -> None:
        """Hold one collected round until it is stored."""
        self._seq += 1
        self._held.append((self._seq, item))
        while len(self._held) > self.limit:
            self._held.popleft()
            self.dropped += 1
            logger.warning("%s: the database hasn't taken a round in %d; dropped the oldest "
                           "(%d dropped so far)", self.name, self.limit, self.dropped)

    async def flush(self, wait_s: float) -> bool:
        """Start the writer if it isn't running, and wait up to `wait_s` for
        it. True when nothing is left held. A writer still waiting on the
        database is left to finish: the next round joins what it stores next."""
        writer = self._writer
        if writer is None or writer.done() or writer.get_loop() is not asyncio.get_running_loop():
            if not self._held:
                return True
            writer = self._writer = asyncio.create_task(self._write())
        done, _ = await asyncio.wait({writer}, timeout=wait_s)
        if not done:
            logger.warning("%s: the database hasn't taken %d held round(s) in %.0f s; "
                           "collecting on", self.name, len(self._held), wait_s)
        return not self._held

    async def _write(self) -> None:
        while self._held:
            batch = list(self._held)
            try:
                await asyncio.wait_for(self.store([item for _, item in batch]), self.deadline_s)
            except TimeoutError:
                logger.error("%s: storing %d held round(s) took over %.0f s; given up, they stay "
                             "held (the database may still finish it)", self.name, len(batch),
                             self.deadline_s)
                return
            except Exception as exc:
                if unreachable(exc):
                    logger.warning("%s: the database can't be reached (%s); %d round(s) stay held",
                                   self.name, exc, len(batch))
                    return
                # A round the database refuses would otherwise be retried, and
                # refused, ahead of every round after it: store each on its own
                # and let go of the ones refused.
                logger.exception("%s: storing %d held round(s) failed; storing each alone",
                                 self.name, len(batch))
                for seq, item in batch:
                    try:
                        await asyncio.wait_for(self.store([item]), self.deadline_s)
                    except Exception as one:
                        if isinstance(one, TimeoutError) or unreachable(one):
                            return
                        logger.error("%s: a round collected was refused and is dropped: %s",
                                     self.name, one)
                    self._drop_through(seq)
                continue
            self._drop_through(batch[-1][0])

    def _drop_through(self, seq: int) -> None:
        while self._held and self._held[0][0] <= seq:
            self._held.popleft()


def unreachable(exc: BaseException) -> bool:
    """Whether a store failed because the database couldn't be reached, as
    opposed to refusing what it was given."""
    from sqlalchemy.exc import DBAPIError, InterfaceError, OperationalError

    from steward.database import DatabaseUnavailable
    if isinstance(exc, (OSError, DatabaseUnavailable, InterfaceError, OperationalError)):
        return True
    return isinstance(exc, DBAPIError) and bool(exc.connection_invalidated)


def holding(app, name: str, store: Callable[[list], Awaitable[None]], limit: int) -> Holding:
    """The app's Holding for one collector, made on first use."""
    holdings = getattr(app, "_holdings", None)
    if holdings is None:
        holdings = app._holdings = {}
    if name not in holdings:
        holdings[name] = Holding(name, store, limit)
    return holdings[name]


async def hold_write(app, name: str, write: Callable[[object], Awaitable[None]],
                     limit: int = 120, wait_s: float = 20.0) -> bool:
    """Hold one collected round as `write(session)` and store it: the held
    writes run oldest first, in one transaction. For a collector whose round
    is easiest said as the writes it makes. True when it was stored within
    `wait_s`."""
    held = holding(app, name, store=lambda writes: _run_writes(app, writes), limit=limit)
    held.hold(write)
    return await held.flush(wait_s)


async def _run_writes(app, writes: list) -> None:
    async with app.db_sessionmaker() as session:
        async with session.begin():
            for write in writes:
                await write(session)