Files
conjurer/conjurer_librarian/conjurer_librarian.py
T
gitea ae1bd67772
CI / compile (pull_request) Successful in 12s
CI / unit (pull_request) Successful in 29s
CI / integration (pull_request) Successful in 31s
build / build (push) Successful in 33s
CI / compile (push) Successful in 15s
CI / unit (push) Successful in 30s
CI / integration (push) Successful in 34s
Librarian: survive transient Crossref failures instead of losing the search
Field report: one httpx ReadTimeout inside habanero surfaced as
'Search <uuid> crashed', and the worker's crash handler then FORGOT the
request - so an expensive search vanished and the user was told it was
eaten, because a public API blinked once.

Two defences:
* Every habanero call goes through _crossref_call, which retries with
  linear backoff (CONJURER_CROSSREF_ATTEMPTS, default 4; backoff
  CONJURER_CROSSREF_BACKOFF, 5s). habanero wraps httpx errors in a plain
  RuntimeError so we can't filter narrowly - retries are simply bounded
  and the last error is re-raised. They now also run via asyncio.to_thread,
  so a slow Crossref no longer blocks the worker's event loop.
* A crashed search is no longer dropped on the first failure: the attempt
  count is persisted with the request and the search is requeued (keeping
  any checkpoint, so a crashed DB scan resumes rather than restarts) until
  CONJURER_SEARCH_MAX_ATTEMPTS (default 3). It stays 'queued' for the
  bot's watchdog while retrying, and only after the cap is it forgotten.

Also: scrape_bot's 'Got blocked' is routine sci-hub behaviour (it backs off
an hour and carries on) - log it as WARNING, not ERROR, so it stops looking
like a fault when scanning for real problems.

Tests: retry-then-succeed, bounded re-raise, no retry on success, the
persisted attempt counter, and forget-on-give-up. Suite: 58 unit + 70
integration green.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-08-04 13:50:47 +02:00

1117 lines
46 KiB
Python

