Payloads¶

def packed(self) -> Payload:
    """This payload with `data` compressed by its codec (itself if it already is)."""
    if self.raw_length is not None:
        return self
    if self.codec not in CODECS:
        raise ValueError(f"unknown codec {self.codec!r}; one of {sorted(CODECS)}")
    return Payload(_compress(self.codec, self.data, self.dictionary), self.kind, self.series, self.codec,
                   None, len(self.data))

def put(self, session: Any, payloads: Iterable[Payload | PayloadGroup]) -> list[int]:
    """Store `payloads` inside the caller's transaction; return their ids, in order: one per
    `Payload`, one per part of a `PayloadGroup` (an empty group stores nothing).

    The bytes are on disk (fdatasync) when this returns, so the caller's commit never
    points to lost bytes. If that transaction does not commit, the volumes go back to
    the size they had before it.
    """
    payloads = list(payloads)
    if not payloads:
        return []
    connection = session.connection()  # begins the caller's transaction
    tx = self._transaction(session)
    cursor = connection.connection.dbapi_connection.cursor()
    try:
        pending: dict[_Volume, list] = {}
        placed = []
        for payload in payloads:
            sizes = None
            if isinstance(payload, PayloadGroup):
                payload, sizes = payload.packed()
                if not sizes:
                    continue
            payload = payload.packed()  # done already when it came from a worker
            data, codec = payload.data, CODECS[payload.codec]
            size = len(data)
            volume = self._volume_for(tx, cursor, payload.series, NEEDLE + size)
            crc = zlib.crc32(data)
            pending.setdefault(volume, [volume.end, []])[1] += (
                _NEEDLE.pack(crc, size, payload.raw_length, codec, 0),
                data,
            )
            row = (volume.id, volume.end, size, payload.raw_length, codec, crc, payload.kind)
            if sizes is None:
                placed.append((*row, None, None))
            else:
                placed += [(*row, start, length) for start, length in zip(accumulate(sizes, initial=0), sizes)]
            volume.end += NEEDLE + size
        for volume, (offset, buffers) in pending.items():
            try:
                _write_all(volume.fd, buffers, offset)
            except OSError as error:  # a full disk: the caller's rollback cuts what got in
                raise OSError(error.errno, f"cannot write payloads to {self.path(volume.id)}: "
                                           f"{error.strerror}") from error
        _fault("appended")
        for volume in pending:
            _datasync(volume.fd)
            cursor.execute("UPDATE ms_volumes SET size = ? WHERE id = ?", (volume.end, volume.id))
        if tx.dir_dirty:  # a new volume file must survive a crash too
            _fsync_dir(self.directory)
            tx.dir_dirty = False
        _fault("synced")
        ids = []
        for row in placed:
            cursor.execute(
                'INSERT INTO ms_payloads (volume_id, "offset", length, raw_length, codec, crc32, kind,'
                " part_offset, part_length) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
                row,
            )
            ids.append(cursor.lastrowid)
        return ids
    finally:
        cursor.close()

def referencing_columns(self) -> list[tuple[str, str]]:
    """(table, column) of every app column that holds payload ids.

    Read from the schema (`PRAGMA foreign_key_list`): an app declares
    `REFERENCES ms_payloads(id)` and registers nothing with ms_flow.
    """
    with self.engine.connect() as conn:
        tables = [r[0] for r in conn.exec_driver_sql("SELECT name FROM sqlite_master WHERE type = 'table'")]
        return [
            (table, fk[3])
            for table in tables
            for fk in conn.exec_driver_sql(f'PRAGMA foreign_key_list("{table}")')
            if fk[2] == "ms_payloads"
        ]

def get_many(self, ids: Iterable[int]) -> dict[int, bytes]:
    """Bytes by id. Read in (volume, offset) order; unknown ids raise KeyError."""
    wanted = list(dict.fromkeys(int(i) for i in ids))
    located = self._locate(wanted)
    missing = [i for i in wanted if i not in located]
    if missing:
        raise KeyError(f"unknown payload ids: {missing[:10]}")
    found, needle, raw = {}, None, b""
    for i in sorted(located, key=lambda i: located[i][:2]):
        location = located[i]
        if location[:2] != needle:  # the parts of a group: one read, one decompression
            needle, raw = location[:2], self._read(i, location)
        found[i] = _part(raw, location)
    return found

def kinds(self, ids: Iterable[int]) -> dict[int, str]:
    """The `kind` each payload was stored with (a format: `sdf`, `pdbqt`...), by id; unknown ids are left out."""
    ids = list(dict.fromkeys(int(i) for i in ids))
    found = {}
    with self.engine.connect() as conn:
        for start in range(0, len(ids), LOCATE_CHUNK):
            chunk = ids[start : start + LOCATE_CHUNK]
            found.update(tuple(row) for row in conn.exec_driver_sql(
                f"SELECT id, kind FROM ms_payloads WHERE id IN ({','.join('?' * len(chunk))})", tuple(chunk)))
    return found

