supervisor.py

#!/usr/bin/env python3

r"""
Usage: supervisor.py <database_file> <input_dir> <worker_socket>

This script has two responsibilities:

  * Watch all files in a directory for tasks to give to workers, and
    store those tasks in a priority queue.

  * Listen for requests from workers and hand out tasks.

The input directory is expected to contain files conforming to the
following:

  * The name of the file must match [^.]+\.[0-9]+\.txt(\.part)?

  * The numeric part of the file name should be an integer that
    increases whenever the file is recreated (such as the current Unix
    time, but the supervisor doesn't require that).

  * The .part suffix should be present if data is still being written
    to that file. When the file is complete, it should be renamed to
    lose the .part suffix.

  * The file should consist of lines that look like

    attribute_path old_ver new_ver optional_url

    (though in reality, the supervisor only cares that there are at
    least two fields; everything past the first space is passed on
    to the workers without the supervisor getting involved).

These files are produced by fetcher processes, and are referred to
as fetcher files or fetcher runs throughout this script.

Tasks are distributed to workers based on how recently a task
corresponding to the same attribute path has been dispatched. Brand
new attribute paths are given highest priority; attribute paths that
have run most recently are given lowest priority.

The worker socket is a Unix stream socket that uses a line-based
protocol for communication. The following commands are expected:

    READY
    DONE <attribute_path> <exit_code>

The possible responses to either command are:

    NOJOBS
    JOB <attribute_path> <rest_of_line>

When a worker receives NOJOBS, it is expected to wait some period of
time before asking for jobs again (with READY).

When a worker receives JOB, it is expected to complete that job and
then send a new DONE request to get the next job.

This code has a unit test suite in supervisor_test.py.
"""

import asyncio
import contextlib
import functools
import hashlib
import os
import pathlib
import sqlite3
import sys
import time
from collections.abc import Callable, Generator, Iterable
from typing import (
    Concatenate,
    ParamSpec,
    Protocol,
    TextIO,
    TypeVar,
    overload,
)

import asyncinotify

P = ParamSpec("P")
S = TypeVar("S")
T = TypeVar("T")
A_co = TypeVar("A_co", covariant=True)
B_co = TypeVar("B_co", covariant=True)
StorageSelf = TypeVar("StorageSelf", bound="Storage")


class Wrapped(Protocol[A_co, B_co]):
    """Protocol for result of `functools.wraps`"""

    @property
    def __wrapped__(self) -> A_co: ...

    @property
    def __call__(self) -> B_co: ...


CurMethod = Callable[Concatenate[S, sqlite3.Cursor, P], T]
WrappedCurMethod = Wrapped[CurMethod[S, P, T], Callable[P, T]]


READY = b"READY\n"
DONE = b"DONE "
JOB = b"JOB "
NOJOBS = b"NOJOBS\n"


def jitter_hash(data: str) -> int:
    """
    Deterministically produce a 32-bit integer from a string, not correlated
    with the sort order of the string.

    The implementation shouldn't matter but if it ever changes, the database
    queue will need to be recreated.
    """
    h = hashlib.md5(data.encode(), usedforsecurity=False)
    return int.from_bytes(h.digest()[:4], byteorder="little", signed=True)


