use a3s_box_core::{
ExecutionGeneration, ExecutionId, ExecutionManagerError, ExecutionManagerResult,
ExecutionSnapshotId, KillExecutionOptions, OperationId, RestartExecutionOptions,
};
use super::record::{
apply_handle, apply_restart_handle, apply_start_handle, clear_live_runtime,
clear_live_runtime_for_cold_pause, execution_id,
};
use super::support::{generation, managed_state};
use super::{LocalExecutionHandle, LocalExecutionManager};
use crate::{
BoxRecord, ManagedExecutionOperation, ManagedExecutionReservation, ManagedExecutionState,
ManagedExecutionStoreError, SnapshotStore,
};
impl LocalExecutionManager {
pub(super) async fn reserve(
&self,
mut record: BoxRecord,
) -> ExecutionManagerResult<ManagedExecutionReservation> {
let store = self.store.clone();
let Some(snapshot_id) = record
.managed_execution
.as_ref()
.and_then(|metadata| metadata.request.rootfs_snapshot_id.clone())
else {
return run_store(move || store.reserve(record)).await;
};
let home_dir = self.home_dir.clone();
tokio::task::spawn_blocking(move || {
let snapshots = SnapshotStore::new(&home_dir.join("snapshots")).map_err(|error| {
ExecutionManagerError::Unavailable(format!(
"failed to open filesystem snapshot store: {error}"
))
})?;
let _snapshot_lock = snapshots.acquire_exclusive_lock().map_err(|error| {
ExecutionManagerError::Unavailable(format!(
"failed to lock filesystem snapshot store: {error}"
))
})?;
let metadata = snapshots
.get(snapshot_id.as_str())
.map_err(|error| {
ExecutionManagerError::Unavailable(format!(
"failed to inspect filesystem snapshot {snapshot_id}: {error}"
))
})?
.ok_or_else(|| {
ExecutionManagerError::Unavailable(format!(
"filesystem snapshot {snapshot_id} is unavailable"
))
})?;
if metadata.id != snapshot_id.as_str()
|| !snapshots.rootfs_path(snapshot_id.as_str()).is_dir()
{
return Err(ExecutionManagerError::Unavailable(format!(
"filesystem snapshot {snapshot_id} is not a valid published snapshot"
)));
}
metadata.require_image_config().map_err(|error| {
ExecutionManagerError::Unavailable(format!(
"filesystem snapshot {snapshot_id} cannot be restored: {error}"
))
})?;
let health_check = record
.health_check
.clone()
.or_else(|| metadata.health_check.clone());
let healthcheck_disabled =
record.healthcheck_disabled || metadata.healthcheck_disabled;
record.health_check = health_check.clone();
record.healthcheck_disabled = healthcheck_disabled;
if let Some(managed) = record.managed_execution.as_mut() {
managed.request.policy.health_check = health_check;
managed.request.policy.healthcheck_disabled = healthcheck_disabled;
}
super::record::validate_record_health(&record)?;
#[cfg(windows)]
if !healthcheck_disabled
&& metadata
.image_config
.as_ref()
.and_then(|config| config.health_check.as_ref())
.is_some_and(a3s_box_core::SnapshotImageHealthCheck::is_enabled)
{
return Err(ExecutionManagerError::InvalidRequest(
"container health checks are not supported on Windows; the snapshot image defines an effective health check"
.to_string(),
));
}
store.reserve(record).map_err(map_store_error)
})
.await
.map_err(|error| {
ExecutionManagerError::Internal(format!("managed state task failed: {error}"))
})?
}
pub(super) async fn get(
&self,
execution_id: &ExecutionId,
) -> ExecutionManagerResult<Option<BoxRecord>> {
let store = self.store.clone();
let execution_id = execution_id.clone();
run_store(move || store.get(&execution_id)).await
}
pub(super) async fn get_by_operation(
&self,
operation_id: &OperationId,
) -> ExecutionManagerResult<Option<BoxRecord>> {
let store = self.store.clone();
let operation_id = operation_id.clone();
run_store(move || store.get_by_operation_id(&operation_id)).await
}
pub(super) async fn transition(
&self,
record: &BoxRecord,
from: ManagedExecutionState,
to: ManagedExecutionState,
update: RuntimeUpdate,
) -> ExecutionManagerResult<BoxRecord> {
let store = self.store.clone();
let execution_id = execution_id(record)?;
let generation = generation(record, &execution_id)?;
run_store(move || {
store.transition_with(&execution_id, generation, from, to, |record| match update {
RuntimeUpdate::None => {}
RuntimeUpdate::Handle(handle) => apply_handle(record, &handle),
RuntimeUpdate::StartHandle(handle) => apply_start_handle(record, &handle),
RuntimeUpdate::Terminal(exit_code) => clear_live_runtime(record, exit_code),
RuntimeUpdate::KillTerminal(exit_code) => {
clear_live_runtime(record, exit_code);
record.stopped_by_user = true;
}
RuntimeUpdate::ColdPause => clear_live_runtime_for_cold_pause(record),
RuntimeUpdate::PauseClaim(keep_memory) => {
if let Some(metadata) = record.managed_execution.as_mut() {
metadata.pending_operation =
Some(ManagedExecutionOperation::Pause { keep_memory });
}
}
RuntimeUpdate::KillClaim(options) => {
if let Some(metadata) = record.managed_execution.as_mut() {
metadata.pending_operation = Some(ManagedExecutionOperation::Kill {
signal: options.signal,
timeout_secs: options.timeout_secs,
});
}
}
RuntimeUpdate::SnapshotClaim {
snapshot_id,
source_state,
} => {
if let Some(metadata) = record.managed_execution.as_mut() {
metadata.pending_operation = Some(ManagedExecutionOperation::Snapshot {
snapshot_id,
source_state,
});
}
}
RuntimeUpdate::RestartClaim {
operation_id,
options,
} => {
if let Some(metadata) = record.managed_execution.as_mut() {
metadata.pending_operation = Some(ManagedExecutionOperation::Restart {
operation_id,
source_generation: metadata.generation,
source_state: from,
stop_timeout_secs: options.stop_timeout_secs,
});
}
}
RuntimeUpdate::RestartAdvance => clear_live_runtime(record, None),
RuntimeUpdate::RestartHandle(handle) => apply_restart_handle(record, &handle),
RuntimeUpdate::RestartFailed(exit_code) => clear_live_runtime(record, exit_code),
})
})
.await
}
pub(super) async fn complete_with_handle(
&self,
record: &BoxRecord,
from: ManagedExecutionState,
to: ManagedExecutionState,
handle: LocalExecutionHandle,
) -> ExecutionManagerResult<BoxRecord> {
let update =
if from == ManagedExecutionState::Starting && to == ManagedExecutionState::Running {
RuntimeUpdate::StartHandle(handle)
} else {
RuntimeUpdate::Handle(handle)
};
self.complete_transition(record, from, to, update).await
}
pub(super) async fn complete_transition(
&self,
record: &BoxRecord,
from: ManagedExecutionState,
to: ManagedExecutionState,
update: RuntimeUpdate,
) -> ExecutionManagerResult<BoxRecord> {
let execution_id = execution_id(record)?;
let current_generation = generation(record, &execution_id)?;
let expected_generation = if matches!(
(from, to),
(
ManagedExecutionState::Pausing,
ManagedExecutionState::Paused
) | (
ManagedExecutionState::Resuming,
ManagedExecutionState::Running
)
) {
ExecutionGeneration::new(current_generation.get().checked_add(1).ok_or_else(|| {
ExecutionManagerError::Internal(format!(
"execution {execution_id} generation is exhausted"
))
})?)?
} else {
current_generation
};
match self.transition(record, from, to, update).await {
Ok(record) => Ok(record),
Err(error @ ExecutionManagerError::Conflict { .. }) => {
let Some(current) = self.get(&execution_id).await? else {
return Err(ExecutionManagerError::NotFound(execution_id));
};
if managed_state(¤t)? == to
&& generation(¤t, &execution_id)? == expected_generation
{
Ok(current)
} else {
Err(error)
}
}
Err(error) => Err(error),
}
}
}
pub(super) enum RuntimeUpdate {
None,
Handle(LocalExecutionHandle),
StartHandle(LocalExecutionHandle),
Terminal(Option<i32>),
KillTerminal(Option<i32>),
ColdPause,
PauseClaim(bool),
KillClaim(KillExecutionOptions),
SnapshotClaim {
snapshot_id: ExecutionSnapshotId,
source_state: ManagedExecutionState,
},
RestartClaim {
operation_id: OperationId,
options: RestartExecutionOptions,
},
RestartAdvance,
RestartHandle(LocalExecutionHandle),
RestartFailed(Option<i32>),
}
pub(super) async fn run_store<T>(
operation: impl FnOnce() -> Result<T, ManagedExecutionStoreError> + Send + 'static,
) -> ExecutionManagerResult<T>
where
T: Send + 'static,
{
tokio::task::spawn_blocking(operation)
.await
.map_err(|error| {
ExecutionManagerError::Internal(format!("managed state task failed: {error}"))
})?
.map_err(map_store_error)
}
pub(super) fn map_store_error(error: ManagedExecutionStoreError) -> ExecutionManagerError {
match error {
ManagedExecutionStoreError::Io(error) => {
ExecutionManagerError::Unavailable(error.to_string())
}
ManagedExecutionStoreError::NotFound(execution_id) => {
ExecutionManagerError::NotFound(execution_id)
}
ManagedExecutionStoreError::Conflict {
execution_id,
message,
} => ExecutionManagerError::Conflict {
execution_id,
message,
},
ManagedExecutionStoreError::Unmanaged(execution_id) => ExecutionManagerError::Internal(
format!("execution record is not managed: {execution_id}"),
),
ManagedExecutionStoreError::InvalidRecord(message) => {
ExecutionManagerError::Internal(message)
}
ManagedExecutionStoreError::InvalidTransition {
execution_id,
from,
to,
} => ExecutionManagerError::Internal(format!(
"invalid managed transition for {execution_id}: {from} -> {to}"
)),
}
}