def iter_bytes(self, ids: Iterable[int], window: int = READ_WINDOW):
    """`(id, bytes)` for each id, `window` ids at a time, each window in (volume, offset) order: to export."""
    ids = iter(ids)
    while batch := list(islice(ids, window)):
        yield from self.get_many(batch).items()

def materialize(self, payload_id: int, directory: str | os.PathLike, suffix: str | None = None):
    """The payload as the file `<directory>/<id><suffix>`, written once: what a job hands its
    tasks as a `File` (a receptor), which executors without the disk carry once per job.
    Without `suffix`, the payload's kind is the extension (`.pdbqt`)."""
    from ms_flow.engine.files import File

    if suffix is None:
        kind = self.kinds([payload_id]).get(int(payload_id))
        suffix = f".{kind}" if kind else ""
    path = Path(directory) / f"{int(payload_id)}{suffix}"
    if not path.exists():
        path.parent.mkdir(parents=True, exist_ok=True)
        part = path.with_name(f".{path.name}.{os.getpid()}.{threading.get_ident()}")
        part.write_bytes(self.get(payload_id))
        os.replace(part, path)
    return File(path)

def train_dictionary(self, series: str, samples: int = DICT_SAMPLES, size: int = DICT_SIZE):
    """Train a zstd dictionary on `samples` payloads spread over `series` (the head of a
    library is not like the rest of it) and make it the series' current one: a `File`."""
    import random

    from ms_flow.engine.files import File

    with self.engine.connect() as conn:
        ids = [r[0] for r in conn.exec_driver_sql(
            "SELECT p.id FROM ms_payloads p JOIN ms_volumes v ON v.id = p.volume_id WHERE v.series = ?",
            (series,))]
    if len(ids) < 100:
        raise ValueError(f"series {series!r} has {len(ids)} payloads; a dictionary needs a few thousand")
    picked = sorted(random.Random(0).sample(ids, min(samples, len(ids))))
    data = [b for _, b in self.iter_bytes(picked)]
    with self.engine.begin() as conn:
        dict_id = conn.exec_driver_sql(
            "INSERT INTO ms_payload_dicts (series, samples, created_at) VALUES (?, ?, ?)",
            (series, len(data), datetime.now().isoformat(" "))).lastrowid
        trained = _zstd().train_dictionary(size, data, dict_id=dict_id)
        path = self.directory / "dicts" / f"{dict_id}.zdict"
        path.parent.mkdir(parents=True, exist_ok=True)
        part = path.with_name(path.name + ".part")
        with open(part, "wb") as file:
            file.write(trained.as_bytes())
            file.flush()
            os.fsync(file.fileno())
        os.replace(part, path)
        _fsync_dir(path.parent)
    return File(path)

def dictionary(self, series: str):
    """The current zstd dictionary of `series` as a `File`, or None if it has none yet."""
    from ms_flow.engine.files import File

    with self.engine.connect() as conn:
        row = conn.exec_driver_sql(
            "SELECT max(id) FROM ms_payload_dicts WHERE series = ?", (series,)).scalar()
    return None if row is None else File(self.directory / "dicts" / f"{row}.zdict")

def check(self, full: bool = False) -> Report:
    """What is wrong or left over; `full` also reads every needle (crc32, codec, dictionary)."""
    report = Report()
    with self.engine.connect() as conn:
        volumes = {r[0]: r[1:] for r in conn.exec_driver_sql("SELECT id, state, size FROM ms_volumes")}
        stats = {r[0]: r[1:] for r in conn.exec_driver_sql(
            f'SELECT volume_id, sum({NEEDLE} + length), max("offset" + {NEEDLE} + length)'
            f" FROM {_NEEDLES} GROUP BY volume_id")}
        report.unreferenced = [r[0] for r in conn.exec_driver_sql(
            f"SELECT id FROM ms_payloads WHERE {self._unreferenced()}")]
        dicts = [r[0] for r in conn.exec_driver_sql("SELECT id FROM ms_payload_dicts")]
    short = {}  # volume id -> its size on disk, when rows point past it
    for volume_id, (state, size) in volumes.items():
        live, end = stats.get(volume_id, (0, HEADER))
        try:
            on_disk = self.path(volume_id).stat().st_size
        except FileNotFoundError:
            if volume_id in stats:
                report.missing_volumes.append(volume_id)
            continue
        if state == "trash" or (state == "sealed" and volume_id not in stats):
            report.empty.append(volume_id)
        elif state == "sealed":
            report.dead[volume_id] = round(1 - live / max(size - HEADER, 1), 4)
        if end > on_disk:
            short[volume_id] = on_disk
    with self.engine.connect() as conn:
        for volume_id in report.missing_volumes:
            for (pid,) in conn.exec_driver_sql("SELECT id FROM ms_payloads WHERE volume_id = ?", (volume_id,)):
                report.broken[pid] = f"volume {volume_id} is missing"
        for volume_id, on_disk in short.items():
            for (pid,) in conn.exec_driver_sql(
                    f'SELECT id FROM ms_payloads WHERE volume_id = ? AND "offset" + {NEEDLE} + length > ?',
                    (volume_id, on_disk)):
                report.broken[pid] = f"past the end of volume {volume_id}"
        if full:
            rows = conn.exec_driver_sql('SELECT id, volume_id, "offset", length, raw_length, codec, crc32'
                                        ' FROM ms_payloads ORDER BY volume_id, "offset"')
            needle, error = None, None
            for pid, *location in rows:
                if tuple(location[:2]) != needle:  # the parts of a group: read once
                    needle, error = tuple(location[:2]), None
                    try:
                        self._read(pid, tuple(location))
                    except Exception as failure:  # a report: whatever stops a read goes in it
                        error = str(failure)
                if error is not None and pid not in report.broken:
                    report.broken[pid] = error
    report.missing_dicts = [d for d in dicts if not (self.directory / "dicts" / f"{d}.zdict").exists()]
    known = set(volumes)
    report.orphans = sorted(volume_id for volume_id in self._volume_files()
                            if volume_id not in known and self._unowned(volume_id))
    return report