class Storage:
    """
    Abstracts over the database that stores the supervisor state.

    This state comprises three tables: `fetcher_runs`, `queue`, and
    `log`. (A fourth table, `fetchers`, is just a simple string
    deduplication table for fetcher names.) The `fetcher_runs` table is
    a record of the distinct fetcher files captured in the database. The
    `log` table stores the times and exit code associated with the last
    job dispatched for each attribute path. The `queue` table is the
    most interesting; it stores the contents of the fetcher files along
    with a bit indicating if that item has been dispatched to a worker.

    The `dequeue` operation has sub-linear performance only because the
    `queue_order_index` index on `queue` aligns with the query used by
    `dequeue`. This index relies on a denormalized column
    `last_started`, which is populated from `log`.`started` via the
    triggers `queue_insert_started`, `log_insert_started`, and
    `log_update_started`.
    """

    def __init__(self, conn: sqlite3.Connection):
        self._conn = conn

    @overload
    @staticmethod
    def _cursor_method(
        fun: CurMethod[StorageSelf, P, T],
    ) -> WrappedCurMethod[StorageSelf, P, T]: ...

    @overload
    @staticmethod
    def _cursor_method(
        *,
        transaction: bool = False,
    ) -> Callable[
        [CurMethod[StorageSelf, P, T]], WrappedCurMethod[StorageSelf, P, T]
    ]: ...

    # NOTE: mypy <1.6 claims this implementation doesn't match the first
    # overload; it's wrong.
    @staticmethod
    def _cursor_method(
        fun: CurMethod[StorageSelf, P, T] | None = None,
        transaction: bool = False,
    ) -> (
        WrappedCurMethod[StorageSelf, P, T]
        | Callable[[CurMethod[StorageSelf, P, T]], WrappedCurMethod[StorageSelf, P, T]]
    ):
        def decorator(
            fun: CurMethod[StorageSelf, P, T],
        ) -> WrappedCurMethod[StorageSelf, P, T]:
            if transaction:

                def wrapper(self: StorageSelf, *args: P.args, **kwargs: P.kwargs) -> T:
                    with self._conn:
                        with contextlib.closing(self._conn.cursor()) as cur:
                            cur.execute("BEGIN")
                            return fun(self, cur, *args, **kwargs)

            else:

                def wrapper(self: StorageSelf, *args: P.args, **kwargs: P.kwargs) -> T:
                    with contextlib.closing(self._conn.cursor()) as cur:
                        return fun(self, cur, *args, **kwargs)

            return functools.wraps(fun)(wrapper)  # type: ignore

        return decorator if fun is None else decorator(fun)

    @_cursor_method
    def upgrade(self, cur: sqlite3.Cursor, version: int = 2) -> None:
        """
        Create or update the database schema.

        The database stores a version number to allow for future schema
        modifications. This function contains scripts for migrating to
        each version number from its immediate predecessor. Every script
        from the current database version to the target version is run.

        For testing purposes, the optional `version` parameter can be
        used to create a database at an earlier version of the schema.

        The `Storage` class documentation explains the current schema.
        """

        cur.execute("PRAGMA user_version")
        user_version = cur.fetchone()[0]
        if user_version < 1 <= version:
            # Here is as good a place as any: we're using backtick quotation
            # extensively in these statements as a way to protect against
            # having a token mistaken for a keyword. The standard SQL quotation
            # style for this purpose is double-quote, which SQLite supports,
            # but SQLite also interprets double-quoted identifiers as strings
            # if they are in a position where an identifier isn't expected;
            # this can suppress or obfuscate error messages when one has made a
            # syntax error. Backticks are nonstandard but used in MySQL and
            # supported by SQLite, and SQLite doesn't try to interpret
            # backtick-quoted tokens as anything other than identifiers.
            cur.executescript(
                """
                BEGIN;

                CREATE TABLE `fetchers` (
                    `fetcher_id` `integer` PRIMARY KEY AUTOINCREMENT,
                    `name` `text` NOT NULL UNIQUE
                ) STRICT;

                CREATE TABLE `fetcher_runs` (
                    `fetcher_id` `integer` NOT NULL REFERENCES `fetchers` (`fetcher_id`),
                    `run_started` `integer` NOT NULL,
                    `is_complete` `integer` NOT NULL DEFAULT 0,
                    PRIMARY KEY (`fetcher_id`, `run_started`)
                ) STRICT;

                CREATE INDEX `fetcher_run_started_index`
                ON `fetcher_runs` (`run_started`, `fetcher_id`);

                CREATE TRIGGER `fetcher_runs_delete`
                AFTER DELETE
                ON `fetcher_runs`
                BEGIN
                    DELETE FROM `fetchers`
                    WHERE `fetcher_id` = OLD.`fetcher_id`
                    AND (
                        SELECT `fetcher_id`
                        FROM `fetcher_runs`
                        WHERE `fetcher_id` = OLD.`fetcher_id`
                    ) IS NULL;
                END;

                CREATE TABLE `queue` (
                    `fetcher_id` `integer` NOT NULL,
                    `fetcher_run_started` `integer` NOT NULL,
                    `attr_path` `text` NOT NULL,
                    `payload` `text` NOT NULL,
                    `is_dequeued` `integer` NOT NULL DEFAULT 0,
                    `last_started` `integer`,
                    PRIMARY KEY (`fetcher_run_started`, `fetcher_id`, `attr_path`),
                    FOREIGN KEY (`fetcher_id`, `fetcher_run_started`)
                        REFERENCES `fetcher_runs` (`fetcher_id`, `run_started`)
                        ON DELETE CASCADE
                ) STRICT;

                CREATE INDEX `queue_order_index`
                ON `queue` (
                    `last_started` ASC,
                    `attr_path`,
                    `fetcher_id`,
                    `fetcher_run_started` DESC
                );

                CREATE INDEX `queue_attr_path_index`
                ON `queue` (`attr_path` ASC);

                CREATE TABLE `log` (
                    `attr_path` `text` PRIMARY KEY,
                    `started` `integer` NOT NULL,
                    `finished` `integer`,
                    `exit_code` `integer`
                ) STRICT, WITHOUT ROWID;

                CREATE TRIGGER `queue_insert_started`
                AFTER INSERT
                ON `queue`
                BEGIN
                    UPDATE `queue` SET
                        `last_started` = `started`
                    FROM `log`
                    WHERE
                        `log`.`attr_path` = NEW.`attr_path` AND
                        `queue`.`rowid` = NEW.`rowid`;
                END;

                CREATE TRIGGER `log_insert_started`
                AFTER INSERT
                ON `log`
                BEGIN
                    UPDATE `queue` SET
                        `last_started` = NEW.`started`
                    WHERE `queue`.`attr_path` = NEW.`attr_path`;
                END;

                CREATE TRIGGER `log_update_started`
                AFTER UPDATE OF `started`
                ON `log`
                BEGIN
                    UPDATE `queue` SET
                        `last_started` = NEW.`started`
                    WHERE `queue`.`attr_path` = NEW.`attr_path`;
                END;

                COMMIT;
                """
            )
        if user_version < 2 <= version:
            # We want to add some disorder to the initial order of packages;
            # dispatching packages that are alphabetically adjacent increases
            # the chances of parallel duplicates or getting stuck on a large
            # block of time-consuming packages.
            #
            # Unfortunately, to apply this jitter retroactively we need to
            # delete most of the rows already in the database.
            cur.executescript(
                """
                BEGIN;

                DELETE FROM `fetcher_runs`;
                DELETE FROM `log`;

                ALTER TABLE `queue`
                ADD COLUMN `order_jitter` `integer`;

                DROP INDEX `queue_order_index`;

                CREATE INDEX `queue_order_index`
                ON `queue` (
                    `last_started` ASC,
                    `order_jitter`,
                    `attr_path`,
                    `fetcher_id`,
                    `fetcher_run_started` DESC
                );

                COMMIT;
                """
            )
        cur.execute(f"PRAGMA user_version = {version}")

    @_cursor_method
    def get_fetcher_runs(self, cur: sqlite3.Cursor) -> dict[tuple[str, int], bool]:
        """Return a set of fetcher runs known to the database."""
        cur.execute(
            """
            SELECT `name`, `run_started`, `is_complete`
            FROM `fetchers`
            JOIN `fetcher_runs` USING (`fetcher_id`)
            """
        )
        return {(r[0], r[1]): r[2] for r in cur}

    @_cursor_method
    def delete_fetcher_run(
        self, cur: sqlite3.Cursor, fetcher: str, run_started: int
    ) -> None:
        """Delete a fetcher run and all of its queue items."""
        cur.execute(
            """
            DELETE FROM `fetcher_runs`
            WHERE `fetcher_id` = (SELECT `fetcher_id` FROM `fetchers` WHERE `name` = ?)
            AND `run_started` = ?
            """,
            (fetcher, run_started),
        )

    @_cursor_method(transaction=True)
    def delete_fetcher_runs(
        self, cur: sqlite3.Cursor, fetchers: Iterable[tuple[str, int]]
    ) -> None:
        """Delete multiple fetcher runs and their queue items."""
        for fetcher, run_started in fetchers:
            self.delete_fetcher_run.__wrapped__(self, cur, fetcher, run_started)

    @_cursor_method(transaction=True)
    def upsert_fetcher_run(
        self,
        cur: sqlite3.Cursor,
        fetcher: str,
        run_started: int,
        is_complete: bool,
    ) -> None:
        """Add or update a fetcher."""
        cur.execute("INSERT OR IGNORE INTO `fetchers` (`name`) VALUES (?)", (fetcher,))
        cur.execute(
            """
            INSERT INTO `fetcher_runs` (`fetcher_id`, `run_started`, `is_complete`)
            VALUES ((SELECT `fetcher_id` FROM `fetchers` WHERE `name` = ?), ?, ?)
            ON CONFLICT DO UPDATE SET `is_complete` = excluded.`is_complete`
            """,
            (fetcher, run_started, is_complete),
        )

    @_cursor_method(transaction=True)
    def enqueue(
        self,
        cur: sqlite3.Cursor,
        fetcher: str,
        run_started: int,
        entries: list[tuple[str, str]],
    ) -> None:
        """
        Add entries for a given fetcher to the queue.

        The same attribute paths can appear multiple times in the queue
        with different payloads, but only once per fetcher run. Fetcher
        files shouldn't contain more than one line for a given attribute
        path, but if they do, the later line overwrites the earlier one.
        """
        cur.executemany(
            """
            INSERT INTO `queue` (`fetcher_id`, `fetcher_run_started`, `attr_path`, `payload`, `order_jitter`)
            SELECT `fetcher_id`, ?, ?, ?, ?
            FROM `fetchers`
            WHERE `name` = ?
            ON CONFLICT DO UPDATE SET `payload` = excluded.`payload`
            """,
            ((run_started, a, p, jitter_hash(a), fetcher) for a, p in entries),
        )

    @_cursor_method(transaction=True)
    def dequeue(self, cur: sqlite3.Cursor, start_time: int) -> tuple[str, str] | None:
        """
        Pull one entry from the top of the queue.

        Returns a tuple (attribute path, payload), or None if nothing is
        currently available in the queue. If an entry is dequeued, a log
        record for this entry will be marked as started as of
        `start_time`.

        Most of the time, if a job for an attribute path was started but
        has not yet finished, any queue entries for that same path will
        be skipped. However, in case a worker dies or fails to report
        back, after 12 hours such entries are eligible again.
        `start_time` is used to determine if this 12-hour exclusion
        period has ended. (These details are only likely to be relevant
        when the queue is very small, like at the beginning or the end
        of a run.)
        """
        cur.execute(
            """
            SELECT
                `fetcher_id`,
                `attr_path`,
                `payload`
            FROM `queue`
            LEFT JOIN `log` USING (`attr_path`)
            GROUP BY `last_started`, `order_jitter`, `attr_path`, `fetcher_id`
            HAVING `is_dequeued` = 0
            AND (`started` IS NULL OR `finished` IS NOT NULL OR `started` + 43200 < ?)
            AND (TRUE OR `max`(`fetcher_run_started`))
            ORDER BY `last_started` ASC, `order_jitter`, `attr_path`, `fetcher_id`
            LIMIT 1
            """,
            (start_time,),
        )
        # NOTE: The `max` call in the above query triggers a nonstandard SQLite
        # behavior; see <https://www.sqlite.org/lang_select.html#bareagg>.
        # This behavior is critical to the correctness of this query. We don't
        # actually need the value of `max`(), though, so we tuck it into the
        # HAVING clause in a position where it can't have any other effect.

        row: tuple[int, str, str] | None = cur.fetchone()
        result = None
        if row is not None:
            fetcher_id, attr_path, payload = row
            cur.execute(
                """
                UPDATE `queue` SET `is_dequeued` = 1
                WHERE `fetcher_id` = ? AND `attr_path` = ?
                """,
                (fetcher_id, attr_path),
            )
            cur.execute(
                """
                INSERT INTO `log` (`attr_path`, `started`) VALUES (?, ?)
                ON CONFLICT DO UPDATE SET
                    `started` = excluded.`started`,
                    `finished` = NULL,
                    `exit_code` = NULL
                """,
                (attr_path, start_time),
            )
            result = attr_path, payload
        return result

    @_cursor_method
    def finish(
        self, cur: sqlite3.Cursor, attr_path: str, finish_time: int, exit_code: int
    ) -> None:
        """Log the completion of a dequeued entry."""
        cur.execute(
            "UPDATE `log` SET `finished` = ?, `exit_code` = ? WHERE `attr_path` = ?",
            (finish_time, exit_code, attr_path),
        )


