Vito
f5856076ce
A browser server inside an --ipc=host container shares the host's /dev/shm but not its PID namespace, so the pid probe saw live servers' scratch folders as orphans and deleted them. A live owner now holds a flock on owner.lock in its folder; the kernel drops it when the owner dies, in any namespace. The cleaner removes only held-named folders whose lock is free and leaves legacy and staging folders alone. Without flock (Windows) the scratch stays on disk.
334 lines
13 KiB
Python
334 lines
13 KiB
Python
"""Run a browser-test fixture seeder off the live database, then publish once (#384).
|
|
|
|
The seeders make hundreds of small commits. On a slow CI disk the live
|
|
database's commits queue behind WAL checkpoint syncs and a reset took over
|
|
30 s. ``ScratchSeeder.run`` instead:
|
|
|
|
1. copies the live database into a scratch database (SQLite backup API) whose
|
|
connections never sync (a crash loses only the copy);
|
|
2. runs the unchanged seeder against an API proxy whose ``db`` is the scratch
|
|
database; everything else (files, thumbnails, caches, command state) is
|
|
the live API's, written exactly as before;
|
|
3. replaces the rows of every table the seed changed in the live database, in
|
|
one transaction on the live connection. Tables the seed left alone are not
|
|
touched, so state another fixture owns stays as it was.
|
|
|
|
A commit by any other live connection while the seed ran fails the reset
|
|
instead of being overwritten (``PRAGMA data_version``).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import os
|
|
import shutil
|
|
import sqlite3
|
|
import tempfile
|
|
import weakref
|
|
from pathlib import Path
|
|
from typing import Any, Callable, TypeVar
|
|
|
|
try:
|
|
import fcntl
|
|
except ImportError: # Windows: no flock, and no /dev/shm to share
|
|
fcntl = None
|
|
|
|
T = TypeVar("T")
|
|
|
|
# Set to 1 to run seeders on the live database as before #384 (the equivalence
|
|
# proof, tests/browser/reset_equivalence.py, compares both ways).
|
|
SEED_ON_LIVE_ENV = "PM_TEST_SEED_ON_LIVE"
|
|
SCRATCH_SCHEMA = "reset_seed"
|
|
# API calls a seed makes that must see the published rows (see _SeedApi).
|
|
DEFERRED_CALLS = frozenset({"invalidate_gallery_cache"})
|
|
# Linux's shared-memory tmpfs (present in the Playwright image and on runners).
|
|
MEMORY_ROOT = Path("/dev/shm")
|
|
# The largest fixtures' scratch copy is about 8 MB (database plus a WAL that
|
|
# autocheckpoints near 4 MB). Docker's default /dev/shm is 64 MB, shared with
|
|
# the browsers and the other workers; with less free than this, the scratch
|
|
# copy goes on disk instead.
|
|
MEMORY_FREE_BYTES = 32 * 1024 * 1024
|
|
# Scratch folders are created as <prefix>new-<owner pid>-<random>, locked, then
|
|
# renamed to <prefix>held-<owner pid>-<random>. The pid is for people only:
|
|
# ownership is the flock on OWNER_LOCK, which every PID namespace sees (#419).
|
|
SCRATCH_PREFIX = "pm-reset-"
|
|
STAGING_PREFIX = f"{SCRATCH_PREFIX}new-"
|
|
HELD_PREFIX = f"{SCRATCH_PREFIX}held-"
|
|
OWNER_LOCK = "owner.lock"
|
|
|
|
|
|
def _quote(name: str) -> str:
|
|
return '"' + name.replace('"', '""') + '"'
|
|
|
|
|
|
def _tables(conn: sqlite3.Connection, schema: str = "main") -> list[str]:
|
|
rows = conn.execute(
|
|
f"SELECT name, sql FROM {schema}.sqlite_master WHERE type = 'table' "
|
|
"ORDER BY name"
|
|
).fetchall()
|
|
for name, sql in rows:
|
|
if sql and sql.lstrip().upper().startswith("CREATE VIRTUAL"):
|
|
raise NotImplementedError(f"virtual table {name} cannot be copied")
|
|
return [name for name, _sql in rows]
|
|
|
|
|
|
def _columns(conn: sqlite3.Connection, table: str, schema: str = "main") -> list[str]:
|
|
"""Writable columns (hidden and generated ones excluded), plus ``rowid``
|
|
when the table has a rowid that no INTEGER PRIMARY KEY column carries."""
|
|
info = conn.execute(f"PRAGMA {schema}.table_xinfo({_quote(table)})").fetchall()
|
|
columns = [row[1] for row in info if row[6] == 0]
|
|
keys = [row for row in info if row[5]]
|
|
sql = conn.execute(
|
|
f"SELECT sql FROM {schema}.sqlite_master WHERE type = 'table' AND name = ?",
|
|
(table,),
|
|
).fetchone()[0]
|
|
without_rowid = sql is not None and "WITHOUT ROWID" in sql.upper()
|
|
integer_key = len(keys) == 1 and str(keys[0][2]).upper() == "INTEGER"
|
|
if not without_rowid and not integer_key and table != "sqlite_sequence":
|
|
columns.insert(0, "rowid")
|
|
return columns
|
|
|
|
|
|
def table_digests(conn: sqlite3.Connection, schema: str = "main") -> dict[str, str]:
|
|
"""A digest of every table's rows (rowids included), by table name."""
|
|
result = {}
|
|
for table in _tables(conn, schema):
|
|
selected = ", ".join(
|
|
_quote(c) if c != "rowid" else "rowid"
|
|
for c in _columns(conn, table, schema)
|
|
)
|
|
digest = hashlib.sha256()
|
|
rows = conn.execute(f"SELECT {selected} FROM {schema}.{_quote(table)}")
|
|
for row in sorted(repr(tuple(row)) for row in rows):
|
|
digest.update(row.encode("utf-8", "surrogatepass"))
|
|
digest.update(b"\n")
|
|
result[table] = digest.hexdigest()
|
|
return result
|
|
|
|
|
|
def _hold(folder: Path, *, create: bool) -> int:
|
|
"""Take ``folder``'s owner lock without waiting; the open descriptor."""
|
|
flags = os.O_RDWR | (os.O_CREAT if create else 0)
|
|
fd = os.open(folder / OWNER_LOCK, flags, 0o600)
|
|
try:
|
|
fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
except OSError:
|
|
os.close(fd)
|
|
raise
|
|
return fd
|
|
|
|
|
|
def remove_orphaned_scratch(root: Path) -> None:
|
|
"""Remove scratch folders of servers that died without cleaning up.
|
|
|
|
A server killed outright (SIGKILL, a crash) never runs ``close``, and
|
|
tmpfs keeps its folder in memory until reboot. A live owner holds the
|
|
flock on the folder's owner lock; the kernel drops it when the owner
|
|
dies, in whatever PID namespace it ran (a container with ``--ipc=host``
|
|
shares /dev/shm but not the pids, #419). A folder whose lock is free is
|
|
removed. Folders without a lock (staging, or named before #419) cannot
|
|
be proven orphaned and are left alone.
|
|
"""
|
|
if fcntl is None:
|
|
return
|
|
for folder in root.glob(f"{HELD_PREFIX}*"):
|
|
try:
|
|
fd = _hold(folder, create=False)
|
|
except OSError:
|
|
continue # held by its owner, no lock to judge by, or gone
|
|
try:
|
|
shutil.rmtree(folder, ignore_errors=True)
|
|
finally:
|
|
os.close(fd)
|
|
|
|
|
|
class _SeedApi:
|
|
"""The live API, except that ``db`` is the scratch database.
|
|
|
|
Cache drops the seed asks for (``DEFERRED_CALLS``) are recorded and made
|
|
after the publish: dropped earlier, a request in between would cache the
|
|
live database's old rows again.
|
|
"""
|
|
|
|
def __init__(self, api: Any, db: Any) -> None:
|
|
object.__setattr__(self, "_api", api)
|
|
object.__setattr__(self, "db", db)
|
|
object.__setattr__(self, "deferred", [])
|
|
|
|
def __getattr__(self, name: str) -> Any:
|
|
if name in DEFERRED_CALLS:
|
|
|
|
def record(*args: Any, **kwargs: Any) -> None:
|
|
self.deferred.append((name, args, kwargs))
|
|
|
|
return record
|
|
return getattr(self._api, name)
|
|
|
|
def __setattr__(self, name: str, value: Any) -> None:
|
|
setattr(self._api, name, value)
|
|
|
|
|
|
class ScratchSeeder:
|
|
"""Seed into a scratch copy and publish the changed tables in one commit."""
|
|
|
|
def __init__(
|
|
self,
|
|
api: Any,
|
|
directory: Path,
|
|
*,
|
|
on_live: bool = False,
|
|
memory_root: Path | None = MEMORY_ROOT,
|
|
) -> None:
|
|
self._api = api
|
|
self._fallback = Path(directory)
|
|
self._memory_root = memory_root
|
|
self._directory: Path | None = None
|
|
self._owner_lock: int | None = None
|
|
self._on_live = on_live
|
|
self._scratch = None
|
|
self._unsynced: "weakref.WeakSet[sqlite3.Connection]" = weakref.WeakSet()
|
|
|
|
def run(self, seed: Callable[[Any], T]) -> T:
|
|
"""Run ``seed(api)`` (blocking; call it in an executor thread)."""
|
|
if self._on_live:
|
|
return seed(self._api)
|
|
live = self._api.db.model.get_connection()
|
|
if live.in_transaction:
|
|
raise RuntimeError("the live connection has an open transaction")
|
|
version = live.execute("PRAGMA data_version").fetchone()[0]
|
|
scratch = self._scratch_database()
|
|
target = self._scratch_connection(scratch)
|
|
with target.using():
|
|
live.backup(target)
|
|
before = table_digests(target)
|
|
target_api = _SeedApi(self._api, scratch)
|
|
result = seed(target_api)
|
|
target = self._scratch_connection(scratch)
|
|
if target.in_transaction:
|
|
raise RuntimeError("the seed left a scratch transaction open")
|
|
after = table_digests(target)
|
|
changed = [table for table in after if after[table] != before.get(table)]
|
|
if set(before) != set(after):
|
|
raise RuntimeError("the seed changed the database schema")
|
|
if live.execute("PRAGMA data_version").fetchone()[0] != version:
|
|
raise RuntimeError("the live database changed while the seed ran")
|
|
self._publish(live, scratch.model.db_path, changed, version)
|
|
for name, args, kwargs in target_api.deferred:
|
|
getattr(self._api, name)(*args, **kwargs)
|
|
return result
|
|
|
|
def _scratch_database(self):
|
|
if self._scratch is None:
|
|
from database.operations import PromptDatabase
|
|
|
|
self._directory = self._scratch_directory()
|
|
self._scratch = PromptDatabase(str(self._directory / "seed.db"))
|
|
return self._scratch
|
|
|
|
def _scratch_connection(self, scratch) -> sqlite3.Connection:
|
|
"""The calling thread's scratch connection, set never to sync."""
|
|
conn = scratch.model.get_connection()
|
|
if conn not in self._unsynced:
|
|
conn.execute("PRAGMA synchronous = OFF")
|
|
self._unsynced.add(conn)
|
|
return conn
|
|
|
|
@staticmethod
|
|
def _publish(
|
|
live: sqlite3.Connection, path: str, tables: list[str], version: int
|
|
) -> None:
|
|
"""Replace ``tables`` with the scratch rows in one live transaction.
|
|
|
|
``version`` is the live ``data_version`` read before the seed. It is
|
|
checked again under the write lock: a writer whose transaction was
|
|
open during the seed commits before ``BEGIN IMMEDIATE`` returns, and
|
|
its rows must fail the reset instead of being overwritten.
|
|
"""
|
|
if not tables:
|
|
return
|
|
foreign_keys = live.execute("PRAGMA foreign_keys").fetchone()[0]
|
|
# Rows are copied as they are; no cascade may reach a table left alone.
|
|
live.execute("PRAGMA foreign_keys = OFF")
|
|
try:
|
|
live.execute(f"ATTACH DATABASE ? AS {SCRATCH_SCHEMA}", (path,))
|
|
try:
|
|
live.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
if live.execute("PRAGMA data_version").fetchone()[0] != version:
|
|
raise RuntimeError(
|
|
"the live database changed while the seed ran"
|
|
)
|
|
for table in tables:
|
|
columns = ", ".join(
|
|
_quote(c) if c != "rowid" else "rowid"
|
|
for c in _columns(live, table)
|
|
)
|
|
name = _quote(table)
|
|
live.execute(f"DELETE FROM main.{name}")
|
|
live.execute(
|
|
f"INSERT INTO main.{name} ({columns}) "
|
|
f"SELECT {columns} FROM {SCRATCH_SCHEMA}.{name}"
|
|
)
|
|
live.commit()
|
|
except BaseException:
|
|
live.rollback()
|
|
raise
|
|
finally:
|
|
live.execute(f"DETACH DATABASE {SCRATCH_SCHEMA}")
|
|
finally:
|
|
live.execute(f"PRAGMA foreign_keys = {int(foreign_keys)}")
|
|
|
|
def close(self) -> None:
|
|
"""Close every scratch connection, then remove the scratch files."""
|
|
if self._on_live:
|
|
return
|
|
if self._scratch is not None:
|
|
self._scratch.close_all()
|
|
self._scratch = None
|
|
if self._directory is not None:
|
|
shutil.rmtree(self._directory, ignore_errors=True)
|
|
self._directory = None
|
|
if self._owner_lock is not None:
|
|
os.close(self._owner_lock) # after the removal: never a free lock
|
|
self._owner_lock = None
|
|
|
|
def _scratch_directory(self) -> Path:
|
|
"""A folder in memory (tmpfs) where there is one, else ``directory``.
|
|
|
|
Even unsynced, scratch writes on the runner's disk wait whenever the
|
|
kernel is writing back the pages they rewrite.
|
|
"""
|
|
root = self._memory_root
|
|
if (
|
|
fcntl is not None
|
|
and root is not None
|
|
and root.is_dir()
|
|
and os.access(root, os.W_OK)
|
|
):
|
|
remove_orphaned_scratch(root)
|
|
if shutil.disk_usage(root).free >= MEMORY_FREE_BYTES:
|
|
held = self._held_memory_directory(root)
|
|
if held is not None:
|
|
return held
|
|
self._fallback.mkdir(parents=True, exist_ok=True)
|
|
return self._fallback
|
|
|
|
def _held_memory_directory(self, root: Path) -> Path | None:
|
|
"""A new folder under ``root`` that this seeder holds the lock of.
|
|
|
|
It is locked under the staging name, which no cleaner judges, and
|
|
only then renamed: no cleaner ever sees it held-named and unlocked.
|
|
"""
|
|
staging = Path(
|
|
tempfile.mkdtemp(prefix=f"{STAGING_PREFIX}{os.getpid()}-", dir=root)
|
|
)
|
|
try:
|
|
fd = _hold(staging, create=True)
|
|
except OSError:
|
|
shutil.rmtree(staging, ignore_errors=True)
|
|
return None # no locks on this filesystem: use the disk
|
|
held = root / (HELD_PREFIX + staging.name[len(STAGING_PREFIX) :])
|
|
os.rename(staging, held)
|
|
self._owner_lock = fd
|
|
return held
|