conjurer/conjurer_librarian/search_bot.py

"""
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, 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()
WORK_Q_SIZE = 35500000
# 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):
    """
    Produces items from the output queue and puts them into the control queue.

    Args:
        out_q (Queue): Output queue.
        control_q (Queue): Control queue.
        filename (str): Name of the file.
        _logger: Logger object for logging.
    """
    try:
        with open(DATABASE_PATH + filename, "r", encoding=ENCODING) as operated_file:
            print(f"Worker {filename} ")
            line_no = 0
            while True:
                line = operated_file.readline()
                line_no += 1
                print(f"\t \t \t \t \t \t W{filename}{line_no}\r", end="")

                if not line:
                    print(f"EOF {filename}")
                    break

                out_q.put(line)
                try:
                    check = control_q.get(block=False)
                except Empty:
                    check = False

                if check is _sentinel:
                    print("TERM signal received")
                    control_q.put(check)
                    break
            print(f"Worker finished: {filename}")
    except OSError as exc:
        # A missing or unreadable chunk must not take the whole search down with
        # it - log and move on. The sentinel below still fires (finally), so the
        # consumers' count stays correct and nothing deadlocks.
        _logger.warning("Chunk %s unreadable, skipping: %s", filename, exc)
        print(f"Worker {filename} failed: {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.
    """
    print(f"Consumer thread started: {no} no")
    empty_counter = 0
    alive_no = 0
    while True:
        done_check = True
        try:
            data = in_q.get(block=True, timeout = 1)
            if data is _sentinel:
                control_dict["sentinels"] += 1
                print(f"Workers finished: {control_dict['sentinels']} reported by consumer {no}")

            else:
                empty_counter = 0
                alive_no += 1
                print(f"C{no}__{alive_no}\r", end="")

                for item in result_list:
                    if item["DOI"] in data and not item["exists"]:
                        print(f"HIT in {no} content {data[0]} line {data[1]} file {data[2]} {item['exists']}")
                        _logger.info(data)
                        _logger.info("HIT")
                        item["exists"] = True
                        live_results.append(item)
                    done_check = done_check and item["exists"]
                if done_check:
                    control_q.put(_sentinel)
        except Empty:
            empty_counter += 1
            time.sleep(1)
            print(f"Consumer {no} empty")
            # 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:
                print(f"Consumer thread finished {no} (idle backstop)")
                break
            if empty_counter > 5:
                print(f"Consumer {no} empty lvl 2")
                time.sleep(2)

        if control_dict["sentinels"] >= expected_sentinels:
            _logger.info(f"All workers finished {no}")
            break



def search_for_doi(doi, live_results, _logger):
    """
    Search for DOI in live_results using _logger for logging.

    Args:
        doi (list): List of DOI to search for.
        live_results (list): List to store the search results.
        _logger: Logger object for logging.
    """
    control_dict = {"sentinels":0}
    result_list = []
    threads = []
    work_q = Queue(maxsize=WORK_Q_SIZE)
    control_q = Queue()

    for item in doi:
        result_list.append({"DOI": item[0], "exists": False, "data": item[1]})

    # 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

    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")
        threads.append(t_cons)
    for filename in chunk_files:
        _logger.info("Creating worker thread for %s", filename)
        threads.append(
            Thread(target=producer, args=(work_q, control_q, filename, _logger))
        )
    for worker in threads:
        worker.start()
    for worker in threads:
        worker.join()
    return result_list


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)