use super::{
align_up_resource, invalid_resource, AdmissionDeferred, AdmissionDemand, AdmissionFitPolicy,
AdmissionPressureAction, AdmissionRejected, AllocationKind, AllocationLifetime, Arc,
BackingClaimCertificate, BackingPrepareDecision, BufferUsage, CapacityEntry, CapacityUnits,
CapacityVector, DeviceBufferRetention, DeviceRuntime, DynamicBackingDeferred,
EvaluatedBackingProjection, EvaluatedBackingRequest, LogicalBackingBufferView,
LogicalBackingSliceAuthority, LogicalBackingSliceEvidence, PhysicalBackingClaimIdentity,
PlanHash, PlanRuntimeResources, ResourceId, TrustedPlanRuntimeBinding, VNextError,
};
use crate::vnext::{
CheckpointAuthorityId, CheckpointCapacityClaimDecision, CheckpointRetentionSkipReason,
LogicalCheckpointLease,
};
use std::collections::BTreeMap;
mod transfer;
use std::sync::Mutex;
mod capture;
pub(crate) use capture::*;
mod maintenance;
pub use maintenance::*;
#[derive(Debug)]
pub(crate) struct CheckpointBackingRequest {
resource_id: ResourceId,
logical_bytes: u64,
}
impl CheckpointBackingRequest {
pub(crate) fn new(resource_id: ResourceId, logical_bytes: u64) -> Result<Self, VNextError> {
if logical_bytes == 0 {
return Err(invalid_resource("checkpoint backing cannot be empty"));
}
Ok(Self {
resource_id,
logical_bytes,
})
}
pub(crate) fn resource_id(&self) -> &ResourceId {
&self.resource_id
}
pub(crate) const fn logical_bytes(&self) -> u64 {
self.logical_bytes
}
}
#[derive(Debug)]
pub(crate) struct CheckpointBackingRequests {
plan_hash: PlanHash,
requests: Vec<CheckpointBackingRequest>,
}
impl CheckpointBackingRequests {
pub(crate) fn new(
plan_hash: PlanHash,
mut requests: Vec<CheckpointBackingRequest>,
) -> Result<Self, VNextError> {
requests.sort_by(|left, right| left.resource_id.cmp(&right.resource_id));
if requests.is_empty()
|| requests
.windows(2)
.any(|pair| pair[0].resource_id == pair[1].resource_id)
{
return Err(invalid_resource(
"checkpoint layout must supply non-empty unique base resources",
));
}
Ok(Self {
plan_hash,
requests,
})
}
pub(crate) fn requests(&self) -> &[CheckpointBackingRequest] {
&self.requests
}
}
pub(crate) enum CheckpointBackingAllocationDecision<R: DeviceRuntime> {
Allocated(Arc<CheckpointBackingOwner<R>>),
Skipped(CheckpointRetentionSkipReason),
Deferred(AdmissionDeferred),
BackingDeferred(DynamicBackingDeferred),
PermanentRejected(AdmissionRejected),
}
#[must_use = "checkpoint backing and its capacity must remain owned through device completion"]
pub(crate) struct CheckpointBackingOwner<R: DeviceRuntime> {
backing_slices: Vec<LogicalBackingSliceAuthority>,
logical_lease: LogicalCheckpointLease,
plan: Arc<PlanRuntimeResources<R>>,
logical_bytes: u64,
extent_bytes: u64,
capture: Mutex<CheckpointCaptureState>,
}
impl<R: DeviceRuntime> CheckpointBackingOwner<R> {
pub(crate) fn authority(&self) -> CheckpointAuthorityId {
self.logical_lease.authority()
}
pub(crate) fn claims(&self) -> &CapacityVector {
self.logical_lease.claims()
}
pub(crate) const fn logical_bytes(&self) -> u64 {
self.logical_bytes
}
pub(crate) const fn extent_bytes(&self) -> u64 {
self.extent_bytes
}
pub(crate) fn backing_evidence(&self) -> impl Iterator<Item = &LogicalBackingSliceEvidence> {
self.backing_slices
.iter()
.map(LogicalBackingSliceAuthority::evidence)
}
pub(in crate::vnext::resource) fn backing_slices(&self) -> &[LogicalBackingSliceAuthority] {
&self.backing_slices
}
pub(in crate::vnext::resource) fn plan_resources(&self) -> &Arc<PlanRuntimeResources<R>> {
&self.plan
}
pub(crate) fn device_buffer_retention(self: &Arc<Self>) -> DeviceBufferRetention {
DeviceBufferRetention::pair(Arc::clone(self), Arc::clone(&self.plan))
}
pub(crate) fn view<'owner>(
self: &'owner Arc<Self>,
resource_id: &ResourceId,
) -> Result<LogicalBackingBufferView<'owner, R::Buffer>, VNextError> {
let index = self
.backing_slices
.binary_search_by(|slice| slice.resource_id().cmp(resource_id))
.map_err(|_| invalid_resource("checkpoint does not own this resource"))?;
let mut view = self.plan.dynamic_pools.view(&self.backing_slices[index])?;
for binding in &mut view.bindings {
binding.retention =
DeviceBufferRetention::pair(Arc::clone(self), Arc::clone(&binding.chunk));
}
Ok(view)
}
}
struct EvaluatedCheckpointBacking<'a> {
slices: Vec<EvaluatedBackingRequest<'a>>,
demand: AdmissionDemand,
logical_bytes: u64,
extent_bytes: u64,
}
impl<R: DeviceRuntime> TrustedPlanRuntimeBinding<R> {
fn evaluate_checkpoint_backing(
&self,
request: &CheckpointBackingRequests,
) -> Result<EvaluatedCheckpointBacking<'_>, VNextError> {
if request.plan_hash != *self.plan_hash() {
return Err(invalid_resource(
"checkpoint layout belongs to another plan",
));
}
let mut requested_slices = Vec::with_capacity(request.requests.len());
let mut domain_bytes = BTreeMap::new();
let mut logical_bytes = 0_u64;
let mut extent_bytes = 0_u64;
for resource in &request.requests {
let (domain, descriptor) = self
.dynamic_pools()
.domains
.iter()
.find_map(|domain| {
domain
.descriptors
.iter()
.find(|descriptor| descriptor.base_resource_id() == &resource.resource_id)
.map(|descriptor| (domain, descriptor))
})
.ok_or_else(|| {
invalid_resource("checkpoint references an unknown plan resource")
})?;
if descriptor.lifetime() != AllocationLifetime::Sequence
|| descriptor.usage() != BufferUsage::State
|| *descriptor.kind() != AllocationKind::Value
|| resource.logical_bytes > descriptor.theoretical_maximum_request_bytes()?
{
return Err(invalid_resource(
"checkpoint allocation must fit the plan's Sequence state descriptor",
));
}
let pool = self
.dynamic_pools()
.pools
.get(domain.pool_id())
.ok_or_else(|| invalid_resource("checkpoint descriptor has no State pool"))?;
let capacity_bytes =
align_up_resource(resource.logical_bytes, pool.allocation_quantum())?;
let total = domain_bytes.entry(domain.domain_id()).or_insert(0_u64);
*total = total
.checked_add(capacity_bytes)
.ok_or_else(|| invalid_resource("checkpoint domain capacity overflows u64"))?;
logical_bytes = logical_bytes
.checked_add(resource.logical_bytes)
.ok_or_else(|| invalid_resource("checkpoint logical bytes overflow u64"))?;
extent_bytes = extent_bytes
.checked_add(capacity_bytes)
.ok_or_else(|| invalid_resource("checkpoint extent bytes overflow u64"))?;
requested_slices.push(EvaluatedBackingRequest {
domain,
claim_identity: PhysicalBackingClaimIdentity::new(
domain.pool_id().clone(),
vec![resource.resource_id.clone()],
)?,
capacity_size_bytes: capacity_bytes,
reusable_execution_bucket_id: None,
projections: vec![EvaluatedBackingProjection {
descriptor,
physical_offset_bytes: 0,
logical_size_bytes: resource.logical_bytes,
capacity_size_bytes: capacity_bytes,
}],
});
}
let capacity = CapacityVector::new(
domain_bytes
.into_iter()
.map(|(domain, bytes)| CapacityEntry::new(domain, CapacityUnits::new(bytes)))
.collect::<Result<Vec<_>, _>>()?,
)?;
let demand = AdmissionDemand::from_plan(
capacity.clone(),
capacity,
AdmissionFitPolicy::ImmediateOnly,
AdmissionPressureAction::WaitForRelease,
)?;
Ok(EvaluatedCheckpointBacking {
slices: requested_slices,
demand,
logical_bytes,
extent_bytes,
})
}
pub(crate) fn try_allocate_checkpoint_backing(
&self,
request: &CheckpointBackingRequests,
) -> Result<CheckpointBackingAllocationDecision<R>, VNextError> {
let _lifecycle = self
.resources
.read_lifecycle("allocate checkpoint backing")?;
let EvaluatedCheckpointBacking {
slices: requested_slices,
demand,
logical_bytes,
extent_bytes,
} = self.evaluate_checkpoint_backing(request)?;
let logical_lease = match self
.logical_admission()
.try_claim_checkpoint(&demand, extent_bytes)?
{
CheckpointCapacityClaimDecision::Claimed(lease) => lease,
CheckpointCapacityClaimDecision::Skipped(reason) => {
return Ok(CheckpointBackingAllocationDecision::Skipped(reason));
}
CheckpointCapacityClaimDecision::Deferred(deferred) => {
return Ok(CheckpointBackingAllocationDecision::Deferred(deferred));
}
CheckpointCapacityClaimDecision::PermanentRejected(rejected) => {
return Ok(CheckpointBackingAllocationDecision::PermanentRejected(
rejected,
));
}
};
if !self
.logical_admission()
.owns_checkpoint_claim(&logical_lease)
{
drop(logical_lease);
return Err(invalid_resource(
"checkpoint claim belongs to another coordinator",
));
}
let prepared = match self
.dynamic_pools()
.prepare_checkpoint_claim(&requested_slices)?
{
BackingPrepareDecision::Prepared(prepared) => prepared,
BackingPrepareDecision::Deferred(deferred) => {
return Ok(CheckpointBackingAllocationDecision::BackingDeferred(
deferred,
));
}
};
let owner = CheckpointBackingOwner {
backing_slices: prepared.commit(),
logical_lease,
plan: Arc::clone(&self.resources),
logical_bytes,
extent_bytes,
capture: Mutex::new(CheckpointCaptureState::default()),
};
BackingClaimCertificate::from_slices(&owner.backing_slices)?
.bind(&owner.backing_slices, &demand)?;
Ok(CheckpointBackingAllocationDecision::Allocated(Arc::new(
owner,
)))
}
}