diff --git a/bot.py b/bot.py index 76dacbe..ddccace 100755 --- a/bot.py +++ b/bot.py @@ -30,17 +30,23 @@ import discord import requests from discord.ext import commands -from communication_subroutine import comm_subroutine +from communication_subroutine import comm_subroutine, librarian_ping from constants import ( ENCODING, FILE_SERVICE_ADDRESS, GET_MP3, + LIBRARIAN_PING, LIBRARIAN_SERVICE_ADDRESS, LOGFILE, RADIO_SERVICE_ADDRESS, TOKEN, + service_headers, ) +# Round-trip health check budget for the librarian ping (bot -> librarian +# internal queue -> pong back). Deliberately short so startup never stalls. +LIBRARIAN_PING_TIMEOUT = 3.0 + logger = logging.getLogger("discord") logger.setLevel(logging.DEBUG) formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") @@ -110,9 +116,11 @@ SERVICE_EXTENSION_GROUPS = { "health_url": f"{RADIO_SERVICE_ADDRESS}/ping", "extensions": ["radio_commands"], }, - # librarian: DOI / Crossref search + # librarian: DOI / Crossref search. NOTE: this one is NOT a plain GET - it + # is gated on a full ping round-trip (see _load_service_groups); the URL + # here is only the label used in the "unreachable" log line. "librarian": { - "health_url": f"{LIBRARIAN_SERVICE_ADDRESS}/", + "health_url": f"{LIBRARIAN_SERVICE_ADDRESS}{LIBRARIAN_PING}", "extensions": ["librarian_commands"], }, } @@ -154,17 +162,37 @@ async def _load_service_groups() -> bool: missing = [e for e in group["extensions"] if e not in client.extensions] if not missing: continue - err = await asyncio.to_thread(_service_health, group["health_url"]) - if err is not None: - logger.warning( - "Service '%s' unreachable (%s) [%s: %s] - cogs stay disabled: %s", - service, - group["health_url"], - type(err).__name__, - err, - ", ".join(missing), + if service == "librarian": + # A plain GET only proves Flask is up. The librarian is only useful + # once its internal queue + worker are flowing, so prove exactly that + # with a ping that must complete the full round-trip (see + # communication_subroutine.librarian_ping). + alive = await asyncio.to_thread( + librarian_ping, + LIBRARIAN_SERVICE_ADDRESS, + LIBRARIAN_PING, + service_headers(), + LIBRARIAN_PING_TIMEOUT, ) - continue + if not alive: + logger.warning( + "Service 'librarian' ping round-trip failed (%s) - cogs stay disabled: %s", + group["health_url"], + ", ".join(missing), + ) + continue + else: + err = await asyncio.to_thread(_service_health, group["health_url"]) + if err is not None: + logger.warning( + "Service '%s' unreachable (%s) [%s: %s] - cogs stay disabled: %s", + service, + group["health_url"], + type(err).__name__, + err, + ", ".join(missing), + ) + continue logger.info("Service '%s' is alive - enabling: %s", service, ", ".join(missing)) for extension in missing: if await _load_extension_safe(extension): diff --git a/communication_subroutine.py b/communication_subroutine.py index dad76fd..49867c1 100644 --- a/communication_subroutine.py +++ b/communication_subroutine.py @@ -4,10 +4,12 @@ import os import re import threading import time +import uuid as uuidlib from queue import Empty, Queue from typing import Optional from urllib import request as urequest +import requests from flask import Flask, abort, jsonify, request from waitress import serve @@ -26,6 +28,11 @@ SRCHTITLE = re.compile(rb"StreamTitle=\\*(?P[^;]*);").search awaiting_q = [] incoming_q = Queue() +# A health-check ping whose pong never comes back (dead/dropped librarian) would +# otherwise leave its record in awaiting_q forever. scan_incoming sweeps ping +# records older than this. Kept well above the ping timeout so a slow-but-alive +# round-trip is never swept out from under a waiter. +PING_TTL_SECONDS = 30 app = Flask(__name__) @@ -232,9 +239,36 @@ def scan_incoming(stop_event: Optional[threading.Event] = None): if stop_event and stop_event.is_set(): logger.info("scan_incoming: stop requested") break + # Sweep stale health-check pings first: if a librarian is dead the pong + # never arrives, so drop ping records past their TTL. scan_incoming is + # the sole remover of awaiting_q, so this needs no lock (scan_queue only + # appends). Snapshot with list() so removal during iteration is safe. + now = time.monotonic() + for record in list(awaiting_q): + if ( + getattr(record, "is_ping", False) + and getattr(record, "answered", None) is not None + and not record.answered.is_set() + and now - getattr(record, "created", now) > PING_TTL_SECONDS + ): + awaiting_q.remove(record) try: answer = incoming_q.get(block=False) logger.info("DATA FOUND") + # Health-check pong (shape: {"__pong__": uuid}). Close the waiter's + # event and drop its record. It must NEVER fall through to the + # result/orphan path below, or the librarian cog would later pull it + # off IN_COMM_Q and post a bogus "no results" message to Discord. + if isinstance(answer, dict) and "__pong__" in answer: + pong_uuid = answer["__pong__"] + for record in list(awaiting_q): + if getattr(record, "uuid", None) == pong_uuid: + event = getattr(record, "answered", None) + if event is not None: + event.set() + awaiting_q.remove(record) + logger.info("PONG matched for %s", pong_uuid) + continue # Collect matched records and drop them from awaiting_q afterwards - # they used to stay forever (awaiting_q only ever grew), leaking # memory over the bot's uptime and letting a reused UUID re-match a @@ -258,6 +292,55 @@ def scan_incoming(stop_event: Optional[threading.Event] = None): time.sleep(1) +def librarian_ping(address: str, endpoint: str, headers: Optional[dict] = None, + timeout: float = 3.0) -> bool: + """Health-check the librarian by round-tripping a ping through the FULL path. + + The ping is a pseudo-query that exercises exactly the same machinery a real + search does, on BOTH sides: + + * bot out: a ``QueryControl`` rides ``OUT_COMM_Q`` -> ``scan_queue`` -> + ``awaiting_q`` just like a real query, + * librarian: it is POSTed to the librarian, which must pull it off its OWN + internal queue and answer WITHOUT running a search, + * bot in: the pong comes back over ``/conjurer`` -> ``incoming_q`` -> + ``scan_incoming``, which matches it by uuid and sets our event. + + Returns True only when that whole loop closes within ``timeout``. Never + blocks longer than roughly ``timeout`` and cannot deadlock: the POST is + bounded, the wait is bounded, and a ping whose pong never arrives is swept + out of ``awaiting_q`` by ``scan_incoming`` (PING_TTL_SECONDS). + """ + logger = logging.getLogger("discord") + ping_uuid = str(uuidlib.uuid4()) + answered = threading.Event() + query = QueryControl("healthcheck", ping_uuid, "__ping__", None) + query.is_ping = True + query.answered = answered + query.created = time.monotonic() + # Enter the bot-side comm queue BEFORE the POST, so the record is already in + # awaiting_q by the time the pong can come back (no lost-wakeup race). + OUT_COMM_Q.put(query) + try: + response = requests.post( + f"{address}{endpoint}", + json={"UUID": ping_uuid}, + headers=headers or {}, + timeout=timeout, + ) + except requests.exceptions.RequestException as exc: + logger.info("Librarian ping POST failed (%s): %s", ping_uuid, exc) + return False # stale record is swept by scan_incoming + if response.status_code != 200: + logger.info("Librarian ping rejected (%s): HTTP %s", ping_uuid, response.status_code) + return False + if answered.wait(timeout): + logger.info("Librarian ping round-trip OK (%s)", ping_uuid) + return True + logger.info("Librarian ping timed out after %ss (%s)", timeout, ping_uuid) + return False + + def get_stream_title(tag: bytes) -> str: title = "" if m := SRCHTITLE(tag): diff --git a/conjurer_librarian/conjurer_librarian.py b/conjurer_librarian/conjurer_librarian.py index 4450fe4..5b27c6c 100644 --- a/conjurer_librarian/conjurer_librarian.py +++ b/conjurer_librarian/conjurer_librarian.py @@ -69,6 +69,19 @@ app = Flask(__name__) librarian_queue = Queue() librarian_list = [] +# Lifecycle of every real search uuid: "queued" (accepted, sitting in +# librarian_queue) -> "processing" (worker pulled it) -> removed (worker +# finished AND attempted to send the result). The bot's per-query watchdog polls +# /query_status against this: a uuid that VANISHES from here without its result +# reaching the bot is a lost result (finished-but-never-delivered) and gets +# flagged in chat. +active_queries: Dict[str, str] = {} +_active_lock = threading.Lock() +# Set while the worker is grinding a real search. A ping arriving during this +# pongs back immediately WITHOUT queueing - being busy is healthy (you can keep +# piling searches on), so "busy" must never look like "dead" to the health check. +worker_busy = threading.Event() + def _service_headers() -> Dict[str, str]: if API_KEY: @@ -81,6 +94,25 @@ def _authorize_request() -> None: abort(401) +def _post_pong(app_logger, ping_uuid) -> None: + """POST a pong for ``ping_uuid`` back to the bot. Non-fatal on failure. + + This is the SAME return path a real result takes (bot's /conjurer), so a + delivered pong proves the librarian->bot leg works - the one thing the ping + needs to establish. Used directly by the /ping route (busy ping, which skips + the queue) and, wrapped in a thread, by the worker (idle ping). Synchronous + so the /ping route can stay a plain (non-async) view.""" + try: + requests.post( + f"{MAIN_BOT_ADDRESS}{SEND_RESULTS}", + json={"__pong__": ping_uuid}, + headers=_service_headers(), + timeout=5, + ) + except requests.exceptions.RequestException as exc: + app_logger.warning("PING pong send failed for %s: %s", ping_uuid, exc) + + # trunk-ignore(pylint/R0902) class Librarian(object): """ @@ -411,89 +443,118 @@ class BackgroundTaskSearch(threading.Thread): while True: database = None ndb_database = None - librarian = librarian_queue.get() - self.app.logger.info("STARTED") - result = await librarian.answer_query(librarian.deep_search) - result = {librarian.uuid: result} - self.app.logger.info("Saving to file") - - # Save results to "not_in_db.json" file - with open(lib_paths.NOT_IN_DB, "r+", encoding="utf-8") as ndb_file: - ndb_database = {} - try: - ndb_database = json.load(ndb_file) - except JSONDecodeError: - pass - if ndb_database: - ndb_database.update(librarian.not_in_db) - else: - ndb_database = librarian.not_in_db - ndb_file.truncate(0) - ndb_file.seek(0) - json.dump(ndb_database, ndb_file) - - # Save results to "s_results.json" file - with open(lib_paths.S_RESULTS, "r+", encoding="utf-8") as s_file: - database = {} - try: - database = json.load(s_file) - except JSONDecodeError: - pass - if database: - self.app.logger.info(database) - self.app.logger.info(result) - database.update(result) - else: - database = result - self.app.logger.info("DUMPING DATA") - s_file.truncate(0) - s_file.seek(0) - json.dump(database, s_file) - self.app.logger.info("FINISHED") - - # Send the result back to the bot. Log EXACTLY what goes out (target, - # uuid, how many DOIs and which) so the librarian log makes it plain a - # result was sent and what was in it. - payload = result # shape: {uuid: {DOI: {"Title": ..., "type": ...}}} - hits = payload.get(librarian.uuid, {}) if isinstance(payload, dict) else {} - target = f"{MAIN_BOT_ADDRESS}{SEND_RESULTS}" - self.app.logger.info( - "SENDING result for %s to %s: %d DOI(s): %s", - librarian.uuid, - target, - len(hits), - list(hits.keys()), - ) - # A failed send must NOT kill this worker - otherwise a bot that is - # momentarily down stalls every future query until the librarian is - # restarted. Log and carry on to the next queued search. - try: - response = await asyncio.to_thread( - requests.post, - target, - json=payload, - headers=_service_headers(), - timeout=360, + item = librarian_queue.get() + # Health-check ping: it has flowed through the internal queue and is + # now pulled off it - that is the whole point. Pong it straight back + # with the same uuid and DO NOT run a search. + if isinstance(item, dict) and "__ping__" in item: + ping_uuid = item["__ping__"] + self.app.logger.info( + "PING %s pulled off internal queue - ponging back (no search)", + ping_uuid, ) - if response.status_code == 200: - self.app.logger.info( - "SENT result for %s -> HTTP 200 (bot accepted)", librarian.uuid - ) - else: - self.app.logger.warning( - "SENT result for %s but bot returned HTTP %s: %s", - librarian.uuid, - response.status_code, - response.text[:500], - ) - except requests.exceptions.RequestException as exc: - self.app.logger.error( - "FAILED to send result for %s to %s: %s", + await asyncio.to_thread(_post_pong, self.app.logger, ping_uuid) + continue + librarian = item + # Mark busy + processing for the whole search, and ALWAYS clear both + # (even on a crash) in finally: worker_busy so a ping doesn't wait + # behind us, and active_queries so the bot's watchdog can tell a + # finished-and-gone query from one still in flight. + worker_busy.set() + with _active_lock: + active_queries[str(librarian.uuid)] = "processing" + try: + self.app.logger.info("STARTED") + result = await librarian.answer_query(librarian.deep_search) + result = {librarian.uuid: result} + self.app.logger.info("Saving to file") + + # Save results to "not_in_db.json" file + with open(lib_paths.NOT_IN_DB, "r+", encoding="utf-8") as ndb_file: + ndb_database = {} + try: + ndb_database = json.load(ndb_file) + except JSONDecodeError: + pass + if ndb_database: + ndb_database.update(librarian.not_in_db) + else: + ndb_database = librarian.not_in_db + ndb_file.truncate(0) + ndb_file.seek(0) + json.dump(ndb_database, ndb_file) + + # Save results to "s_results.json" file + with open(lib_paths.S_RESULTS, "r+", encoding="utf-8") as s_file: + database = {} + try: + database = json.load(s_file) + except JSONDecodeError: + pass + if database: + self.app.logger.info(database) + self.app.logger.info(result) + database.update(result) + else: + database = result + self.app.logger.info("DUMPING DATA") + s_file.truncate(0) + s_file.seek(0) + json.dump(database, s_file) + self.app.logger.info("FINISHED") + + # Send the result back to the bot. Log EXACTLY what goes out + # (target, uuid, how many DOIs and which) so the librarian log + # makes it plain a result was sent and what was in it. + payload = result # shape: {uuid: {DOI: {"Title": ..., "type": ...}}} + hits = payload.get(librarian.uuid, {}) if isinstance(payload, dict) else {} + target = f"{MAIN_BOT_ADDRESS}{SEND_RESULTS}" + self.app.logger.info( + "SENDING result for %s to %s: %d DOI(s): %s", librarian.uuid, target, - exc, + len(hits), + list(hits.keys()), ) - await asyncio.sleep(1) + # A failed send must NOT kill this worker - otherwise a bot that + # is momentarily down stalls every future query until the + # librarian is restarted. Log and carry on to the next search. + try: + response = await asyncio.to_thread( + requests.post, + target, + json=payload, + headers=_service_headers(), + timeout=360, + ) + if response.status_code == 200: + self.app.logger.info( + "SENT result for %s -> HTTP 200 (bot accepted)", librarian.uuid + ) + else: + self.app.logger.warning( + "SENT result for %s but bot returned HTTP %s: %s", + librarian.uuid, + response.status_code, + response.text[:500], + ) + except requests.exceptions.RequestException as exc: + self.app.logger.error( + "FAILED to send result for %s to %s: %s", + librarian.uuid, + target, + exc, + ) + except Exception as exc: # pylint: disable=broad-exception-caught + # A crashing search must not kill the worker thread (which would + # freeze the whole queue). Log and move on; finally still clears + # busy/active so the query is correctly seen as "gone". + self.app.logger.exception("Search %s crashed: %s", librarian.uuid, exc) + finally: + worker_busy.clear() + with _active_lock: + active_queries.pop(str(librarian.uuid), None) + await asyncio.sleep(1) # ==================================SERVER ROUTES========================================== @@ -522,6 +583,10 @@ async def query_database(): cl = Librarian(app, record["query"], uuid, deep_search) librarian_queue.put(cl) librarian_list.append(cl) + # The bot's per-query watchdog polls /query_status for this uuid; mark it + # "queued" now so it counts as known the moment we accept it. + with _active_lock: + active_queries[str(uuid)] = "queued" answer_data = (record["query"], record["UUID"], librarian_queue.qsize()) return_data = ( jsonify(isError=False, message="Success", statusCode=200, data=answer_data), @@ -530,6 +595,66 @@ async def query_database(): return return_data +@app.route("/ping", methods=["POST"]) +def ping_roundtrip(): + _authorize_request() + """ + Health-check round-trip. + + Two cases, one guarantee - the pong always comes back over the librarian->bot + return path (the only thing the ping must prove): + + * IDLE: put a ping marker onto the SAME internal ``librarian_queue`` real + searches use and return 200. The worker pulls it off and pongs it back, + so a successful pong proves the whole pipeline flows (queue + worker + the + return leg), not just that Flask is up. + * BUSY (a search is grinding): DO NOT queue - the ping would just wait behind + a possibly hours-long search and time out, making a perfectly healthy busy + librarian look dead. Pong back immediately instead. Being busy is fine; you + can keep piling searches on. The ping only needs to catch a BROKEN return + path, and the direct pong exercises exactly that. + """ + record = json.loads(request.data) + ping_uuid = record["UUID"] + if worker_busy.is_set(): + app.logger.info("PING %s while busy grinding - direct pong (skip queue)", ping_uuid) + _post_pong(app.logger, ping_uuid) + else: + app.logger.info("PING received %s - queued for round-trip", ping_uuid) + librarian_queue.put({"__ping__": ping_uuid}) + return ( + jsonify(isError=False, message="ping-queued", statusCode=200, data=ping_uuid), + 200, + ) + + +@app.route("/query_status", methods=["POST"]) +def query_status(): + _authorize_request() + """ + Per-query watchdog probe. + + Returns whether ``UUID`` is still known to the librarian (queued or being + processed). The bot polls this after dispatching a search: while the uuid is + known the search is progressing; once it VANISHES here without the result + ever reaching the bot, the result was lost in transit and the bot tells the + user. A busy/queued search is never mistaken for a lost one. + """ + record = json.loads(request.data) + uuid = str(record["UUID"]) + with _active_lock: + state = active_queries.get(uuid, "unknown") + return ( + jsonify( + isError=False, + message="Success", + statusCode=200, + data={"uuid": uuid, "known": state != "unknown", "state": state}, + ), + 200, + ) + + @app.route("/get_partial_result", methods=["POST"]) async def get_partial(): _authorize_request() diff --git a/constants.py b/constants.py index 83a3986..686c21d 100644 --- a/constants.py +++ b/constants.py @@ -90,6 +90,12 @@ CREATE_PRIO_PLAYLIST = "/create_priority_playlist" REQUEST_MUSIC = "/request_radio_file" CLEAR_PRIO = "/clear_pr_pls" SEND_QUERY = "/query" +# Health-check round-trip: a pseudo-query that the librarian must pull off its +# own internal queue and answer (same uuid) WITHOUT running a real search. +LIBRARIAN_PING = "/ping" +# Per-query watchdog probe: "do you still know this uuid?" (queued/processing). +# A uuid that vanishes here without its result reaching the bot was lost. +QUERY_STATUS = "/query_status" TIME_BETWEEN_CALLS = 100000 LAST_SPONTANEOUS_CALL = datetime.now() diff --git a/librarian_commands.py b/librarian_commands.py index 05580a8..9cf195c 100644 --- a/librarian_commands.py +++ b/librarian_commands.py @@ -3,6 +3,7 @@ import io import logging import os import random +import time import uuid from queue import Empty @@ -15,15 +16,40 @@ from discord.ext import commands, tasks from ai_functions import handle_response from communication_subroutine import IN_COMM_Q, OUT_COMM_Q, QueryControl, submit_ai_query -from constants import DIR_PATH_SADOX, LIBRARIAN_SERVICE_ADDRESS, SEND_QUERY, service_headers +from constants import ( + DIR_PATH_SADOX, + LIBRARIAN_SERVICE_ADDRESS, + QUERY_STATUS, + SEND_QUERY, + service_headers, +) +from librarian_watchdog import FLAG, pending_verdict SERVICE_HEADERS = service_headers() +# Per-query watchdog tuning. +PENDING_WATCH_SECONDS = 30 # how often to ask the librarian about a uuid +PENDING_GRACE_SECONDS = 45 # unknown-but-pending must persist this long +PENDING_HARD_TTL = 60 * 60 * 24 * 3 # drop tracking after 3 days no matter what + class DataModule(commands.Cog): def __init__(self, bot, logger_name): self.bot = bot self.logger = logging.getLogger(logger_name) + # uuid -> {"ctx", "query", "created", "unknown_since"} for every search + # dispatched but not yet answered. watch_pending polls the librarian for + # each; check_data_q removes an entry the moment its result is rendered. + self.pending = {} + + def _track_pending(self, query_uuid, query, ctx): + """Start watching a dispatched search so a lost result can be caught.""" + self.pending[str(query_uuid)] = { + "ctx": ctx, + "query": query, + "created": time.monotonic(), + "unknown_since": None, + } @commands.hybrid_command( nsfw=True, @@ -100,6 +126,9 @@ class DataModule(commands.Cog): fresh_data = IN_COMM_Q.get(block=False) entries = [] if fresh_data.stop: + # The result arrived and is about to be rendered - stop the + # watchdog from ever flagging this uuid as lost. + self.pending.pop(str(fresh_data.uuid), None) searcher = fresh_data.author query = fresh_data.content # ai_lines is a clean, plain rendering of the SAME list in the @@ -175,6 +204,77 @@ class DataModule(commands.Cog): except Empty: pass + @tasks.loop(seconds=PENDING_WATCH_SECONDS) + async def watch_pending(self): + """Per-query safety net for lost results (case a). + + For each dispatched-but-unanswered search, ask the librarian whether it + still knows the uuid (queued or processing). While it does, the search is + progressing - leave it alone (a busy librarian is fine). The moment a + uuid VANISHES on the librarian while still pending here, its result was + computed but never reached us: after a short grace window (to rule out a + result that is merely in flight) we tell the channel - but ONLY then. + + A normally-delivered result is popped from self.pending by check_data_q, + so it never reaches the flag path. + """ + now = time.monotonic() + for query_uuid in list(self.pending.keys()): + info = self.pending.get(query_uuid) + if info is None: + continue + # Hard cap so a permanently-unreachable librarian can't leak entries. + if now - info["created"] > PENDING_HARD_TTL: + self.logger.warning("Dropping stale pending query %s (hard TTL)", query_uuid) + self.pending.pop(query_uuid, None) + continue + try: + response = await asyncio.to_thread( + requests.post, + f"{LIBRARIAN_SERVICE_ADDRESS}{QUERY_STATUS}", + json={"UUID": query_uuid}, + headers=SERVICE_HEADERS, + timeout=5, + ) + known = ( + response.status_code == 200 + and response.json().get("data", {}).get("known", False) + ) + except ( + requests.exceptions.RequestException, + ValueError, + AttributeError, + KeyError, + TypeError, + ) as exc: + # Librarian unreachable / garbled or unexpected answer: we can't + # judge, so don't cry wolf, and don't let one bad poll kill the + # loop. Reset the clock and try again next tick. + self.logger.info("Pending check for %s inconclusive: %s", query_uuid, exc) + info["unknown_since"] = None + continue + action, info["unknown_since"] = pending_verdict( + known, info["unknown_since"], now, PENDING_GRACE_SECONDS + ) + # Re-check membership: the await above yields, so check_data_q may + # have just delivered (and popped) this result. + if action == FLAG and query_uuid in self.pending: + await self._flag_lost(query_uuid, info) + self.pending.pop(query_uuid, None) + + async def _flag_lost(self, query_uuid, info): + """Tell the querent their finished search never made it back.""" + message = ( + "*Winda na książki z hukiem wraca z podziemi PUSTA. Z głośnika trzeszczy:* " + f"Twoje zapytanie {query_uuid} (\"{info['query']}\") przemieliło się w " + "bibliotece do końca, ale wynik przepadł gdzieś w drodze do baru - nic nie " + "dotarło. Zawołaj szefa albo puść jeszcze raz." + ) + try: + await info["ctx"].send(message) + except Exception: # pylint: disable=broad-exception-caught + self.logger.exception("Failed to post lost-result notice for %s", query_uuid) + @commands.hybrid_command( name="wyszukaj_linki_do_dokumentow", description="Szuka linkow doi w bazie crossref i podaje linki do scihuba", @@ -233,6 +333,7 @@ class DataModule(commands.Cog): username = ctx.message.author.name query_object = QueryControl(username, query_uuid, query, ctx) OUT_COMM_Q.put(query_object) + self._track_pending(query_uuid, query, ctx) await ctx.send( f"No dobra poszło. Zapamiętaj proszę {query_uuid} - to identyfikator twojego zapytania. Jesteś {queue_size} w kolejce." + " Zapytania obsługuje algorytm zasilany czterema chomikami zapierdalającymi w kołowrotku - więc wyniki najwcześniej za kilka godzi - ale mogą być też dni." @@ -291,6 +392,7 @@ class DataModule(commands.Cog): username = ctx.message.author.name query_object = QueryControl(username, query_uuid, query, ctx, ai_review=True) OUT_COMM_Q.put(query_object) + self._track_pending(query_uuid, query, ctx) await ctx.send( f"Poszło z recenzją AI. Identyfikator: {query_uuid}. Jesteś {queue_size} w kolejce." + " Najpierw dojadą surowe wyniki, a zaraz po nich przesortowanie i recenzja od AI." @@ -387,6 +489,7 @@ class DataModule(commands.Cog): username = ctx.message.author.name query_object = QueryControl(username, query_uuid, query, ctx) OUT_COMM_Q.put(query_object) + self._track_pending(query_uuid, query, ctx) await ctx.send( f"Wypij wypi {query_uuid} - to identyfikator twojego zapytania. Jesteś {queue_size} w kolejce niestety kolejka obowiazuje zawsze." + " Zapytania obsługuje algorytm zasilany czterema chomikami zapierdalającymi w kołowrotku - więc wyniki najwcześniej za kilka godzin - ale mogą być też dni. Głębokie zabawy trwają dłużej. *Znowu mruga*" @@ -397,5 +500,6 @@ async def setup(bot): logger = logging.getLogger("discord") dm = DataModule(bot, "discord") dm.check_data_q.start() + dm.watch_pending.start() await bot.add_cog(dm) logger.info("Loading data sharing commands module done") diff --git a/librarian_watchdog.py b/librarian_watchdog.py new file mode 100644 index 0000000..b095d4f --- /dev/null +++ b/librarian_watchdog.py @@ -0,0 +1,51 @@ +"""Pure decision logic for the librarian per-query watchdog. + +Split out of ``librarian_commands`` (which pulls in discord / pdf libs, so it is +not importable in the pytest-only unit job) so the one subtle part - the grace +window that stops a just-delivered result from being falsely flagged as lost - +can be unit-tested in isolation. + +The watchdog polls the librarian's /query_status for each dispatched query: + +* the librarian reports the uuid ``known`` while it is queued or processing, +* once the search finishes it is dropped there, so the uuid goes ``unknown``. + +A result that arrives normally is removed from the pending set by the result +handler, so the watchdog never even sees it. Only a uuid that goes ``unknown`` +on the librarian *and is still pending on the bot* is a lost result - but we +require it to stay that way for a grace window first, because there is always a +brief moment where the librarian has finished (uuid gone) yet the result is +still in flight / not yet rendered. +""" + +WAIT = "wait" +FLAG = "flag" + + +def pending_verdict(known, unknown_since, now, grace_seconds): + """Decide what to do this tick for one pending query. + + Args: + known: did the librarian report the uuid as still known this tick? + unknown_since: monotonic timestamp the uuid was first seen unknown, or + None if it was known last tick. + now: current monotonic time. + grace_seconds: how long a uuid must stay unknown-but-pending before it + is declared lost. + + Returns: + (action, unknown_since) where action is WAIT or FLAG and the returned + ``unknown_since`` is what the caller should store for the next tick. + """ + if known: + # Still queued/processing (or freshly back to known) - reset the clock. + return WAIT, None + if unknown_since is None: + # First tick we see it gone: start the grace clock, don't flag yet - the + # result may simply be in flight. + return WAIT, now + if now - unknown_since >= grace_seconds: + # Gone for the whole grace window and still pending: the result was lost. + return FLAG, unknown_since + # Gone, but not long enough yet - keep waiting. + return WAIT, unknown_since diff --git a/tests/integration/test_librarian_ping.py b/tests/integration/test_librarian_ping.py new file mode 100644 index 0000000..36ad6de --- /dev/null +++ b/tests/integration/test_librarian_ping.py @@ -0,0 +1,121 @@ +"""Integration: the librarian health check is a full comm round-trip, not a GET. + +``librarian_ping`` injects a pseudo-query into the SAME machinery a real search +uses - it rides ``OUT_COMM_Q`` -> ``scan_queue`` -> ``awaiting_q``, the librarian +is expected to pull it off its own queue and pong it back over ``/conjurer`` -> +``incoming_q`` -> ``scan_incoming``, which matches it by uuid and wakes the +waiter. These tests stand in for the librarian with a stubbed ``requests.post`` +and assert the loop closes (and, crucially, that a pong never leaks into +``IN_COMM_Q`` where the librarian cog would mistake it for a real result). +""" +import threading +import time + +import pytest + +import communication_subroutine as cs + + +def _drain(queue): + while not queue.empty(): + queue.get() + + +@pytest.fixture +def comm_threads(): + """Run scan_queue + scan_incoming (the two workers librarian_ping relies on) + for the duration of a test, on cleared shared state.""" + cs.awaiting_q.clear() + _drain(cs.incoming_q) + _drain(cs.OUT_COMM_Q) + _drain(cs.IN_COMM_Q) + stop = threading.Event() + workers = [ + threading.Thread(target=cs.scan_queue, kwargs={"stop_event": stop}, daemon=True), + threading.Thread(target=cs.scan_incoming, kwargs={"stop_event": stop}, daemon=True), + ] + for worker in workers: + worker.start() + yield + stop.set() + for worker in workers: + worker.join(timeout=3) + + +def _await_in_awaiting(ping_uuid, timeout=2): + """Block until scan_queue has moved the ping into awaiting_q.""" + deadline = time.time() + timeout + while time.time() < deadline: + if any(getattr(r, "uuid", None) == ping_uuid for r in list(cs.awaiting_q)): + return True + time.sleep(0.01) + return False + + +class _Resp: + def __init__(self, status_code=200): + self.status_code = status_code + + +def test_ping_round_trip_ok(comm_threads, monkeypatch): + # Stub librarian: only pong AFTER the record is in awaiting_q, mirroring the + # real network latency that always lets scan_queue win. + def fake_post(url, json=None, headers=None, timeout=None): + ping_uuid = json["UUID"] + _await_in_awaiting(ping_uuid) + cs.incoming_q.put({"__pong__": ping_uuid}) + return _Resp(200) + + monkeypatch.setattr(cs.requests, "post", fake_post) + + assert cs.librarian_ping("http://lib", "/ping", {}, timeout=3.0) is True + # A pong must NEVER reach the cog's inbound queue... + assert cs.IN_COMM_Q.empty() + # ...and the ping record must be cleaned out of awaiting_q. + assert not any(getattr(r, "is_ping", False) for r in list(cs.awaiting_q)) + + +def test_ping_times_out_when_librarian_accepts_but_never_pongs(comm_threads, monkeypatch): + monkeypatch.setattr(cs.requests, "post", lambda *a, **k: _Resp(200)) + start = time.monotonic() + assert cs.librarian_ping("http://lib", "/ping", {}, timeout=0.3) is False + # Bounded: it must not block much beyond the timeout. + assert time.monotonic() - start < 2.0 + + +def test_ping_false_when_librarian_unreachable(comm_threads, monkeypatch): + def boom(*a, **k): + raise cs.requests.exceptions.RequestException("no route to host") + + monkeypatch.setattr(cs.requests, "post", boom) + assert cs.librarian_ping("http://lib", "/ping", {}, timeout=0.3) is False + + +def test_ping_false_on_non_200(comm_threads, monkeypatch): + monkeypatch.setattr(cs.requests, "post", lambda *a, **k: _Resp(503)) + assert cs.librarian_ping("http://lib", "/ping", {}, timeout=0.3) is False + + +def test_orphan_pong_is_dropped_not_enqueued(comm_threads): + # A pong with no matching waiter (e.g. after the ping already timed out) must + # be silently dropped - never turned into an "Orphaned" IN_COMM_Q record that + # makes the librarian cog post a bogus "no results" message. + cs.incoming_q.put({"__pong__": "no-such-uuid"}) + deadline = time.time() + 3 + while time.time() < deadline and not cs.incoming_q.empty(): + time.sleep(0.02) + time.sleep(0.2) # give scan_incoming a beat to (not) enqueue anything + assert cs.IN_COMM_Q.empty() + + +def test_real_result_still_reaches_in_comm_q(comm_threads): + # Guard the existing path: a normal {uuid: {...}} result must still match its + # QueryControl and land in IN_COMM_Q for the cog to render. + query = cs.QueryControl("user", "real-uuid", "jakieś zapytanie", None) + cs.OUT_COMM_Q.put(query) + assert _await_in_awaiting("real-uuid") + cs.incoming_q.put({"real-uuid": {"10.1/x": {"Title": ["Tytuł"], "type": "article"}}}) + got = cs.IN_COMM_Q.get(timeout=3) + assert got.uuid == "real-uuid" + assert got.stop is True + assert "10.1/x" in got.entries diff --git a/tests/integration/test_librarian_query_lifecycle.py b/tests/integration/test_librarian_query_lifecycle.py new file mode 100644 index 0000000..03b27b9 --- /dev/null +++ b/tests/integration/test_librarian_query_lifecycle.py @@ -0,0 +1,103 @@ +"""Integration: the librarian's health/liveness surface. + +Two behaviours, both proven against the real Flask app: + +* /ping is busy-aware - while a search is grinding it pongs back immediately + WITHOUT queueing (busy is healthy); when idle it routes the ping through the + internal queue for the worker to answer. +* /query_status reports whether a uuid is still known (queued/processing), which + is what the bot's per-query watchdog polls to catch a lost result. + +Only the SYNC routes (/ping, /query_status) are exercised - the async /query +route needs flask[async], which the integration job doesn't install, so query +state is seeded directly on the module. +""" +import sys +import types + +# conjurer_librarian does `from habanero import Crossref` at import time and +# habanero isn't installed in the integration job. Stub it before importing the +# service (we never build a real Librarian here, so Crossref is just a name). +if "habanero" not in sys.modules: + _habanero = types.ModuleType("habanero") + _habanero.Crossref = object + sys.modules["habanero"] = _habanero + +import conjurer_librarian as lib # noqa: E402 + + +def _client(key=None): + lib.API_KEY = key + return lib.app.test_client() + + +def _reset(): + with lib._active_lock: + lib.active_queries.clear() + lib.worker_busy.clear() + while not lib.librarian_queue.empty(): + lib.librarian_queue.get() + + +def test_query_status_known_vs_unknown(): + _reset() + client = _client() + with lib._active_lock: + lib.active_queries["abc"] = "queued" + + known = client.post("/query_status", json={"UUID": "abc"}).get_json()["data"] + assert known == {"uuid": "abc", "known": True, "state": "queued"} + + unknown = client.post("/query_status", json={"UUID": "nope"}).get_json()["data"] + assert unknown == {"uuid": "nope", "known": False, "state": "unknown"} + + +def test_ping_idle_routes_through_internal_queue(monkeypatch): + _reset() + posted = [] + monkeypatch.setattr(lib.requests, "post", lambda *a, **k: posted.append((a, k))) + client = _client() + + resp = client.post("/ping", json={"UUID": "ping-idle"}) + + assert resp.status_code == 200 + # Idle => it went onto the internal queue for the worker, NOT posted directly. + assert posted == [] + assert lib.librarian_queue.get_nowait() == {"__ping__": "ping-idle"} + + +def test_ping_while_busy_pongs_directly_without_queue(monkeypatch): + _reset() + lib.worker_busy.set() # a search is grinding + posted = [] + + class _Resp: + status_code = 200 + + def fake_post(url, json=None, headers=None, timeout=None): + posted.append({"url": url, "json": json}) + return _Resp() + + monkeypatch.setattr(lib.requests, "post", fake_post) + client = _client() + + resp = client.post("/ping", json={"UUID": "ping-busy"}) + + assert resp.status_code == 200 + # Busy => direct pong back to the bot, and NOTHING queued (it would only wait + # behind the long search). + assert lib.librarian_queue.empty() + assert len(posted) == 1 + assert posted[0]["json"] == {"__pong__": "ping-busy"} + assert posted[0]["url"].endswith(lib.SEND_RESULTS) + + +def test_query_status_enforces_api_key(): + _reset() + client = _client(key="secret") + denied = client.post("/query_status", json={"UUID": "x"}) + assert denied.status_code == 401 + ok = client.post( + "/query_status", json={"UUID": "x"}, headers={"X-Conjurer-Api-Key": "secret"} + ) + assert ok.status_code == 200 diff --git a/tests/unit/test_librarian_watchdog.py b/tests/unit/test_librarian_watchdog.py new file mode 100644 index 0000000..0875d1f --- /dev/null +++ b/tests/unit/test_librarian_watchdog.py @@ -0,0 +1,51 @@ +"""Unit tests for the per-query watchdog verdict logic. + +The grace window is the whole point: a search that has just finished is briefly +'unknown' on the librarian while its result is still in flight, and we must NOT +flag that as lost. Only a uuid that stays unknown-but-pending past the grace +window is a genuinely lost result. +""" +from librarian_watchdog import FLAG, WAIT, pending_verdict + +GRACE = 45 + + +def test_known_resets_clock_and_waits(): + action, unknown_since = pending_verdict( + known=True, unknown_since=100.0, now=200.0, grace_seconds=GRACE + ) + assert action == WAIT + assert unknown_since is None # clock reset while it's still known + + +def test_first_unknown_starts_grace_but_does_not_flag(): + action, unknown_since = pending_verdict( + known=False, unknown_since=None, now=1000.0, grace_seconds=GRACE + ) + assert action == WAIT + assert unknown_since == 1000.0 # clock started now + + +def test_unknown_within_grace_keeps_waiting(): + action, unknown_since = pending_verdict( + known=False, unknown_since=1000.0, now=1000.0 + GRACE - 1, grace_seconds=GRACE + ) + assert action == WAIT + assert unknown_since == 1000.0 # unchanged, still counting + + +def test_unknown_past_grace_flags_lost(): + action, unknown_since = pending_verdict( + known=False, unknown_since=1000.0, now=1000.0 + GRACE, grace_seconds=GRACE + ) + assert action == FLAG + assert unknown_since == 1000.0 + + +def test_recovered_to_known_after_being_unknown_resets(): + # It reappeared (e.g. requeued / status flapped): do not flag, reset. + action, unknown_since = pending_verdict( + known=True, unknown_since=1000.0, now=1000.0 + GRACE + 10, grace_seconds=GRACE + ) + assert action == WAIT + assert unknown_since is None