gitea
ac16b77f56
CI / compile (pull_request) Successful in 10s
CI / unit (pull_request) Successful in 28s
CI / integration (pull_request) Successful in 27s
build / build (push) Successful in 42s
CI / compile (push) Successful in 10s
CI / unit (push) Successful in 26s
CI / integration (push) Successful in 25s
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>
345 lines
15 KiB
Python
345 lines
15 KiB
Python
"""
|
|
This module contains functions for searching for DOI (Digital Object Identifier)
|
|
in a list of live results.
|
|
It includes a producer-consumer pattern implementation using threads and queues.
|
|
|
|
Functions:
|
|
- producer: Reads lines from a file and puts them into an output queue.
|
|
- consumer: Consumes items from an input queue and checks if DOI exists in the live results list.
|
|
- search_for_doi: Searches for DOI in live results using the producer-consumer pattern.
|
|
|
|
Global Variables:
|
|
- MAXTHREADS: Maximum number of worker threads.
|
|
- DATABASE_PATH: Path to the database files.
|
|
- ENCODING: Encoding of the database files.
|
|
- CHUNK: File name pattern for the database files.
|
|
- _sentinel: Sentinel object used to signal termination.
|
|
- result_list: List to store the search results.
|
|
- WORK_Q_SIZE: Maximum size of the work queue.
|
|
"""
|
|
|
|
# TODO: Wpiemdolić to wszystko w klasę z loggerem przysłanym z góry
|
|
import os
|
|
import re
|
|
from queue import Empty, Full, Queue
|
|
from threading import Thread
|
|
import time
|
|
q = Queue()
|
|
#TODO: Count number of lines in files and print to approximate on which part of the file search is
|
|
|
|
# Deployment data is environment-overridable so the local DOI database can live
|
|
# on a mounted volume (Docker/Linux) instead of the hardcoded Windows path.
|
|
DATABASE_PATH = os.getenv("CONJURER_LIBRARIAN_DB_PATH", r"C:\\Database\\chunks\\")
|
|
|
|
ENCODING = os.getenv("CONJURER_ENCODING", "utf-8")
|
|
CHUNK = os.getenv("CONJURER_LIBRARIAN_CHUNK", "_chunk.txt")
|
|
|
|
# DEPRECATED. The chunk files are now auto-discovered from DATABASE_PATH, so the
|
|
# thread count and the termination threshold both derive from what is actually
|
|
# on disk. This used to be BOTH "how many chunk files to read" AND "how many
|
|
# sentinels to wait for", which had to match exactly: set too low it silently
|
|
# skipped trailing chunks, set too high it referenced a nonexistent file whose
|
|
# producer crashed, starving the sentinel count and hanging the search forever.
|
|
# Kept only so old env files / references don't break; it no longer gates logic.
|
|
MAXTHREADS = int(os.getenv("CONJURER_LIBRARIAN_MAXTHREADS", "0"))
|
|
|
|
_sentinel = object()
|
|
# BOUNDED work queue. The producers stream the WHOLE DOI database (potentially
|
|
# tens of millions of lines across chunks) into this queue; the previous cap of
|
|
# 35_500_000 items was effectively unbounded (~3.5 GB of buffered lines), which
|
|
# OOM-killed the 1 GiB container mid-search. A small bound makes the producers
|
|
# backpressure to the consumers, keeping RAM to a few MB. The producer put below
|
|
# stays responsive to the TERM sentinel so a full queue can never deadlock it.
|
|
WORK_Q_SIZE = int(os.getenv("CONJURER_LIBRARIAN_WORKQ_SIZE", "100000"))
|
|
# Idle backstop: after this many consecutive empty seconds a consumer assumes
|
|
# the producers are done (or dead) and exits, so the search can never hang even
|
|
# if a sentinel were somehow lost. The primary, correct termination is still the
|
|
# sentinel count reaching the number of producers actually started.
|
|
EMPTY_LIMIT = 30
|
|
|
|
# Chunk files are named "<n>_chunk.txt" (suffix from CHUNK).
|
|
_CHUNK_RE = re.compile(r"^(\d+)" + re.escape(CHUNK) + r"$")
|
|
|
|
|
|
def discover_chunk_files(_logger):
|
|
"""Return the ``<n>_chunk.txt`` files present in DATABASE_PATH, numeric order.
|
|
|
|
Reading what exists (rather than files 0..MAXTHREADS-1) removes both historic
|
|
failure modes at once: no trailing chunk is ever silently skipped, and no
|
|
producer is ever pointed at a missing file, so it cannot crash before
|
|
emitting its sentinel and deadlock the consumers.
|
|
"""
|
|
try:
|
|
names = os.listdir(DATABASE_PATH)
|
|
except OSError as exc:
|
|
_logger.error("Cannot list DOI database dir %s: %s", DATABASE_PATH, exc)
|
|
return []
|
|
indexed = []
|
|
for name in names:
|
|
match = _CHUNK_RE.match(name)
|
|
if match:
|
|
indexed.append((int(match.group(1)), name))
|
|
indexed.sort()
|
|
ordered = [name for _, name in indexed]
|
|
_logger.info("Discovered %d chunk files in %s", len(ordered), DATABASE_PATH)
|
|
return ordered
|
|
|
|
|
|
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.
|
|
|
|
``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
|
|
# scraped DOI dumps) becomes U+FFFD instead of raising UnicodeDecodeError
|
|
# mid-file. Without it the readline() below would blow up, killing the
|
|
# 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:
|
|
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
|
|
# print flooded stdout / the log file with millions of lines.
|
|
if line_no % 500000 == 0:
|
|
_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 stop_event / the TERM sentinel so a full queue whose
|
|
# consumers have already finished can never deadlock us here.
|
|
stopped = False
|
|
while True:
|
|
try:
|
|
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)
|
|
stopped = True
|
|
break
|
|
except Empty:
|
|
pass
|
|
if stopped:
|
|
_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)
|
|
except Empty:
|
|
check = False
|
|
|
|
if check is _sentinel:
|
|
_logger.debug("Producer %s: TERM signal", filename)
|
|
control_q.put(check)
|
|
break
|
|
_logger.debug("Producer finished: %s", filename)
|
|
except Exception as exc: # pylint: disable=broad-except
|
|
# No per-file error (missing/unreadable chunk, a decode edge case that
|
|
# slips past errors="replace", anything unforeseen) may take the whole
|
|
# search down or crash the thread with a traceback. Log it and move on;
|
|
# the sentinel below still fires (finally), so the consumers' count stays
|
|
# correct and nothing deadlocks or silently loses a producer.
|
|
_logger.warning("Chunk %s failed, skipping rest of it: %s", filename, exc)
|
|
finally:
|
|
# ALWAYS emit exactly one sentinel per producer, on every exit path (EOF,
|
|
# early TERM, or crash). This is what lets the consumers count producers
|
|
# deterministically instead of hanging on a lost sentinel.
|
|
out_q.put(_sentinel)
|
|
|
|
|
|
# IMPORTANT!!! ONLY ONE CONSUMER THREAD AS WE ARE NOT PUTTING SENTINELS BACK
|
|
def consumer(in_q, control_q, doi, live_results, result_list, control_dict, expected_sentinels, no, _logger):
|
|
"""
|
|
Consumes items from an input queue and checks if DOI exists in the live results list.
|
|
|
|
Args:
|
|
in_q (Queue): Input queue.
|
|
control_q (Queue): Control queue.
|
|
doi (list): List of DOI to search for.
|
|
live_results (list): List to store the search results.
|
|
_logger: Logger object for logging.
|
|
"""
|
|
_logger.debug("Consumer %s started", no)
|
|
empty_counter = 0
|
|
alive_no = 0
|
|
# DOI -> result item, so a line is matched with one O(1) dict lookup instead
|
|
# of scanning every queried DOI. Items are shared with result_list, so
|
|
# setting exists here is seen by everyone.
|
|
doi_index = {item["DOI"]: item for item in result_list}
|
|
while True:
|
|
try:
|
|
data = in_q.get(block=True, timeout = 1)
|
|
if data is _sentinel:
|
|
control_dict["sentinels"] += 1
|
|
_logger.debug(
|
|
"Consumer %s: producer done (%d/%d)",
|
|
no, control_dict["sentinels"], expected_sentinels,
|
|
)
|
|
|
|
else:
|
|
empty_counter = 0
|
|
alive_no += 1
|
|
|
|
# Each DB line is a DOI (optionally followed by metadata). Match
|
|
# the WHOLE first token exactly - the old `item["DOI"] in data`
|
|
# was a substring test, so a DOI that is a prefix of a longer one
|
|
# (10.1/1 vs 10.1/12) produced a false 'exists' hit.
|
|
parts = data.split()
|
|
line_doi = parts[0] if parts else ""
|
|
item = doi_index.get(line_doi)
|
|
if item is not None and not item["exists"]:
|
|
# HIT can fire thousands of times for a deep search -> DEBUG.
|
|
_logger.debug("HIT %s (consumer %s)", line_doi, no)
|
|
item["exists"] = True
|
|
live_results.append(item)
|
|
# All found? Signal producers to stop early (rare -> cheap).
|
|
if all(it["exists"] for it in result_list):
|
|
control_q.put(_sentinel)
|
|
except Empty:
|
|
empty_counter += 1
|
|
time.sleep(1)
|
|
# Order matters: the >EMPTY_LIMIT break must be checked BEFORE the
|
|
# lesser threshold, otherwise (as in the original) the first branch
|
|
# always wins and the break is dead code, leaving the sentinel count
|
|
# as the only exit - which is exactly what used to hang the search.
|
|
if empty_counter > EMPTY_LIMIT:
|
|
_logger.debug("Consumer %s finished (idle backstop)", no)
|
|
break
|
|
if empty_counter > 5:
|
|
time.sleep(2)
|
|
|
|
if control_dict["sentinels"] >= expected_sentinels:
|
|
_logger.debug("Consumer %s: all producers finished", no)
|
|
break
|
|
|
|
|
|
|
|
def search_for_doi(doi, live_results, _logger, stop_event=None, resume=None):
|
|
"""Search for DOI in live_results, resumably.
|
|
|
|
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:
|
|
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.
|
|
chunk_files = discover_chunk_files(_logger)
|
|
expected = len(chunk_files)
|
|
if expected == 0:
|
|
_logger.error(
|
|
"No '<n>%s' chunk files in %s - DOI search cannot run", CHUNK, DATABASE_PATH
|
|
)
|
|
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.debug("Consumer thread created")
|
|
threads.append(t_cons)
|
|
for filename in chunk_files:
|
|
_logger.debug("Creating worker thread for %s", filename)
|
|
threads.append(
|
|
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()
|
|
interrupted = bool(stop_event and stop_event.is_set())
|
|
return result_list, positions, interrupted
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import logging
|
|
|
|
logger = logging.getLogger()
|
|
logger.setLevel("DEBUG")
|
|
h1 = logging.StreamHandler()
|
|
logger.addHandler(h1)
|
|
logger.info("TEST RUN")
|
|
live_result = []
|
|
logger.info(
|
|
search_for_doi(
|
|
[
|
|
("10.1002/9781118786352.wbieg0998.pub2", "DATA"),
|
|
("10.1002/j.2050-0416.2002.tb00563.x", "DATA"),
|
|
("10.1111/j.1365-2958.1994.tb00448.x", "DATA"),
|
|
("10.2165/00128415-200309690-00017", "DATA"),
|
|
("10.5772/48313", "DATA"),
|
|
("10.15803/ijnc.7.2_419", "DATA"),
|
|
("10.1051/0004-6361/201321596e", "DATA"),
|
|
("10.2307/40835941", "DATA"),
|
|
],
|
|
live_result,
|
|
logger,
|
|
)
|
|
)
|
|
logger.info(live_result)
|