diff --git a/gateway/hosted_room_driver.py b/gateway/hosted_room_driver.py index eeca73ec39..3d0213af83 100644 --- a/gateway/hosted_room_driver.py +++ b/gateway/hosted_room_driver.py @@ -13,6 +13,7 @@ import hashlib import json import math import sqlite3 +from contextlib import closing from dataclasses import dataclass from functools import partial from typing import Any, Callable, Literal, get_args @@ -375,17 +376,15 @@ def _load_active_room(conn: sqlite3.Connection, room_id: str) -> sqlite3.Row: return row -def _require_room_authority( - conn: sqlite3.Connection, *, room_id: str, gateway_id: str, authority_epoch: int) -> sqlite3.Row: +def _require_room_authority(conn: sqlite3.Connection, room_id: str, gateway_id: str, epoch: int) -> sqlite3.Row: room = _load_active_room(conn, room_id) - if room["authority_gateway_id"] != gateway_id or int(room["authority_epoch"]) != authority_epoch: + if room["authority_gateway_id"] != gateway_id or int(room["authority_epoch"]) != epoch: raise StaleLeaseError("hosted room authority changed") return room def _require_lease_authority(conn: sqlite3.Connection, lease: DriverLease) -> sqlite3.Row: - return _require_room_authority( - conn, room_id=lease.room_id, gateway_id=lease.gateway_id, authority_epoch=lease.authority_epoch) + return _require_room_authority(conn, lease.room_id, lease.gateway_id, lease.authority_epoch) def _lease_row_matches(row: sqlite3.Row | None, lease: DriverLease) -> bool: @@ -429,15 +428,23 @@ def _terminal_settlement_id(settlement_id: Any, status: Any) -> str: return settlement_id -def _settlement_replay(settlement_id: str, status: str, result_json: str) -> Callable[[sqlite3.Row], Any]: - """Replay predicate: identical settlement already committed; a different one is a conflict.""" +def _settlement( + settlement_id: Any, status: Any, result: Any, clock: Clock +) -> tuple[float, Callable[[sqlite3.Row], Any], tuple[Any, ...]]: + """Validate one terminal settlement -> (now, replay predicate, ``_SETTLE_SET`` params). + + Replay: an identical settlement already committed is idempotent; a different one is a conflict. + """ + settlement_id = _terminal_settlement_id(settlement_id, status) + result_json = _canonical_json(result) + now = _timestamp(clock) def replay(row: sqlite3.Row) -> dict[str, Any] | None: if row["settlement_id"] is None: return None if (row["settlement_id"], row["settlement_status"], row["result_json"]) == (settlement_id, status, result_json): return _task_from_row(row, idempotent=True) raise TaskConflictError("task already has a different terminal settlement") - return replay + return now, replay, (status, settlement_id, status, result_json, now, now) def _cancel_replay(cancel_id: str, status: str = "cancelled") -> Callable[[sqlite3.Row], Any]: @@ -472,10 +479,8 @@ def _transition( if lease is not None and lease_first: _require_active_lease(conn, lease, now=now) row = _load_task(conn, identity) - if replay is not None: - replayed = replay(row) - if replayed is not None: - return replayed + if replay is not None and (replayed := replay(row)) is not None: + return replayed if lease is not None and not lease_first: _require_active_lease(conn, lease, now=now) if guard is not None: @@ -533,7 +538,7 @@ def acquire_lease( now = _timestamp(clock) expires_at = _expiry(now, ttl_seconds) with _transaction(db_path) as conn: - _require_room_authority(conn, room_id=room_id, gateway_id=gateway_id, authority_epoch=authority_epoch) + _require_room_authority(conn, room_id, gateway_id, authority_epoch) row = conn.execute(_SELECT_LEASE, (room_id,)).fetchone() if row is None: conn.execute("""INSERT INTO hosted_room_driver_leases ( @@ -602,8 +607,9 @@ def release_lease(db_path: DbPath, lease: DriverLease, *, clock: Clock) -> dict[ conn.execute("""UPDATE hosted_room_driver_leases SET expires_at=?, updated_at=?, released_at=? WHERE room_id=? AND lease_generation=?""", (now, now, now, lease.room_id, lease.lease_generation)) - current = {**dict(row), "expires_at": now, "updated_at": now, "released_at": now} - return {"lease": _lease_from_row(current), "idempotent": False} + return { + "lease": _lease_from_row({**dict(row), "expires_at": now, "updated_at": now, "released_at": now}), + "idempotent": False} def admit_task(db_path: DbPath, identity: TaskIdentity, *, payload: Any, clock: Clock) -> dict[str, Any]: @@ -677,13 +683,10 @@ def settle_task( db_path: DbPath, attempt: TaskAttempt, *, settlement_id: Any, status: TerminalStatus, result: Any, clock: Clock ) -> dict[str, Any]: """Commit one terminal result if every lease and task fence still matches.""" - settlement_id = _terminal_settlement_id(settlement_id, status) - result_json = _canonical_json(result) - now = _timestamp(clock) + now, replay, set_params = _settlement(settlement_id, status, result, clock) return _run_fence_transition( - db_path, attempt, guard_stale="task attempt is stale or cancelled", lease_first=False, now=now, - replay=_settlement_replay(settlement_id, status, result_json), sql=_SETTLE_RUNNING_SQL, - set_params=(status, settlement_id, status, result_json, now, now), stale="task changed during settlement") + db_path, attempt, guard_stale="task attempt is stale or cancelled", lease_first=False, now=now, replay=replay, + sql=_SETTLE_RUNNING_SQL, set_params=set_params, stale="task changed during settlement") def settle_stopping_task( @@ -691,16 +694,13 @@ def settle_stopping_task( expected_cancel_generation: int, settlement_id: Any, status: TerminalStatus, result: Any, clock: Clock ) -> dict[str, Any]: """Commit a completion that won the race with an unacknowledged Stop.""" - settlement_id = _terminal_settlement_id(settlement_id, status) + _terminal_settlement_id(settlement_id, status) # settlement errors take precedence over generation errors if expected_execution_generation < 1 or expected_cancel_generation < 1: raise DriverValidationError("stopping settlement generations are invalid") - result_json = _canonical_json(result) - now = _timestamp(clock) + now, replay, set_params = _settlement(settlement_id, status, result, clock) return _transition( - db_path, identity, lease=lease, lease_first=False, now=now, - replay=_settlement_replay(settlement_id, status, result_json), sql=_SETTLE_STOPPING_SQL, - set_params=(status, settlement_id, status, result_json, now, now), - fence_params=(expected_execution_generation, expected_cancel_generation), + db_path, identity, lease=lease, lease_first=False, now=now, replay=replay, sql=_SETTLE_STOPPING_SQL, + set_params=set_params, fence_params=(expected_execution_generation, expected_cancel_generation), stale="task completion lost the stop race") @@ -710,13 +710,10 @@ def resolve_indeterminate_task( ) -> dict[str, Any]: """Commit a verified historical receipt under the current room lease.""" _expected_generations(lease, identity, expected_execution_generation, expected_cancel_generation) - settlement_id = _terminal_settlement_id(settlement_id, status) - result_json = _canonical_json(result) - now = _timestamp(clock) + now, replay, set_params = _settlement(settlement_id, status, result, clock) return _generation_transition( db_path, identity, lease, "resolve", expected_execution_generation, expected_cancel_generation, now=now, - replay=_settlement_replay(settlement_id, status, result_json), - set_params=(status, settlement_id, status, result_json, now, now)) + replay=replay, set_params=set_params) def resolve_indeterminate_cancellation( @@ -856,17 +853,10 @@ def recover_room(db_path: DbPath, lease: DriverLease, *, clock: Clock) -> dict[s for status in ("queued", "indeterminate")} -def _read(db_path: DbPath, query: Callable[[sqlite3.Connection], Any]) -> Any: - conn = _connect(db_path) - try: - return query(conn) - finally: - conn.close() - - def get_task(db_path: DbPath, identity: TaskIdentity) -> dict[str, Any]: """Read one task without mutating its state.""" - return _read(db_path, lambda conn: _task_from_row(_load_task(conn, identity))) + with closing(_connect(db_path)) as conn: + return _task_from_row(_load_task(conn, identity)) def list_tasks(db_path: DbPath, *, room_id: Any, status: TaskStatus | None = None) -> list[dict[str, Any]]: @@ -874,7 +864,8 @@ def list_tasks(db_path: DbPath, *, room_id: Any, status: TaskStatus | None = Non room_id = _identifier(room_id, label="room_id") if status is not None and status not in TASK_STATUSES: raise DriverValidationError("invalid task status") - return _read(db_path, lambda conn: [_task_from_row(row) for row in _tasks_in_order(conn, room_id, status)]) + with closing(_connect(db_path)) as conn: + return [_task_from_row(row) for row in _tasks_in_order(conn, room_id, status)] def prune_published_terminal_tasks( @@ -899,9 +890,7 @@ def prune_published_terminal_tasks( ORDER BY t.terminal_at DESC, t.task_id ASC""", (room_id,)).fetchall() cutoff = now - float(retention_seconds) candidates = [ - str(row["task_id"]) - for index, - row in enumerate(rows) + str(row["task_id"]) for index, row in enumerate(rows) if index >= retain or (row["terminal_at"] is not None and float(row["terminal_at"]) <= cutoff) ][:MAX_TASK_PRUNE_BATCH] if not candidates: