use core::num::NonZeroUsize;
use super::{
ByteLease, ByteLeaseReservation, BytePool, BytePoolMaterializeOutcome, BytePoolReserveOutcome,
LeaseStatus, RetainedByteLease,
};
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct BytePoolStatistics {
pub byte_capacity: usize,
pub free_byte_capacity: usize,
pub maximum_lease_count: usize,
pub active_lease_count: usize,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum BytePoolInconsistency {
LeaseUnavailable { slot: usize },
StorageUnavailable { start: usize, end: usize },
}
#[must_use = "byte-pool lease outcomes must be handled"]
pub enum BytePoolLeaseOutcome<'pool> {
Leased(ByteLease<'pool>),
Exhausted {
requested_maximum: NonZeroUsize,
},
LeaseCapacityReached {
requested_maximum: NonZeroUsize,
},
LeaseIdentitiesExhausted {
requested_maximum: NonZeroUsize,
},
Inconsistent {
requested_maximum: NonZeroUsize,
reason: BytePoolInconsistency,
},
}
#[must_use = "retained byte-pool lease outcomes must be handled"]
pub enum BytePoolResumeOutcome<'pool> {
Resumed(ByteLease<'pool>),
Unavailable {
retained: RetainedByteLease,
},
Inconsistent {
retained: RetainedByteLease,
reason: BytePoolInconsistency,
},
}
#[derive(Debug, Eq, PartialEq)]
#[must_use = "byte-pool recycle outcomes must be handled"]
pub enum BytePoolRecycleOutcome {
Recycled,
Unavailable { retained: RetainedByteLease },
}
impl BytePool {
pub fn byte_capacity(&self) -> usize {
self.storage.len()
}
pub fn reserve_up_to(&mut self, requested_maximum: NonZeroUsize) -> BytePoolReserveOutcome<'_> {
self.prepare_reservation(requested_maximum, false)
}
#[cfg(feature = "performance-evidence")]
pub const fn evidence(&self) -> super::BytePoolEvidence {
self.evidence
}
#[cfg(feature = "performance-evidence")]
pub fn take_evidence(&mut self) -> super::BytePoolEvidence {
self.collect_released();
let active_leases = self.evidence.active_leases().current();
core::mem::replace(
&mut self.evidence,
super::BytePoolEvidence::starting_with_active_leases(active_leases),
)
}
pub fn statistics(&mut self) -> BytePoolStatistics {
self.collect_released();
BytePoolStatistics {
byte_capacity: self.storage.len(),
free_byte_capacity: self.free_extents.iter().map(core::ops::Range::len).sum(),
maximum_lease_count: self.leases.len(),
active_lease_count: self
.leases
.iter()
.filter(|lease| lease.status.get() == LeaseStatus::Leased)
.count(),
}
}
pub fn lease(&mut self, requested_maximum: NonZeroUsize) -> BytePoolLeaseOutcome<'_> {
match self.prepare_reservation(requested_maximum, false) {
BytePoolReserveOutcome::Reserved(reservation) => {
let reserved_maximum = reservation.reserved_maximum();
lease_outcome(reservation.materialize(reserved_maximum), requested_maximum)
}
BytePoolReserveOutcome::Exhausted { requested_maximum } => {
BytePoolLeaseOutcome::Exhausted { requested_maximum }
}
BytePoolReserveOutcome::LeaseCapacityReached { requested_maximum } => {
BytePoolLeaseOutcome::LeaseCapacityReached { requested_maximum }
}
BytePoolReserveOutcome::LeaseIdentitiesExhausted { requested_maximum } => {
BytePoolLeaseOutcome::LeaseIdentitiesExhausted { requested_maximum }
}
BytePoolReserveOutcome::Inconsistent {
requested_maximum,
reason,
} => BytePoolLeaseOutcome::Inconsistent {
requested_maximum,
reason,
},
}
}
pub fn reserve_exact(
&mut self,
requested_byte_count: NonZeroUsize,
) -> BytePoolReserveOutcome<'_> {
self.prepare_reservation(requested_byte_count, true)
}
fn prepare_reservation(
&mut self,
requested_maximum: NonZeroUsize,
exact: bool,
) -> BytePoolReserveOutcome<'_> {
self.collect_released();
#[cfg(feature = "performance-evidence")]
self.evidence
.record_reservation_attempt(requested_maximum.get(), self.free_extents.len());
let Some(slot) = self.available_lease_head else {
if self.exhausted_lease_count > 0 {
#[cfg(feature = "performance-evidence")]
self.evidence.record_lease_identity_refusal();
return BytePoolReserveOutcome::LeaseIdentitiesExhausted { requested_maximum };
}
#[cfg(feature = "performance-evidence")]
self.evidence.record_lease_capacity_refusal();
return BytePoolReserveOutcome::LeaseCapacityReached { requested_maximum };
};
let Some(lease) = self.leases.get(slot) else {
#[cfg(feature = "performance-evidence")]
self.evidence.record_inconsistency();
return BytePoolReserveOutcome::Inconsistent {
requested_maximum,
reason: BytePoolInconsistency::LeaseUnavailable { slot },
};
};
if lease.status.get() != LeaseStatus::Vacant || lease.generation == u64::MAX {
#[cfg(feature = "performance-evidence")]
self.evidence.record_inconsistency();
return BytePoolReserveOutcome::Inconsistent {
requested_maximum,
reason: BytePoolInconsistency::LeaseUnavailable { slot },
};
}
let Some((extent_index, extent)) = self
.free_extents
.iter()
.enumerate()
.max_by_key(|(_, extent)| extent.len())
else {
#[cfg(feature = "performance-evidence")]
self.evidence.record_exhaustion();
return BytePoolReserveOutcome::Exhausted { requested_maximum };
};
if exact && extent.len() < requested_maximum.get() {
#[cfg(feature = "performance-evidence")]
self.evidence.record_exhaustion();
return BytePoolReserveOutcome::Exhausted { requested_maximum };
}
let Some(reserved_maximum) = NonZeroUsize::new(extent.len().min(requested_maximum.get()))
else {
#[cfg(feature = "performance-evidence")]
self.evidence.record_exhaustion();
return BytePoolReserveOutcome::Exhausted { requested_maximum };
};
self.reserve_extent(
slot,
extent_index,
extent.start,
reserved_maximum,
requested_maximum,
)
}
fn reserve_extent(
&mut self,
slot: usize,
extent_index: usize,
start: usize,
reserved_maximum: NonZeroUsize,
requested_maximum: NonZeroUsize,
) -> BytePoolReserveOutcome<'_> {
let Some(end) = start.checked_add(reserved_maximum.get()) else {
#[cfg(feature = "performance-evidence")]
self.evidence.record_inconsistency();
return BytePoolReserveOutcome::Inconsistent {
requested_maximum,
reason: BytePoolInconsistency::StorageUnavailable {
start,
end: usize::MAX,
},
};
};
if self.storage.get(start..end).is_none() {
#[cfg(feature = "performance-evidence")]
self.evidence.record_inconsistency();
return BytePoolReserveOutcome::Inconsistent {
requested_maximum,
reason: BytePoolInconsistency::StorageUnavailable { start, end },
};
}
#[cfg(feature = "performance-evidence")]
self.evidence.record_reservation();
BytePoolReserveOutcome::Reserved(ByteLeaseReservation {
pool: self,
slot,
extent_index,
reserved_maximum,
})
}
pub fn resume(&mut self, retained: RetainedByteLease) -> BytePoolResumeOutcome<'_> {
#[cfg(feature = "performance-evidence")]
self.evidence.record_resume_attempt();
if self.matching_lease(&retained).is_none() {
#[cfg(feature = "performance-evidence")]
self.evidence.record_resume_refusal();
return BytePoolResumeOutcome::Unavailable { retained };
}
#[cfg(feature = "performance-evidence")]
self.evidence.record_resumed();
match self.active_lease(retained.slot) {
Ok(lease) => {
lease.identity.set(retained.identity);
BytePoolResumeOutcome::Resumed(lease)
}
Err(reason) => BytePoolResumeOutcome::Inconsistent { retained, reason },
}
}
pub fn recycle(&mut self, retained: RetainedByteLease) -> BytePoolRecycleOutcome {
#[cfg(feature = "performance-evidence")]
self.evidence.record_recycle_attempt();
let Some(lease) = self.matching_lease(&retained) else {
#[cfg(feature = "performance-evidence")]
self.evidence.record_recycle_refusal();
return BytePoolRecycleOutcome::Unavailable { retained };
};
lease.identity.set(retained.identity);
self.release(retained.slot);
#[cfg(feature = "performance-evidence")]
self.evidence.record_recycled();
BytePoolRecycleOutcome::Recycled
}
pub(super) fn materialize(
&mut self,
slot: usize,
extent_index: usize,
requested: NonZeroUsize,
) -> Result<ByteLease<'_>, BytePoolInconsistency> {
let extent = self.free_extents.get(extent_index).cloned().ok_or(
BytePoolInconsistency::StorageUnavailable {
start: 0,
end: requested.get(),
},
)?;
let end = extent.start.checked_add(requested.get()).ok_or(
BytePoolInconsistency::StorageUnavailable {
start: extent.start,
end: usize::MAX,
},
)?;
if end > extent.end {
return Err(BytePoolInconsistency::StorageUnavailable {
start: extent.start,
end,
});
}
let lease = self
.leases
.get(slot)
.ok_or(BytePoolInconsistency::LeaseUnavailable { slot })?;
if lease.status.get() != LeaseStatus::Vacant || self.available_lease_head != Some(slot) {
return Err(BytePoolInconsistency::LeaseUnavailable { slot });
}
let generation = lease
.generation
.checked_add(1)
.ok_or(BytePoolInconsistency::LeaseUnavailable { slot })?;
let range = extent.start..end;
if self.storage.get(range.clone()).is_none() {
return Err(BytePoolInconsistency::StorageUnavailable {
start: range.start,
end: range.end,
});
}
if end == extent.end {
self.free_extents.remove(extent_index);
} else if let Some(remaining) = self.free_extents.get_mut(extent_index) {
remaining.start = end;
}
let lease = self
.leases
.get_mut(slot)
.ok_or(BytePoolInconsistency::LeaseUnavailable { slot })?;
self.available_lease_head = lease.next.get();
lease.next.set(None);
lease.generation = generation;
lease.range = range;
lease.status.set(LeaseStatus::Leased);
#[cfg(feature = "performance-evidence")]
self.evidence.record_materialization(requested.get());
self.active_lease(slot)
}
fn matching_lease(&self, retained: &RetainedByteLease) -> Option<&super::LeaseRecord> {
if self.storage.as_ptr().addr() != retained.pool_address {
return None;
}
self.leases.get(retained.slot).filter(|lease| {
lease.identity_address == retained.identity.as_ptr().addr()
&& lease.generation == retained.generation
&& lease.range == retained.range
&& lease.status.get() == LeaseStatus::Leased
})
}
fn active_lease(&mut self, slot: usize) -> Result<ByteLease<'_>, BytePoolInconsistency> {
let Self {
storage,
leases,
released_lease_head,
..
} = self;
let Some(lease) = leases.get(slot) else {
return Err(BytePoolInconsistency::LeaseUnavailable { slot });
};
let range = lease.range.clone();
let pool_address = storage.as_ptr().addr();
let Some(bytes) = storage.get_mut(range.clone()) else {
return Err(BytePoolInconsistency::StorageUnavailable {
start: range.start,
end: range.end,
});
};
Ok(ByteLease {
pool_address,
identity: &lease.identity,
slot,
generation: lease.generation,
range,
status: &lease.status,
next: &lease.next,
released_lease_head,
storage: bytes,
recycle_on_drop: true,
})
}
fn release(&self, slot: usize) {
if let Some(lease) = self.leases.get(slot)
&& lease.status.get() == LeaseStatus::Leased
{
lease.next.set(self.released_lease_head.get());
lease.status.set(LeaseStatus::Released);
self.released_lease_head.set(Some(slot));
}
}
fn collect_released(&mut self) {
while let Some(slot) = self.released_lease_head.get() {
let Some(lease) = self.leases.get_mut(slot) else {
self.released_lease_head.set(None);
break;
};
self.released_lease_head.set(lease.next.get());
lease.next.set(None);
let range = lease.range.clone();
let generation = lease.generation;
Self::insert_free_extent(&mut self.free_extents, range);
#[cfg(feature = "performance-evidence")]
self.evidence.record_released_lease();
if generation == u64::MAX {
lease.status.set(LeaseStatus::Exhausted);
self.exhausted_lease_count = self.exhausted_lease_count.saturating_add(1);
} else {
lease.status.set(LeaseStatus::Vacant);
lease.next.set(self.available_lease_head);
self.available_lease_head = Some(slot);
}
}
}
fn insert_free_extent(
free_extents: &mut alloc::vec::Vec<core::ops::Range<usize>>,
mut released: core::ops::Range<usize>,
) {
let index = free_extents.partition_point(|extent| extent.end < released.start);
while let Some(following) = free_extents.get(index)
&& following.start <= released.end
{
released.start = released.start.min(following.start);
released.end = released.end.max(following.end);
free_extents.remove(index);
}
free_extents.insert(index, released);
}
}
fn lease_outcome(
materialized: BytePoolMaterializeOutcome<'_>,
requested_maximum: NonZeroUsize,
) -> BytePoolLeaseOutcome<'_> {
match materialized {
BytePoolMaterializeOutcome::Leased(lease) => BytePoolLeaseOutcome::Leased(lease),
BytePoolMaterializeOutcome::ReservationExceeded { .. } => {
BytePoolLeaseOutcome::Exhausted { requested_maximum }
}
BytePoolMaterializeOutcome::Inconsistent { reason, .. } => {
BytePoolLeaseOutcome::Inconsistent {
requested_maximum,
reason,
}
}
}
}
#[cfg(test)]
mod tests;