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