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 <noreply@anthropic.com>

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,
+        )

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 '<n>%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__":