use super::{
expected_lease_transition, invalid_resource, AdmissionFitPolicy, AdmissionPressureAction, Arc,
BTreeSet, BufferDescriptor, CoreOwnedAllocation, DeviceRuntime, DynamicResourceShape,
ResourceAllocation, ResourceId, ResourceLeaseAction, ResourceLeaseEntry, ResourceLeaseState,
ResourceOwnedBuffer, ResourceReservation, ResourceReservationBatch,
ResourceTransactionIdentity, ResourceWorkShape, StaticProvisioningBinding, VNextError,
};
pub(super) struct OwnedLeaseSlot<B> {
pub(super) entry: ResourceLeaseEntry,
pub(super) actual_resource_id: Option<ResourceId>,
pub(super) actual_generation: Option<u64>,
pub(super) descriptor: Option<BufferDescriptor>,
pub(super) buffer: Option<B>,
}
impl<B> OwnedLeaseSlot<B> {
pub(super) fn new(reservation: &ResourceReservation) -> Self {
Self {
entry: ResourceLeaseEntry::from_reservation(reservation, ResourceLeaseState::Active),
actual_resource_id: None,
actual_generation: None,
descriptor: None,
buffer: None,
}
}
pub(super) fn install(&mut self, allocation: CoreOwnedAllocation<B>) {
self.actual_resource_id = Some(allocation.resource_id);
self.actual_generation = Some(allocation.generation);
self.descriptor = Some(allocation.descriptor);
self.buffer = Some(allocation.buffer);
}
pub(super) fn clear(&mut self) {
drop(self.buffer.take());
self.descriptor.take();
self.actual_resource_id.take();
self.actual_generation.take();
}
pub(super) fn take_allocation(&mut self) -> Option<CoreOwnedAllocation<B>> {
Some(CoreOwnedAllocation {
resource_id: self.actual_resource_id.take()?,
generation: self.actual_generation.take()?,
descriptor: self.descriptor.take()?,
buffer: self.buffer.take()?,
})
}
pub(super) fn restore_allocation(&mut self, allocation: CoreOwnedAllocation<B>) {
debug_assert!(self.buffer.is_none());
self.install(allocation);
}
}
pub struct LeasedBufferView<'a, B> {
pub(super) identity: &'a ResourceTransactionIdentity,
pub(super) admission: &'a StaticProvisioningBinding,
pub(super) resource_id: &'a ResourceId,
pub(super) generation: u64,
pub(super) descriptor: &'a BufferDescriptor,
pub(super) buffer: &'a B,
}
impl<'a, B> LeasedBufferView<'a, B> {
pub fn identity(&self) -> &ResourceTransactionIdentity {
self.identity
}
pub fn admission(&self) -> &StaticProvisioningBinding {
self.admission
}
pub fn resource_id(&self) -> &ResourceId {
self.resource_id
}
pub const fn generation(&self) -> u64 {
self.generation
}
pub fn committed_descriptor(&self) -> &BufferDescriptor {
self.descriptor
}
pub fn buffer(&self) -> &B {
self.buffer
}
}
#[must_use = "a resource lease is the batch owner of committed buffers"]
pub struct StaticProvisioningLease<R>
where
R: DeviceRuntime,
{
pub(super) slots: Vec<OwnedLeaseSlot<R::Buffer>>,
pub(super) identity: ResourceTransactionIdentity,
pub(super) admission: StaticProvisioningBinding,
pub(super) runtime: Arc<R>,
}
impl<R> StaticProvisioningLease<R>
where
R: DeviceRuntime,
{
pub(crate) fn runtime(&self) -> &Arc<R> {
&self.runtime
}
pub(super) fn new(
runtime: Arc<R>,
identity: &ResourceTransactionIdentity,
admission: &StaticProvisioningBinding,
reservations: &ResourceReservationBatch,
) -> Self {
Self {
slots: reservations
.reservations()
.iter()
.map(OwnedLeaseSlot::new)
.collect(),
identity: identity.clone(),
admission: admission.clone(),
runtime,
}
}
pub fn identity(&self) -> &ResourceTransactionIdentity {
&self.identity
}
pub fn admission(&self) -> &StaticProvisioningBinding {
&self.admission
}
pub fn state(&self) -> ResourceLeaseState {
let mut states = self
.slots
.iter()
.filter(|slot| slot.buffer.is_some())
.map(|slot| slot.entry.state);
let Some(first) = states.next() else {
return ResourceLeaseState::Cancelled;
};
if states.all(|state| state == first) {
first
} else {
ResourceLeaseState::Mixed
}
}
pub fn entries(&self) -> impl Iterator<Item = &ResourceLeaseEntry> {
self.slots
.iter()
.filter(|slot| slot.buffer.is_some())
.map(|slot| &slot.entry)
}
pub fn plan_static_entries(&self) -> impl Iterator<Item = &ResourceLeaseEntry> {
self.slots
.iter()
.filter(|slot| slot.buffer.is_some())
.map(|slot| &slot.entry)
}
pub(crate) fn view(
&self,
resource_id: &ResourceId,
generation: u64,
) -> Result<LeasedBufferView<'_, R::Buffer>, VNextError> {
let slot = self
.slots
.iter()
.find(|slot| {
slot.entry.resource_id == *resource_id && slot.entry.generation == generation
})
.ok_or_else(|| invalid_resource("lease does not contain that resource generation"))?;
if slot.entry.state != ResourceLeaseState::Active {
return Err(VNextError::InvalidLeaseTransition {
lease_id: self.identity.transaction_id.to_string(),
from: slot.entry.state.as_str(),
action: "borrow_live_buffer",
});
}
let descriptor = slot
.descriptor
.as_ref()
.ok_or_else(|| invalid_resource("lease resource is not committed"))?;
let buffer = slot
.buffer
.as_ref()
.ok_or_else(|| invalid_resource("lease resource buffer is not live"))?;
Ok(LeasedBufferView {
identity: &self.identity,
admission: &self.admission,
resource_id: &slot.entry.resource_id,
generation: slot.entry.generation,
descriptor,
buffer,
})
}
pub(crate) fn plan_static_view(
&self,
slot_index: usize,
allocation: &ResourceAllocation,
) -> Result<LeasedBufferView<'_, R::Buffer>, VNextError> {
let slot = self
.slots
.get(slot_index)
.ok_or_else(|| invalid_resource("plan-static slot index is out of range"))?;
let descriptor = slot
.descriptor
.as_ref()
.ok_or_else(|| invalid_resource("plan-static resource is not committed"))?;
let buffer = slot
.buffer
.as_ref()
.ok_or_else(|| invalid_resource("plan-static resource buffer is not live"))?;
if slot.entry.resource_id() != allocation.resource_id()
|| slot.entry.size_bytes() != allocation.size_bytes()
|| slot.entry.alignment_bytes() != allocation.alignment_bytes()
|| slot.entry.usage() != allocation.usage()
|| slot.entry.element_type() != allocation.element_type()
|| slot.entry.generation() == 0
|| slot.entry.state() != ResourceLeaseState::Active
|| slot.actual_resource_id.as_ref() != Some(allocation.resource_id())
|| slot.actual_generation != Some(slot.entry.generation())
|| descriptor.resource_id != *allocation.resource_id()
|| descriptor.size_bytes != allocation.size_bytes()
|| descriptor.alignment_bytes != allocation.alignment_bytes()
|| descriptor.usage != allocation.usage()
|| descriptor.element_type != allocation.element_type()
{
return Err(invalid_resource(
"plan-static slot differs from its immutable allocation",
));
}
Ok(LeasedBufferView {
identity: &self.identity,
admission: &self.admission,
resource_id: &slot.entry.resource_id,
generation: slot.entry.generation,
descriptor,
buffer,
})
}
pub(super) fn buffer(&self, order: usize) -> Option<&R::Buffer> {
self.slots.get(order).and_then(|slot| slot.buffer.as_ref())
}
pub(super) fn install(&mut self, order: usize, allocation: CoreOwnedAllocation<R::Buffer>) {
self.slots[order].install(allocation);
}
pub(super) fn clear(&mut self, order: usize) {
self.slots[order].clear();
}
pub(super) fn transition_subset(
&mut self,
orders: &[usize],
action: ResourceLeaseAction,
) -> Result<
(
ResourceLeaseState,
ResourceLeaseState,
Vec<ResourceLeaseEntry>,
),
VNextError,
> {
if orders.is_empty() {
return Err(invalid_resource(
"lease transition subset must not be empty",
));
}
let mut unique = BTreeSet::new();
let mut common_before = None;
for &order in orders {
if !unique.insert(order) {
return Err(invalid_resource("lease transition subset is duplicated"));
}
let slot = self
.slots
.get(order)
.ok_or_else(|| invalid_resource("lease transition order is out of bounds"))?;
if slot.buffer.is_none() {
return Err(invalid_resource(
"lease transition targets a non-live buffer",
));
}
let before = slot.entry.state;
if common_before.is_some_and(|common| common != before) {
return Err(invalid_resource(
"one lease receipt cannot hide heterogeneous before states",
));
}
if expected_lease_transition(action, before).is_none() {
return Err(VNextError::InvalidLeaseTransition {
lease_id: self.identity.transaction_id.to_string(),
from: before.as_str(),
action: action.as_str(),
});
}
common_before = Some(before);
}
let before = common_before.expect("non-empty subset has a before state");
let after = expected_lease_transition(action, before)
.expect("lease subset was preflight validated");
for &order in orders {
self.slots[order].entry.state = after;
}
Ok((
before,
after,
orders
.iter()
.map(|&order| self.slots[order].entry.clone())
.collect(),
))
}
pub(super) fn take_owned_buffers(
&mut self,
reservations: &ResourceReservationBatch,
) -> Vec<ResourceOwnedBuffer<R::Buffer>> {
self.slots
.iter_mut()
.zip(reservations.reservations())
.enumerate()
.filter_map(|(order, (slot, reservation))| {
let allocation = slot.take_allocation()?;
Some(ResourceOwnedBuffer {
order,
expected_resource_id: reservation.resource_id.clone(),
actual_resource_id: allocation.resource_id,
expected_generation: reservation.generation,
actual_generation: allocation.generation,
expected_descriptor: BufferDescriptor {
resource_id: reservation.resource_id.clone(),
size_bytes: reservation.size_bytes,
alignment_bytes: reservation.alignment_bytes,
usage: reservation.usage,
element_type: reservation.element_type,
},
actual_descriptor: allocation.descriptor,
buffer: allocation.buffer,
})
})
.collect()
}
pub(super) fn restore_owned_buffers(&mut self, buffers: Vec<ResourceOwnedBuffer<R::Buffer>>) {
for buffer in buffers {
let (order, allocation) = buffer.into_allocation();
self.slots[order].restore_allocation(allocation);
}
}
}
macro_rules! scoped_resource_admission_request {
($name:ident, $single_sequence:literal) => {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct $name {
pub(super) work_shape: ResourceWorkShape,
pub(super) fit_policy: AdmissionFitPolicy,
pub(super) pressure_action: AdmissionPressureAction,
}
impl $name {
pub fn new(
work_shape: ResourceWorkShape,
fit_policy: AdmissionFitPolicy,
pressure_action: AdmissionPressureAction,
) -> Result<Self, VNextError> {
if $single_sequence
&& (work_shape.immediate_sequences() != 1 || work_shape.fit_sequences() != 1)
{
return Err(invalid_resource(
"sequence resource admission requires a single-sequence shape",
));
}
Ok(Self {
work_shape,
fit_policy,
pressure_action,
})
}
pub fn work_shape(&self) -> &ResourceWorkShape {
&self.work_shape
}
pub(crate) const fn immediate_shape(&self) -> DynamicResourceShape {
self.work_shape.immediate_shape()
}
pub(crate) const fn fit_shape(&self) -> DynamicResourceShape {
match self.fit_policy {
AdmissionFitPolicy::ImmediateOnly => self.work_shape.immediate_shape(),
AdmissionFitPolicy::FullInputMustFit => self.work_shape.fit_shape(),
}
}
pub const fn fit_policy(&self) -> AdmissionFitPolicy {
self.fit_policy
}
pub const fn pressure_action(&self) -> AdmissionPressureAction {
self.pressure_action
}
}
};
}
scoped_resource_admission_request!(RequestResourceAdmissionRequest, false);
scoped_resource_admission_request!(SequenceResourceAdmissionRequest, true);