"""
This module contains the implementation of the Librarian class and
related functions for searching and refining queries.
Classes:
- Librarian: Represents a librarian object that performs search and
refinement operations on queries
Functions:
- flask_debug: Starts a Flask application in debug mode without using the reloader.
- waitress_run: Serves the Flask application using the Waitress WSGI server.
- BackgroundTaskSearch: Represents a background task for running the
Librarian object asynchronously
"""
import asyncio
import hashlib
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 Empty, Queue
from typing import Dict, Optional
import requests
import lib_paths
import scrape_bot
import search_bot
# import search_bot2 as search_bot
from durable_queue import DiskQueue
from flask import Flask, jsonify, request, abort
from habanero import Crossref
from waitress import serve
try:
import netrc
except ImportError: # pragma: no cover
netrc = None
# Constants
def _env(name: str, default: str) -> str:
return os.getenv(name, default)
def _env_path(name: str, default: str) -> Path:
return Path(os.getenv(name, default)).expanduser().resolve()
BASE_DIR = Path(
os.getenv("CONJURER_LIBRARIAN_BASE", str(Path(__file__).resolve().parent))
)
NETRC_FILE = _env_path("CONJURER_NETRC_FILE", str(Path.home() / ".netrc"))
HOST_ADDRESS = _env("CONJURER_LIBRARIAN_HOST", "0.0.0.0")
PORT_ADDRESS = int(_env("CONJURER_LIBRARIAN_PORT", "5001"))
MAIN_BOT_ADDRESS = _env("CONJURER_MAIN_BOT", "http://127.0.0.1:5000")
SEND_RESULTS = _env("CONJURER_LIBRARIAN_RESULTS_ENDPOINT", "/conjurer")
MAX_CR_RESULTS = int(_env("CONJURER_LIBRARIAN_MAX_RESULTS", "500"))
ENCODING = _env("CONJURER_ENCODING", "utf-8")
API_KEY = os.getenv("CONJURER_API_KEY")
LOGFILE_PATH = _env_path(
"CONJURER_LIBRARIAN_LOG", str(BASE_DIR / "librarian.log")
)
# Durable OUTBOX for finished results. A result is expensive (hours of compute),
# so it is written here and only removed once the bot ACKs it (HTTP 200). Lives
# on the librarian's persistent state volume, so it survives a librarian restart
# and a transient bot outage; the resender thread keeps retrying until delivered.
OUTBOX_DIR = _env("CONJURER_LIBRARIAN_OUTBOX", os.path.join(lib_paths.STATE_DIR, "outbox"))
RESULT_SEND_ATTEMPTS = int(_env("CONJURER_RESULT_SEND_ATTEMPTS", "3"))
RESULT_SEND_BACKOFF = float(_env("CONJURER_RESULT_SEND_BACKOFF", "2"))
OUTBOX_RESEND_SECONDS = int(_env("CONJURER_OUTBOX_RESEND_SECONDS", "60"))
_outbox = DiskQueue(OUTBOX_DIR)
# cr_results/rr_results/s_results.json are write-only debug dumps (nothing reads
# them). They used to accumulate EVERY search forever AND json.load the whole
# growing file on each write - unbounded RAM + disk, and for a deep search the
# raw dump is hundreds of MB. Off by default now; when explicitly enabled they
# are overwritten with just the latest search (no load, no accumulation).
DEBUG_DUMPS = _env("CONJURER_LIBRARIAN_DEBUG_DUMPS", "0").lower() in ("1", "true", "yes")
# Log level: INFO keeps normal runs readable (the desktop-era per-line/per-file
# chatter is now DEBUG); set DEBUG to get the full verbosity back.
LOG_LEVEL = _env("CONJURER_LIBRARIAN_LOG_LEVEL", "INFO").upper()
def _dump_debug(path, uuid, data) -> None:
"""Optionally dump the latest search's data for debugging.
Overwrites (never accumulates) and does nothing unless DEBUG_DUMPS is on, so
it can't grow RAM or the state volume in normal operation."""
if not DEBUG_DUMPS:
return
try:
with open(path, "w", encoding="utf-8") as handle:
json.dump({uuid: data}, handle)
except OSError as exc:
logging.getLogger("conjurer_librarian").warning("Debug dump to %s failed: %s", path, exc)
app = Flask(__name__)
librarian_queue = Queue()
librarian_list = []
# Lifecycle of every real search uuid: "queued" (accepted, sitting in
# librarian_queue) -> "processing" (worker pulled it) -> removed (worker
# finished AND attempted to send the result). The bot's per-query watchdog polls
# /query_status against this: a uuid that VANISHES from here without its result
# reaching the bot is a lost result (finished-but-never-delivered) and gets
# flagged in chat.
active_queries: Dict[str, str] = {}
_active_lock = threading.Lock()
# Set while the worker is grinding a real search. A ping arriving during this
# pongs back immediately WITHOUT queueing - being busy is healthy (you can keep
# 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)
# Simple result cache: a repeat of the same query (normalised) skips the whole
# Crossref + DB-scan and returns the stored hits. Disk-backed so it survives a
# restart, TTL'd, and size-bounded. Set CONJURER_LIBRARIAN_CACHE_TTL=0 to disable.
CACHE_DIR = _env("CONJURER_LIBRARIAN_CACHE", os.path.join(lib_paths.STATE_DIR, "cache"))
CACHE_TTL_SECONDS = int(_env("CONJURER_LIBRARIAN_CACHE_TTL", str(7 * 24 * 3600)))
CACHE_MAX_ENTRIES = int(_env("CONJURER_LIBRARIAN_CACHE_MAX", "500"))
_cache = DiskQueue(CACHE_DIR)
def _cache_key(query, deep_search) -> str:
"""Stable key for a query: whitespace-normalised, case-insensitive, and
scoped by deep vs shallow (they return different result sets)."""
normalised = " ".join(str(query).lower().split())
return hashlib.sha256(f"{int(bool(deep_search))}:{normalised}".encode("utf-8")).hexdigest()
def _cache_get(query, deep_search):
"""Return the cached final_result for this query, or None on miss/expiry."""
if CACHE_TTL_SECONDS <= 0:
return None
entry = _cache.get(_cache_key(query, deep_search))
if not entry or entry.get("expires", 0) < time.time():
return None
return entry.get("final_result")
def _cache_put(query, deep_search, final_result) -> None:
"""Store a completed search's hits, with a TTL, and bound the cache size."""
if CACHE_TTL_SECONDS <= 0:
return
_cache.put(
_cache_key(query, deep_search),
{"query": str(query), "final_result": final_result, "expires": time.time() + CACHE_TTL_SECONDS},
)
_cache.prune(CACHE_MAX_ENTRIES)
# ---- "Still alive" heartbeat for the running search ------------------------
# A deep scan runs for hours with nothing in the log between start and finish.
# Every HEARTBEAT_SECONDS the running search says it is still going, with its
# uuid, the phrase, and a ROUGH how-far-along. The estimate is deliberately
# cheap: producers already record a byte offset per chunk file, and the total
# size is stat()'d once at search start - so it costs a sum over ~40 ints.
HEARTBEAT_SECONDS = int(_env("CONJURER_LIBRARIAN_HEARTBEAT_SECONDS", "1200")) # 20 min
_current_search: Dict[str, object] = {}
_current_lock = threading.Lock()
def _set_current_search(uuid, query, progress, live_results) -> None:
with _current_lock:
_current_search.clear()
_current_search.update({
"uuid": str(uuid), "query": str(query), "started": time.monotonic(),
"progress": progress, "live": live_results,
})
def _clear_current_search() -> None:
with _current_lock:
_current_search.clear()
def _progress_summary(progress):
"""(done_bytes, total_bytes, percent) from a live progress dict. Cheap: a
sum over one int per chunk file. Percent is 0.0 when the total is unknown."""
progress = progress or {}
positions = progress.get("positions") or {}
total = progress.get("total_bytes") or 0
done = sum(positions.values())
if total > 0:
done = min(done, total) # a partially-buffered tail can nudge past 100%
return done, total, 100.0 * done / total
return done, total, 0.0
def search_heartbeat(app_logger) -> None:
"""Log a 'still searching' line every HEARTBEAT_SECONDS while one runs."""
while not SHUTDOWN_EVENT.wait(HEARTBEAT_SECONDS):
try:
with _current_lock:
snapshot = dict(_current_search) if _current_search else None
if not snapshot:
continue # nothing running - stay quiet
done, total, percent = _progress_summary(snapshot.get("progress"))
app_logger.info(
"SEARCH ALIVE %s | '%s' | ~%.1f%% przeskanowane (%.2f/%.2f GB, "
"~%.2f GB do końca) | %d trafień | %.0f min",
snapshot["uuid"], snapshot["query"], percent,
done / 1e9, total / 1e9, max(0, total - done) / 1e9,
len(snapshot.get("live") or []),
(time.monotonic() - snapshot["started"]) / 60.0,
)
except Exception as exc: # pylint: disable=broad-exception-caught
app_logger.warning("Heartbeat tick failed: %s", exc)
# Crossref is a public service that times out / rate-limits under load. A single
# transient ReadTimeout used to blow up the whole (expensive) search, so every
# habanero call is retried with backoff, and a search that still fails is retried
# as a whole a few times before being given up on.
CROSSREF_ATTEMPTS = int(_env("CONJURER_CROSSREF_ATTEMPTS", "4"))
CROSSREF_BACKOFF = float(_env("CONJURER_CROSSREF_BACKOFF", "5"))
SEARCH_MAX_ATTEMPTS = int(_env("CONJURER_SEARCH_MAX_ATTEMPTS", "3"))
def _crossref_call(app_logger, what, func, *args, **kwargs):
"""Run one habanero call, retrying transient failures with linear backoff.
habanero wraps httpx errors (ReadTimeout, connection resets, 5xx) in a plain
RuntimeError, so we cannot filter narrowly - we retry a BOUNDED number of
times on any failure and re-raise the last error if none succeed. Blocking
on purpose: callers invoke it via asyncio.to_thread, which also keeps the
worker's event loop free while Crossref is slow."""
last_exc = None
for attempt in range(1, max(1, CROSSREF_ATTEMPTS) + 1):
try:
return func(*args, **kwargs)
except Exception as exc: # pylint: disable=broad-exception-caught
last_exc = exc
app_logger.warning(
"Crossref %s failed (attempt %d/%d): %s", what, attempt, CROSSREF_ATTEMPTS, exc
)
if attempt < CROSSREF_ATTEMPTS:
time.sleep(CROSSREF_BACKOFF * attempt)
raise last_exc
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:
return {"X-Conjurer-Api-Key": API_KEY}
return {}
def _authorize_request() -> None:
if API_KEY and request.headers.get("X-Conjurer-Api-Key") != API_KEY:
abort(401)
def _post_pong(app_logger, ping_uuid, callback="") -> None:
"""POST a pong for ``ping_uuid`` back to the pinging bot. Non-fatal.
``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(
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 _outbox_target_payload(entry):
"""Unpack an OUTBOX entry into (target, payload).
New shape: {"target": <bot address>, "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(
url, json=payload, headers=_service_headers(), timeout=60
)
if response.status_code == 200:
app_logger.info("Result %s delivered to %s (HTTP 200) attempt %d", uuid, url, attempt)
return True
app_logger.warning(
"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 to %s failed (attempt %d/%d): %s",
uuid, url, attempt, attempts, exc,
)
if attempt < attempts:
time.sleep(RESULT_SEND_BACKOFF * attempt)
return False
def _resend_once(app_logger) -> None:
"""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)
def outbox_resender(app_logger) -> None:
"""Background loop: periodically flush the OUTBOX until the bot is reachable.
This is what makes an expensive result survive a transient bot outage or a
librarian restart - on restart the persisted OUTBOX is simply resent."""
pending = len(_outbox)
if pending:
app_logger.info("OUTBOX has %d un-acked result(s) on startup - will resend", pending)
while True:
try:
_resend_once(app_logger)
except Exception as exc: # pylint: disable=broad-exception-caught
app_logger.exception("OUTBOX resend sweep failed: %s", exc)
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), 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)
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):
"""
Represents a librarian object that performs search and refinement operations on queries.
"""
def __init__(self, _app, query, uuid, _deep_search, callback="") -> None:
"""
Initializes a Librarian object.
Args:
- _app: The Flask application object.
- query: The query to be searched.
- uuid: The unique identifier for the search.
Attributes:
- cr: The Crossref object for performing the search.
- query: The query to be searched.
- uuid: The unique identifier for the search.
- limit: The maximum number of search results to fetch.
- fetched: The number of search results fetched so far.
- hit: The number of search results that match the refinement criteria.
- total: The total number of search results.
- app: The Flask application object.
- live_results: A list to store live search results.
- final_result: A list to store the final refined search results.
- not_in_db: A list to store search results that are not in the local database.
- search_result_from_cr: A dictionary to store the search results from Crossref.
- done: A flag indicating if the search is done.
"""
# Crossref only needs a contact mailto. It can come from
# CONJURER_CROSSREF_MAILTO (the usual container setup) OR from a
# "crossref" entry in the netrc; netrc takes precedence when present.
mailto_contact: Optional[str] = os.getenv("CONJURER_CROSSREF_MAILTO")
if netrc:
try:
netrc_mod = netrc.netrc(str(NETRC_FILE))
auth_tokens = netrc_mod.authenticators("crossref")
if auth_tokens:
mailto_contact = auth_tokens[0]
except (FileNotFoundError, netrc.NetrcParseError) as exc:
# A missing/unreadable netrc is NORMAL when the mailto is set via
# env - don't cry wolf on every single search. Only warn when we
# genuinely have no contact from either source.
_log = logging.getLogger("conjurer_librarian")
if mailto_contact:
_log.debug(
"netrc %s not used (%s) - using CONJURER_CROSSREF_MAILTO",
NETRC_FILE, exc,
)
else:
_log.warning(
"Crossref contact not configured: netrc %s unreadable (%s) "
"and CONJURER_CROSSREF_MAILTO unset",
NETRC_FILE, exc,
)
if not mailto_contact:
raise RuntimeError(
"Crossref credentials not configured. Set CONJURER_CROSSREF_MAILTO or add to netrc."
)
self.cr = Crossref(
mailto=mailto_contact,
ua_string=f"Conjurer project. mailto:{mailto_contact}"
)
self.query = query
self.uuid = str(uuid)
self.limit = MAX_CR_RESULTS
self.fetched = 0
self.hit = 0
self.total = 0
self.app = _app
self.live_results = []
self.final_result = {}
self.not_in_db = {}
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.
self.interrupted = False
async def search_crossref(self, query, deep_search=False):
"""
Performs a search on Crossref for the given query.
Args:
- query: The query to be searched.
Returns:
- result: The search result from Crossref.
Raises:
- None.
"""
self.app.logger.info("STARTED SEARCH")
if not deep_search:
query_limit = MAX_CR_RESULTS if MAX_CR_RESULTS < 1000 else 1000
cr_result = await asyncio.to_thread(
_crossref_call, self.app.logger, "works",
self.cr.works, query=query, limit=query_limit,
)
self.search_result_from_cr.update(cr_result)
self.total = cr_result["message"]["total-results"]
self.fetched += len(cr_result["message"]["items"])
self.app.logger.info(self.total)
self.app.logger.info(self.fetched)
while self.total > self.fetched and self.limit > self.fetched:
tmp_result = await asyncio.to_thread(
_crossref_call, self.app.logger, "works(offset)",
self.cr.works, query=query, limit=query_limit, offset=self.fetched,
)
cr_result["message"]["items"].extend(tmp_result["message"]["items"])
self.total = tmp_result["message"]["total-results"]
self.fetched = len(cr_result["message"]["items"])
self.app.logger.info(self.total)
self.app.logger.info(self.fetched)
await asyncio.sleep(0.1)
else:
cr_result = await asyncio.to_thread(
_crossref_call, self.app.logger, "works(deep cursor)",
self.cr.works, query=query, cursor_max=15000, cursor='*', progress_bar=True,
)
result = cr_result[0]
for item in cr_result[1:]:
result["message"]["items"].extend(item["message"]["items"])
self.total = item["message"]["total-results"]
self.fetched = len(result["message"]["items"])
self.app.logger.info(self.total)
self.app.logger.info(self.fetched)
cr_result = result
self.app.logger.info("Total, fetched:")
self.app.logger.info(self.total)
self.app.logger.info(self.fetched)
self.search_result_from_cr.update(cr_result)
self.total = cr_result["message"]["total-results"]
self.app.logger.info("CROSSREF DONE")
self.app.logger.info("CROSSREF DONE")
_dump_debug(lib_paths.CR_RESULTS, self.uuid, self.search_result_from_cr)
return cr_result
async def refine_search(self, unrefined_result):
"""
Refines the search query based on the unrefined search result.
Args:
- unrefined_result: The unrefined search result.
Returns:
- refined_result: The refined search result.
Raises:
- None.
"""
summarized_results = []
self.app.logger.info("REFINE: Removing all derived works from the list")
for item in unrefined_result["message"]["items"]:
summarized_results.append(
{
"DOI": item["DOI"],
"title": item["title"] if "title" in item else None,
"type": item["type"] if "type" in item else None,
}
)
partial_result = {
self.uuid: {
"total_results": self.search_result_from_cr["message"][
"total-results"
],
"on_page": 1,
"summary": summarized_results,
"results": self.search_result_from_cr["message"]["items"],
}
}
self.app.logger.info("REFINE: Dumping to file")
temp = []
refined_result = {}
for key in partial_result:
self.app.logger.info("KEY:")
self.app.logger.info(key)
for item in partial_result[self.uuid]["summary"]:
if item["title"]:
temp.append(item)
for item in temp:
refined_result[item["DOI"]]= item
_dump_debug(lib_paths.RR_RESULTS, self.uuid, refined_result)
return refined_result
async def check_if_exists(self, refined_result, resume=None):
"""
Checks if the given DOI exists.
Args:
- doi: The DOI to be checked.
- brute_force: A flag indicating if brute force method should be used.
Returns:
- result: The search result.
Raises:
- None.
"""
self.app.logger.info("REFINE: Running search in the backend app")
dois = []
for item, value in refined_result.items():
dois.append([item, value])
# Publish this scan as "the running search" so the heartbeat can report
# it; cleared in finally so a finished/crashed scan never lingers there.
progress = {}
_set_current_search(self.uuid, self.query, progress, self.live_results)
try:
result, positions, interrupted = await asyncio.to_thread(
search_bot.search_for_doi,
dois, self.live_results, self.app.logger, SHUTDOWN_EVENT, resume,
progress,
)
finally:
_clear_current_search()
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:
if item["exists"]:
result_list.append(item)
else:
result_no_db.append(item)
self.hit = len(result)
return result_list, result_no_db
async def answer_query(self, deep_search=False):
"""
Answers the search query.
Args:
- None.
Returns:
- result: The search result.
Raises:
- None.
"""
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:
# Cache hit: an identical (normalised) query ran recently - return its
# stored hits and skip Crossref + the whole DB scan entirely.
cached = _cache_get(self.query, deep_search)
if cached is not None:
self.final_result = cached
self.app.logger.info(
"Cache HIT for %s (%d hits): %s", self.uuid, len(cached), self.query
)
return self.final_result
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
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("Search %s produced %d hits", self.uuid, len(self.final_result))
# Cache the completed result so a repeat of this query is instant.
_cache_put(self.query, deep_search, self.final_result)
return self.final_result
# ============================= FLASK INTERNALS===============================
def flask_debug():
"""
Starts a Flask application in debug mode without using the reloader.
Args:
- None.
Returns:
- None.
Raises:
- None.
"""
# trunk-ignore(bandit/B201)
app.run(debug=True, use_reloader=False, host=HOST_ADDRESS, port=PORT_ADDRESS)
def waitress_run():
"""
Serves the Flask application using the Waitress WSGI server.
Args:
- None.
Returns:
- None.
Raises:
- None.
"""
serve(app, host=HOST_ADDRESS, port=PORT_ADDRESS)
class BackgroundTaskSearch(threading.Thread):
"""
A background task for searching and saving results to files.
This class extends the `threading.Thread` class and is responsible for running
the search task in the background. It retrieves queries from a queue, performs
the search, and saves the results to files.
Attributes:
app (App): The application instance.
"""
def run(self):
"""
Run the background task.
This method is called when the thread is started. It creates a new event loop,
runs the `_run` method, and closes the event loop.
"""
loop = asyncio.new_event_loop()
loop.run_until_complete(self._run())
loop.close()
async def _run(self):
"""
Perform the search task.
This method is an asynchronous coroutine that runs in a loop. It retrieves a
librarian from the queue, answers the query, and saves the results to files.
It also sends the results to a remote server.
The search task continues running indefinitely until the thread is stopped.
"""
while not SHUTDOWN_EVENT.is_set():
database = None
ndb_database = None
# 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.
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,
)
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
# (even on a crash) in finally: worker_busy so a ping doesn't wait
# behind us, and active_queries so the bot's watchdog can tell a
# finished-and-gone query from one still in flight.
worker_busy.set()
requeued = False # set when a crash schedules another attempt
with _active_lock:
active_queries[str(librarian.uuid)] = "processing"
try:
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}
# Save results to "not_in_db.json" file
with open(lib_paths.NOT_IN_DB, "r+", encoding="utf-8") as ndb_file:
ndb_database = {}
try:
ndb_database = json.load(ndb_file)
except JSONDecodeError:
pass
if ndb_database:
ndb_database.update(librarian.not_in_db)
else:
ndb_database = librarian.not_in_db
ndb_file.truncate(0)
ndb_file.seek(0)
json.dump(ndb_database, ndb_file)
# Optional debug dump of the final result (off by default).
_dump_debug(lib_paths.S_RESULTS, librarian.uuid, result[librarian.uuid])
self.app.logger.info("Search %s finished", librarian.uuid)
# Persist the result to the durable OUTBOX FIRST, then try to
# deliver it. Writing to disk before sending is the whole point:
# an expensive (hours-long) result now survives a failed send, a
# bot outage, or a librarian restart - the resender keeps
# 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 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 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, target, uuid, payload, self.app.logger):
_outbox.remove(uuid)
else:
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 (that would
# freeze the whole queue). It also must not silently vanish just
# because Crossref timed out once: retry the whole search a
# bounded number of times (attempt count persisted with the
# request, so it can't loop forever), keeping any checkpoint so a
# crashed DB scan resumes rather than restarts. Only after
# SEARCH_MAX_ATTEMPTS do we give up and let the bot's watchdog
# tell the user it vanished.
uuid = str(librarian.uuid)
self.app.logger.exception("Search %s crashed: %s", uuid, exc)
stored = _requests.get(uuid) or {}
attempts = int(stored.get("attempts", 0)) + 1
if attempts < max(1, SEARCH_MAX_ATTEMPTS):
stored.update({
"query": librarian.query,
"deep_search": librarian.deep_search,
"callback": librarian.callback,
"attempts": attempts,
})
_requests.put(uuid, stored)
librarian_queue.put(
Librarian(self.app, librarian.query, uuid,
librarian.deep_search, librarian.callback)
)
requeued = True
self.app.logger.warning(
"Search %s requeued after crash (attempt %d/%d)",
uuid, attempts, SEARCH_MAX_ATTEMPTS,
)
else:
self.app.logger.error(
"Search %s failed %d times - giving up", uuid, attempts
)
_forget_search(uuid)
finally:
worker_busy.clear()
with _active_lock:
if requeued:
# Still known to the bot's watchdog - it's going round again.
active_queries[str(librarian.uuid)] = "queued"
else:
active_queries.pop(str(librarian.uuid), None)
await asyncio.sleep(1)
SHUTDOWN_DONE.set()
self.app.logger.info("Search worker stopped cleanly")
# ==================================SERVER ROUTES==========================================
@app.route("/query", methods=["POST"])
async def query_database():
_authorize_request()
"""
Endpoint for querying the database.
This function receives a POST request containing a JSON payload with a query and a UUID.
It creates a Librarian object with the query and UUID,
and adds it to the librarian_queue and librarian_list.
Finally, it returns a JSON response indicating the success
of the operation, along with the query, UUID,
and the current size of the librarian_queue.
Returns:
tuple: A tuple containing a JSON response and a status code.
"""
record = json.loads(request.data)
uuid = record["UUID"]
deep_search = record["deep_search"]
# 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, "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
# "queued" now so it counts as known the moment we accept it.
with _active_lock:
active_queries[str(uuid)] = "queued"
answer_data = (record["query"], record["UUID"], librarian_queue.qsize())
return_data = (
jsonify(isError=False, message="Success", statusCode=200, data=answer_data),
200,
)
return return_data
@app.route("/ping", methods=["POST"])
def ping_roundtrip():
_authorize_request()
"""
Health-check round-trip.
Two cases, one guarantee - the pong always comes back over the librarian->bot
return path (the only thing the ping must prove):
* IDLE: put a ping marker onto the SAME internal ``librarian_queue`` real
searches use and return 200. The worker pulls it off and pongs it back,
so a successful pong proves the whole pipeline flows (queue + worker + the
return leg), not just that Flask is up.
* BUSY (a search is grinding): DO NOT queue - the ping would just wait behind
a possibly hours-long search and time out, making a perfectly healthy busy
librarian look dead. Pong back immediately instead. Being busy is fine; you
can keep piling searches on. The ping only needs to catch a BROKEN return
path, and the direct pong exercises exactly that.
"""
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, callback)
else:
app.logger.info("PING received %s - queued for round-trip", ping_uuid)
librarian_queue.put({"__ping__": ping_uuid, "callback": callback})
return (
jsonify(isError=False, message="ping-queued", statusCode=200, data=ping_uuid),
200,
)
@app.route("/query_status", methods=["POST"])
def query_status():
_authorize_request()
"""
Per-query watchdog probe.
Returns whether ``UUID`` is still known to the librarian (queued or being
processed). The bot polls this after dispatching a search: while the uuid is
known the search is progressing; once it VANISHES here without the result
ever reaching the bot, the result was lost in transit and the bot tells the
user. A busy/queued search is never mistaken for a lost one.
"""
record = json.loads(request.data)
uuid = str(record["UUID"])
with _active_lock:
state = active_queries.get(uuid, "unknown")
return (
jsonify(
isError=False,
message="Success",
statusCode=200,
data={"uuid": uuid, "known": state != "unknown", "state": state},
),
200,
)
@app.route("/get_partial_result", methods=["POST"])
async def get_partial():
_authorize_request()
"""
Retrieves the partial result for a given UUID.
Returns:
A JSON response containing the partial result.
"""
record = json.loads(request.data)
app.logger.info(record)
app.logger.info(record["UUID"])
for lib in librarian_list:
if lib.uuid == record["UUID"]:
answer_data = lib.live_results
break
return_data = (
jsonify(isError=False, message="Success", statusCode=200, data=answer_data),
200,
)
return return_data
# =======================================MAIN===================================================
if __name__ == "__main__":
# Default INFO (readable). Set CONJURER_LIBRARIAN_LOG_LEVEL=DEBUG for the
# full per-file / per-line search chatter.
app.logger.setLevel(LOG_LEVEL)
_fmt = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
LOGFILE_PATH.parent.mkdir(parents=True, exist_ok=True)
h1 = handlers.RotatingFileHandler(
filename=str(LOGFILE_PATH),
encoding=ENCODING,
mode="a",
maxBytes=6 * 1024 * 1024,
backupCount=6,
)
h1.setFormatter(_fmt)
app.logger.addHandler(h1)
# Console handler so `kubectl logs` shows what's happening (k8s reads stdout);
# the search internals no longer print() straight to stdout.
_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))
bgtask = BackgroundTaskSearch()
bgtask.app = app
bgtask.daemon = True
threads.append(bgtask)
threads.append(
threading.Thread(
target=scrape_bot.scraper, args=(app.logger,), daemon=True
)
)
# Durable delivery: keep flushing the OUTBOX so any result not yet acked by
# the bot (transient outage, or left over from before a restart) is resent.
threads.append(
threading.Thread(target=outbox_resender, args=(app.logger,), daemon=True)
)
# "Still searching" heartbeat, so an hours-long scan isn't radio silence.
threads.append(
threading.Thread(target=search_heartbeat, args=(app.logger,), daemon=True)
)
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,
)