use super::{
invalid_resource, Arc, CheckpointAuthorityId, CheckpointBackingOwner, DeviceRuntime, VNextError,
};
use crate::vnext::SuccessfulCheckpointCaptureSeal;
use serde::Serialize;
use std::num::NonZeroU64;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize)]
pub(crate) struct CheckpointCaptureAttemptId {
checkpoint: CheckpointAuthorityId,
serial: NonZeroU64,
}
impl CheckpointCaptureAttemptId {
pub(crate) fn checkpoint_authority(self) -> CheckpointAuthorityId {
self.checkpoint
}
pub(crate) fn serial(self) -> u64 {
self.serial.get()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CapturePhase {
Fresh,
Reserved(CheckpointCaptureAttemptId),
PossiblySubmitted(CheckpointCaptureAttemptId),
Captured(CheckpointCaptureAttemptId),
Poisoned,
}
pub(super) struct CheckpointCaptureState {
next_serial: Option<NonZeroU64>,
phase: CapturePhase,
}
impl Default for CheckpointCaptureState {
fn default() -> Self {
Self {
next_serial: Some(NonZeroU64::MIN),
phase: CapturePhase::Fresh,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PermitPhase {
Reserved,
PossiblySubmitted,
}
#[must_use = "capture destination ownership must reach a proven terminal or stay retained"]
pub(crate) struct CheckpointCapturePermit<R: DeviceRuntime> {
backing: Option<Arc<CheckpointBackingOwner<R>>>,
attempt: CheckpointCaptureAttemptId,
phase: PermitPhase,
}
#[must_use = "captured storage retains its physical and logical budget"]
pub(crate) struct CapturedCheckpointBacking<R: DeviceRuntime> {
backing: Arc<CheckpointBackingOwner<R>>,
attempt: CheckpointCaptureAttemptId,
}
impl<R: DeviceRuntime> std::fmt::Debug for CheckpointCapturePermit<R> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("CheckpointCapturePermit")
.field("attempt", &self.attempt)
.field("phase", &self.phase)
.finish_non_exhaustive()
}
}
impl<R: DeviceRuntime> std::fmt::Debug for CapturedCheckpointBacking<R> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("CapturedCheckpointBacking")
.field("attempt", &self.attempt)
.finish_non_exhaustive()
}
}
impl<R: DeviceRuntime> CapturedCheckpointBacking<R> {
pub(crate) fn backing(&self) -> &Arc<CheckpointBackingOwner<R>> {
&self.backing
}
pub(crate) fn attempt_id(&self) -> CheckpointCaptureAttemptId {
self.attempt
}
}
pub(crate) struct CheckpointCaptureFinishFailure<R: DeviceRuntime> {
error: VNextError,
permit: CheckpointCapturePermit<R>,
}
impl<R: DeviceRuntime> std::fmt::Debug for CheckpointCaptureFinishFailure<R> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("CheckpointCaptureFinishFailure")
.field("error", &self.error)
.field("attempt", &self.permit.attempt)
.finish()
}
}
impl<R: DeviceRuntime> CheckpointCaptureFinishFailure<R> {
pub(crate) fn error(&self) -> &VNextError {
&self.error
}
pub(crate) fn into_parts(self) -> (VNextError, CheckpointCapturePermit<R>) {
(self.error, self.permit)
}
}
impl<R: DeviceRuntime> CheckpointBackingOwner<R> {
pub(crate) fn try_reserve_capture(
self: &Arc<Self>,
) -> Result<CheckpointCapturePermit<R>, VNextError> {
let _lifecycle = self
.plan
.read_lifecycle("reserve checkpoint capture destination")?;
let mut state = self
.capture
.lock()
.map_err(|_| invalid_resource("checkpoint capture mutex is poisoned"))?;
if state.phase != CapturePhase::Fresh {
return Err(invalid_resource(
"checkpoint storage is not a fresh capture destination",
));
}
let serial = state.next_serial.ok_or_else(|| {
invalid_resource("checkpoint capture attempt identities are exhausted")
})?;
let attempt = CheckpointCaptureAttemptId {
checkpoint: self.authority(),
serial,
};
state.next_serial = serial.get().checked_add(1).and_then(NonZeroU64::new);
state.phase = CapturePhase::Reserved(attempt);
Ok(CheckpointCapturePermit {
backing: Some(Arc::clone(self)),
attempt,
phase: PermitPhase::Reserved,
})
}
}
impl<R: DeviceRuntime> CheckpointCapturePermit<R> {
pub(crate) fn backing(&self) -> &Arc<CheckpointBackingOwner<R>> {
self.backing
.as_ref()
.expect("live capture permit retains its backing")
}
pub(crate) fn attempt_id(&self) -> CheckpointCaptureAttemptId {
self.attempt
}
pub(crate) fn mark_possibly_submitted(&mut self) -> Result<(), VNextError> {
let mut state = self
.backing()
.capture
.lock()
.map_err(|_| invalid_resource("checkpoint capture mutex is poisoned"))?;
if self.phase != PermitPhase::Reserved
|| state.phase != CapturePhase::Reserved(self.attempt)
{
return Err(invalid_resource(
"capture submission does not own its exact reserved attempt",
));
}
state.phase = CapturePhase::PossiblySubmitted(self.attempt);
drop(state);
self.phase = PermitPhase::PossiblySubmitted;
Ok(())
}
pub(crate) fn definitely_not_submitted(self) -> Result<(), CheckpointCaptureFinishFailure<R>> {
self.finish_without_capture(CapturePhase::Fresh, false)
}
pub(crate) fn finish_failed_but_quiescent(
self,
) -> Result<(), CheckpointCaptureFinishFailure<R>> {
self.finish_without_capture(CapturePhase::Poisoned, true)
}
fn finish_without_capture(
mut self,
next: CapturePhase,
recover_poison: bool,
) -> Result<(), CheckpointCaptureFinishFailure<R>> {
let transitioned = (|| {
let mut state = match self.backing().capture.lock() {
Ok(state) => state,
Err(poisoned) if recover_poison => poisoned.into_inner(),
Err(_) => return Err(invalid_resource("checkpoint capture mutex is poisoned")),
};
if self.phase != PermitPhase::PossiblySubmitted
|| state.phase != CapturePhase::PossiblySubmitted(self.attempt)
{
return Err(invalid_resource(
"capture terminal does not own its exact submitted attempt",
));
}
state.phase = next;
Ok(())
})();
if let Err(error) = transitioned {
return Err(CheckpointCaptureFinishFailure {
error,
permit: self,
});
}
drop(self.backing.take());
Ok(())
}
pub(crate) fn finish_succeeded(
mut self,
seal: &SuccessfulCheckpointCaptureSeal,
) -> Result<CapturedCheckpointBacking<R>, CheckpointCaptureFinishFailure<R>> {
let transitioned = (|| {
if seal.attempt_id() != self.attempt {
return Err(invalid_resource(
"capture success belongs to another destination attempt",
));
}
let mut state = self
.backing()
.capture
.lock()
.map_err(|_| invalid_resource("checkpoint capture mutex is poisoned"))?;
if self.phase != PermitPhase::PossiblySubmitted
|| state.phase != CapturePhase::PossiblySubmitted(self.attempt)
{
return Err(invalid_resource(
"capture success lacks its exact submitted destination permit",
));
}
state.phase = CapturePhase::Captured(self.attempt);
Ok(())
})();
if let Err(error) = transitioned {
return Err(CheckpointCaptureFinishFailure {
error,
permit: self,
});
}
Ok(CapturedCheckpointBacking {
backing: self
.backing
.take()
.expect("successful capture retains its backing"),
attempt: self.attempt,
})
}
}
impl<R: DeviceRuntime> Drop for CheckpointCapturePermit<R> {
fn drop(&mut self) {
let Some(backing) = self.backing.take() else {
return;
};
let (mut state, was_poisoned) = match backing.capture.lock() {
Ok(state) => (state, false),
Err(poisoned) => (poisoned.into_inner(), true),
};
if !was_poisoned
&& self.phase == PermitPhase::Reserved
&& state.phase == CapturePhase::Reserved(self.attempt)
{
state.phase = CapturePhase::Fresh;
} else {
state.phase = CapturePhase::Poisoned;
}
drop(state);
if self.phase == PermitPhase::PossiblySubmitted {
std::mem::forget(backing);
}
}
}
#[cfg(test)]
impl<R: DeviceRuntime> CheckpointBackingOwner<R> {
pub(in crate::vnext::resource) fn test_only_exhaust_capture_serial(&self) {
self.capture.lock().unwrap().next_serial = Some(NonZeroU64::MAX);
}
pub(in crate::vnext::resource) fn test_only_poison_capture_mutex(&self) {
let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let _state = self.capture.lock().unwrap();
panic!("injected capture state panic");
}));
}
}