async def listen_for_workers(storage: Storage, socket_path: pathlib.Path) -> None:
    """Open a Unix stream socket and handle requests from workers."""

    async def worker_connected(
        reader: asyncio.StreamReader, writer: asyncio.StreamWriter
    ) -> None:
        try:
            while True:
                line = await reader.readline()
                if line == b"":
                    break
                now = int(time.time())
                do_dequeue = False
                if line == READY:
                    do_dequeue = True
                elif line.startswith(DONE):
                    parts = line.split(b" ")
                    attr_path = parts[1].decode()
                    exit_code = int(parts[2])
                    storage.finish(attr_path, now, exit_code)
                    do_dequeue = True
                else:
                    print(f"Unexpected command from worker: {line!r}")
                    break

                if do_dequeue:
                    entry = storage.dequeue(now)
                    if entry:
                        writer.write(
                            b"".join(
                                [
                                    JOB,
                                    entry[0].encode(),
                                    b" ",
                                    entry[1].encode(),
                                    b"\n",
                                ]
                            )
                        )
                    else:
                        writer.write(NOJOBS)
                    await writer.drain()
        finally:
            writer.close()
            await writer.wait_closed()

    server = await asyncio.start_unix_server(worker_connected, socket_path)
    await server.serve_forever()


class FetcherDataWatcher:
    """
    Monitors a directory containing fetcher files and syncs them to
    storage.
    """

    _dir_events = asyncinotify.Mask.CREATE | asyncinotify.Mask.DELETE
    _file_events = asyncinotify.Mask.MODIFY | asyncinotify.Mask.MOVE_SELF

    class _FileDeleted(Exception):
        """A fetcher file was deleted."""

    class _FileMoved(Exception):
        """A fetcher file was moved."""

    def __init__(
        self,
        storage: Storage,
        fetcher_data_path: pathlib.Path,
        inotify: asyncinotify.Inotify,
    ) -> None:
        self._storage = storage
        self._fetcher_data_path = fetcher_data_path
        self._inotify = inotify
        self._gtors: dict[tuple[str, int], Generator[None, None, None]] = {}

    async def watch(self) -> None:
        """Start the watcher."""

        self._inotify.add_watch(self._fetcher_data_path, self._dir_events)

        try:
            known_fetcher_runs = self._storage.get_fetcher_runs()
            for path in self._fetcher_data_path.iterdir():
                if (that := self._parse_fetcher_filename(path.name)) is None:
                    continue
                name, run_started, is_complete = that
                if not known_fetcher_runs.pop((name, run_started), False):
                    self._on_fetcher(path, name, run_started, is_complete)
            self._storage.delete_fetcher_runs(known_fetcher_runs.keys())

            async for event in self._inotify:
                if event.path is None:
                    continue
                if (that := self._parse_fetcher_filename(event.path.name)) is None:
                    continue
                name, run_started, is_complete = that
                with contextlib.suppress(KeyError):
                    match event.mask:
                        case asyncinotify.Mask.CREATE:
                            self._on_fetcher(event.path, name, run_started, is_complete)
                        case asyncinotify.Mask.DELETE:
                            if not is_complete:
                                self._close_fetcher(
                                    name, run_started, self._FileDeleted()
                                )
                            self._storage.delete_fetcher_run(name, run_started)
                        case asyncinotify.Mask.MODIFY:
                            self._gtors[(name, run_started)].send(None)
                        case asyncinotify.Mask.MOVE_SELF:
                            self._close_fetcher(name, run_started, self._FileMoved())
        finally:
            with contextlib.suppress(KeyError):
                while True:
                    self._gtors.popitem()[1].close()

    def _on_fetcher(
        self,
        path: pathlib.Path,
        name: str,
        run_started: int,
        is_complete: bool,
    ) -> None:
        watch = None
        try:
            if not is_complete:
                watch = self._inotify.add_watch(path, self._file_events)
            file = path.open(encoding="utf-8")
        except FileNotFoundError:
            return
        self._storage.upsert_fetcher_run(name, run_started, is_complete)
        gtor = self._read_fetcher_file(file, watch, name, run_started)
        gtor.send(None)
        if is_complete:
            gtor.close()
        else:
            self._gtors[(name, run_started)] = gtor

    def _close_fetcher(self, name: str, run_started: int, ex: Exception) -> None:
        with contextlib.suppress(StopIteration):
            self._gtors.pop((name, run_started)).throw(ex)

    def _parse_fetcher_filename(self, name: str) -> tuple[str, int, bool] | None:
        match name.split("."):
            case [stem, run_started, "txt", "part"]:
                return stem, int(run_started), False
            case [stem, run_started, "txt"]:
                return stem, int(run_started), True
        return None

    def _read_fetcher_file(
        self,
        file: TextIO,
        watch: asyncinotify.Watch | None,
        name: str,
        run_started: int,
    ) -> Generator[None, None, None]:
        with file:
            try:
                while True:
                    self._storage.enqueue(
                        name,
                        run_started,
                        (yield from self._read_fetcher_lines(file)),
                    )
            except self._FileDeleted:
                pass
            except self._FileMoved:
                try:
                    target_path_stat = (
                        self._fetcher_data_path / f"{name}.{run_started}.txt"
                    ).stat()
                except FileNotFoundError:
                    pass
                else:
                    if target_path_stat.st_ino == os.stat(file.fileno()).st_ino:
                        for entries in self._read_fetcher_lines(file):
                            if entries is not None:
                                self._storage.enqueue(name, run_started, entries)
                            break
                        self._storage.upsert_fetcher_run(name, run_started, True)
                        assert watch is not None
                        self._inotify.rm_watch(watch)
                        return
                self._storage.delete_fetcher_run(name, run_started)

    def _read_fetcher_lines(
        self, file: TextIO
    ) -> Generator[None, None, list[tuple[str, str]]]:
        """
        Read all available complete lines from an open fetcher file.

        This is a generator, but not one that yields each line. It will
        *return* all lines as a non-empty list. If no complete lines are
        available, however, it will yield. Calling code should reenter
        the generator when more content becomes available, or use `yield
        from` to pass that responsibility outward.
        """
        entries: list[tuple[str, str]] = []
        while True:
            cookie = file.tell()
            line = file.readline()
            if line == "" or line[-1] != "\n":
                file.seek(cookie)
                if entries:
                    return entries
                yield
                continue
            match line.strip().split(" ", maxsplit=1):
                case [attr_path, payload]:
                    entries.append((attr_path, payload))
                case _:
                    print(f"Unexpected line in {file.name!r}: {line!r}")


async def main(
    db_path: pathlib.Path,
    fetcher_data_path: pathlib.Path,
    socket_path: pathlib.Path,
) -> None:
    """Run all supervisor responsibilities."""
    fetcher_data_path.mkdir(parents=True, exist_ok=True)
    with contextlib.closing(sqlite3.connect(db_path, isolation_level=None)) as conn:
        conn.execute("PRAGMA foreign_keys = ON")
        storage = Storage(conn)
        storage.upgrade()
        with asyncinotify.Inotify() as inotify:
            watcher = FetcherDataWatcher(storage, fetcher_data_path, inotify)
            await asyncio.gather(
                listen_for_workers(storage, socket_path),
                watcher.watch(),
            )


if __name__ == "__main__":
    asyncio.run(main(*[pathlib.Path(arg) for arg in sys.argv[1:4]]))