Vito Sansevero
94d50d15f2
History writes deliberately give up after a bounded 250 ms wait so they never stall ComfyUI; the existing bounded-timeout test still proves that. The concurrent replay tests assert that every peer succeeds, which is a serialization and first-capture property, so they now raise the budget to the shared 30 s busy timeout while the first writer is held open. The budget is one named constant, EXECUTION_BUSY_TIMEOUT_MS, shared by the transaction helper and the background writer. Before: 20/20 'database is locked' with the held writer. After: 0/20.
199 lines
7.2 KiB
Python
199 lines
7.2 KiB
Python
"""Transactional attempt writes and process-local restore invalidation."""
|
|
|
|
import os
|
|
import sqlite3
|
|
import threading
|
|
from contextlib import contextmanager
|
|
from typing import Callable, Iterator
|
|
|
|
try:
|
|
from ..utils.execution_history_events import (
|
|
RunFinish,
|
|
RunStart,
|
|
classify_evidence,
|
|
encode_command,
|
|
)
|
|
from ..utils.logging_config import get_logger
|
|
except ImportError:
|
|
from utils.execution_history_events import (
|
|
RunFinish,
|
|
RunStart,
|
|
classify_evidence,
|
|
encode_command,
|
|
)
|
|
from utils.logging_config import get_logger
|
|
|
|
|
|
_logger = get_logger("prompt_manager.database.execution_history")
|
|
# Bounded wait for a history write lock. History is best effort and must
|
|
# never stall ComfyUI's queue: a write that cannot start in time raises
|
|
# "database is locked" and the writer records a gap instead of waiting.
|
|
EXECUTION_BUSY_TIMEOUT_MS = 250
|
|
_restore_lock = threading.Lock()
|
|
_restore_observers: dict[str, dict[object, Callable[[], None]]] = {}
|
|
|
|
|
|
def _database_identity(db_path: str) -> str:
|
|
return os.path.normcase(os.path.realpath(db_path))
|
|
|
|
|
|
def subscribe_history_restore(
|
|
db_path: str, callback: Callable[[], None]
|
|
) -> Callable[[], None]:
|
|
"""Observe successful restores of the canonical database path."""
|
|
identity = _database_identity(db_path)
|
|
token = object()
|
|
with _restore_lock:
|
|
_restore_observers.setdefault(identity, {})[token] = callback
|
|
|
|
def unsubscribe() -> None:
|
|
with _restore_lock:
|
|
observers = _restore_observers.get(identity)
|
|
if observers is not None:
|
|
observers.pop(token, None)
|
|
if not observers:
|
|
_restore_observers.pop(identity, None)
|
|
|
|
return unsubscribe
|
|
|
|
|
|
def notify_history_restore(db_path: str) -> None:
|
|
"""Notify outside the registry lock; observers cannot fail a restore."""
|
|
with _restore_lock:
|
|
callbacks = tuple(
|
|
_restore_observers.get(_database_identity(db_path), {}).values()
|
|
)
|
|
for callback in callbacks:
|
|
try:
|
|
callback()
|
|
except Exception:
|
|
_logger.warning("Execution history restore observer failed", exc_info=True)
|
|
|
|
|
|
class HistoryEpochChanged(Exception):
|
|
"""A command belongs to storage replaced by a completed restore."""
|
|
|
|
def __init__(self, storage_epoch: str):
|
|
self.storage_epoch = storage_epoch
|
|
super().__init__("Execution history storage epoch changed")
|
|
|
|
|
|
class ExecutionHistoryOperations:
|
|
"""Mixin using the calling thread's existing managed model connection."""
|
|
|
|
def get_execution_epoch(self) -> str:
|
|
row = (
|
|
self.model.get_connection()
|
|
.execute(
|
|
"SELECT storage_epoch FROM execution_history_metadata "
|
|
"WHERE singleton_id=1"
|
|
)
|
|
.fetchone()
|
|
)
|
|
if row is None:
|
|
raise sqlite3.DatabaseError("Missing execution history metadata")
|
|
return row[0]
|
|
|
|
@contextmanager
|
|
def _execution_transaction(
|
|
self, expected_epoch: str
|
|
) -> Iterator[sqlite3.Connection]:
|
|
conn = self.model.get_connection()
|
|
if conn.in_transaction:
|
|
raise sqlite3.ProgrammingError(
|
|
"Execution history requires its own transaction"
|
|
)
|
|
timeout = conn.execute("PRAGMA busy_timeout").fetchone()[0]
|
|
conn.execute(f"PRAGMA busy_timeout={int(EXECUTION_BUSY_TIMEOUT_MS)}")
|
|
try:
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
epoch = self.get_execution_epoch()
|
|
if epoch != expected_epoch:
|
|
raise HistoryEpochChanged(epoch)
|
|
yield conn
|
|
conn.commit()
|
|
except BaseException:
|
|
conn.rollback()
|
|
raise
|
|
finally:
|
|
conn.execute(f"PRAGMA busy_timeout={timeout}")
|
|
|
|
def write_execution_start(self, command: RunStart) -> bool:
|
|
encode_command(command)
|
|
if type(command) is not RunStart:
|
|
raise ValueError("Expected execution start command")
|
|
with self._execution_transaction(command.storage_epoch) as conn:
|
|
conn.execute(
|
|
"INSERT INTO execution_runs "
|
|
"(run_id,session_id,queue_item_id,job_id,storage_epoch,record_version,"
|
|
"observed_dequeue_at,outcome,diagnostics_omitted) "
|
|
"VALUES(?,?,?,?,?,1,?,'observed',0) "
|
|
"ON CONFLICT(session_id,queue_item_id) DO NOTHING",
|
|
(
|
|
command.run_id,
|
|
command.session_id,
|
|
command.queue_item_id,
|
|
command.job_id,
|
|
command.storage_epoch,
|
|
command.observed_dequeue_at,
|
|
),
|
|
)
|
|
return True
|
|
|
|
def write_execution_finish(self, command: RunFinish) -> bool:
|
|
encode_command(command)
|
|
if type(command) is not RunFinish:
|
|
raise ValueError("Expected execution finish command")
|
|
with self._execution_transaction(command.storage_epoch) as conn:
|
|
row = conn.execute(
|
|
"SELECT * FROM execution_runs WHERE session_id=? AND queue_item_id=? "
|
|
"AND storage_epoch=?",
|
|
(command.session_id, command.queue_item_id, command.storage_epoch),
|
|
).fetchone()
|
|
if row is None:
|
|
return False
|
|
completion = classify_evidence(row["job_id"], command.evidence)
|
|
if row["outcome"] == "unknown":
|
|
return True
|
|
if row["outcome"] != "observed":
|
|
contradictory = completion.reason == "conflicting_terminal" or (
|
|
completion.outcome in ("succeeded", "failed", "interrupted")
|
|
and completion.outcome != row["outcome"]
|
|
)
|
|
if contradictory:
|
|
conn.execute(
|
|
"UPDATE execution_runs SET outcome='unknown',"
|
|
"reason='conflicting_terminal' WHERE run_id=?",
|
|
(row["run_id"],),
|
|
)
|
|
return True
|
|
conn.execute(
|
|
"UPDATE execution_runs SET observed_callback_at=?,"
|
|
"executor_started_at=?,"
|
|
"terminal_event_at=?,outcome=?,reason=?,diagnostics_omitted=?,"
|
|
"node_id=?,"
|
|
"node_type=?,exception_type=? WHERE run_id=?",
|
|
(
|
|
command.observed_callback_at,
|
|
completion.executor_started_at,
|
|
completion.terminal_event_at,
|
|
completion.outcome,
|
|
completion.reason,
|
|
int(completion.diagnostics_omitted),
|
|
completion.node_id,
|
|
completion.node_type,
|
|
completion.exception_type,
|
|
row["run_id"],
|
|
),
|
|
)
|
|
return True
|
|
|
|
def get_execution_run(self, run_id: str) -> dict | None:
|
|
row = (
|
|
self.model.get_connection()
|
|
.execute("SELECT * FROM execution_runs WHERE run_id=?", (run_id,))
|
|
.fetchone()
|
|
)
|
|
return dict(row) if row is not None else None
|