"""Snapshot decisions and timers owned by execution-notified policies."""
from __future__ import annotations
import asyncio
import threading
from abc import ABC, abstractmethod
from collections.abc import Callable
from concurrent.futures import Future, InvalidStateError
from contextlib import suppress
from enum import Enum
from math import isfinite
from time import monotonic
from typing import Protocol
[docs]
class SnapshotEvent(Enum):
"""Notifications for one policy instance owned by one lazy session.
RESET starts a new VM cycle. VM_STARTING precedes the VM spawn request.
COMMAND_STARTED reserves an accepted command
before its subprocess is created. COMMAND_FINISHED confirms its exit event,
including nonzero exits. INSPECTION_STARTED/FINISHED bracket a live filesystem
request, including its stream. They affect idle time, not command counts.
CLOSE releases resources after completion or failure.
"""
RESET = "reset"
VM_STARTING = "vm_starting"
COMMAND_STARTED = "command_started"
COMMAND_FINISHED = "command_finished"
INSPECTION_STARTED = "inspection_started"
INSPECTION_FINISHED = "inspection_finished"
CLOSE = "close"
class SnapshotTimer(Protocol):
"""Cancellation handle returned by a policy's timer factory."""
def cancel(self) -> None: ...
TimerFactory = Callable[[float, Callable[[], None]], SnapshotTimer]
def _start_timer(seconds: float, callback: Callable[[], None]) -> SnapshotTimer:
try:
loop = asyncio.get_running_loop()
except RuntimeError:
timer = threading.Timer(seconds, callback)
timer.daemon = True
timer.start()
return timer
return loop.call_later(seconds, callback)
def _resolve(future: Future[None], error: BaseException | None = None) -> None:
# A previous cycle can be cancelled while a timer or child policy completes.
with suppress(InvalidStateError):
if error is None:
future.set_result(None)
else:
future.set_exception(error)
[docs]
class AbstractSnapshotPolicy(ABC):
"""Own the decision state and expose one future per VM cycle.
Use one policy instance per lazy session. Subclasses call super().__init__()
and super().notify(event) to reset or cancel the cycle future. Resolve the
future to request a snapshot, or fail it to report a policy error. The session
still waits for all accepted commands and inspections before stopping its VM.
"""
def __init__(self) -> None:
self.future: Future[None] = Future()
[docs]
@abstractmethod
def notify(self, event: SnapshotEvent) -> None:
"""Handle an execution event without blocking or performing network I/O."""
if event is SnapshotEvent.RESET:
self.future.cancel()
self.future = Future()
elif event is SnapshotEvent.CLOSE:
self.future.cancel()
[docs]
def should_snapshot(self) -> Future[None]:
"""Return this cycle's future; successful completion requests a snapshot.
The caller must not cancel or resolve it. To await it from asyncio, use
shield(wrap_future(...)) so cancellation does not affect the policy.
Returns:
A future shared by all readers of the current policy decision.
"""
return self.future
[docs]
class IdleSnapshotPolicy(AbstractSnapshotPolicy):
"""Request a snapshot after commands and inspections have been idle for seconds.
Timers start when the last command or inspection finishes. Admission,
reset, or close cancels the timer. Async callers use their event loop's timer; sync callers
use a daemon threading.Timer. Inject timer_factory for a custom scheduler.
"""
def __init__(self, seconds: float, *, timer_factory: TimerFactory = _start_timer) -> None:
if not isfinite(seconds) or seconds <= 0:
raise ValueError("idle seconds must be finite and positive")
super().__init__()
self.seconds = seconds
self.timer_factory = timer_factory
self._lock = threading.RLock()
self._active = 0
self._generation = 0
self._timer: SnapshotTimer | None = None
def _cancel_timer(self) -> None:
self._generation += 1
if self._timer is not None:
self._timer.cancel()
self._timer = None
def _expired(self, generation: int) -> None:
with self._lock:
if generation == self._generation and self._active == 0:
self._timer = None
_resolve(self.future)
[docs]
def notify(self, event: SnapshotEvent) -> None:
with self._lock:
super().notify(event)
if event in {SnapshotEvent.RESET, SnapshotEvent.CLOSE}:
self._cancel_timer()
self._active = 0
elif event in {SnapshotEvent.COMMAND_STARTED, SnapshotEvent.INSPECTION_STARTED}:
self._cancel_timer()
self._active += 1
elif event in {SnapshotEvent.COMMAND_FINISHED, SnapshotEvent.INSPECTION_FINISHED}:
if self._active <= 0:
raise ValueError("activity finish has no matching admission")
self._active -= 1
if self._active == 0 and not self.future.done():
generation = self._generation
self._timer = self.timer_factory(self.seconds, lambda: self._expired(generation))
[docs]
class DurationSnapshotPolicy(AbstractSnapshotPolicy):
"""Request a snapshot after a fixed VM lifetime, including startup time."""
def __init__(
self,
seconds: float,
*,
timer_factory: TimerFactory = _start_timer,
clock: Callable[[], float] = monotonic,
) -> None:
if not isfinite(seconds) or seconds <= 0:
raise ValueError("duration seconds must be finite and positive")
super().__init__()
self.seconds = seconds
self.timer_factory = timer_factory
self.clock = clock
self.lock = threading.RLock()
self.timer: SnapshotTimer | None = None
self.started_at: float | None = None
self.generation = 0
[docs]
def start_deferred(self) -> None:
"""Record startup time; set_duration later arms the effective deadline."""
with self.lock:
if self.started_at is None and not self.future.done():
self.started_at = self.clock()
[docs]
def set_duration(self, seconds: float) -> None:
"""Update the deadline relative to the original VM startup time.
Raises:
ValueError: The duration is not finite and positive.
"""
if not isfinite(seconds) or seconds <= 0:
raise ValueError("duration seconds must be finite and positive")
with self.lock:
self.seconds = seconds
self.schedule_timer()
[docs]
def schedule_timer(self) -> None:
self.generation += 1
if self.timer is not None:
self.timer.cancel()
self.timer = None
if self.started_at is None or self.future.done():
return
remaining = self.seconds - (self.clock() - self.started_at)
if remaining <= 0:
_resolve(self.future)
else:
generation = self.generation
self.timer = self.timer_factory(remaining, lambda: self.expired(generation))
[docs]
def expired(self, generation: int) -> None:
with self.lock:
if generation == self.generation:
_resolve(self.future)
[docs]
def notify(self, event: SnapshotEvent) -> None:
with self.lock:
super().notify(event)
if event in {SnapshotEvent.RESET, SnapshotEvent.CLOSE}:
self.started_at = None
self.schedule_timer()
elif event is SnapshotEvent.VM_STARTING and self.started_at is None and not self.future.done():
self.start_deferred()
self.schedule_timer()
[docs]
class CommandCountSnapshotPolicy(AbstractSnapshotPolicy):
"""Request a snapshot after exactly limit command admissions in one VM."""
def __init__(self, limit: int) -> None:
if isinstance(limit, bool) or not isinstance(limit, int) or limit <= 0:
raise ValueError("command limit must be a positive integer")
super().__init__()
self.limit = limit
self.commands_started = 0
[docs]
def notify(self, event: SnapshotEvent) -> None:
super().notify(event)
if event in {SnapshotEvent.RESET, SnapshotEvent.CLOSE}:
self.commands_started = 0
elif event is SnapshotEvent.COMMAND_STARTED:
self.commands_started += 1
if self.commands_started >= self.limit:
_resolve(self.future)
[docs]
class CompositeSnapshotPolicy(AbstractSnapshotPolicy):
"""Request a snapshot when any child policy's future completes."""
def __init__(self, *policies: AbstractSnapshotPolicy) -> None:
if not policies:
raise ValueError("at least one snapshot policy is required")
super().__init__()
self.policies = policies
self._watch_children()
def _watch_children(self) -> None:
target = self.future
def completed(child: Future[None]) -> None:
if child.cancelled():
target.cancel()
else:
_resolve(target, child.exception())
for policy in self.policies:
policy.should_snapshot().add_done_callback(completed)
[docs]
def notify(self, event: SnapshotEvent) -> None:
super().notify(event)
for policy in self.policies:
policy.notify(event)
if event is SnapshotEvent.RESET:
self._watch_children()