393 lines
15 KiB
Python
393 lines
15 KiB
Python
"""A bounded, single-worker transport for best-effort execution history."""
|
|
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from functools import partial
|
|
import math
|
|
import queue
|
|
import threading
|
|
import time
|
|
from typing import Callable, Iterator
|
|
|
|
try:
|
|
from ..database import PromptDatabase
|
|
from ..database.execution_history import (
|
|
HistoryEpochChanged,
|
|
subscribe_history_restore,
|
|
)
|
|
except ImportError:
|
|
from database import PromptDatabase
|
|
from database.execution_history import (
|
|
HistoryEpochChanged,
|
|
subscribe_history_restore,
|
|
)
|
|
|
|
from .execution_context_budget import (
|
|
SHARED_CONTEXT_BUDGET,
|
|
ContextBudget,
|
|
PayloadLease,
|
|
SealedContext,
|
|
)
|
|
from .execution_contexts import ContextBundle, bundle_payload_size
|
|
from .execution_history_events import RunCommand, RunFinish, RunStart, encode_command
|
|
from .logging_config import get_logger
|
|
|
|
_logger = get_logger("prompt_manager.execution_history_writer")
|
|
_MAX_PENDING = 1024
|
|
_PROBE_INTERVAL = 5.0
|
|
_MAX_GAPS = 2**63 - 1
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class HistoryWriterStatus:
|
|
epoch: str | None
|
|
accepting: bool
|
|
pending: int
|
|
recording_gaps: int
|
|
alive: bool
|
|
capture_generation: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _Finalization:
|
|
command: RunFinish
|
|
context: SealedContext
|
|
|
|
|
|
class HistoryWriter:
|
|
"""Keep database initialization and all writes off producer threads."""
|
|
|
|
def __init__(
|
|
self,
|
|
db_factory: Callable[[], PromptDatabase],
|
|
*,
|
|
clock: Callable[[], float] = time.monotonic,
|
|
context_budget: ContextBudget | None = None,
|
|
) -> None:
|
|
self.context_budget = (
|
|
SHARED_CONTEXT_BUDGET if context_budget is None else context_budget
|
|
)
|
|
self._db_factory = db_factory
|
|
self._clock = clock
|
|
self._context_commit_gate = threading.Lock()
|
|
self._lock = threading.Lock()
|
|
self._wake = threading.Event()
|
|
self._queue: queue.Queue[RunCommand | _Finalization] = queue.Queue(
|
|
maxsize=_MAX_PENDING
|
|
)
|
|
self._thread: threading.Thread | None = None
|
|
self._epoch: str | None = None
|
|
self._accepting = False
|
|
self._pending = 0
|
|
self._gaps = 0
|
|
self._generation = 0
|
|
self._closing = False
|
|
self._close_deadline = 0.0
|
|
self._next_probe = 0.0
|
|
|
|
def start(self) -> None:
|
|
with self._lock:
|
|
if self._thread is not None or self._closing:
|
|
return
|
|
self._thread = threading.Thread(
|
|
target=self._run, name="prompt-manager-history", daemon=True
|
|
)
|
|
self._thread.start()
|
|
|
|
def submit(self, command: RunCommand) -> bool:
|
|
try:
|
|
encode_command(command)
|
|
except ValueError:
|
|
with self._lock:
|
|
self._record_gap_locked()
|
|
self._wake.set()
|
|
return False
|
|
with self._lock:
|
|
if (
|
|
not self._accepting
|
|
or command.storage_epoch != self._epoch
|
|
or self._pending >= _MAX_PENDING
|
|
):
|
|
self._record_gap_locked()
|
|
accepted = False
|
|
else:
|
|
self._queue.put_nowait(command)
|
|
self._pending += 1
|
|
accepted = True
|
|
self._wake.set()
|
|
return accepted
|
|
|
|
@staticmethod
|
|
def _release_context(context: SealedContext) -> None:
|
|
if isinstance(context, SealedContext) and type(context.lease) is PayloadLease:
|
|
context.lease.release()
|
|
|
|
def _consume_context_locked(self, context: SealedContext) -> SealedContext:
|
|
try:
|
|
if type(context) is not SealedContext:
|
|
raise ValueError("Invalid sealed context")
|
|
context.validate_header()
|
|
if context.bundle is None:
|
|
if context.lease is not None:
|
|
raise ValueError("Unexpected payload lease")
|
|
return context
|
|
size = bundle_payload_size(context.bundle)
|
|
if (
|
|
context.capture_generation != self._generation
|
|
or context.capture_generation is None
|
|
or type(context.lease) is not PayloadLease
|
|
or context.lease._pool is not self.context_budget
|
|
or context.lease.byte_count < size
|
|
or context.state != context.bundle.state
|
|
or context.reason != context.bundle.reason
|
|
):
|
|
raise ValueError("Invalid context ownership or generation")
|
|
owner = context.lease.transfer()
|
|
return SealedContext(
|
|
context.bundle,
|
|
owner,
|
|
context.state,
|
|
context.reason,
|
|
context.capture_generation,
|
|
)
|
|
except ValueError:
|
|
self._release_context(context)
|
|
return SealedContext(None, None, "unavailable", "invalid_evidence")
|
|
|
|
def submit_finalization(self, command: RunFinish, context: SealedContext) -> bool:
|
|
try:
|
|
if type(command) is not RunFinish:
|
|
raise ValueError("Finalization requires finish")
|
|
encode_command(command)
|
|
except ValueError:
|
|
self._release_context(context)
|
|
self.record_gap()
|
|
return False
|
|
with self._lock:
|
|
if (
|
|
not self._accepting
|
|
or command.storage_epoch != self._epoch
|
|
or self._pending >= _MAX_PENDING
|
|
):
|
|
self._release_context(context)
|
|
self._record_gap_locked()
|
|
accepted = False
|
|
else:
|
|
owned = self._consume_context_locked(context)
|
|
self._queue.put_nowait(_Finalization(command, owned))
|
|
self._pending += 1
|
|
accepted = True
|
|
self._wake.set()
|
|
return accepted
|
|
|
|
def _record_gap_locked(self) -> None:
|
|
self._gaps = min(_MAX_GAPS, self._gaps + 1)
|
|
|
|
def record_gap(self) -> None:
|
|
"""Account for unavailable capture without constructing a fake command."""
|
|
with self._lock:
|
|
self._record_gap_locked()
|
|
self._wake.set()
|
|
|
|
@contextmanager
|
|
def _context_generation_guard(
|
|
self, generation: int, epoch: str, capture_generation: int | None
|
|
) -> Iterator[bool]:
|
|
with self._context_commit_gate:
|
|
with self._lock:
|
|
current = (
|
|
generation == self._generation
|
|
and epoch == self._epoch
|
|
and (capture_generation is None or capture_generation == generation)
|
|
)
|
|
yield current
|
|
|
|
def invalidate(self) -> None:
|
|
with self._context_commit_gate:
|
|
with self._lock:
|
|
self._accepting = False
|
|
self._epoch = None
|
|
self._generation += 1
|
|
self._next_probe = self._clock() + _PROBE_INTERVAL
|
|
self._discard_locked()
|
|
self._wake.set()
|
|
|
|
def status(self) -> HistoryWriterStatus:
|
|
with self._lock:
|
|
return HistoryWriterStatus(
|
|
self._epoch,
|
|
self._accepting,
|
|
self._pending,
|
|
self._gaps,
|
|
self._thread is not None and self._thread.is_alive(),
|
|
self._generation,
|
|
)
|
|
|
|
def close(self, timeout: float = 2.0) -> bool:
|
|
if not math.isfinite(timeout) or timeout < 0:
|
|
raise ValueError("Shutdown timeout must be finite and nonnegative")
|
|
with self._lock:
|
|
deadline = time.monotonic() + timeout
|
|
if not self._closing:
|
|
self._closing = True
|
|
self._close_deadline = deadline
|
|
self._accepting = False
|
|
thread = self._thread
|
|
self._wake.set()
|
|
if thread is not None:
|
|
thread.join(max(0.0, deadline - time.monotonic()))
|
|
with self._lock:
|
|
if self._closing and time.monotonic() >= self._close_deadline:
|
|
self._discard_locked()
|
|
return thread is None or not thread.is_alive()
|
|
|
|
def _discard_locked(self) -> None:
|
|
while True:
|
|
try:
|
|
item = self._queue.get_nowait()
|
|
except queue.Empty:
|
|
return
|
|
if isinstance(item, _Finalization):
|
|
self._release_context(item.context)
|
|
self._pending -= 1
|
|
self._record_gap_locked()
|
|
|
|
def _run(self) -> None:
|
|
db: PromptDatabase | None = None
|
|
unsubscribe: Callable[[], None] | None = None
|
|
reported_gaps = 0
|
|
next_diagnostic = 0.0
|
|
try:
|
|
while True:
|
|
self._wake.clear()
|
|
with self._lock:
|
|
if self._closing and (
|
|
self._pending == 0 or time.monotonic() >= self._close_deadline
|
|
):
|
|
self._discard_locked()
|
|
break
|
|
generation = self._generation
|
|
probe = (
|
|
self._epoch is None
|
|
and not self._closing
|
|
and self._clock() >= self._next_probe
|
|
)
|
|
if probe:
|
|
self._next_probe = self._clock() + _PROBE_INTERVAL
|
|
command = None
|
|
if self._epoch is not None:
|
|
try:
|
|
command = self._queue.get_nowait()
|
|
except queue.Empty:
|
|
pass
|
|
if probe:
|
|
try:
|
|
if db is None:
|
|
db = self._db_factory()
|
|
with self._lock:
|
|
closing = self._closing
|
|
if closing:
|
|
continue
|
|
if unsubscribe is None:
|
|
unsubscribe = subscribe_history_restore(
|
|
db.model.db_path, self.invalidate
|
|
)
|
|
db.model.get_connection().execute("PRAGMA busy_timeout=250")
|
|
epoch = db.get_execution_epoch()
|
|
with self._lock:
|
|
if not self._closing and generation == self._generation:
|
|
self._epoch = epoch
|
|
self._accepting = True
|
|
except Exception:
|
|
# The next scheduled probe opens a clean managed connection.
|
|
if db is not None:
|
|
db.model.close()
|
|
with self._lock:
|
|
self._record_gap_locked()
|
|
if command is not None:
|
|
failed = False
|
|
context = None
|
|
if isinstance(command, _Finalization):
|
|
context = command.context
|
|
command = command.command
|
|
try:
|
|
# Restore may have reopened this thread's managed connection.
|
|
db.model.get_connection().execute("PRAGMA busy_timeout=250")
|
|
if isinstance(command, RunStart):
|
|
db.write_execution_start(command)
|
|
else:
|
|
db.write_execution_finish(command)
|
|
if context is not None:
|
|
with self._lock:
|
|
current = (
|
|
generation == self._generation
|
|
and command.storage_epoch == self._epoch
|
|
and (
|
|
context.capture_generation is None
|
|
or context.capture_generation
|
|
== self._generation
|
|
)
|
|
)
|
|
if current:
|
|
try:
|
|
db.write_execution_contexts_for_finish(
|
|
command,
|
|
(
|
|
context.bundle
|
|
if context.bundle is not None
|
|
else ContextBundle(
|
|
(), context.state, context.reason
|
|
)
|
|
),
|
|
generation_guard=partial(
|
|
self._context_generation_guard,
|
|
generation,
|
|
command.storage_epoch,
|
|
context.capture_generation,
|
|
),
|
|
)
|
|
except ValueError:
|
|
# Validation/collision cannot revoke the outcome.
|
|
with self._lock:
|
|
self._record_gap_locked()
|
|
except HistoryEpochChanged:
|
|
# The exception's storage_epoch is only a hint; readiness
|
|
# comes from a later probe and never replays this command.
|
|
failed = True
|
|
except Exception:
|
|
failed = True
|
|
finally:
|
|
if context is not None:
|
|
self._release_context(context)
|
|
context = None
|
|
with self._lock:
|
|
self._pending -= 1
|
|
if failed:
|
|
self._record_gap_locked()
|
|
self._accepting = False
|
|
self._epoch = None
|
|
self._generation += 1
|
|
self._next_probe = self._clock() + _PROBE_INTERVAL
|
|
self._discard_locked()
|
|
with self._lock:
|
|
gaps = self._gaps
|
|
has_work = self._epoch is not None and not self._queue.empty()
|
|
now = self._clock()
|
|
if gaps != reported_gaps and now >= next_diagnostic:
|
|
_logger.warning(
|
|
"Execution history recording gaps: %s", gaps - reported_gaps
|
|
)
|
|
reported_gaps = gaps
|
|
next_diagnostic = now + _PROBE_INTERVAL
|
|
if not has_work:
|
|
self._wake.wait(0.05)
|
|
finally:
|
|
with self._lock:
|
|
self._accepting = False
|
|
self._epoch = None
|
|
self._discard_locked()
|
|
if unsubscribe is not None:
|
|
unsubscribe()
|
|
if db is not None:
|
|
db.model.close()
|