From ac16b77f56cf851fa96fa0e470a47b8159656948 Mon Sep 17 00:00:00 2001 From: Polish Hammer Date: Sun, 2 Aug 2026 22:09:49 +0200 Subject: [PATCH] Librarian: graceful shutdown with resumable search state A restart of the librarian used to throw away an in-flight search (and any searches still queued). Now search state survives a restart: * Resumable DB scan (search_bot): each producer records a tell()-cookie watermark per chunk file as it goes (safe because search_for_doi drains the work queue before returning), and can seek back to it. search_for_doi now takes stop_event + resume and returns (result_list, positions, interrupted). * Persisted requests: /query writes the accepted request to a disk queue before enqueuing; replay_requests re-enqueues unfinished ones on startup. So even a search still waiting in the queue survives a restart. * Checkpoints: when a graceful shutdown interrupts a scan, the librarian writes {dois, found-so-far, per-file offsets}. On restart answer_query loads it, skips the (already done) Crossref+refine, and continues the scan from the saved offsets with the found DOIs pre-marked - no line is read twice and none is missed. A finished or crashed search forgets its request+checkpoint (no poison-pill replay). * Graceful shutdown: SIGTERM/SIGINT set a shutdown event; the running scan checkpoints and the worker stops. The main thread then exits within a BOUNDED window (CONJURER_LIBRARIAN_GRACEFUL_TIMEOUT, default 45s) so the pod can never become an un-killable zombie. Needs terminationGracePeriod >= that in the deploy (separate PR). Tests: search_bot resume correctness (seek past scanned, don't miss/re-scan; stop_event -> interrupted) and librarian state mechanics (request replay, forget, checkpoint round-trip, poison-pill drop). Suite: 58 unit + 49 integration green. Co-Authored-By: Claude Opus 4.8 --- conjurer_librarian/conjurer_librarian.py | 195 +++++++++++++++--- conjurer_librarian/search_bot.py | 92 ++++++--- durable_queue.py | 8 + .../test_librarian_resume_state.py | 90 ++++++++ tests/unit/test_search_bot.py | 78 ++++++- 5 files changed, 397 insertions(+), 66 deletions(-) create mode 100644 tests/integration/test_librarian_resume_state.py diff --git a/conjurer_librarian/conjurer_librarian.py b/conjurer_librarian/conjurer_librarian.py index 376fa87..c5bde98 100644 --- a/conjurer_librarian/conjurer_librarian.py +++ b/conjurer_librarian/conjurer_librarian.py @@ -17,12 +17,13 @@ import asyncio import json import logging import os +import signal import threading import time from json.decoder import JSONDecodeError from logging import handlers from pathlib import Path -from queue import Queue +from queue import Empty, Queue from typing import Dict, Optional import requests @@ -118,6 +119,29 @@ _active_lock = threading.Lock() # piling searches on), so "busy" must never look like "dead" to the health check. worker_busy = threading.Event() +# ---- Graceful shutdown + resumable search state ---------------------------- +# SHUTDOWN_EVENT is set by the SIGTERM/SIGINT handler; the running search checks +# it (via search_bot) and checkpoints itself. SHUTDOWN_DONE is set by the worker +# once it has stopped cleanly, so the main thread can exit promptly - bounded by +# GRACEFUL_TIMEOUT so we never become an un-killable zombie pod. +SHUTDOWN_EVENT = threading.Event() +SHUTDOWN_DONE = threading.Event() +GRACEFUL_TIMEOUT = float(_env("CONJURER_LIBRARIAN_GRACEFUL_TIMEOUT", "45")) +# Persisted, per-uuid: accepted-but-unfinished search REQUESTS (so a restart +# re-runs them) and in-progress CHECKPOINTS (found-so-far + per-file resume +# offset, so a restart CONTINUES a long scan instead of restarting it). +REQUESTS_DIR = _env("CONJURER_LIBRARIAN_REQUESTS", os.path.join(lib_paths.STATE_DIR, "requests")) +CHECKPOINT_DIR = _env("CONJURER_LIBRARIAN_CHECKPOINTS", os.path.join(lib_paths.STATE_DIR, "checkpoints")) +_requests = DiskQueue(REQUESTS_DIR) +_checkpoints = DiskQueue(CHECKPOINT_DIR) + + +def _forget_search(uuid) -> None: + """A search is fully done (or abandoned): drop its persisted request and any + checkpoint so it is never replayed or resumed again.""" + _requests.remove(uuid) + _checkpoints.remove(uuid) + def _service_headers() -> Dict[str, str]: if API_KEY: @@ -206,6 +230,39 @@ def outbox_resender(app_logger) -> None: time.sleep(OUTBOX_RESEND_SECONDS) +def replay_requests(app_logger) -> None: + """Re-enqueue accepted-but-unfinished searches after a restart. + + Requests persisted by /query but never completed are put back on the internal + queue. Those with a checkpoint resume mid-scan (answer_query loads it); the + rest simply re-run. Marked 'queued' so the bot's watchdog sees them as known + again.""" + pending = _requests.items() + if not pending: + return + 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)) + except Exception as exc: # pylint: disable=broad-exception-caught + app_logger.warning("Cannot replay request %s (dropping): %s", uuid, exc) + _forget_search(uuid) + continue + with _active_lock: + active_queries[str(uuid)] = "queued" + librarian_queue.put(cl) + librarian_list.append(cl) + + +def _handle_shutdown(signum, _frame) -> None: + """SIGTERM/SIGINT: ask the running search to checkpoint and stop. The main + thread then waits (bounded) for it to finish - see __main__.""" + logging.getLogger("conjurer_librarian").warning( + "Signal %s received - beginning graceful shutdown", signum + ) + SHUTDOWN_EVENT.set() + + # trunk-ignore(pylint/R0902) class Librarian(object): """ @@ -283,6 +340,10 @@ class Librarian(object): self.search_result_from_cr = {} self.done = False self.deep_search = _deep_search + # 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. + self.interrupted = False async def search_crossref(self, query, deep_search=False): """ @@ -389,7 +450,7 @@ class Librarian(object): _dump_debug(lib_paths.RR_RESULTS, self.uuid, refined_result) return refined_result - async def check_if_exists(self, refined_result): + async def check_if_exists(self, refined_result, resume=None): """ Checks if the given DOI exists. @@ -403,15 +464,28 @@ class Librarian(object): Raises: - None. """ - result = {} self.app.logger.info("REFINE: Running search in the backend app") dois = [] for item, value in refined_result.items(): dois.append([item, value]) - coro = asyncio.to_thread( - search_bot.search_for_doi, dois, self.live_results, self.app.logger + result, positions, interrupted = await asyncio.to_thread( + search_bot.search_for_doi, + dois, self.live_results, self.app.logger, SHUTDOWN_EVENT, resume, ) - result = await coro + if interrupted: + # Graceful shutdown hit mid-scan: checkpoint found-so-far + per-file + # resume offsets + the DOI list, so a restart continues instead of + # restarting. The worker sees self.interrupted and does NOT deliver. + found = [item["DOI"] for item in result if item["exists"]] + _checkpoints.put( + self.uuid, + {"dois": refined_result, "found": found, "positions": positions}, + ) + self.interrupted = True + self.app.logger.info( + "Search %s checkpointed (%d found so far) for resume", self.uuid, len(found) + ) + return [], [] result_list = [] result_no_db = [] for item in result: @@ -436,21 +510,37 @@ class Librarian(object): Raises: - None. """ - self.app.logger.info(f"Search started {self.uuid}") - cr_result = await self.search_crossref(query=self.query, deep_search=deep_search) - refined_result = await self.refine_search(cr_result) - answer, negative_answer = await self.check_if_exists(refined_result) + checkpoint = _checkpoints.get(self.uuid) + if checkpoint is not None: + # Resume a search interrupted by a previous shutdown: the expensive + # Crossref + refine work is already captured in the checkpoint, so go + # straight to the DB scan with the saved offsets + found-so-far. + self.app.logger.info( + "Resuming search %s from checkpoint (%d found so far)", + self.uuid, len(checkpoint.get("found", [])), + ) + refined_result = checkpoint["dois"] + resume = { + "found": checkpoint.get("found", []), + "positions": checkpoint.get("positions", {}), + } + else: + self.app.logger.info(f"Search started {self.uuid}") + cr_result = await self.search_crossref(query=self.query, deep_search=deep_search) + refined_result = await self.refine_search(cr_result) + resume = None - self.app.logger.info("Returning result") - self.app.logger.info(answer) - self.app.logger.info(negative_answer) + answer, negative_answer = await self.check_if_exists(refined_result, resume=resume) + if self.interrupted: + # Graceful shutdown mid-scan: checkpoint is written, request stays. + # Signal the worker (None) NOT to deliver - a restart resumes this. + return None for item in answer: self.final_result[item["DOI"]] = {"Title": item["data"]["title"], "type": item["data"]["type"]} for item in negative_answer: self.not_in_db[item["DOI"]] = {"Title": item["data"]["title"], "type": item["data"]["type"]} - self.app.logger.info("Returning result case2") - self.app.logger.info(self.final_result) + self.app.logger.info("Search %s produced %d hits", self.uuid, len(self.final_result)) return self.final_result # ============================= FLASK INTERNALS=============================== @@ -522,10 +612,15 @@ class BackgroundTaskSearch(threading.Thread): The search task continues running indefinitely until the thread is stopped. """ - while True: + while not SHUTDOWN_EVENT.is_set(): database = None ndb_database = None - item = librarian_queue.get() + # Bounded get so the loop can observe SHUTDOWN_EVENT while idle + # (blocked on a plain get() it would never notice a shutdown). + try: + item = librarian_queue.get(timeout=1) + except Empty: + continue # 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. @@ -546,10 +641,18 @@ class BackgroundTaskSearch(threading.Thread): with _active_lock: active_queries[str(librarian.uuid)] = "processing" try: - self.app.logger.info("STARTED") + self.app.logger.info("Processing search %s", librarian.uuid) result = await librarian.answer_query(librarian.deep_search) + if result is None: + # Graceful shutdown interrupted this search mid-scan. Its + # checkpoint + persisted request stay in place, so a restart + # picks it up and RESUMES from where it stopped. + self.app.logger.info( + "Search %s interrupted by shutdown - will resume on restart", + librarian.uuid, + ) + break 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: @@ -589,16 +692,23 @@ class BackgroundTaskSearch(threading.Thread): self.app.logger.warning( "Result %s not acked yet - left in OUTBOX for the resender", uuid ) + # Computed + handed to the durable OUTBOX: the search is done, so + # forget its request + checkpoint (never replay/resume it again). + _forget_search(uuid) 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". + # freeze the whole queue). Log and give up on it - forget the + # request/checkpoint so it isn't retried forever as a poison pill; + # the bot's watchdog tells the user it vanished. self.app.logger.exception("Search %s crashed: %s", librarian.uuid, exc) + _forget_search(str(librarian.uuid)) finally: worker_busy.clear() with _active_lock: active_queries.pop(str(librarian.uuid), None) await asyncio.sleep(1) + SHUTDOWN_DONE.set() + self.app.logger.info("Search worker stopped cleanly") # ==================================SERVER ROUTES========================================== @@ -619,11 +729,12 @@ async def query_database(): tuple: A tuple containing a JSON response and a status code. """ record = json.loads(request.data) - app.logger.info(record) - app.logger.info(record["query"]) - app.logger.info(record["UUID"]) uuid = record["UUID"] deep_search = record["deep_search"] + app.logger.info("Query accepted %s: %s", uuid, 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) librarian_queue.put(cl) librarian_list.append(cl) @@ -743,6 +854,13 @@ if __name__ == "__main__": _console = logging.StreamHandler() _console.setFormatter(_fmt) app.logger.addHandler(_console) + + # Graceful shutdown: on SIGTERM (k8s) / SIGINT the running search checkpoints + # itself and the worker stops; the main thread then exits within a bounded + # window so we never linger as an un-killable zombie pod. + signal.signal(signal.SIGTERM, _handle_shutdown) + signal.signal(signal.SIGINT, _handle_shutdown) + threads = [] threads.append(threading.Thread(target=waitress_run, daemon=True)) # threads.append(threading.Thread(target=flask_debug)) @@ -760,13 +878,22 @@ if __name__ == "__main__": threads.append( threading.Thread(target=outbox_resender, args=(app.logger,), daemon=True) ) - i = 0 - try: - for worker in threads: - app.logger.info("App number: %s", i) - i += 1 - worker.start() - for worker in threads: - worker.join() - except KeyboardInterrupt: - app.logger.info("Shutdown requested - exiting librarian service") + for worker in threads: + worker.start() + # Re-enqueue searches that were accepted/in-progress before the last stop. + replay_requests(app.logger) + app.logger.info("Librarian ready (graceful-shutdown timeout %ss)", GRACEFUL_TIMEOUT) + + # Main thread parks until a shutdown signal, then gives the worker a BOUNDED + # window to checkpoint. sleep() (not Event.wait) so the signal is delivered + # promptly to this thread on every platform. + while not SHUTDOWN_EVENT.is_set(): + time.sleep(0.5) + app.logger.info("Waiting up to %ss for the search to checkpoint...", GRACEFUL_TIMEOUT) + if SHUTDOWN_DONE.wait(GRACEFUL_TIMEOUT): + app.logger.info("Graceful shutdown complete - state saved") + else: + app.logger.warning( + "Graceful shutdown TIMED OUT after %ss - exiting anyway (no zombie)", + GRACEFUL_TIMEOUT, + ) diff --git a/conjurer_librarian/search_bot.py b/conjurer_librarian/search_bot.py index 214e455..ba75e79 100644 --- a/conjurer_librarian/search_bot.py +++ b/conjurer_librarian/search_bot.py @@ -85,15 +85,17 @@ def discover_chunk_files(_logger): return ordered -def producer(out_q, control_q, filename, _logger): - """ - Produces items from the output queue and puts them into the control queue. +def producer(out_q, control_q, filename, _logger, stop_event=None, positions=None, + start_offsets=None): + """Stream a chunk file's lines onto the work queue, resumably. - Args: - out_q (Queue): Output queue. - control_q (Queue): Control queue. - filename (str): Name of the file. - _logger: Logger object for logging. + ``start_offsets[filename]`` (a tell() cookie) is where to RESUME reading from + - so a search continued after a restart skips the part already scanned. + ``positions[filename]`` is updated to the tell() cookie just PAST each line + successfully enqueued; because search_for_doi drains the queue before it + returns, that cookie is a safe "everything up to here is processed" watermark + to checkpoint. ``stop_event`` (graceful shutdown) makes the producer stop + reading and record its watermark, mirroring the early-TERM path. """ try: # errors="replace" so a stray non-UTF-8 byte in a chunk (they happen in @@ -102,9 +104,16 @@ def producer(out_q, control_q, filename, _logger): # producer partway and leaving every DOI after the bad byte unsearched. # DOIs are ASCII, so a replaced byte can only affect junk, never a match. with open(DATABASE_PATH + filename, "r", encoding=ENCODING, errors="replace") as operated_file: - _logger.debug("Producer started: %s", filename) + if start_offsets and filename in start_offsets: + operated_file.seek(start_offsets[filename]) + _logger.debug("Producer %s: resuming at offset %s", filename, start_offsets[filename]) + else: + _logger.debug("Producer started: %s", filename) line_no = 0 while True: + if stop_event is not None and stop_event.is_set(): + _logger.debug("Producer %s: stop requested (graceful)", filename) + break line = operated_file.readline() line_no += 1 # Coarse progress at DEBUG only - the old per-line carriage-return @@ -113,11 +122,15 @@ def producer(out_q, control_q, filename, _logger): _logger.debug("Producer %s: %d lines read", filename, line_no) if not line: + # EOF: record the end offset so a resume seeks here and stops + # immediately (the file is fully scanned). + if positions is not None: + positions[filename] = operated_file.tell() _logger.debug("Producer %s: EOF at %d lines", filename, line_no) break - # Backpressure-safe put onto the BOUNDED queue: wait for room, - # but keep polling the TERM sentinel so a full queue whose + # Backpressure-safe put onto the BOUNDED queue: wait for room, but + # keep polling stop_event / the TERM sentinel so a full queue whose # consumers have already finished can never deadlock us here. stopped = False while True: @@ -125,6 +138,9 @@ def producer(out_q, control_q, filename, _logger): out_q.put(line, timeout=1) break except Full: + if stop_event is not None and stop_event.is_set(): + stopped = True + break try: if control_q.get(block=False) is _sentinel: control_q.put(_sentinel) @@ -133,8 +149,12 @@ def producer(out_q, control_q, filename, _logger): except Empty: pass if stopped: - _logger.debug("Producer %s: TERM (queue full)", filename) + _logger.debug("Producer %s: stop while enqueuing", filename) break + # Watermark AFTER a successful enqueue: safe to resume past here + # once the queue drains (which it does before search_for_doi ends). + if positions is not None: + positions[filename] = operated_file.tell() try: check = control_q.get(block=False) @@ -227,23 +247,41 @@ def consumer(in_q, control_q, doi, live_results, result_list, control_dict, expe -def search_for_doi(doi, live_results, _logger): - """ - Search for DOI in live_results using _logger for logging. +def search_for_doi(doi, live_results, _logger, stop_event=None, resume=None): + """Search for DOI in live_results, resumably. - Args: - doi (list): List of DOI to search for. - live_results (list): List to store the search results. - _logger: Logger object for logging. + Returns ``(result_list, positions, interrupted)``: + * ``result_list`` - the queried DOIs with their ``exists`` flag, + * ``positions`` - ``{filename: tell()-cookie}`` safe-to-resume watermarks + (the queue is drained before return, so everything up to each cookie is + processed), + * ``interrupted`` - True if ``stop_event`` fired (the scan is PARTIAL; check + point ``positions`` + the found DOIs and call again with ``resume=`` to + continue where it left off). + + ``resume`` is ``{"positions": {...}, "found": [doi, ...]}`` from a previous + interrupted run: already-found DOIs are pre-marked and each producer seeks to + its saved offset, so no already-scanned line is read twice. """ control_dict = {"sentinels":0} result_list = [] threads = [] work_q = Queue(maxsize=WORK_Q_SIZE) control_q = Queue() + positions = {} + + resume = resume or {} + already_found = set(resume.get("found", [])) + start_offsets = resume.get("positions", {}) for item in doi: - result_list.append({"DOI": item[0], "exists": False, "data": item[1]}) + entry = {"DOI": item[0], "exists": False, "data": item[1]} + if item[0] in already_found: + # Pre-mark hits from the previous (interrupted) run so we neither + # re-scan for them nor drop them from live_results. + entry["exists"] = True + live_results.append(entry) + result_list.append(entry) # One producer per chunk file that actually exists; the sentinel threshold is # that same count, so the two can never drift apart the way MAXTHREADS did. @@ -253,25 +291,29 @@ def search_for_doi(doi, live_results, _logger): _logger.error( "No '%s' chunk files in %s - DOI search cannot run", CHUNK, DATABASE_PATH ) - return result_list + return result_list, positions, bool(stop_event and stop_event.is_set()) for i in range (0, (len(doi)//1000)+2): t_cons = Thread( target=consumer, args=(work_q, control_q, doi, live_results, result_list, control_dict, expected, i, _logger), ) - _logger.info("Consumer thread created") + _logger.debug("Consumer thread created") threads.append(t_cons) for filename in chunk_files: - _logger.info("Creating worker thread for %s", filename) + _logger.debug("Creating worker thread for %s", filename) threads.append( - Thread(target=producer, args=(work_q, control_q, filename, _logger)) + Thread( + target=producer, + args=(work_q, control_q, filename, _logger, stop_event, positions, start_offsets), + ) ) for worker in threads: worker.start() for worker in threads: worker.join() - return result_list + interrupted = bool(stop_event and stop_event.is_set()) + return result_list, positions, interrupted if __name__ == "__main__": diff --git a/durable_queue.py b/durable_queue.py index 671a85e..725f3f8 100644 --- a/durable_queue.py +++ b/durable_queue.py @@ -52,6 +52,14 @@ class DiskQueue: def contains(self, key) -> bool: return os.path.exists(self._path(key)) + def get(self, key): + """Return the payload stored for ``key``, or None if absent/unreadable.""" + try: + with open(self._path(key), encoding="utf-8") as handle: + return json.load(handle)["payload"] + except (OSError, ValueError, KeyError, TypeError): + return None + def remove(self, key) -> None: try: os.remove(self._path(key)) diff --git a/tests/integration/test_librarian_resume_state.py b/tests/integration/test_librarian_resume_state.py new file mode 100644 index 0000000..cb89075 --- /dev/null +++ b/tests/integration/test_librarian_resume_state.py @@ -0,0 +1,90 @@ +"""Integration: the librarian's persisted search state (requests + checkpoints). + +These pin the durable-state mechanics that let a search survive a restart: +* accepted requests are replayed (re-enqueued) after a restart, +* a finished/abandoned search is forgotten (request + checkpoint dropped), +* a checkpoint round-trips through disk intact. + +The RESUME correctness itself (seek past scanned, don't miss, don't re-scan) +lives in tests/unit/test_search_bot.py. +""" +import logging +import sys +import types + +import pytest + +if "habanero" not in sys.modules: + _habanero = types.ModuleType("habanero") + _habanero.Crossref = object + sys.modules["habanero"] = _habanero + +import conjurer_librarian as lib # noqa: E402 +from durable_queue import DiskQueue # noqa: E402 + +_LOG = logging.getLogger("test-resume-state") +_LOG.addHandler(logging.NullHandler()) + + +class _DummyCrossref: + def __init__(self, **kwargs): + self.kwargs = kwargs + + +@pytest.fixture +def state(tmp_path, monkeypatch): + monkeypatch.setattr(lib, "_requests", DiskQueue(str(tmp_path / "req"))) + monkeypatch.setattr(lib, "_checkpoints", DiskQueue(str(tmp_path / "cp"))) + monkeypatch.setattr(lib, "Crossref", _DummyCrossref) + monkeypatch.setenv("CONJURER_CROSSREF_MAILTO", "test@example.com") + monkeypatch.setattr(lib, "NETRC_FILE", "/nonexistent/conjurer/.netrc") + while not lib.librarian_queue.empty(): + lib.librarian_queue.get() + lib.librarian_list.clear() + with lib._active_lock: + lib.active_queries.clear() + return None + + +def test_accepted_request_is_replayed_after_restart(state): + lib._requests.put("u1", {"query": "kwas foliowy", "deep_search": False}) + + lib.replay_requests(_LOG) + + item = lib.librarian_queue.get_nowait() + assert isinstance(item, lib.Librarian) + assert item.uuid == "u1" + assert item.query == "kwas foliowy" + assert lib.active_queries["u1"] == "queued" # known again to the watchdog + + +def test_forget_search_drops_request_and_checkpoint(state): + lib._requests.put("u2", {"query": "x", "deep_search": False}) + lib._checkpoints.put("u2", {"dois": {}, "found": [], "positions": {}}) + + lib._forget_search("u2") + + assert not lib._requests.contains("u2") + assert not lib._checkpoints.contains("u2") + + +def test_checkpoint_round_trips_through_disk(state): + checkpoint = { + "dois": {"10.1/x": {"DOI": "10.1/x", "title": ["T"], "type": "article"}}, + "found": ["10.1/already"], + "positions": {"0_chunk.txt": 4096}, + } + lib._checkpoints.put("u3", checkpoint) + assert lib._checkpoints.get("u3") == checkpoint + + +def test_unreadable_replayed_request_is_dropped_not_looped(state, monkeypatch): + # A request that can't be reconstructed (e.g. missing Crossref contact) must + # be dropped, not retried forever. + lib._requests.put("u4", {"query": "x", "deep_search": False}) + monkeypatch.delenv("CONJURER_CROSSREF_MAILTO", raising=False) + + lib.replay_requests(_LOG) + + assert lib.librarian_queue.empty() + assert not lib._requests.contains("u4") # forgotten, not left to loop diff --git a/tests/unit/test_search_bot.py b/tests/unit/test_search_bot.py index 9b3d006..545b3c6 100644 --- a/tests/unit/test_search_bot.py +++ b/tests/unit/test_search_bot.py @@ -27,17 +27,32 @@ def _write_chunks(directory, count, target=None, target_index=None): (directory / f"{n}_chunk.txt").write_text("".join(lines), encoding="utf-8") -def _run_bounded(dois, timeout=20): - """Run search_for_doi in a thread; return (finished_in_time, result).""" +def _run_full(dois, timeout=20, stop_event=None, resume=None): + """Run search_for_doi in a thread; return the whole result box. + + search_for_doi now returns (result_list, positions, interrupted); the box + exposes all three (plus 'finished' and 'live') for the resume tests. + """ box = {} live = [] - worker = threading.Thread( - target=lambda: box.update(result=search_bot.search_for_doi(dois, live, _LOG)), - daemon=True, - ) + + def _run(): + result_list, positions, interrupted = search_bot.search_for_doi( + dois, live, _LOG, stop_event=stop_event, resume=resume + ) + box.update(result=result_list, positions=positions, interrupted=interrupted, live=live) + + worker = threading.Thread(target=_run, daemon=True) worker.start() worker.join(timeout) - return (not worker.is_alive()), box.get("result") + box["finished"] = not worker.is_alive() + return box + + +def _run_bounded(dois, timeout=20, stop_event=None, resume=None): + """Back-compat wrapper: return (finished_in_time, result_list).""" + box = _run_full(dois, timeout, stop_event, resume) + return box.get("finished"), box.get("result") def test_finds_doi_in_trailing_chunk(tmp_path, monkeypatch): @@ -133,6 +148,55 @@ def test_discover_chunk_files_sorted_numerically(tmp_path, monkeypatch): assert found == ["0_chunk.txt", "1_chunk.txt", "2_chunk.txt", "10_chunk.txt"] +def _offset_after(path, marker): + """Byte-cookie (tell) just past the line equal to `marker` in `path`.""" + with open(path, "r", encoding="utf-8") as handle: + while True: + line = handle.readline() + if not line: + raise AssertionError(f"marker {marker!r} not found") + if line.strip() == marker: + return handle.tell() + + +def test_resume_seeks_past_scanned_part_and_continues(tmp_path, monkeypatch): + # Chunk: early | first-half decoy | MIDDLE | late. Resume from just past + # MIDDLE with 'early' pre-found. The scan must: keep 'early' (pre-marked), + # find 'late' (after the resume point), and NOT find the first-half decoy + # (proving it seeked past it instead of re-reading from the top). + monkeypatch.setattr(search_bot, "DATABASE_PATH", str(tmp_path) + "/") + path = tmp_path / "0_chunk.txt" + path.write_text( + "10.1/early\n10.1/only-first-half\nMIDDLE\n10.1/late\n", encoding="utf-8" + ) + offset = _offset_after(str(path), "MIDDLE") + + resume = {"found": ["10.1/early"], "positions": {"0_chunk.txt": offset}} + box = _run_full( + [("10.1/early", "D"), ("10.1/late", "D"), ("10.1/only-first-half", "D")], + resume=resume, + ) + + assert box["finished"] + by_doi = {r["DOI"]: r["exists"] for r in box["result"]} + assert by_doi["10.1/early"] is True # carried over from the checkpoint + assert by_doi["10.1/late"] is True # found after the resume offset + assert by_doi["10.1/only-first-half"] is False # skipped - not re-scanned + + +def test_stop_event_interrupts_and_reports_positions(tmp_path, monkeypatch): + monkeypatch.setattr(search_bot, "DATABASE_PATH", str(tmp_path) + "/") + _write_chunks(tmp_path, count=2) + stop = __import__("threading").Event() + stop.set() # already asked to stop before it starts + + box = _run_full([("10.0000/decoy-0-a", "D")], stop_event=stop) + + assert box["finished"], "an already-set stop must not hang the search" + assert box["interrupted"] is True + assert isinstance(box["positions"], dict) + + def test_bounded_queue_does_not_deadlock_on_early_termination(tmp_path, monkeypatch): # The OOM fix bounds the work queue. That means a producer can block on a # FULL queue - and if the consumers have already finished (all DOIs found)