From 4af32f8bd52edb16605e3c168d3e0646dfe39c79 Mon Sep 17 00:00:00 2001 From: Krzysztof kuhy Rudnicki Date: Fri, 4 Sep 2026 21:35:34 +0200 Subject: [PATCH] Handle frames side by side, and budget the bot's own share of a reply Bot.listen awaited every frame to completion, so one Commons fetch or one attachment upload held up every text reply behind it (DOCS-latency.md). Each frame now runs in its own task, never more than LIMIT at once; the reader waits for a slot before spawning, so a burst applies back-pressure at the websocket instead of fanning out upstream. One person's commands stay serialised (user_lock.py): two quick !kot would otherwise both pass the cooldown gate before either recorded a use. Every served command reports its timings split into the bot's own time and everybody else's (served.py); /_selftest carries that block and tool/latency.py holds own_ms to 100 ms at p95, with the command's cooldown switched off for the run and put back afterwards. Co-Authored-By: Claude Opus 4.8 --- DOCS-latency.md | 41 ++++++- DOCS-selftest.md | 9 ++ signal_bot/analytics/probe.py | 11 +- signal_bot/app.py | 14 ++- signal_bot/bot.py | 46 ++++++-- signal_bot/dispatcher.py | 72 ++++++++++--- signal_bot/inflight.py | 94 ++++++++++++++++ signal_bot/latency.py | 150 ++++++++++++++++++++++++++ signal_bot/selftest.py | 14 ++- signal_bot/served.py | 56 ++++++++++ signal_bot/user_lock.py | 50 +++++++++ tests/test_bot_listen.py | 152 ++++++++++++++++++++++++++ tests/test_dispatcher_concurrency.py | 146 +++++++++++++++++++++++++ tests/test_inflight.py | 129 ++++++++++++++++++++++ tests/test_latency.py | 155 +++++++++++++++++++++++++++ tests/test_selftest.py | 49 ++++++++- tests/test_served.py | 54 ++++++++++ tests/test_user_lock.py | 123 +++++++++++++++++++++ tool/latency.py | 66 ++++++++++++ 19 files changed, 1390 insertions(+), 41 deletions(-) create mode 100644 signal_bot/inflight.py create mode 100644 signal_bot/latency.py create mode 100644 signal_bot/served.py create mode 100644 signal_bot/user_lock.py create mode 100644 tests/test_bot_listen.py create mode 100644 tests/test_dispatcher_concurrency.py create mode 100644 tests/test_inflight.py create mode 100644 tests/test_latency.py create mode 100644 tests/test_served.py create mode 100644 tests/test_user_lock.py create mode 100755 tool/latency.py diff --git a/DOCS-latency.md b/DOCS-latency.md index 32a706c..8fac352 100644 --- a/DOCS-latency.md +++ b/DOCS-latency.md @@ -34,10 +34,43 @@ building anything, measure these two separately: that is the real win and it makes the prefetch pool far more valuable, because the pool could hold ids instead of bytes. -Head-of-line blocking is a separate, confirmed problem and does not need -re-measuring: `bot.py` awaits each frame to completion, so one slow upstream -fetch delays every other message behind it. Dispatching each frame as its own -task under a bounded semaphore fixes that regardless of what the send costs. +## Head-of-line blocking: fixed + +This one was confirmed without re-measuring: `bot.py` awaited each frame to +completion, so one slow upstream fetch delayed every message behind it. Now +every frame runs in a task of its own (`signal_bot/inflight.py`), never more +than `LIMIT` at once. The reader waits for a slot *before* spawning, so a +burst applies back-pressure at the websocket instead of fanning out into a +hundred upstream connections. The one thing that must not run side by side +is one person's own commands -- two quick `!kot` would both pass the cooldown +gate before either had recorded a use -- so `signal_bot/user_lock.py` +serialises per sender, and nobody else waits. + +## The bot's own share, measured and budgeted + +Every measurement above is a measurement of somebody else. To say anything +about *this* code, each served command now reports its timings split into +the bot's own time and everybody else's (`signal_bot/served.py`): + +| field | what it covers | +|---|---| +| `queue_ms` | frame read → a slot free: waiting behind other frames | +| `dispatch_ms` | slot → handler: parse, scope, tally, resolve, sender lock | +| `handler_ms` | the handler, wall clock | +| `send_ms` | of which: `POST /v2/send` and friends | +| `upstream_ms` | of which: Commons, the cat API, and the like | +| `own_ms` | `queue + dispatch + handler - send - upstream` | + +`POST /_selftest` answers with that block, and `tool/latency.py` holds it to +the budget: **`own_ms` ≤ `OWN_BUDGET_MS` (100 ms) at p95**, over N replies of +a text-only command in a test group, with the command's cooldown switched off +for the run and put back afterwards. Exit 1 when over. The whole reply is +bounded by the container and reported (`round_trip_ms`), not budgeted. + +```sh +tool/latency.py testbot # 20 x !pomoc +tool/latency.py testbot --text '!ja' --times 50 +``` ## Rule diff --git a/DOCS-selftest.md b/DOCS-selftest.md index 26f467c..f579c3f 100644 --- a/DOCS-selftest.md +++ b/DOCS-selftest.md @@ -29,6 +29,15 @@ scripts/selftest.sh testbot '!papież' 3 the self-test is always the same author, so a tighter loop just exercises the cooldown instead of the command. +## What it answers with + +`delivered`, the group id, the text -- and, when the text was a command, +`served`: what that one reply cost, split into the bot's own time and the +time spent on signal-cli and upstream (`signal_bot/served.py`, and the table +in `DOCS-latency.md`). A plain message or an unknown word is delivered and +counted but has no `served`, because nothing ran. `tool/latency.py` reads +that block repeatedly to hold the bot to its budget. + ## What it will not do It **posts for real**, into a real group. Two things keep that from becoming a diff --git a/signal_bot/analytics/probe.py b/signal_bot/analytics/probe.py index 7cef26a..2fbb7b8 100644 --- a/signal_bot/analytics/probe.py +++ b/signal_bot/analytics/probe.py @@ -10,11 +10,12 @@ A ``ContextVar`` avoids all of it. The dispatcher opens a probe around the handler call; anything running underneath finds it and notes what it did. ``contextvars`` is stdlib, so the runtime stays at exactly six dependencies. -``Bot.listen`` awaits ``dispatcher.handle`` directly rather than spawning a -task per frame, so set/reset in a ``try/finally`` cannot interleave. A -handler that does spawn a task gets a *copy* of the context, which is the -behaviour wanted: the copy's notes land in the probe, and the reset in the -parent cannot be undone by the child. +``Bot.listen`` runs every frame in a task of its own (``inflight.py``), and +a task starts with a *copy* of the context, so frames handled side by side +cannot see each other's probe: each ``probing()`` sets the variable in its +own copy, and the reset in one task cannot be undone by another. The same +holds for a handler that spawns a task of its own -- the copy's notes land +in that handler's probe, nowhere else. Every ``note_*`` is a no-op when nothing is probing. That branch lives here, once, so no caller needs to ask whether it is being measured -- which is diff --git a/signal_bot/app.py b/signal_bot/app.py index 3b07d24..ed307aa 100644 --- a/signal_bot/app.py +++ b/signal_bot/app.py @@ -12,7 +12,7 @@ from __future__ import annotations import asyncio import contextlib import logging -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import httpx from fastapi import FastAPI, Response, status @@ -124,11 +124,14 @@ async def healthz(response: Response) -> dict[str, str]: @app.post("/_selftest") -async def selftest(request: SelfTestRequest, response: Response) -> dict[str, str]: +async def selftest(request: SelfTestRequest, response: Response) -> dict[str, Any]: """Type a command at the bot on its own behalf, for verification. Bound to loopback with the rest of the app, and refuses any group not named in SELFTEST_GROUPS. See signal_bot/selftest.py for why this exists. + When the text was a command, ``served`` carries what it cost, split into + the bot's own time and the time spent waiting on Commons and signal-cli + (``served.py``); ``tool/latency.py`` reads it to hold the bot to a budget. """ starter: Starter = app.state.starter runtime = starter.runtime @@ -136,7 +139,7 @@ async def selftest(request: SelfTestRequest, response: Response) -> dict[str, st response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE return {"status": "connecting", "reason": starter.reason} try: - group_id = await run_selftest( + group_id, served = await run_selftest( runtime.bot, app.state.client, request.group, @@ -146,7 +149,10 @@ async def selftest(request: SelfTestRequest, response: Response) -> dict[str, st except SelfTestError as refusal: response.status_code = status.HTTP_403_FORBIDDEN return {"status": "refused", "reason": str(refusal)} - return {"status": "delivered", "group": group_id, "text": request.text} + answer: dict[str, Any] = {"status": "delivered", "group": group_id, "text": request.text} + if served is not None: + answer["served"] = served.as_dict() + return answer @app.post("/_reload") diff --git a/signal_bot/bot.py b/signal_bot/bot.py index 7ced04a..dd804b7 100644 --- a/signal_bot/bot.py +++ b/signal_bot/bot.py @@ -8,16 +8,19 @@ half is "what the bot does"; startup.py is "how it comes up". from __future__ import annotations +import asyncio import datetime as dt import logging +import time from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING from signal_bot.announce import announce_all from signal_bot.changelog import load as load_changelog from signal_bot.dispatcher import Dispatcher from signal_bot.http_client import ImageFetcher +from signal_bot.inflight import Inflight from signal_bot.mute import is_muted from signal_bot.portraits import PortraitCache from signal_bot.registry_holder import RegistryHolder @@ -33,6 +36,7 @@ if TYPE_CHECKING: from signal_bot.config import Config from signal_bot.discovery import GroupDirectory from signal_bot.liveness import Liveness + from signal_bot.served import Served from signal_bot.store import Store from signal_bot.transport_base import Transport @@ -73,6 +77,8 @@ class Bot: # "which day is it" cannot differ between writing an event and reporting # on it -- which is precisely the drift the old reset-after-post had. clock: Clock = now_local + # The frames being handled right now, bounded: see inflight. + inflight: Inflight = field(default_factory=Inflight) async def publish_report(self) -> None: """Post each group **one** message: today, plus anything still owed. @@ -149,13 +155,37 @@ class Bot: ) async def listen(self) -> None: - """Consume frames forever, handing each to the dispatcher.""" - async for frame in self.transport.receive(): - # Every frame counts, including receipts and typing indicators: - # the question this answers is "can anything still reach us?", - # not "did anybody address the bot?". - self.liveness.note_envelope() - await self.dispatcher.handle(frame) + """Consume frames forever, handing each to the dispatcher. + + Each frame runs in a task of its own, so a slow reply -- a picture + fetched from Commons, an attachment upload -- never holds up the + text reply behind it. The reader only pauses when every slot is + busy. When the stream ends the frames in flight are finished; when + the bot is stopped they are stopped with it. + """ + try: + async for frame in self.transport.receive(): + # Every frame counts, including receipts and typing indicators: + # the question this answers is "can anything still reach us?", + # not "did anybody address the bot?". + self.liveness.note_envelope() + await self.inflight.start( + self.dispatcher.handle(frame, received_at=time.perf_counter()) + ) + await self.inflight.drain() + except asyncio.CancelledError: + self.inflight.cancel() + raise + + async def serve(self, frame: str) -> Served | None: + """Handle one frame here and now, through the same slots ``listen`` uses. + + For the self-test: it wants the timings back, which a spawned task + cannot give it, and it must queue behind real frames exactly as a + real frame would, or its numbers are not the bot's. + """ + async with self.inflight.slot(): + return await self.dispatcher.handle(frame, received_at=time.perf_counter()) def reload_commands(self) -> bool: """Rebuild the command table from the data on disk, keeping the old one on failure.""" diff --git a/signal_bot/dispatcher.py b/signal_bot/dispatcher.py index 8c7e690..5ec9259 100644 --- a/signal_bot/dispatcher.py +++ b/signal_bot/dispatcher.py @@ -5,8 +5,9 @@ from __future__ import annotations import datetime as dt import logging +import time from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING from signal_bot import mute @@ -16,14 +17,16 @@ from signal_bot.context import CommandContext from signal_bot.counting import should_count, watched_group from signal_bot.envelope_shape import shape_of from signal_bot.models import parse_envelope +from signal_bot.served import Served from signal_bot.stats.recorder import record_message, sent_at_ms_of from signal_bot.unknown import answer_unknown +from signal_bot.user_lock import UserLocks if TYPE_CHECKING: from signal_bot.config import Config from signal_bot.discovery import GroupDirectory from signal_bot.models import IncomingMessage - from signal_bot.registry import CommandRegistry + from signal_bot.registry import CommandRegistry, Resolved from signal_bot.registry_holder import RegistryHolder from signal_bot.store import Store from signal_bot.transport_base import Transport @@ -31,6 +34,11 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) Clock = Callable[[], dt.datetime] +MS_PER_SECOND = 1000 + + +def _ms_since(started: float) -> int: + return round((time.perf_counter() - started) * MS_PER_SECOND) @dataclass(frozen=True, slots=True) @@ -46,6 +54,9 @@ class Dispatcher: store: Store transport: Transport clock: Clock + # Frames run concurrently (bot.py), so one person's commands are + # serialised here: see user_lock for the cooldown race that closes. + users: UserLocks = field(default_factory=UserLocks) @property def registry(self) -> CommandRegistry: @@ -72,15 +83,23 @@ class Dispatcher: except Exception: logger.exception("could not answer an unrecognised command in %s", group) - async def handle(self, frame: str) -> None: - """Process one raw websocket frame end to end.""" + async def handle(self, frame: str, *, received_at: float | None = None) -> Served | None: + """Process one raw websocket frame end to end. + + ``received_at`` is when the reader took the frame off the socket, + as ``time.perf_counter()``; the gap to here is time spent waiting + for a slot behind other frames. Returns the timings when a command + ran -- refused or not -- and None for every other frame. + """ + entered = time.perf_counter() + queue_ms = 0 if received_at is None else round((entered - received_at) * MS_PER_SECOND) # Keys and types only, never values -- see envelope_shape. This is # how "what does a frame with an attachment look like?" gets answered # without a code change and a lucky moment. logger.debug("envelope shape: %s", shape_of(frame)) message = parse_envelope(frame) if message is None: - return + return None # Before the scoping gate: a group-update frame arrives while its # group is still unknown, and it is what vouches for that group. @@ -104,7 +123,7 @@ class Dispatcher: message.group_id, list(self.directory.groups), ) - return + return None # Resolved before recording, not after: is_command is part of the # row, so "was this a command" has to be known at write time. It is @@ -136,16 +155,32 @@ class Dispatcher: now = self.clock() if not mute.allows(self.store, group, now.timestamp(), command): logger.debug("muted in %s; not answering", group) - return + return None - if command is None: + if resolved is None: # is_command above stays False, and this stays below the # recording rather than in front of it: an unrecognised word is # not a command just because the bot now says something about it. if counted: await self._answer_unknown(message, group) - return + return None + # Under this person's lock, gate included: the second of two quick + # !kot must see the first one's recorded use, or both are served. + async with self.users.holding(message.source_uuid): + return await self._run(message, group, resolved, now, queue_ms, entered) + + async def _run( + self, + message: IncomingMessage, + group: str, + resolved: Resolved, + now: dt.datetime, + queue_ms: int, + entered: float, + ) -> Served: + """Run one resolved command and record every trace of the attempt.""" + command = resolved.command logger.info("running command %s for %s", command.name, message.source_name) ctx = CommandContext( message=message, @@ -156,7 +191,7 @@ class Dispatcher: command_name=command.name, now=now.timestamp(), day=now.date().isoformat(), - args=resolved.args if resolved is not None else (), + args=resolved.args, ) # The envelope's own timestamp, not the clock, and the same value # record_message keyed its row on -- that shared key is what lets @@ -170,6 +205,7 @@ class Dispatcher: # cooldown -- recording a refusal there would push the wait N seconds # further away every time somebody was told to wait N seconds. # See analytics.recorder for the invariant and the test that holds it. + dispatch_ms = _ms_since(entered) with probe.probing() as notes, recorder.attempting(command.name) as attempt: await run_unless_off(self.store, command, ctx) recorder.record_attempt( @@ -184,7 +220,15 @@ class Dispatcher: attempt, notes, ) - if not attempt.ok: - return - # Recorded once, here, so stacked gates cannot double-count a use. - self.store.record_use(group, message.source_uuid, command.name, ctx.now, ctx.day) + if attempt.ok: + # Recorded once, here, so stacked gates cannot double-count a use. + self.store.record_use(group, message.source_uuid, command.name, ctx.now, ctx.day) + return Served( + command=command.name, + outcome=attempt.outcome, + queue_ms=queue_ms, + dispatch_ms=dispatch_ms, + handler_ms=attempt.handler_ms, + send_ms=notes.send_ms, + upstream_ms=sum(call.ms for call in notes.calls), + ) diff --git a/signal_bot/inflight.py b/signal_bot/inflight.py new file mode 100644 index 0000000..cf0d1d1 --- /dev/null +++ b/signal_bot/inflight.py @@ -0,0 +1,94 @@ +# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence. +"""Frames in flight: each one its own task, never more than a few at once. + +``Bot.listen`` used to await every frame to completion before reading the +next, so one slow upstream fetch -- a picture from Commons, a 2.6 s +attachment upload -- delayed every message behind it (``DOCS-latency.md``). +Now a frame is handed to a task of its own and the reader carries on. + +The count is bounded on purpose. Unbounded tasks would let a burst of frames +open a hundred upstream connections at once, and the reader would sail on +while memory filled with work nothing had got to. Waiting for a slot +*before* spawning applies the pressure where it belongs: the websocket +reader pauses, the socket buffers, and signal-cli holds the rest. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +from typing import TYPE_CHECKING, Final + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Coroutine + +logger = logging.getLogger(__name__) + +# Handlers running at once. Enough that a slow picture never queues a text +# reply behind it; few enough that a burst cannot fan out into a flood of +# upstream requests. Not configuration: nothing about a group changes it. +LIMIT: Final = 8 + + +class Inflight: + """A bounded set of running frame tasks.""" + + def __init__(self, limit: int = LIMIT) -> None: + """Allow at most ``limit`` frames to be handled at the same time.""" + self._slots = asyncio.Semaphore(limit) + self._tasks: set[asyncio.Task[object]] = set() + # Set whenever the last task has been *accounted for*, which is a + # moment after it has finished: a done callback runs on a later loop + # turn, and gather() on an already-finished task may not yield at + # all, so "no tasks left" cannot be read off the task set alone. + self._idle = asyncio.Event() + self._idle.set() + + @property + def count(self) -> int: + """How many spawned frames have not finished yet.""" + return len(self._tasks) + + @contextlib.asynccontextmanager + async def slot(self) -> AsyncIterator[None]: + """Hold one of the slots for the duration of the block.""" + async with self._slots: + yield + + async def start(self, work: Coroutine[object, object, object]) -> None: + """Run ``work`` in its own task, once a slot is free. + + Returns as soon as the task exists, not when it finishes -- that is + the whole point. A task that raises is logged rather than lost: + nothing awaits it, so an unhandled exception would otherwise be + reported only when the task is garbage-collected, if at all. + """ + await self._slots.acquire() + task = asyncio.create_task(work) + self._tasks.add(task) + self._idle.clear() + task.add_done_callback(self._finished) + + def _finished(self, task: asyncio.Task[object]) -> None: + # The slot goes back here, not in a wrapper around the work: a task + # cancelled before its first step never runs a wrapper's `finally`, + # and a slot lost that way would be lost until restart. + self._slots.release() + self._tasks.discard(task) + if not self._tasks: + self._idle.set() + if task.cancelled(): + return + error = task.exception() + if error is not None: + logger.error("a frame's task failed", exc_info=error) + + def cancel(self) -> None: + """Stop everything in flight; for shutdown, where finishing is not wanted.""" + for task in self._tasks: + task.cancel() + + async def drain(self) -> None: + """Wait for everything in flight to finish, whichever way it ends.""" + await self._idle.wait() diff --git a/signal_bot/latency.py b/signal_bot/latency.py new file mode 100644 index 0000000..facf017 --- /dev/null +++ b/signal_bot/latency.py @@ -0,0 +1,150 @@ +# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence. +"""Hold the bot's own share of a reply to its budget, against the live bot. + +``DOCS-latency.md`` measured a reply from the outside and found the bot's +code was a rounding error next to the container and the upstream fetch. That +is a claim about one measurement of one command; this is the check that keeps +it true. It types one text command at the bot through ``/_selftest`` many +times over, reads the ``served`` timings each answer carries (``served.py``), +and fails if the 95th percentile of ``own_ms`` -- everything that is not a +send or an upstream request -- is over ``OWN_BUDGET_MS``. + +The self-test is always the same person, so the second attempt would be +refused by the cooldown. The command's cooldown is overridden to zero for the +run through the same switch the management page uses, and the row is put +back exactly as it was, refusal or crash included. +""" + +from __future__ import annotations + +import math +import time +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final + +import httpx + +from signal_bot.served import OWN_BUDGET_MS + +if TYPE_CHECKING: + from collections.abc import Sequence + +MS_PER_SECOND: Final = 1000 +# What the report prints and what the budget is held against, in that order. +COLUMNS: Final = ("own_ms", "queue_ms", "dispatch_ms", "handler_ms", "send_ms", "round_trip_ms") +PERCENTILES: Final = (50, 95) + + +class LatencyError(RuntimeError): + """The run could not measure what it was asked to.""" + + +@dataclass(frozen=True, slots=True) +class Sample: + """One reply's timings plus how long the whole self-test call took.""" + + served: dict[str, Any] + round_trip_ms: int + + def __getitem__(self, column: str) -> int: + """One column of the report, whichever side of the wire it was measured on.""" + if column == "round_trip_ms": + return self.round_trip_ms + return int(self.served[column]) + + +def percentile(values: Sequence[int], p: int) -> int: + """Nearest-rank percentile: the value at least ``p``% of the samples sit at or under.""" + ordered = sorted(values) + rank = max(1, math.ceil(p / 100 * len(ordered))) + return ordered[rank - 1] + + +@dataclass(frozen=True, slots=True) +class Report: + """Every sample of one run, and the verdict.""" + + command: str + samples: tuple[Sample, ...] + budget_ms: int = OWN_BUDGET_MS + + def at(self, column: str, p: int) -> int: + """One percentile of one column.""" + return percentile([sample[column] for sample in self.samples], p) + + @property + def within_budget(self) -> bool: + """Whether the bot's own p95 is at or under the budget.""" + return self.at("own_ms", 95) <= self.budget_ms + + def render(self) -> str: + """A table a human reads once and a CI log keeps.""" + lines = [f"{self.command}: {len(self.samples)} replies"] + lines.extend( + f" {column:<14}" + "".join(f"p{p}={self.at(column, p):>6} ms" for p in PERCENTILES) + for column in COLUMNS + ) + verdict = "within" if self.within_budget else "OVER" + lines.append(f" own_ms p95 {verdict} the {self.budget_ms} ms budget") + return "\n".join(lines) + + +async def _selftest(client: httpx.AsyncClient, base: str, group: str, text: str) -> Sample: + started = time.perf_counter() + response = await client.post(f"{base}/_selftest", json={"group": group, "text": text}) + round_trip_ms = round((time.perf_counter() - started) * MS_PER_SECOND) + body = response.json() + if response.status_code != httpx.codes.OK or body.get("status") != "delivered": + message = f"self-test refused: {body.get('reason', response.text)}" + raise LatencyError(message) + if "served" not in body: + message = f"{text!r} is not a command the bot ran; nothing to time" + raise LatencyError(message) + return Sample(served=body["served"], round_trip_ms=round_trip_ms) + + +async def _setting(client: httpx.AsyncClient, base: str, name: str) -> dict[str, Any]: + """The command's switch row as it stands, so it can be put back.""" + response = await client.get(f"{base}/_commands") + response.raise_for_status() + for row in response.json()["commands"]: + if row["name"] == name: + return { + "enabled": row["enabled"], + "disabled_until": row["disabled_until"], + "cooldown_seconds": row["cooldown_seconds"], + } + message = f"the bot ran {name!r} but does not list it" + raise LatencyError(message) + + +async def _put_setting( + client: httpx.AsyncClient, base: str, name: str, setting: dict[str, Any] +) -> None: + response = await client.put(f"{base}/_commands/{name}", json=setting) + response.raise_for_status() + + +async def measure( + client: httpx.AsyncClient, + base: str, + group: str, + text: str, + times: int, + budget_ms: int = OWN_BUDGET_MS, +) -> Report: + """Type ``text`` at the bot ``times`` times and report what each reply cost. + + Sequential on purpose: the point is what one reply costs when nothing + else is going on, and a burst would measure the queue behind ourselves. + """ + first = await _selftest(client, base, group, text) + command = str(first.served["command"]) + original = await _setting(client, base, command) + await _put_setting(client, base, command, {**original, "cooldown_seconds": 0}) + samples = [first] + try: + samples.extend([await _selftest(client, base, group, text) for _ in range(times - 1)]) + finally: + await _put_setting(client, base, command, original) + return Report(command=command, samples=tuple(samples), budget_ms=budget_ms) diff --git a/signal_bot/selftest.py b/signal_bot/selftest.py index 3cea93f..83f1c5f 100644 --- a/signal_bot/selftest.py +++ b/signal_bot/selftest.py @@ -41,6 +41,7 @@ if TYPE_CHECKING: from signal_bot.bot import Bot from signal_bot.config import Config + from signal_bot.served import Served GROUPS_ENV: Final = "SELFTEST_GROUPS" @@ -130,11 +131,16 @@ async def run( text: str, env: Mapping[str, str] | None = None, attachments: Sequence[str] = (), -) -> str: - """Feed one synthetic message, optionally carrying media, to the dispatcher.""" +) -> tuple[str, Served | None]: + """Feed one synthetic message, optionally carrying media, to the dispatcher. + + Returns the group posted into and, when the text was a command, what + it cost -- through ``Bot.serve``, so the frame queues for a slot the way + a real one does and the timings are the bot's own. + """ if group not in allowed_groups(os.environ if env is None else env): msg = f"{group!r} is not in {GROUPS_ENV}; refusing to post there" raise SelfTestError(msg) group_id, author = await resolve_group(client, bot.config, group) - await bot.dispatcher.handle(build_frame(group_id, author, text, attachments)) - return group_id + served = await bot.serve(build_frame(group_id, author, text, attachments)) + return group_id, served diff --git a/signal_bot/served.py b/signal_bot/served.py new file mode 100644 index 0000000..2a8964f --- /dev/null +++ b/signal_bot/served.py @@ -0,0 +1,56 @@ +# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence. +"""What one command cost, split into the bot's own time and everybody else's. + +``DOCS-latency.md`` found 72% of a ``!kot`` reply inside the signal-cli +container and most of the rest on Commons -- none of it this repository's +code. So "the bot is slow" needs a number that excludes both, or every +measurement is a measurement of somebody else. ``own_ms`` is that number: +the wall clock from the frame being read to the handler returning, minus +every upstream request and every send. It is what ``/_selftest`` reports +and what ``tool/latency.py`` holds to a budget. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final + +# The bot's own share of one reply, at the 95th percentile. The whole +# reply is bounded by the container, and is reported, not budgeted. +OWN_BUDGET_MS: Final = 100 + + +@dataclass(frozen=True, slots=True) +class Served: + """The timings of one command attempt, in milliseconds.""" + + command: str + outcome: str + # Frame read to a slot free: the queue behind other people's frames. + queue_ms: int + # Slot to handler: parse, scope, tally, resolve, and this person's lock. + dispatch_ms: int + # The handler, wall clock, everything included. + handler_ms: int + # Of which: outbound sends to signal-cli... + send_ms: int + # ...and requests to Commons, the cat API, and the like. + upstream_ms: int + + @property + def own_ms(self) -> int: + """Time spent in this repository's code, and nowhere else.""" + return self.queue_ms + self.dispatch_ms + self.handler_ms - self.send_ms - self.upstream_ms + + def as_dict(self) -> dict[str, int | str]: + """The shape ``/_selftest`` answers with.""" + return { + "command": self.command, + "outcome": self.outcome, + "queue_ms": self.queue_ms, + "dispatch_ms": self.dispatch_ms, + "handler_ms": self.handler_ms, + "send_ms": self.send_ms, + "upstream_ms": self.upstream_ms, + "own_ms": self.own_ms, + } diff --git a/signal_bot/user_lock.py b/signal_bot/user_lock.py new file mode 100644 index 0000000..82582f9 --- /dev/null +++ b/signal_bot/user_lock.py @@ -0,0 +1,50 @@ +# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence. +"""One person's commands run one at a time; everybody else's interleave. + +Handling frames concurrently (``inflight.py``) opens a race the sequential +loop never had: two ``!kot`` from the same person in the same second would +both pass the cooldown gate before either had recorded a use, and both be +served. Serialising per person closes it without serialising the group -- +the gate and the write it guards happen under that person's lock, so the +second attempt sees the first one's use and is refused, exactly as before. + +Locks are dropped when nobody holds or wants them, so the table cannot grow +with every person who ever typed a command. +""" + +from __future__ import annotations + +import asyncio +import contextlib +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + +class UserLocks: + """A lock per person, created on first use and forgotten when idle.""" + + def __init__(self) -> None: + """Start with nobody waiting on anything.""" + self._locks: dict[str, asyncio.Lock] = {} + self._waiting: dict[str, int] = {} + + @property + def held(self) -> frozenset[str]: + """Who has a lock right now -- held or queued for. For tests.""" + return frozenset(self._locks) + + @contextlib.asynccontextmanager + async def holding(self, user: str) -> AsyncIterator[None]: + """Run the block as the only thing this person has going.""" + lock = self._locks.setdefault(user, asyncio.Lock()) + self._waiting[user] = self._waiting.get(user, 0) + 1 + try: + async with lock: + yield + finally: + self._waiting[user] -= 1 + if self._waiting[user] == 0: + del self._waiting[user] + del self._locks[user] diff --git a/tests/test_bot_listen.py b/tests/test_bot_listen.py new file mode 100644 index 0000000..35013c0 --- /dev/null +++ b/tests/test_bot_listen.py @@ -0,0 +1,152 @@ +# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence. +"""The listen loop hands frames to tasks: slow ones no longer block quick ones.""" + +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING + +import httpx + +from signal_bot.bot import Bot, build_bot +from signal_bot.liveness import Liveness +from signal_bot.transport_fake import FakeTransport +from tests.conftest import WATCHED_GROUP, envelope, mock_client +from tests.test_bot import NOW + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + from signal_bot.config import Config + from signal_bot.context import CommandContext + from signal_bot.discovery import GroupDirectory + from signal_bot.store import Store + +BOB = ("uuid-bob", "Bob") + + +async def _settled() -> None: + for _ in range(5): + await asyncio.sleep(0) + + +def _bot_with_slow_command( + config: Config, + store: Store, + transport: FakeTransport, + directory: GroupDirectory, + gate: asyncio.Event, +) -> Bot: + """The real bot, plus a ``!slow`` that answers only once ``gate`` is set.""" + client = mock_client(lambda r: httpx.Response(200)) + bot = build_bot( + config, client, store, transport, directory, Liveness(config, client), clock=lambda: NOW + ) + + async def slow(ctx: CommandContext) -> None: + await gate.wait() + await ctx.reply("finally") + + bot.dispatcher.registry.register("slow", ["!slow"], slow, "waits") + return bot + + +async def test_a_quick_command_is_answered_while_a_slow_one_is_still_running( + config: Config, store: Store, transport: FakeTransport, directory: GroupDirectory +) -> None: + """The bug ``DOCS-latency.md`` describes: one picture used to hold up every text.""" + gate = asyncio.Event() + bot = _bot_with_slow_command(config, store, transport, directory, gate) + transport.queue( + [ + envelope(text="!slow", group_id=WATCHED_GROUP), + envelope(text="!pomoc", group_id=WATCHED_GROUP, sender=BOB), + ] + ) + + listening = asyncio.ensure_future(bot.listen()) + await _settled() + # The stream has ended and !pomoc is answered, while !slow still waits. + assert not listening.done() + assert len(transport.sent) == 1 + assert transport.sent[0].text != "finally" + + gate.set() + await asyncio.wait_for(listening, timeout=1) + assert transport.sent[-1].text == "finally" + assert bot.inflight.count == 0 + + +async def test_listen_finishes_what_is_in_flight_when_the_stream_ends( + config: Config, store: Store, transport: FakeTransport, directory: GroupDirectory +) -> None: + gate = asyncio.Event() + gate.set() + bot = _bot_with_slow_command(config, store, transport, directory, gate) + transport.queue([envelope(text="!slow", group_id=WATCHED_GROUP)]) + + await bot.listen() + + # listen() returned only once the reply had gone out, not merely once the + # frame had been handed off -- the row and the send are both there. + assert [m.text for m in transport.sent] == ["finally"] + assert bot.inflight.count == 0 + + +class _OpenStream(FakeTransport): + """A transport whose stream, like the real websocket, never ends on its own.""" + + async def receive(self) -> AsyncIterator[str]: + for frame in list(self.frames): + yield frame + await asyncio.Event().wait() + + +async def test_stopping_the_bot_stops_the_frames_in_flight( + config: Config, store: Store, directory: GroupDirectory +) -> None: + gate = asyncio.Event() + transport = _OpenStream() + bot = _bot_with_slow_command(config, store, transport, directory, gate) + transport.queue([envelope(text="!slow", group_id=WATCHED_GROUP)]) + + listening = asyncio.ensure_future(bot.listen()) + await _settled() + assert bot.inflight.count == 1 + listening.cancel() + await asyncio.gather(listening, return_exceptions=True) + await _settled() + assert bot.inflight.count == 0 + assert transport.sent == [] + + +async def test_stopping_the_bot_while_it_drains_stops_the_drain_too( + config: Config, store: Store, transport: FakeTransport, directory: GroupDirectory +) -> None: + gate = asyncio.Event() + bot = _bot_with_slow_command(config, store, transport, directory, gate) + transport.queue([envelope(text="!slow", group_id=WATCHED_GROUP)]) + + # The fake stream ends at once, so listen() is already waiting on the + # slow frame when the cancel arrives. + listening = asyncio.ensure_future(bot.listen()) + await _settled() + assert bot.inflight.count == 1 + listening.cancel() + await asyncio.gather(listening, return_exceptions=True) + await _settled() + assert bot.inflight.count == 0 + assert transport.sent == [] + + +async def test_serve_returns_the_timings_and_takes_a_slot_like_any_frame( + config: Config, store: Store, transport: FakeTransport, directory: GroupDirectory +) -> None: + gate = asyncio.Event() + bot = _bot_with_slow_command(config, store, transport, directory, gate) + + served = await bot.serve(envelope(text="!pomoc", group_id=WATCHED_GROUP)) + assert served is not None + assert (served.command, served.outcome) == ("help", "ok") + assert served.queue_ms >= 0 + assert bot.inflight.count == 0 diff --git a/tests/test_dispatcher_concurrency.py b/tests/test_dispatcher_concurrency.py new file mode 100644 index 0000000..2411a32 --- /dev/null +++ b/tests/test_dispatcher_concurrency.py @@ -0,0 +1,146 @@ +# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence. +"""Frames handled side by side: what still holds, and what each one costs.""" + +from __future__ import annotations + +import asyncio +import datetime as dt +import time +from typing import TYPE_CHECKING + +from signal_bot.decorators import cooldown +from signal_bot.dispatcher import Dispatcher +from signal_bot.registry import CommandRegistry +from signal_bot.registry_holder import RegistryHolder +from tests.conftest import envelope + +if TYPE_CHECKING: + from signal_bot.config import Config + from signal_bot.context import CommandContext + from signal_bot.discovery import GroupDirectory + from signal_bot.store import Store + from signal_bot.transport_fake import FakeTransport + +NOW = dt.datetime(2026, 8, 29, 12, 0, tzinfo=dt.UTC) +BOB = ("uuid-bob", "Bob") + + +async def _settled() -> None: + for _ in range(3): + await asyncio.sleep(0) + + +def _dispatcher( + config: Config, + directory: GroupDirectory, + store: Store, + transport: FakeTransport, + gate: asyncio.Event, + runs: list[str], +) -> Dispatcher: + """A ``!cat`` that blocks on ``gate`` before answering, with a cooldown.""" + + async def handler(ctx: CommandContext) -> None: + runs.append(ctx.message.source_uuid) + await gate.wait() + await ctx.reply("meow") + + registry = CommandRegistry() + registry.register("cat", ["!cat"], cooldown(10, "wait {seconds}")(handler), "kot") + return Dispatcher( + config=config, + directory=directory, + holder=RegistryHolder(lambda: registry), + store=store, + transport=transport, + clock=lambda: NOW, + ) + + +async def test_the_same_persons_second_command_waits_and_is_then_refused( + config: Config, directory: GroupDirectory, store: Store, transport: FakeTransport +) -> None: + """Two ``!cat`` in one second used to both pass the gate; now one does. + + Without the per-person lock, both would read "never used" before either + had recorded a use, and the cooldown would be a decoration. + """ + gate = asyncio.Event() + runs: list[str] = [] + dispatcher = _dispatcher(config, directory, store, transport, gate, runs) + + first = asyncio.ensure_future(dispatcher.handle(envelope(text="!cat"))) + second = asyncio.ensure_future(dispatcher.handle(envelope(text="!cat"))) + await _settled() + assert runs == ["uuid-alice"] + assert not second.done() + + gate.set() + served_first, served_second = await asyncio.gather(first, second) + assert runs == ["uuid-alice"] + assert served_first is not None + assert served_second is not None + assert (served_first.outcome, served_second.outcome) == ("ok", "refused_cooldown") + assert transport.texts[0] == "meow" + assert transport.texts[1].startswith("wait ") + + +async def test_another_persons_command_is_answered_while_the_first_still_runs( + config: Config, directory: GroupDirectory, store: Store, transport: FakeTransport +) -> None: + gate = asyncio.Event() + runs: list[str] = [] + dispatcher = _dispatcher(config, directory, store, transport, gate, runs) + + alice = asyncio.ensure_future(dispatcher.handle(envelope(text="!cat"))) + bob = asyncio.ensure_future(dispatcher.handle(envelope(text="!cat", sender=BOB))) + await _settled() + # Both are inside the handler at once: Bob did not queue behind Alice. + assert runs == ["uuid-alice", "uuid-bob"] + gate.set() + await asyncio.gather(alice, bob) + assert len(transport.sent) == 2 + assert dispatcher.users.held == frozenset() + + +async def test_what_a_served_command_reports( + config: Config, directory: GroupDirectory, store: Store, transport: FakeTransport +) -> None: + gate = asyncio.Event() + gate.set() + dispatcher = _dispatcher(config, directory, store, transport, gate, []) + + served = await dispatcher.handle(envelope(text="!cat"), received_at=time.perf_counter() - 0.25) + + assert served is not None + assert (served.command, served.outcome) == ("cat", "ok") + # A quarter of a second sat in the queue before this frame got a slot. + assert 200 <= served.queue_ms <= 400 + assert served.dispatch_ms >= 0 + assert served.handler_ms >= served.send_ms + assert served.upstream_ms == 0 + + +async def test_a_frame_with_no_command_in_it_costs_nothing_to_report( + config: Config, directory: GroupDirectory, store: Store, transport: FakeTransport +) -> None: + dispatcher = _dispatcher(config, directory, store, transport, asyncio.Event(), []) + assert await dispatcher.handle(envelope(text="hej"), received_at=time.perf_counter()) is None + assert await dispatcher.handle(envelope(text="!nic")) is None + assert await dispatcher.handle(envelope(text="!cat", group_id="elsewhere")) is None + assert dispatcher.users.held == frozenset() + + +async def test_frames_that_never_reach_a_command_take_no_lock( + config: Config, directory: GroupDirectory, store: Store, transport: FakeTransport +) -> None: + """A plain message from Alice must not queue behind Alice's slow ``!cat``.""" + gate = asyncio.Event() + dispatcher = _dispatcher(config, directory, store, transport, gate, []) + slow = asyncio.ensure_future(dispatcher.handle(envelope(text="!cat"))) + await _settled() + assert dispatcher.users.held == {"uuid-alice"} + plain = await asyncio.wait_for(dispatcher.handle(envelope(text="hello")), timeout=1) + assert plain is None + gate.set() + await slow diff --git a/tests/test_inflight.py b/tests/test_inflight.py new file mode 100644 index 0000000..c71f4fd --- /dev/null +++ b/tests/test_inflight.py @@ -0,0 +1,129 @@ +# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence. +"""Frames as tasks: a few at once, the rest waiting at the reader.""" + +from __future__ import annotations + +import asyncio +import logging +from typing import TYPE_CHECKING + +from signal_bot.inflight import LIMIT, Inflight + +if TYPE_CHECKING: + import pytest + + +async def _settled() -> None: + """Let every task that is ready to run take a turn.""" + for _ in range(3): + await asyncio.sleep(0) + + +async def test_start_returns_before_the_work_is_done() -> None: + gate = asyncio.Event() + done: list[str] = [] + + async def work() -> None: + await gate.wait() + done.append("x") + + inflight = Inflight() + await inflight.start(work()) + assert inflight.count == 1 + assert done == [] + gate.set() + await inflight.drain() + assert done == ["x"] + assert inflight.count == 0 + + +async def test_the_reader_waits_only_when_every_slot_is_busy() -> None: + gate = asyncio.Event() + + async def work() -> None: + await gate.wait() + + inflight = Inflight(limit=2) + await inflight.start(work()) + await inflight.start(work()) + third = asyncio.ensure_future(inflight.start(work())) + await _settled() + # Two slots, two frames in them: the third start is still waiting. + assert not third.done() + assert inflight.count == 2 + gate.set() + await third + await inflight.drain() + assert inflight.count == 0 + + +async def test_a_task_that_raises_is_logged_and_frees_its_slot( + caplog: pytest.LogCaptureFixture, +) -> None: + async def boom() -> None: + msg = "handler broke" + raise RuntimeError(msg) + + async def fine() -> None: + return None + + inflight = Inflight(limit=1) + with caplog.at_level(logging.ERROR, logger="signal_bot.inflight"): + await inflight.start(boom()) + await inflight.drain() + # The slot came back, or this would hang. + await asyncio.wait_for(inflight.start(fine()), timeout=1) + await inflight.drain() + assert "a frame's task failed" in caplog.text + assert "handler broke" in caplog.text + + +async def test_cancel_stops_what_is_in_flight_without_a_log_line( + caplog: pytest.LogCaptureFixture, +) -> None: + async def forever() -> None: + await asyncio.Event().wait() + + async def fine() -> None: + return None + + inflight = Inflight(limit=1) + await inflight.start(forever()) + with caplog.at_level(logging.ERROR, logger="signal_bot.inflight"): + # Cancelled before its first step, which is the case a wrapper's + # `finally` would miss: the slot must come back regardless. + inflight.cancel() + await inflight.drain() + assert inflight.count == 0 + assert caplog.text == "" + await asyncio.wait_for(inflight.start(fine()), timeout=1) + await inflight.drain() + + +async def test_a_slot_held_inline_counts_against_the_same_limit() -> None: + gate = asyncio.Event() + + async def work() -> None: + await gate.wait() + + inflight = Inflight(limit=1) + await inflight.start(work()) + + async def inline() -> str: + async with inflight.slot(): + return "ran" + + held = asyncio.ensure_future(inline()) + await _settled() + assert not held.done() + gate.set() + assert await held == "ran" + await inflight.drain() + + +async def test_drain_with_nothing_in_flight_returns_at_once() -> None: + await asyncio.wait_for(Inflight().drain(), timeout=1) + + +def test_the_default_limit_is_a_handful() -> None: + assert 1 < LIMIT <= 16 diff --git a/tests/test_latency.py b/tests/test_latency.py new file mode 100644 index 0000000..730e74d --- /dev/null +++ b/tests/test_latency.py @@ -0,0 +1,155 @@ +# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence. +"""The latency run: what it types, what it restores, and how it judges.""" + +from __future__ import annotations + +import json +from typing import Any + +import httpx +import pytest + +from signal_bot.latency import ( + LatencyError, + Report, + Sample, + measure, + percentile, +) +from tests.conftest import mock_client + +BASE = "http://bot" +# The first sample, the row read, and the row overridden: the second sample +# is the first one made after the cooldown is off. +SETUP_CALLS = 3 + + +def _served(own_ms: int, **overrides: int | str) -> dict[str, Any]: + base: dict[str, Any] = { + "command": "help", + "outcome": "ok", + "queue_ms": 0, + "dispatch_ms": 1, + "handler_ms": own_ms + 40, + "send_ms": 40, + "upstream_ms": 0, + "own_ms": own_ms + 1, + } + return {**base, **overrides} + + +class _Bot: + """Enough of the bot's HTTP face for the run: selftest and the switchboard.""" + + def __init__(self, own_ms: list[int], *, listed: bool = True) -> None: + self.own_ms = own_ms + self.listed = listed + self.requests: list[tuple[str, str, Any]] = [] + self.setting: dict[str, Any] = { + "enabled": True, + "disabled_until": None, + "cooldown_seconds": 30.0, + } + + def __call__(self, request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) if request.content else None + self.requests.append((request.method, request.url.path, body)) + if request.url.path == "/_selftest": + served = _served(self.own_ms.pop(0)) + return httpx.Response( + 200, json={"status": "delivered", "group": "g", "text": "!pomoc", "served": served} + ) + if request.url.path == "/_commands": + # Another command first, so the lookup has to walk past it. + rows = [{"name": "cat", **self.setting}] + if self.listed: + rows.append({"name": "help", **self.setting}) + return httpx.Response(200, json={"commands": rows}) + assert request.url.path == "/_commands/help" + self.setting = body + return httpx.Response(200, json={"name": "help", **body}) + + +async def test_a_run_types_the_command_with_the_cooldown_off_and_puts_it_back() -> None: + bot = _Bot([5, 7, 6]) + report = await measure(mock_client(bot), BASE, "testbot", "!pomoc", times=3) + + assert report.command == "help" + assert [s["own_ms"] for s in report.samples] == [6, 8, 7] + assert report.within_budget + methods = [(m, p) for m, p, _ in bot.requests] + assert methods == [ + ("POST", "/_selftest"), + ("GET", "/_commands"), + ("PUT", "/_commands/help"), + ("POST", "/_selftest"), + ("POST", "/_selftest"), + ("PUT", "/_commands/help"), + ] + # Zero for the run, then exactly the row that was there before. + assert bot.requests[2][2]["cooldown_seconds"] == 0 + assert bot.setting == {"enabled": True, "disabled_until": None, "cooldown_seconds": 30.0} + + +async def test_the_row_is_put_back_even_when_a_later_reply_fails() -> None: + class _Breaks(_Bot): + def __call__(self, request: httpx.Request) -> httpx.Response: + if request.url.path == "/_selftest" and len(self.requests) >= SETUP_CALLS: + self.requests.append((request.method, request.url.path, None)) + return httpx.Response(503, json={"status": "connecting", "reason": "down"}) + return super().__call__(request) + + bot = _Breaks([5, 7, 6]) + with pytest.raises(LatencyError, match="down"): + await measure(mock_client(bot), BASE, "testbot", "!pomoc", times=3) + assert bot.setting["cooldown_seconds"] == 30.0 + + +async def test_a_refused_group_is_reported_with_the_bots_reason() -> None: + client = mock_client( + lambda _: httpx.Response( + 403, json={"status": "refused", "reason": "not in SELFTEST_GROUPS"} + ) + ) + with pytest.raises(LatencyError, match="not in SELFTEST_GROUPS"): + await measure(client, BASE, "prod", "!pomoc", times=2) + + +async def test_text_that_is_not_a_command_is_refused_rather_than_timed() -> None: + client = mock_client( + lambda _: httpx.Response(200, json={"status": "delivered", "group": "g", "text": "hej"}) + ) + with pytest.raises(LatencyError, match="not a command"): + await measure(client, BASE, "testbot", "hej", times=2) + + +async def test_a_command_the_bot_ran_but_does_not_list_is_an_error() -> None: + bot = _Bot([5], listed=False) + with pytest.raises(LatencyError, match="does not list"): + await measure(mock_client(bot), BASE, "testbot", "!pomoc", times=1) + + +def test_percentiles_are_nearest_rank() -> None: + values = list(range(1, 21)) + assert percentile(values, 50) == 10 + assert percentile(values, 95) == 19 + assert percentile([42], 95) == 42 + + +def test_the_report_judges_the_p95_of_own_time_and_renders_every_column() -> None: + samples = tuple( + Sample(served=_served(own_ms), round_trip_ms=own_ms + 300) + for own_ms in [3, 4, 5, 6, 7, 8, 9, 10, 11, 250] + ) + report = Report(command="help", samples=samples, budget_ms=100) + + assert not report.within_budget + text = report.render() + assert text.startswith("help: 10 replies") + for column in ("own_ms", "queue_ms", "dispatch_ms", "handler_ms", "send_ms", "round_trip_ms"): + assert column in text + assert "own_ms p95 OVER the 100 ms budget" in text + assert report.at("round_trip_ms", 50) == 307 + + assert Report(command="help", samples=samples[:-1]).within_budget + assert "within the 100 ms budget" in Report(command="help", samples=samples[:-1]).render() diff --git a/tests/test_selftest.py b/tests/test_selftest.py index 88e177a..00d38a8 100644 --- a/tests/test_selftest.py +++ b/tests/test_selftest.py @@ -115,13 +115,31 @@ async def test_run_feeds_the_live_dispatcher( {"name": "testbot", "internal_id": WATCHED_GROUP, "members": [config.phone_number, OTHER]} ] - delivered = await run( + delivered, served = await run( bot, _client(entries), "testbot", "!papiez", env={"SELFTEST_GROUPS": "testbot"} ) assert delivered == WATCHED_GROUP assert transport.sent[0].image_base64 is not None assert "pl.wikipedia.org" in transport.sent[0].caption + # The cost came back with it, and it is the bot's own path that was timed. + assert served is not None + assert (served.command, served.outcome) == ("pope", "ok") + + +async def test_run_reports_no_cost_for_a_frame_that_was_not_a_command( + config: Config, store: Store, transport: FakeTransport, directory: GroupDirectory +) -> None: + bot = _bot(config, store, transport, directory) + entries = [ + {"name": "testbot", "internal_id": WATCHED_GROUP, "members": [config.phone_number, OTHER]} + ] + + env = {"SELFTEST_GROUPS": "testbot"} + _, served = await run(bot, _client(entries), "testbot", "hej", env=env) + + assert served is None + assert transport.sent == [] async def test_the_endpoint_reports_connecting_before_the_bot_is_up() -> None: @@ -167,10 +185,37 @@ async def test_the_endpoint_delivers_and_says_where( body = await selftest_route(SelfTestRequest(group="testbot", text="!papiez"), response) - assert body == {"status": "delivered", "group": WATCHED_GROUP, "text": "!papiez"} + assert body["served"]["command"] == "pope" + assert body["served"]["outcome"] == "ok" + assert body["served"]["own_ms"] >= 0 + assert {k: v for k, v in body.items() if k != "served"} == { + "status": "delivered", + "group": WATCHED_GROUP, + "text": "!papiez", + } assert transport.sent[0].image_base64 is not None +async def test_the_endpoint_carries_no_cost_for_text_that_ran_nothing( + config: Config, + store: Store, + transport: FakeTransport, + directory: GroupDirectory, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("SELFTEST_GROUPS", "testbot") + bot = _bot(config, store, transport, directory) + app_module.app.state.starter = SimpleNamespace(runtime=SimpleNamespace(bot=bot), reason="") + app_module.app.state.client = _client( + [{"name": "testbot", "internal_id": WATCHED_GROUP, "members": [config.phone_number, OTHER]}] + ) + + body = await selftest_route(SelfTestRequest(group="testbot", text="hej"), Response()) + + assert body == {"status": "delivered", "group": WATCHED_GROUP, "text": "hej"} + assert transport.sent == [] + + async def test_resolve_group_rejects_an_entry_with_no_id(config: Config) -> None: entries = [{"name": "testbot", "members": [config.phone_number, OTHER]}] with pytest.raises(SelfTestError, match="no usable id or no other member"): diff --git a/tests/test_served.py b/tests/test_served.py new file mode 100644 index 0000000..68fc4ac --- /dev/null +++ b/tests/test_served.py @@ -0,0 +1,54 @@ +# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence. +"""The bot's own share of one reply is everything that is not somebody else's.""" + +from __future__ import annotations + +from signal_bot.served import OWN_BUDGET_MS, Served + + +def _kot() -> Served: + return Served( + command="cat", + outcome="ok", + queue_ms=3, + dispatch_ms=2, + handler_ms=1400, + send_ms=900, + upstream_ms=480, + ) + + +def test_own_time_excludes_sends_and_upstream_requests() -> None: + # 3 + 2 + 1400 of wall clock, of which 900 was signal-cli and 480 the + # cat API: 25 ms is what this repository's code cost. + assert _kot().own_ms == 25 + + +def test_as_dict_is_the_selftest_shape_with_the_derived_number_included() -> None: + assert _kot().as_dict() == { + "command": "cat", + "outcome": "ok", + "queue_ms": 3, + "dispatch_ms": 2, + "handler_ms": 1400, + "send_ms": 900, + "upstream_ms": 480, + "own_ms": 25, + } + + +def test_a_refusal_has_no_handler_time_and_costs_only_the_gate() -> None: + refused = Served( + command="cat", + outcome="refused_cooldown", + queue_ms=0, + dispatch_ms=4, + handler_ms=0, + send_ms=0, + upstream_ms=0, + ) + assert refused.own_ms == 4 + + +def test_the_budget_is_the_one_the_latency_note_settled_on() -> None: + assert OWN_BUDGET_MS == 100 diff --git a/tests/test_user_lock.py b/tests/test_user_lock.py new file mode 100644 index 0000000..c6b08d3 --- /dev/null +++ b/tests/test_user_lock.py @@ -0,0 +1,123 @@ +# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence. +"""One lock per person, none for anybody who is not in the middle of something.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from signal_bot.user_lock import UserLocks + + +async def _settled() -> None: + for _ in range(3): + await asyncio.sleep(0) + + +async def test_the_same_person_runs_one_thing_at_a_time() -> None: + locks = UserLocks() + gate = asyncio.Event() + order: list[str] = [] + + async def first() -> None: + async with locks.holding("alice"): + order.append("first in") + await gate.wait() + order.append("first out") + + async def second() -> None: + async with locks.holding("alice"): + order.append("second in") + + one = asyncio.ensure_future(first()) + two = asyncio.ensure_future(second()) + await _settled() + assert order == ["first in"] + assert not two.done() + gate.set() + await asyncio.gather(one, two) + assert order == ["first in", "first out", "second in"] + + +async def test_different_people_do_not_wait_for_each_other() -> None: + locks = UserLocks() + gate = asyncio.Event() + order: list[str] = [] + + async def slow() -> None: + async with locks.holding("alice"): + order.append("alice in") + await gate.wait() + + async def quick() -> None: + async with locks.holding("bob"): + order.append("bob in") + + one = asyncio.ensure_future(slow()) + two = asyncio.ensure_future(quick()) + await _settled() + assert order == ["alice in", "bob in"] + assert two.done() + gate.set() + await one + + +async def test_a_lock_is_forgotten_once_nobody_wants_it() -> None: + locks = UserLocks() + async with locks.holding("alice"): + assert locks.held == {"alice"} + assert locks.held == frozenset() + + +async def test_a_newcomer_waits_for_a_holder_who_was_once_a_waiter() -> None: + """Releasing must not drop a lock somebody else is queued on. + + Otherwise the next person to arrive would get a fresh lock and run + alongside the one who had just been let through. + """ + locks = UserLocks() + first_gate = asyncio.Event() + second_gate = asyncio.Event() + order: list[str] = [] + + async def holding(name: str, gate: asyncio.Event) -> None: + async with locks.holding("alice"): + order.append(name) + await gate.wait() + + one = asyncio.ensure_future(holding("first", first_gate)) + two = asyncio.ensure_future(holding("second", second_gate)) + await _settled() + first_gate.set() + await one + await _settled() + assert order == ["first", "second"] + assert locks.held == {"alice"} + + three = asyncio.ensure_future(holding("third", asyncio.Event())) + await _settled() + assert order == ["first", "second"] + assert not three.done() + second_gate.set() + await two + await _settled() + assert order == ["first", "second", "third"] + three.cancel() + await asyncio.gather(three, return_exceptions=True) + assert locks.held == frozenset() + + +async def test_a_block_that_raises_still_releases_and_forgets() -> None: + locks = UserLocks() + + async def breaks() -> None: + async with locks.holding("alice"): + msg = "boom" + raise RuntimeError(msg) + + with pytest.raises(RuntimeError, match="boom"): + await breaks() + assert locks.held == frozenset() + async with locks.holding("alice"): + assert locks.held == {"alice"} diff --git a/tool/latency.py b/tool/latency.py new file mode 100755 index 0000000..2f95c7f --- /dev/null +++ b/tool/latency.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence. +"""Measure what one text reply costs the bot's own code, against the live bot. + + tool/latency.py testbot # 20 x !pomoc, judged at 100 ms p95 + tool/latency.py testbot --text '!ja' --times 50 + tool/latency.py testbot --budget 50 + +Posts for real into the named group, which must be in SELFTEST_GROUPS -- the +same fail-closed rule as scripts/selftest.sh. The command's cooldown is +switched off for the run and put back afterwards. Exit 1 when the 95th +percentile of the bot's own time is over budget; see signal_bot/latency.py +for what "own" excludes and why. +""" + +from __future__ import annotations + +import argparse +import asyncio +import os +import sys +from typing import Final + +import httpx +from dotenv import load_dotenv + +from signal_bot.app import HTTP_TIMEOUT_SECONDS, USER_AGENT +from signal_bot.latency import LatencyError, measure +from signal_bot.served import OWN_BUDGET_MS + +EXIT_OVER_BUDGET: Final = 1 +EXIT_REFUSED: Final = 2 +DEFAULT_TIMES: Final = 20 + + +async def run(group: str, text: str, times: int, budget_ms: int, base: str) -> int: + """One measured run; the exit code is the verdict.""" + async with httpx.AsyncClient( + timeout=HTTP_TIMEOUT_SECONDS, headers={"User-Agent": USER_AGENT} + ) as client: + report = await measure(client, base, group, text, times, budget_ms) + sys.stdout.write(report.render() + "\n") + return 0 if report.within_budget else EXIT_OVER_BUDGET + + +def main() -> int: + """Parse arguments and run one measurement.""" + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("group", help="test group name; must be in SELFTEST_GROUPS") + parser.add_argument("--text", default="!pomoc", help="the command to type (default: !pomoc)") + parser.add_argument("--times", type=int, default=DEFAULT_TIMES, help="replies to sample") + parser.add_argument("--budget", type=int, default=OWN_BUDGET_MS, help="p95 own_ms allowed") + args = parser.parse_args() + load_dotenv() + base = f"http://127.0.0.1:{os.environ.get('BIND_PORT', '8137')}" + try: + return asyncio.run(run(args.group, args.text, args.times, args.budget, base)) + except (LatencyError, httpx.HTTPError) as error: + sys.stderr.write(f"latency: {error}\n") + return EXIT_REFUSED + + +if __name__ == "__main__": + raise SystemExit(main())