diff --git a/bot.py b/bot.py index c6cbf1f..1b4047c 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"], }, } @@ -152,7 +160,20 @@ async def _load_service_groups() -> bool: missing = [e for e in group["extensions"] if e not in client.extensions] if not missing: continue - alive = await asyncio.to_thread(_service_alive, group["health_url"]) + 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, + ) + else: + alive = await asyncio.to_thread(_service_alive, group["health_url"]) if not alive: logger.warning( "Service '%s' unreachable (%s) - cogs stay disabled: %s", diff --git a/communication_subroutine.py b/communication_subroutine.py index 31fd0e9..1332bd1 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 record_stored = False for record in awaiting_q: if record.uuid in answer.keys(): @@ -252,6 +286,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..c2dc0ef 100644 --- a/conjurer_librarian/conjurer_librarian.py +++ b/conjurer_librarian/conjurer_librarian.py @@ -411,7 +411,30 @@ class BackgroundTaskSearch(threading.Thread): while True: database = None ndb_database = None - librarian = librarian_queue.get() + 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, + ) + try: + await asyncio.to_thread( + requests.post, + f"{MAIN_BOT_ADDRESS}{SEND_RESULTS}", + json={"__pong__": ping_uuid}, + headers=_service_headers(), + timeout=5, + ) + except requests.exceptions.RequestException as exc: + self.app.logger.warning( + "PING pong send failed for %s: %s", ping_uuid, exc + ) + continue + librarian = item self.app.logger.info("STARTED") result = await librarian.answer_query(librarian.deep_search) result = {librarian.uuid: result} @@ -530,6 +553,29 @@ async def query_database(): return return_data +@app.route("/ping", methods=["POST"]) +async def ping_roundtrip(): + _authorize_request() + """ + Health-check round-trip. + + Puts a lightweight ping marker onto the SAME internal ``librarian_queue`` + that real searches go through and returns 200 immediately. The background + worker pulls it off the queue and pongs it back to the bot with the same + uuid, WITHOUT running any Crossref/DOI search. A successful pong therefore + proves the whole pipeline (HTTP in -> internal queue -> worker -> HTTP out) + is flowing, not just that Flask is up. + """ + record = json.loads(request.data) + ping_uuid = record["UUID"] + 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("/get_partial_result", methods=["POST"]) async def get_partial(): _authorize_request() diff --git a/constants.py b/constants.py index 83a3986..1fc67e1 100644 --- a/constants.py +++ b/constants.py @@ -90,6 +90,9 @@ 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" TIME_BETWEEN_CALLS = 100000 LAST_SPONTANEOUS_CALL = datetime.now() 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