refactor(gateway/hosted_room_driver): fold settlement validation+replay into _settlement (3 sites), drop _read wrapper for contextlib.closing, walrus replay

This commit is contained in:
Teknium
2026-09-02 22:47:28 -07:00
parent f07062f9a7
commit 380df1bfcb
+35 -46
View File
@@ -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: