diff --git a/bot.py b/bot.py index ddccace..509711a 100755 --- a/bot.py +++ b/bot.py @@ -39,6 +39,7 @@ from constants import ( LIBRARIAN_SERVICE_ADDRESS, LOGFILE, RADIO_SERVICE_ADDRESS, + SELF_CALLBACK, TOKEN, service_headers, ) @@ -173,6 +174,7 @@ async def _load_service_groups() -> bool: LIBRARIAN_PING, service_headers(), LIBRARIAN_PING_TIMEOUT, + SELF_CALLBACK, ) if not alive: logger.warning( diff --git a/communication_subroutine.py b/communication_subroutine.py index 0f7827b..ad77eff 100644 --- a/communication_subroutine.py +++ b/communication_subroutine.py @@ -349,9 +349,13 @@ def scan_incoming(stop_event: Optional[threading.Event] = None): def librarian_ping(address: str, endpoint: str, headers: Optional[dict] = None, - timeout: float = 3.0) -> bool: + timeout: float = 3.0, callback: str = "") -> bool: """Health-check the librarian by round-tripping a ping through the FULL path. + ``callback`` is where THIS bot wants the pong sent back to, so a shared + librarian pongs each bot at its own address (else a second bot's ping would + be ponged to the first and its health check would always time out). + The ping is a pseudo-query that exercises exactly the same machinery a real search does, on BOTH sides: @@ -380,7 +384,7 @@ def librarian_ping(address: str, endpoint: str, headers: Optional[dict] = None, try: response = requests.post( f"{address}{endpoint}", - json={"UUID": ping_uuid}, + json={"UUID": ping_uuid, "callback": callback}, headers=headers or {}, timeout=timeout, ) diff --git a/conjurer_librarian/conjurer_librarian.py b/conjurer_librarian/conjurer_librarian.py index c5bde98..002a31c 100644 --- a/conjurer_librarian/conjurer_librarian.py +++ b/conjurer_librarian/conjurer_librarian.py @@ -154,49 +154,55 @@ 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. +def _post_pong(app_logger, ping_uuid, callback="") -> None: + """POST a pong for ``ping_uuid`` back to the pinging bot. Non-fatal. - 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.""" + ``callback`` is the address of the bot that sent the ping, so a librarian + shared by several bots pongs each at its OWN address (empty => the static + MAIN_BOT_ADDRESS). Same return path a real result takes (bot's /conjurer), + so a delivered pong proves the librarian->that-bot leg works.""" + target = f"{callback or MAIN_BOT_ADDRESS}{SEND_RESULTS}" try: requests.post( - f"{MAIN_BOT_ADDRESS}{SEND_RESULTS}", - json={"__pong__": ping_uuid}, - headers=_service_headers(), - timeout=5, + target, 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) -def _deliver_result(uuid, payload, app_logger, attempts=RESULT_SEND_ATTEMPTS) -> bool: - """POST one result to the bot, retrying with backoff. True only on HTTP 200. +def _outbox_target_payload(entry): + """Unpack an OUTBOX entry into (target, payload). - The bot's /conjurer is idempotent (dedups by uuid), so re-POSTing a result - it already has is safe - it just answers 200 again. That is what lets the - OUTBOX keep retrying until the result is truly acknowledged, without ever - double-delivering to the user. - """ - target = f"{MAIN_BOT_ADDRESS}{SEND_RESULTS}" + New shape: {"target": , "payload": {uuid: result}}. Old shape + (from before per-origin callbacks) is the raw payload - delivered to the + default bot - so an upgrade doesn't strand results already on disk.""" + if isinstance(entry, dict) and "target" in entry and "payload" in entry: + return entry["target"], entry["payload"] + return "", entry + + +def _deliver_result(target, uuid, payload, app_logger, attempts=RESULT_SEND_ATTEMPTS) -> bool: + """POST one result to ``target``'s /conjurer, retrying with backoff. True only + on HTTP 200. ``target`` is the origin bot's address so a shared librarian + answers each bot where its query came from (empty => MAIN_BOT_ADDRESS). The + bot's /conjurer is idempotent, so re-POSTing what it already has is safe.""" + url = f"{target or MAIN_BOT_ADDRESS}{SEND_RESULTS}" for attempt in range(1, max(1, attempts) + 1): try: response = requests.post( - target, json=payload, headers=_service_headers(), timeout=60 + url, json=payload, headers=_service_headers(), timeout=60 ) if response.status_code == 200: - app_logger.info("Result %s delivered (HTTP 200) on attempt %d", uuid, attempt) + app_logger.info("Result %s delivered to %s (HTTP 200) attempt %d", uuid, url, attempt) return True app_logger.warning( - "Result %s: bot returned HTTP %s (attempt %d/%d): %s", - uuid, response.status_code, attempt, attempts, response.text[:300], + "Result %s: %s returned HTTP %s (attempt %d/%d): %s", + uuid, url, response.status_code, attempt, attempts, response.text[:300], ) except requests.exceptions.RequestException as exc: app_logger.warning( - "Result %s delivery failed (attempt %d/%d): %s", uuid, attempt, attempts, exc + "Result %s delivery to %s failed (attempt %d/%d): %s", + uuid, url, attempt, attempts, exc, ) if attempt < attempts: time.sleep(RESULT_SEND_BACKOFF * attempt) @@ -204,13 +210,12 @@ def _deliver_result(uuid, payload, app_logger, attempts=RESULT_SEND_ATTEMPTS) -> def _resend_once(app_logger) -> None: - """One sweep of the OUTBOX: try to deliver every un-acked result, once each. - - Removes each entry only after a positive ACK, so nothing is dropped until - the bot has it. Corrupt/unreadable entries are skipped by DiskQueue.items(). - """ - for uuid, payload, _ts in _outbox.items(): - if _deliver_result(uuid, payload, app_logger, attempts=1): + """One sweep of the OUTBOX: try to deliver every un-acked result, once each, + to the ORIGIN bot recorded with it. Removes each entry only after a positive + ACK. Corrupt/unreadable entries are skipped by DiskQueue.items().""" + for uuid, entry, _ts in _outbox.items(): + target, payload = _outbox_target_payload(entry) + if _deliver_result(target, uuid, payload, app_logger, attempts=1): _outbox.remove(uuid) @@ -243,7 +248,10 @@ def replay_requests(app_logger) -> None: app_logger.info("Replaying %d unfinished search request(s) after restart", len(pending)) for uuid, payload, _ts in pending: try: - cl = Librarian(app, payload["query"], uuid, payload.get("deep_search", False)) + cl = Librarian( + app, payload["query"], uuid, + payload.get("deep_search", False), payload.get("callback", ""), + ) except Exception as exc: # pylint: disable=broad-exception-caught app_logger.warning("Cannot replay request %s (dropping): %s", uuid, exc) _forget_search(uuid) @@ -269,7 +277,7 @@ class Librarian(object): Represents a librarian object that performs search and refinement operations on queries. """ - def __init__(self, _app, query, uuid, _deep_search) -> None: + def __init__(self, _app, query, uuid, _deep_search, callback="") -> None: """ Initializes a Librarian object. @@ -340,6 +348,9 @@ class Librarian(object): self.search_result_from_cr = {} self.done = False self.deep_search = _deep_search + # Where to send this search's result back to (the bot that asked). Lets + # one librarian serve several bots; empty => static MAIN_BOT_ADDRESS. + self.callback = callback or "" # Set True when a graceful shutdown interrupts this search mid-scan; the # worker then leaves the request + checkpoint in place instead of # delivering, so a restart resumes it. @@ -630,7 +641,9 @@ class BackgroundTaskSearch(threading.Thread): "PING %s pulled off internal queue - ponging back (no search)", ping_uuid, ) - await asyncio.to_thread(_post_pong, self.app.logger, ping_uuid) + await asyncio.to_thread( + _post_pong, self.app.logger, ping_uuid, item.get("callback", "") + ) continue librarian = item # Mark busy + processing for the whole search, and ALWAYS clear both @@ -680,13 +693,16 @@ class BackgroundTaskSearch(threading.Thread): # retrying until the bot ACKs, and only then is it removed. payload = result # shape: {uuid: {DOI: {"Title": ..., "type": ...}}} uuid = str(librarian.uuid) + target = librarian.callback # answer the bot that asked hits = payload.get(librarian.uuid, {}) if isinstance(payload, dict) else {} - _outbox.put(uuid, payload) + # OUTBOX entry carries the origin bot's address so the resender + # delivers it to the right bot even after a librarian restart. + _outbox.put(uuid, {"target": target, "payload": payload}) self.app.logger.info( - "SENDING result for %s: %d DOI(s): %s (queued to OUTBOX)", - uuid, len(hits), list(hits.keys()), + "SENDING result for %s to %s: %d DOI(s): %s (queued to OUTBOX)", + uuid, target or "default", len(hits), list(hits.keys()), ) - if await asyncio.to_thread(_deliver_result, uuid, payload, self.app.logger): + if await asyncio.to_thread(_deliver_result, target, uuid, payload, self.app.logger): _outbox.remove(uuid) else: self.app.logger.warning( @@ -731,11 +747,17 @@ async def query_database(): record = json.loads(request.data) uuid = record["UUID"] deep_search = record["deep_search"] - app.logger.info("Query accepted %s: %s", uuid, record["query"]) + # Where to answer THIS query - the bot that sent it. Persisted with the + # request so a replay after restart still answers the right bot. + callback = record.get("callback", "") + app.logger.info("Query accepted %s (callback %s): %s", uuid, callback or "default", record["query"]) # Persist the request BEFORE enqueuing, so an accepted search survives a # restart (it is replayed on startup) - not just an in-progress one. - _requests.put(str(uuid), {"query": record["query"], "deep_search": deep_search}) - cl = Librarian(app, record["query"], uuid, deep_search) + _requests.put( + str(uuid), + {"query": record["query"], "deep_search": deep_search, "callback": callback}, + ) + cl = Librarian(app, record["query"], uuid, deep_search, callback) librarian_queue.put(cl) librarian_list.append(cl) # The bot's per-query watchdog polls /query_status for this uuid; mark it @@ -771,12 +793,15 @@ def ping_roundtrip(): """ record = json.loads(request.data) ping_uuid = record["UUID"] + # Pong goes back to the bot that pinged (carried through the queue when idle), + # so a shared librarian health-checks correctly for every bot. + callback = record.get("callback", "") 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) + _post_pong(app.logger, ping_uuid, callback) else: app.logger.info("PING received %s - queued for round-trip", ping_uuid) - librarian_queue.put({"__ping__": ping_uuid}) + librarian_queue.put({"__ping__": ping_uuid, "callback": callback}) return ( jsonify(isError=False, message="ping-queued", statusCode=200, data=ping_uuid), 200, diff --git a/constants.py b/constants.py index 9602697..7c135f4 100644 --- a/constants.py +++ b/constants.py @@ -254,6 +254,10 @@ SKIP_TRACK = os.getenv("CONJURER_SKIP_ENDPOINT", "/skip") LIBRARIAN_SERVICE_ADDRESS = os.getenv( "CONJURER_LIBRARIAN_SERVICE", "http://192.168.1.192:5001" ) +# The address the librarian (and its pongs) should send results BACK to for THIS +# bot - so one librarian can serve several bots (test + deploy), each getting its +# own answers. Empty => the librarian falls back to its static CONJURER_MAIN_BOT. +SELF_CALLBACK = os.getenv("CONJURER_SELF_CALLBACK", "") HOST_ADDRESS = os.getenv("CONJURER_DISCORD_HOST", "192.168.1.191") PORT_ADDRESS = int(os.getenv("CONJURER_DISCORD_PORT", "5000")) diff --git a/librarian_commands.py b/librarian_commands.py index 85fd80d..57aaae8 100644 --- a/librarian_commands.py +++ b/librarian_commands.py @@ -26,6 +26,7 @@ from constants import ( DIR_PATH_SADOX, LIBRARIAN_SERVICE_ADDRESS, QUERY_STATUS, + SELF_CALLBACK, SEND_QUERY, service_headers, ) @@ -313,6 +314,7 @@ class DataModule(commands.Cog): "query": str(query), "page": 1, "deep_search": False, + "callback": SELF_CALLBACK, } coroutine = asyncio.to_thread( requests.post, @@ -372,6 +374,7 @@ class DataModule(commands.Cog): "query": str(query), "page": 1, "deep_search": False, + "callback": SELF_CALLBACK, } coroutine = asyncio.to_thread( requests.post, @@ -469,6 +472,7 @@ class DataModule(commands.Cog): "query": str(query), "page": 1, "deep_search": True, + "callback": SELF_CALLBACK, } coroutine = asyncio.to_thread( requests.post, diff --git a/tests/integration/test_librarian_outbox.py b/tests/integration/test_librarian_outbox.py index dc472a8..c9d272d 100644 --- a/tests/integration/test_librarian_outbox.py +++ b/tests/integration/test_librarian_outbox.py @@ -46,16 +46,24 @@ def outbox(tmp_path, monkeypatch): def test_deliver_succeeds_first_try(monkeypatch): - calls = [] - monkeypatch.setattr(lib.requests, "post", lambda *a, **k: calls.append(1) or _Resp(200)) - assert lib._deliver_result("u1", {"u1": {}}, _LOG, attempts=3) is True - assert len(calls) == 1 # no needless retries after a 200 + urls = [] + monkeypatch.setattr(lib.requests, "post", lambda url, **k: urls.append(url) or _Resp(200)) + assert lib._deliver_result("http://bot-a:5000", "u1", {"u1": {}}, _LOG, attempts=3) is True + assert len(urls) == 1 # no needless retries after a 200 + assert urls[0] == "http://bot-a:5000" + lib.SEND_RESULTS # to the origin bot + + +def test_deliver_falls_back_to_main_bot_when_no_target(monkeypatch): + urls = [] + monkeypatch.setattr(lib.requests, "post", lambda url, **k: urls.append(url) or _Resp(200)) + assert lib._deliver_result("", "u1b", {"u1b": {}}, _LOG, attempts=1) is True + assert urls[0] == lib.MAIN_BOT_ADDRESS + lib.SEND_RESULTS # empty target -> default def test_deliver_retries_then_succeeds(monkeypatch): responses = iter([_Resp(503), _Resp(500), _Resp(200)]) monkeypatch.setattr(lib.requests, "post", lambda *a, **k: next(responses)) - assert lib._deliver_result("u2", {"u2": {}}, _LOG, attempts=3) is True + assert lib._deliver_result("http://bot", "u2", {"u2": {}}, _LOG, attempts=3) is True def test_deliver_returns_false_when_all_attempts_fail(monkeypatch): @@ -63,27 +71,47 @@ def test_deliver_returns_false_when_all_attempts_fail(monkeypatch): raise lib.requests.exceptions.RequestException("bot down") monkeypatch.setattr(lib.requests, "post", boom) - assert lib._deliver_result("u3", {"u3": {}}, _LOG, attempts=2) is False + assert lib._deliver_result("http://bot", "u3", {"u3": {}}, _LOG, attempts=2) is False -def test_resend_once_removes_only_acked_entries(outbox, monkeypatch): - outbox.put("ok", {"ok": {}}) - outbox.put("bad", {"bad": {}}) +def test_resend_delivers_each_result_to_its_own_origin_bot(outbox, monkeypatch): + outbox.put("ok", {"target": "http://bot-a:5000", "payload": {"ok": {}}}) + outbox.put("bad", {"target": "http://bot-b:5000", "payload": {"bad": {}}}) + seen = [] - def fake_deliver(query_uuid, _payload, _logger, attempts=1): + def fake_deliver(target, query_uuid, _payload, _logger, attempts=1): + seen.append((target, query_uuid)) return query_uuid == "ok" monkeypatch.setattr(lib, "_deliver_result", fake_deliver) lib._resend_once(_LOG) + assert ("http://bot-a:5000", "ok") in seen # delivered to A's address + assert ("http://bot-b:5000", "bad") in seen # attempted to B's address assert not outbox.contains("ok") # acked -> dropped assert outbox.contains("bad") # not acked -> kept for the next sweep +def test_resend_handles_legacy_entry_shape(outbox, monkeypatch): + # An OUTBOX entry from before per-origin callbacks (raw payload, no target) + # must still be delivered - to the default bot. + outbox.put("old", {"old": {"10.1/x": {"Title": ["P"], "type": "a"}}}) + seen = [] + + def fake_deliver(target, query_uuid, _payload, _logger, attempts=1): + seen.append((target, query_uuid)) + return True + + monkeypatch.setattr(lib, "_deliver_result", fake_deliver) + lib._resend_once(_LOG) + assert seen == [("", "old")] # empty target -> _deliver_result uses MAIN_BOT + assert not outbox.contains("old") + + def test_resend_keeps_result_until_bot_recovers(outbox, monkeypatch): # Simulate: bot down for the first sweep, up for the second. The result must # survive the outage and be delivered on recovery. - outbox.put("u9", {"u9": {"10.1/x": {"Title": ["P"], "type": "article"}}}) + outbox.put("u9", {"target": "http://bot", "payload": {"u9": {"10.1/x": {"Title": ["P"], "type": "article"}}}}) state = {"up": False} def flaky_post(*_a, **_k): diff --git a/tests/integration/test_librarian_query_lifecycle.py b/tests/integration/test_librarian_query_lifecycle.py index 03b27b9..0c28c6e 100644 --- a/tests/integration/test_librarian_query_lifecycle.py +++ b/tests/integration/test_librarian_query_lifecycle.py @@ -58,15 +58,19 @@ def test_ping_idle_routes_through_internal_queue(monkeypatch): monkeypatch.setattr(lib.requests, "post", lambda *a, **k: posted.append((a, k))) client = _client() - resp = client.post("/ping", json={"UUID": "ping-idle"}) + resp = client.post("/ping", json={"UUID": "ping-idle", "callback": "http://bot-a:5000"}) assert resp.status_code == 200 - # Idle => it went onto the internal queue for the worker, NOT posted directly. + # Idle => it went onto the internal queue for the worker (carrying the + # callback so the worker pongs the right bot), NOT posted directly. assert posted == [] - assert lib.librarian_queue.get_nowait() == {"__ping__": "ping-idle"} + assert lib.librarian_queue.get_nowait() == { + "__ping__": "ping-idle", + "callback": "http://bot-a:5000", + } -def test_ping_while_busy_pongs_directly_without_queue(monkeypatch): +def test_ping_while_busy_pongs_directly_to_the_pinging_bot(monkeypatch): _reset() lib.worker_busy.set() # a search is grinding posted = [] @@ -81,15 +85,33 @@ def test_ping_while_busy_pongs_directly_without_queue(monkeypatch): monkeypatch.setattr(lib.requests, "post", fake_post) client = _client() - resp = client.post("/ping", json={"UUID": "ping-busy"}) + resp = client.post("/ping", json={"UUID": "ping-busy", "callback": "http://bot-b:5000"}) assert resp.status_code == 200 - # Busy => direct pong back to the bot, and NOTHING queued (it would only wait - # behind the long search). + # Busy => direct pong, NOTHING queued, and it goes to the CALLBACK bot (not + # the static default) so a shared librarian health-checks each bot correctly. assert lib.librarian_queue.empty() assert len(posted) == 1 assert posted[0]["json"] == {"__pong__": "ping-busy"} - assert posted[0]["url"].endswith(lib.SEND_RESULTS) + assert posted[0]["url"] == "http://bot-b:5000" + lib.SEND_RESULTS + + +def test_ping_without_callback_pongs_to_default_bot(monkeypatch): + _reset() + lib.worker_busy.set() + posted = [] + + class _Resp: + status_code = 200 + + monkeypatch.setattr( + lib.requests, "post", + lambda url, **k: posted.append(url) or _Resp(), + ) + client = _client() + client.post("/ping", json={"UUID": "ping-nocb"}) # no callback + + assert posted == [lib.MAIN_BOT_ADDRESS + lib.SEND_RESULTS] def test_query_status_enforces_api_key():