def release(self, ids: Iterable[int]) -> int:
    """Delete the payload rows of `ids` that no app row references (after the app deleted its
    rows); a sealed volume left without rows goes, unread. Returns how many rows went."""
    ids = list(dict.fromkeys(int(i) for i in ids))
    guard = self._unreferenced()
    released = 0
    with self.engine.begin() as conn:
        for start in range(0, len(ids), LOCATE_CHUNK):
            chunk = ids[start : start + LOCATE_CHUNK]
            released += conn.exec_driver_sql(
                f"DELETE FROM ms_payloads WHERE id IN ({','.join('?' * len(chunk))}) AND {guard}",
                tuple(chunk)).rowcount
        self._trash_empty(conn)
    self._empty_trash()
    return released

def compact(self, min_dead: float = 0.5, small: int | None = None, order_by: str | None = None) -> dict:
    """Copy the live needles of sealed volumes with at least `min_dead` dead bytes (or of two or
    more volumes of a series under `small` bytes, a quarter of VOLUME_SIZE by default) to new volumes, move their rows there and
    delete the old ones. Needles are recompressed with the series' current dictionary.

    `order_by` is SQL that gives payload ids in the order wanted next to each other on disk
    (`SELECT current_payload_id FROM molecules ORDER BY score`); the rest keep theirs.
    Never touches an active volume, so it can run next to writers and readers (I6, I7).
    """
    small = VOLUME_SIZE // 4 if small is None else small
    self._empty_trash()
    with self.engine.connect() as conn:
        rows = conn.exec_driver_sql(
            f"SELECT v.id, v.series, v.size, coalesce(sum({NEEDLE} + p.length), 0) FROM ms_volumes v"
            f" LEFT JOIN {_NEEDLES} p ON p.volume_id = v.id WHERE v.state = 'sealed' GROUP BY v.id").fetchall()
    picked: dict[str, list[int]] = {}
    for series in {r[1] for r in rows}:
        mine = [(vid, 1 - live / max(size - HEADER, 1), size) for vid, s, size, live in rows if s == series]
        chosen = [vid for vid, dead, _ in mine if dead >= min_dead]
        smalls = [vid for vid, dead, size in mine if dead < min_dead and size < small]
        chosen += smalls if len(smalls) > 1 else []  # one small volume alone gains nothing
        if chosen:
            picked[series] = chosen
    stats = {"volumes": 0, "moved": 0, "read_mb": 0.0, "written_mb": 0.0}
    for series, volume_ids in picked.items():
        self._compact_series(series, volume_ids, order_by, stats)
    self._empty_trash()
    return stats

def repair(self, report: Report | None = None) -> Report:
    """Clear what `check` reports as garbage: cut the tails of dead writers, delete unreferenced
    rows, empty volumes and orphan files. A referenced row is never deleted. Returns a new check."""
    report = report or self.check()
    self._recover()
    self.release(report.unreferenced)
    for volume_id in report.orphans:
        if self._unowned(volume_id):
            self.path(volume_id).unlink(missing_ok=True)
    with self.engine.begin() as conn:
        self._trash_empty(conn)
    self._empty_trash()
    return self.check()

def close(self) -> None:
    """Close files. Call it with no transaction of this store open."""
    with self._lock:
        volumes = [v for free in self._free.values() for v in free]
        self._free.clear()
    for volume in volumes:
        os.close(volume.fd)
    with self._fd_lock:
        for entry in self._fds.values():
            entry[2] = True
            if entry[1] == 0:
                os.close(entry[0])
        self._fds.clear()