conjurer_librarian/search_bot.py

@@ -21,7 +21,7 @@ Global Variables:
 # TODO: Wpiemdolić to wszystko w klasę z loggerem przysłanym z góry
 import os
 import re
-from queue import Empty, Queue
+from queue import Empty, Full, Queue
 from threading import Thread
 import time
 q = Queue()
@@ -44,7 +44,13 @@ CHUNK = os.getenv("CONJURER_LIBRARIAN_CHUNK", "_chunk.txt")
 MAXTHREADS = int(os.getenv("CONJURER_LIBRARIAN_MAXTHREADS", "0"))
 
 _sentinel = object()
-WORK_Q_SIZE = 35500000
+# 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
@@ -107,7 +113,26 @@ def producer(out_q, control_q, filename, _logger):
                     print(f"EOF {filename}")
                     break
 
-                out_q.put(line)
+                # Backpressure-safe put onto the BOUNDED queue: wait for room,
+                # but keep polling 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:
+                        try:
+                            if control_q.get(block=False) is _sentinel:
+                                control_q.put(_sentinel)
+                                stopped = True
+                                break
+                        except Empty:
+                            pass
+                if stopped:
+                    print("TERM signal received (queue full)")
+                    break
+
                 try:
                     check = control_q.get(block=False)
                 except Empty:

tests/unit/test_search_bot.py

@@ -131,3 +131,22 @@ def test_discover_chunk_files_sorted_numerically(tmp_path, monkeypatch):
 
     # Numeric order (10 after 2, not lexicographic), and non-chunk files ignored.
     assert found == ["0_chunk.txt", "1_chunk.txt", "2_chunk.txt", "10_chunk.txt"]
+
+
+def test_bounded_queue_does_not_deadlock_on_early_termination(tmp_path, monkeypatch):
+    # The OOM fix bounds the work queue. That means a producer can block on a
+    # FULL queue - and if the consumers have already finished (all DOIs found)
+    # it must notice the TERM sentinel instead of hanging forever. Tiny queue +
+    # target on the first line + thousands of trailing decoys the producer still
+    # holds is exactly that situation.
+    monkeypatch.setattr(search_bot, "DATABASE_PATH", str(tmp_path) + "/")
+    monkeypatch.setattr(search_bot, "WORK_Q_SIZE", 3)  # force the producer to block
+    target = "10.1234/found.on.line.one"
+    lines = [target + "\n"] + [f"10.0000/decoy-{i}\n" for i in range(5000)]
+    (tmp_path / "0_chunk.txt").write_text("".join(lines), encoding="utf-8")
+
+    finished, result = _run_bounded([(target, "DATA")], timeout=20)
+
+    assert finished, "a full bounded queue deadlocked the producer on early termination"
+    hit = [r for r in result if r["DOI"] == target and r["exists"]]
+    assert hit, "the target on the first line should have been found"