use std::ops::Deref;
use std::sync::{Arc, Mutex};
use super::budget::Budgets;
use super::{Store, StoreError, StoreState};
#[cfg(feature = "graph-cypher")]
mod allocation_interval;
#[cfg(feature = "graph-cypher")]
pub use allocation_interval::{
AllocationInterval, AllocationIntervalError, AllocationIntervalSnapshot,
};
#[cfg(feature = "graph-cypher")]
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct GraphWorkLedger {
pub storage_lookups: u64,
pub storage_scans: u64,
pub storage_adjacency_entries: u64,
pub storage_copied_bytes: u64,
pub storage_pages_decoded: u64,
pub storage_pages_copied: u64,
pub storage_property_values: u64,
pub storage_property_bytes: u64,
pub storage_adjacency_physical_entries: u64,
pub storage_adjacency_merged_visits: u64,
pub storage_adjacency_merge_runs: u64,
pub canonical_comparison_bytes: u64,
pub canonical_encoding_bytes: u64,
pub wal_codec_units: u64,
pub encoded_wal_bytes: u64,
pub artifact_bytes_written: u64,
pub artifact_writes: u64,
pub wal_bytes_appended: u64,
pub wal_appends: u64,
pub full_sync_attempts: u64,
pub full_sync_successes: u64,
pub directory_sync_attempts: u64,
pub directory_sync_successes: u64,
}
#[cfg(feature = "graph-cypher")]
#[derive(Clone, Copy)]
pub(crate) enum GraphWorkKind {
StorageLookups,
StorageScans,
StorageAdjacencyEntries,
StorageCopiedBytes,
StoragePagesDecoded,
StoragePagesCopied,
StoragePropertyValues,
StoragePropertyBytes,
StorageAdjacencyPhysicalEntries,
StorageAdjacencyMergedVisits,
StorageAdjacencyMergeRuns,
CanonicalComparisonBytes,
CanonicalEncodingBytes,
WalCodecUnits,
EncodedWalBytes,
ArtifactBytesWritten,
ArtifactWrites,
WalBytesAppended,
WalAppends,
FullSyncAttempts,
FullSyncSuccesses,
DirectorySyncAttempts,
DirectorySyncSuccesses,
}
#[cfg(feature = "graph-cypher")]
impl GraphWorkLedger {
pub(crate) fn add(&mut self, kind: GraphWorkKind, units: u64) {
let slot = match kind {
GraphWorkKind::StorageLookups => &mut self.storage_lookups,
GraphWorkKind::StorageScans => &mut self.storage_scans,
GraphWorkKind::StorageAdjacencyEntries => &mut self.storage_adjacency_entries,
GraphWorkKind::StorageCopiedBytes => &mut self.storage_copied_bytes,
GraphWorkKind::StoragePagesDecoded => &mut self.storage_pages_decoded,
GraphWorkKind::StoragePagesCopied => &mut self.storage_pages_copied,
GraphWorkKind::StoragePropertyValues => &mut self.storage_property_values,
GraphWorkKind::StoragePropertyBytes => &mut self.storage_property_bytes,
GraphWorkKind::StorageAdjacencyPhysicalEntries => {
&mut self.storage_adjacency_physical_entries
}
GraphWorkKind::StorageAdjacencyMergedVisits => {
&mut self.storage_adjacency_merged_visits
}
GraphWorkKind::StorageAdjacencyMergeRuns => &mut self.storage_adjacency_merge_runs,
GraphWorkKind::CanonicalComparisonBytes => &mut self.canonical_comparison_bytes,
GraphWorkKind::CanonicalEncodingBytes => &mut self.canonical_encoding_bytes,
GraphWorkKind::WalCodecUnits => &mut self.wal_codec_units,
GraphWorkKind::EncodedWalBytes => &mut self.encoded_wal_bytes,
GraphWorkKind::ArtifactBytesWritten => &mut self.artifact_bytes_written,
GraphWorkKind::ArtifactWrites => &mut self.artifact_writes,
GraphWorkKind::WalBytesAppended => &mut self.wal_bytes_appended,
GraphWorkKind::WalAppends => &mut self.wal_appends,
GraphWorkKind::FullSyncAttempts => &mut self.full_sync_attempts,
GraphWorkKind::FullSyncSuccesses => &mut self.full_sync_successes,
GraphWorkKind::DirectorySyncAttempts => &mut self.directory_sync_attempts,
GraphWorkKind::DirectorySyncSuccesses => &mut self.directory_sync_successes,
};
*slot = slot.saturating_add(units);
}
}
#[cfg(feature = "graph-cypher")]
impl GraphWorkLedger {
pub(crate) fn merge(&mut self, other: Self) {
self.storage_lookups = self.storage_lookups.saturating_add(other.storage_lookups);
self.storage_scans = self.storage_scans.saturating_add(other.storage_scans);
self.storage_adjacency_entries = self
.storage_adjacency_entries
.saturating_add(other.storage_adjacency_entries);
self.storage_copied_bytes = self
.storage_copied_bytes
.saturating_add(other.storage_copied_bytes);
self.storage_pages_decoded = self
.storage_pages_decoded
.saturating_add(other.storage_pages_decoded);
self.storage_pages_copied = self
.storage_pages_copied
.saturating_add(other.storage_pages_copied);
self.storage_property_values = self
.storage_property_values
.saturating_add(other.storage_property_values);
self.storage_property_bytes = self
.storage_property_bytes
.saturating_add(other.storage_property_bytes);
self.storage_adjacency_physical_entries = self
.storage_adjacency_physical_entries
.saturating_add(other.storage_adjacency_physical_entries);
self.storage_adjacency_merged_visits = self
.storage_adjacency_merged_visits
.saturating_add(other.storage_adjacency_merged_visits);
self.storage_adjacency_merge_runs = self
.storage_adjacency_merge_runs
.saturating_add(other.storage_adjacency_merge_runs);
self.canonical_comparison_bytes = self
.canonical_comparison_bytes
.saturating_add(other.canonical_comparison_bytes);
self.canonical_encoding_bytes = self
.canonical_encoding_bytes
.saturating_add(other.canonical_encoding_bytes);
self.wal_codec_units = self.wal_codec_units.saturating_add(other.wal_codec_units);
self.encoded_wal_bytes = self
.encoded_wal_bytes
.saturating_add(other.encoded_wal_bytes);
self.artifact_bytes_written = self
.artifact_bytes_written
.saturating_add(other.artifact_bytes_written);
self.artifact_writes = self.artifact_writes.saturating_add(other.artifact_writes);
self.wal_bytes_appended = self
.wal_bytes_appended
.saturating_add(other.wal_bytes_appended);
self.wal_appends = self.wal_appends.saturating_add(other.wal_appends);
self.full_sync_attempts = self
.full_sync_attempts
.saturating_add(other.full_sync_attempts);
self.full_sync_successes = self
.full_sync_successes
.saturating_add(other.full_sync_successes);
self.directory_sync_attempts = self
.directory_sync_attempts
.saturating_add(other.directory_sync_attempts);
self.directory_sync_successes = self
.directory_sync_successes
.saturating_add(other.directory_sync_successes);
}
}
#[cfg_attr(not(feature = "graph-cypher"), allow(dead_code))]
#[derive(Clone, Copy)]
pub(crate) enum AllocationComponent {
Snapshot,
Active,
Wal,
Cache,
Temporary,
QueryPool,
NativeGraph,
}
impl AllocationComponent {
const fn name(self) -> &'static str {
match self {
Self::Snapshot => "snapshot",
Self::Active => "active segment",
Self::Wal => "wal",
Self::Cache => "cache",
Self::Temporary => "temporary",
Self::QueryPool => "query pool",
Self::NativeGraph => "native graph registry",
}
}
const fn is_temporary(self) -> bool {
matches!(self, Self::Temporary)
}
}
#[derive(Default)]
struct AccountingState {
resident_owned_bytes: u64,
resident_peak_bytes: u64,
#[cfg(feature = "graph-cypher")]
application_bytes: u64,
#[cfg(feature = "graph-cypher")]
application_peak_bytes: u64,
wal_bytes: u64,
cache_bytes: u64,
temporary_bytes: u64,
snapshot_bytes: u64,
active_bytes: u64,
query_pool_bytes: u64,
native_graph_bytes: u64,
mapped_bytes: u64,
scans_by_reason: [u64; crate::planner::SCAN_REASON_COUNTER_COUNT],
graph_segments_served: u64,
#[cfg(feature = "graph-cypher")]
allocation_interval_owner: u64,
#[cfg(feature = "graph-cypher")]
allocation_interval: Option<allocation_interval::AllocationIntervalState>,
}
pub(crate) struct Accounting {
budgets: Budgets,
#[cfg(feature = "graph-cypher")]
resident_ceiling: std::sync::atomic::AtomicU64,
state: Mutex<AccountingState>,
#[cfg(feature = "graph-cypher")]
graph_work: Mutex<GraphWorkLedger>,
#[cfg(all(test, feature = "graph-cypher"))]
graph_work_merges: std::sync::atomic::AtomicU64,
}
impl Accounting {
pub(crate) const fn new(max_resident_bytes: u64, max_temp_bytes: u64) -> Self {
Self {
budgets: Budgets::new(max_resident_bytes, max_temp_bytes),
#[cfg(feature = "graph-cypher")]
resident_ceiling: std::sync::atomic::AtomicU64::new(max_resident_bytes),
#[cfg(all(test, feature = "graph-cypher"))]
graph_work_merges: std::sync::atomic::AtomicU64::new(0),
#[cfg(feature = "graph-cypher")]
graph_work: Mutex::new(GraphWorkLedger {
storage_lookups: 0,
storage_scans: 0,
storage_adjacency_entries: 0,
storage_copied_bytes: 0,
storage_pages_decoded: 0,
storage_pages_copied: 0,
storage_property_values: 0,
storage_property_bytes: 0,
storage_adjacency_physical_entries: 0,
storage_adjacency_merged_visits: 0,
storage_adjacency_merge_runs: 0,
canonical_comparison_bytes: 0,
canonical_encoding_bytes: 0,
wal_codec_units: 0,
encoded_wal_bytes: 0,
artifact_bytes_written: 0,
artifact_writes: 0,
wal_bytes_appended: 0,
wal_appends: 0,
full_sync_attempts: 0,
full_sync_successes: 0,
directory_sync_attempts: 0,
directory_sync_successes: 0,
}),
state: Mutex::new(AccountingState {
resident_owned_bytes: 0,
resident_peak_bytes: 0,
#[cfg(feature = "graph-cypher")]
application_bytes: 0,
#[cfg(feature = "graph-cypher")]
application_peak_bytes: 0,
wal_bytes: 0,
cache_bytes: 0,
temporary_bytes: 0,
snapshot_bytes: 0,
active_bytes: 0,
query_pool_bytes: 0,
native_graph_bytes: 0,
mapped_bytes: 0,
scans_by_reason: [0; crate::planner::SCAN_REASON_COUNTER_COUNT],
graph_segments_served: 0,
#[cfg(feature = "graph-cypher")]
allocation_interval_owner: 0,
#[cfg(feature = "graph-cypher")]
allocation_interval: None,
}),
}
}
#[cfg(feature = "graph-cypher")]
pub(crate) fn resident_limit(&self) -> u64 {
self.resident_ceiling
.load(std::sync::atomic::Ordering::Relaxed)
}
#[cfg(feature = "graph-cypher")]
pub(crate) fn enable_graph_ceiling(&self) -> Result<(), StoreError> {
let state = self.state.lock().map_err(|_| StoreError::Synchronization {
component: "memory accounting",
})?;
let limit = self
.resident_limit()
.min(crate::property_graph::resources::MAX_GRAPH_RESIDENT_BYTES);
if state.resident_owned_bytes > limit {
return Err(StoreError::BudgetExceeded {
needed: state.resident_owned_bytes,
budget: limit,
component: "graph aggregate configuration",
});
}
self.resident_ceiling
.store(limit, std::sync::atomic::Ordering::Relaxed);
Ok(())
}
#[cfg(feature = "graph-cypher")]
pub(crate) fn resident_peak_bytes(&self) -> Result<u64, StoreError> {
Ok(self
.state
.lock()
.map_err(|_| StoreError::Synchronization {
component: "memory accounting",
})?
.resident_peak_bytes)
}
fn reserve(
self: &Arc<Self>,
bytes: u64,
component: AllocationComponent,
) -> Result<Reservation, StoreError> {
self.add(bytes, component)?;
Ok(Reservation {
accounting: Arc::clone(self),
bytes,
component,
})
}
#[cfg(feature = "graph-cypher")]
pub(crate) fn check_temporary_capacity(
&self,
bytes: usize,
component: &'static str,
) -> Result<(), StoreError> {
let state = self.state.lock().map_err(|_| StoreError::Synchronization {
component: "memory accounting",
})?;
let (resident, _) = self.budgets.check(
state.resident_owned_bytes,
state.temporary_bytes,
bytes as u64,
component,
true,
)?;
#[cfg(feature = "graph-cypher")]
if resident > self.resident_limit() {
return Err(StoreError::BudgetExceeded {
needed: resident,
budget: self.resident_limit(),
component,
});
}
Ok(())
}
fn add(&self, bytes: u64, component: AllocationComponent) -> Result<(), StoreError> {
let mut state = self.state.lock().map_err(|_| StoreError::Synchronization {
component: "memory accounting",
})?;
let (resident, temporary) = self.budgets.check(
state.resident_owned_bytes,
state.temporary_bytes,
bytes,
component.name(),
component.is_temporary(),
)?;
#[cfg(feature = "graph-cypher")]
if resident > self.resident_limit() {
return Err(StoreError::BudgetExceeded {
needed: resident,
budget: self.resident_limit(),
component: component.name(),
});
}
state.resident_owned_bytes = resident;
state.resident_peak_bytes = state.resident_peak_bytes.max(resident);
#[cfg(feature = "graph-cypher")]
if let Some(interval) = state.allocation_interval.as_mut() {
interval.peak_reserved_bytes = interval.peak_reserved_bytes.max(resident);
}
match component {
AllocationComponent::Snapshot => {
state.snapshot_bytes = state.snapshot_bytes.saturating_add(bytes);
}
AllocationComponent::Active => {
state.active_bytes = state.active_bytes.saturating_add(bytes);
}
AllocationComponent::Wal => state.wal_bytes = state.wal_bytes.saturating_add(bytes),
AllocationComponent::Cache => {
state.cache_bytes = state.cache_bytes.saturating_add(bytes);
}
AllocationComponent::Temporary => state.temporary_bytes = temporary,
AllocationComponent::QueryPool => {
state.query_pool_bytes = state.query_pool_bytes.saturating_add(bytes);
}
AllocationComponent::NativeGraph => {
state.native_graph_bytes = state.native_graph_bytes.saturating_add(bytes);
}
}
drop(state);
Ok(())
}
pub(crate) fn audit(&self) -> Result<AccountingAudit, StoreError> {
let state = self.state.lock().map_err(|_| StoreError::Synchronization {
component: "memory accounting",
})?;
Ok(AccountingAudit {
resident_owned_bytes: state.resident_owned_bytes,
wal_bytes: state.wal_bytes,
cache_bytes: state.cache_bytes,
temporary_bytes: state.temporary_bytes,
snapshot_bytes: state.snapshot_bytes,
active_bytes: state.active_bytes,
query_pool_bytes: state.query_pool_bytes,
native_graph_bytes: state.native_graph_bytes,
mapped_bytes: state.mapped_bytes,
scans_by_reason: state.scans_by_reason,
graph_segments_served: state.graph_segments_served,
})
}
#[cfg(feature = "graph-cypher")]
pub(crate) fn graph_resource_bytes(&self) -> Result<(u64, u64, u64, u64), StoreError> {
let state = self.state.lock().map_err(|_| StoreError::Synchronization {
component: "graph resources",
})?;
Ok((
state.resident_owned_bytes,
state.resident_peak_bytes,
state.application_bytes,
state.application_peak_bytes,
))
}
#[cfg(feature = "graph-cypher")]
pub(crate) fn reserve_application(
self: &Arc<Self>,
bytes: usize,
) -> Result<ApplicationReservation, StoreError> {
let bytes = bytes as u64;
let mut state = self.state.lock().map_err(|_| StoreError::Synchronization {
component: "application retention",
})?;
let total =
state
.application_bytes
.checked_add(bytes)
.ok_or(StoreError::BudgetExceeded {
needed: u64::MAX,
budget: u64::MAX,
component: "application retention arithmetic",
})?;
state.application_bytes = total;
state.application_peak_bytes = state.application_peak_bytes.max(total);
Ok(ApplicationReservation {
accounting: Arc::clone(self),
bytes,
})
}
#[cfg(all(
feature = "graph-cypher",
any(test, feature = "test-seams", feature = "benchmark-kernels")
))]
pub(crate) fn graph_work(&self) -> GraphWorkLedger {
*self
.graph_work
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
#[cfg(feature = "graph-cypher")]
pub(crate) fn merge_graph_work(&self, delta: GraphWorkLedger) {
#[cfg(test)]
self.graph_work_merges
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
self.graph_work
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.merge(delta);
}
#[cfg(all(test, feature = "graph-cypher"))]
pub(crate) fn ze76_work_merges(&self) -> u64 {
self.graph_work_merges
.load(std::sync::atomic::Ordering::Relaxed)
}
#[cfg(all(test, feature = "graph-cypher"))]
#[allow(clippy::unwrap_used)]
pub(crate) fn ze76_overflow_work(&self) {
self.graph_work.lock().unwrap().wal_appends = u64::MAX;
}
#[cfg(all(test, feature = "graph-cypher"))]
#[allow(clippy::unwrap_used, clippy::panic)]
pub(crate) fn ze76_poison_work(&self) {
let _ = std::panic::catch_unwind(|| {
let _guard = self.graph_work.lock().unwrap();
panic!("test-only diagnostic lock poison");
});
}
pub(crate) fn record_plans(
&self,
plans: &[crate::planner::SegmentPlan],
) -> Result<(), StoreError> {
let mut state = self.state.lock().map_err(|_| StoreError::Synchronization {
component: "query plan accounting",
})?;
let mut scans_by_reason = state.scans_by_reason;
let mut graph_segments_served = state.graph_segments_served;
for plan in plans {
if let Some(reason) = plan.scan_reason {
let slot = reason.counter().index();
let counter =
scans_by_reason
.get_mut(slot)
.ok_or_else(|| StoreError::Statistics {
component: "scan reason counter",
source: std::io::Error::other(
"scan reason slot is outside the registry",
),
})?;
*counter = counter
.checked_add(1)
.ok_or_else(|| StoreError::Statistics {
component: "scan reason counter",
source: std::io::Error::other("scan reason counter overflow"),
})?;
}
if matches!(
plan.branch,
crate::planner::SegmentBranch::Graph | crate::planner::SegmentBranch::FilteredGraph
) {
graph_segments_served =
graph_segments_served
.checked_add(1)
.ok_or_else(|| StoreError::Statistics {
component: "graph segment counter",
source: std::io::Error::other("graph segment counter overflow"),
})?;
}
}
state.scans_by_reason = scans_by_reason;
state.graph_segments_served = graph_segments_served;
drop(state);
Ok(())
}
pub(crate) fn track_mapping(
self: &Arc<Self>,
bytes: u64,
) -> Result<MappingReservation, StoreError> {
let mut state = self.state.lock().map_err(|_| StoreError::Synchronization {
component: "memory accounting",
})?;
state.mapped_bytes =
state
.mapped_bytes
.checked_add(bytes)
.ok_or_else(|| StoreError::Statistics {
component: "mapped bytes",
source: std::io::Error::other("mapped byte count overflow"),
})?;
drop(state);
Ok(MappingReservation {
accounting: Arc::clone(self),
bytes,
})
}
}
#[derive(Clone, Copy)]
pub(crate) struct AccountingAudit {
pub(crate) resident_owned_bytes: u64,
pub(crate) wal_bytes: u64,
pub(crate) cache_bytes: u64,
pub(crate) temporary_bytes: u64,
pub(crate) snapshot_bytes: u64,
pub(crate) active_bytes: u64,
pub(crate) query_pool_bytes: u64,
pub(crate) native_graph_bytes: u64,
pub(crate) mapped_bytes: u64,
pub(crate) scans_by_reason: [u64; crate::planner::SCAN_REASON_COUNTER_COUNT],
pub(crate) graph_segments_served: u64,
}
impl AccountingAudit {
pub(crate) const fn component_sum(self) -> u64 {
self.wal_bytes
.saturating_add(self.cache_bytes)
.saturating_add(self.temporary_bytes)
.saturating_add(self.snapshot_bytes)
.saturating_add(self.active_bytes)
.saturating_add(self.query_pool_bytes)
.saturating_add(self.native_graph_bytes)
}
}
pub(crate) struct MappingReservation {
accounting: Arc<Accounting>,
bytes: u64,
}
impl Drop for MappingReservation {
fn drop(&mut self) {
let mut state = match self.accounting.state.lock() {
Ok(state) => state,
Err(poisoned) => poisoned.into_inner(),
};
state.mapped_bytes = state.mapped_bytes.saturating_sub(self.bytes);
}
}
struct Reservation {
accounting: Arc<Accounting>,
bytes: u64,
component: AllocationComponent,
}
impl Reservation {
fn grow(&mut self, additional_bytes: u64) -> Result<(), StoreError> {
self.accounting.add(additional_bytes, self.component)?;
self.bytes =
self.bytes
.checked_add(additional_bytes)
.ok_or(StoreError::BudgetExceeded {
needed: u64::MAX,
budget: u64::MAX,
component: self.component.name(),
})?;
Ok(())
}
fn shrink(&mut self, released_bytes: u64) {
let released = released_bytes.min(self.bytes);
self.bytes = self.bytes.saturating_sub(released);
let mut state = match self.accounting.state.lock() {
Ok(state) => state,
Err(poisoned) => poisoned.into_inner(),
};
state.resident_owned_bytes = state.resident_owned_bytes.saturating_sub(released);
match self.component {
AllocationComponent::Snapshot => {
state.snapshot_bytes = state.snapshot_bytes.saturating_sub(released);
}
AllocationComponent::Active => {
state.active_bytes = state.active_bytes.saturating_sub(released);
}
AllocationComponent::Wal => {
state.wal_bytes = state.wal_bytes.saturating_sub(released);
}
AllocationComponent::Cache => {
state.cache_bytes = state.cache_bytes.saturating_sub(released);
}
AllocationComponent::Temporary => {
state.temporary_bytes = state.temporary_bytes.saturating_sub(released);
}
AllocationComponent::QueryPool => {
state.query_pool_bytes = state.query_pool_bytes.saturating_sub(released);
}
AllocationComponent::NativeGraph => {
state.native_graph_bytes = state.native_graph_bytes.saturating_sub(released);
}
}
}
}
impl Drop for Reservation {
fn drop(&mut self) {
let mut state = match self.accounting.state.lock() {
Ok(state) => state,
Err(poisoned) => poisoned.into_inner(),
};
state.resident_owned_bytes = state.resident_owned_bytes.saturating_sub(self.bytes);
match self.component {
AllocationComponent::Snapshot => {
state.snapshot_bytes = state.snapshot_bytes.saturating_sub(self.bytes);
}
AllocationComponent::Active => {
state.active_bytes = state.active_bytes.saturating_sub(self.bytes);
}
AllocationComponent::Wal => {
state.wal_bytes = state.wal_bytes.saturating_sub(self.bytes);
}
AllocationComponent::Cache => {
state.cache_bytes = state.cache_bytes.saturating_sub(self.bytes);
}
AllocationComponent::Temporary => {
state.temporary_bytes = state.temporary_bytes.saturating_sub(self.bytes);
}
AllocationComponent::QueryPool => {
state.query_pool_bytes = state.query_pool_bytes.saturating_sub(self.bytes);
}
AllocationComponent::NativeGraph => {
state.native_graph_bytes = state.native_graph_bytes.saturating_sub(self.bytes);
}
}
}
}
pub(crate) struct AccountedCounter {
reservation: Reservation,
}
impl AccountedCounter {
pub(crate) fn new(
accounting: &Arc<Accounting>,
component: AllocationComponent,
) -> Result<Self, StoreError> {
Ok(Self {
reservation: accounting.reserve(0, component)?,
})
}
pub(crate) fn set(&mut self, bytes: usize) -> Result<(), StoreError> {
let bytes = u64::try_from(bytes).map_err(|_| StoreError::BudgetExceeded {
needed: u64::MAX,
budget: u64::MAX,
component: self.reservation.component.name(),
})?;
if bytes >= self.reservation.bytes {
return self
.reservation
.grow(bytes.saturating_sub(self.reservation.bytes));
}
self.reservation
.shrink(self.reservation.bytes.saturating_sub(bytes));
Ok(())
}
pub(crate) const fn bytes(&self) -> u64 {
self.reservation.bytes
}
#[cfg(feature = "graph-cypher")]
pub(crate) fn release(&mut self, bytes: u64) {
self.reservation.shrink(bytes);
}
#[cfg(feature = "graph-cypher")]
pub(crate) fn belongs_to(&self, accounting: &Arc<Accounting>) -> bool {
Arc::ptr_eq(&self.reservation.accounting, accounting)
}
}
pub(crate) struct Accounted<T> {
value: T,
_reservation: Option<Reservation>,
element_limit: Option<usize>,
}
impl<T> Deref for Accounted<T> {
type Target = T;
fn deref(&self) -> &Self::Target {
&self.value
}
}
impl<T> Accounted<Vec<T>> {
pub(crate) fn unaccounted_empty() -> Self {
Self {
value: Vec::new(),
_reservation: None,
element_limit: None,
}
}
pub(crate) fn try_with_capacity(
accounting: &Arc<Accounting>,
capacity: usize,
component: AllocationComponent,
) -> Result<Self, StoreError> {
let requested =
capacity
.checked_mul(std::mem::size_of::<T>())
.ok_or(StoreError::BudgetExceeded {
needed: u64::MAX,
budget: u64::MAX,
component: component.name(),
})?;
let bytes = u64::try_from(requested).map_err(|_| StoreError::BudgetExceeded {
needed: u64::MAX,
budget: u64::MAX,
component: component.name(),
})?;
let reservation = accounting.reserve(bytes, component)?;
let mut value = Vec::new();
#[cfg(feature = "allocation-audit")]
let allocation = crate::allocation_audit::attributed(|| value.try_reserve_exact(capacity));
#[cfg(not(feature = "allocation-audit"))]
let allocation = value.try_reserve_exact(capacity);
allocation.map_err(|_| StoreError::AllocationFailed {
needed: bytes,
component: component.name(),
})?;
Ok(Self {
value,
_reservation: Some(reservation),
element_limit: Some(capacity),
})
}
pub(crate) fn try_from_vec(
accounting: &Arc<Accounting>,
value: Vec<T>,
component: AllocationComponent,
) -> Result<Self, StoreError> {
let capacity = value.capacity();
let requested =
capacity
.checked_mul(std::mem::size_of::<T>())
.ok_or(StoreError::BudgetExceeded {
needed: u64::MAX,
budget: u64::MAX,
component: component.name(),
})?;
let bytes = u64::try_from(requested).map_err(|_| StoreError::BudgetExceeded {
needed: u64::MAX,
budget: u64::MAX,
component: component.name(),
})?;
let reservation = accounting.reserve(bytes, component)?;
Ok(Self {
value,
_reservation: Some(reservation),
element_limit: Some(capacity),
})
}
pub(crate) fn resident_bytes(&self) -> u64 {
self._reservation
.as_ref()
.map_or(0, |reservation| reservation.bytes)
}
pub(crate) fn try_reserve_total(
&mut self,
accounting: &Arc<Accounting>,
capacity: usize,
component: AllocationComponent,
) -> Result<(), StoreError> {
if capacity <= self.element_limit.unwrap_or(0) {
return Ok(());
}
let mut replacement = Self::try_with_capacity(accounting, capacity, component)?;
if self.value.len() > capacity {
return Err(StoreError::AllocationFailed {
needed: self.value.len() as u64,
component: "accounted vector capacity",
});
}
replacement.value.append(&mut self.value);
*self = replacement;
Ok(())
}
pub(crate) fn push(&mut self, value: T) -> Result<(), StoreError> {
let accounted_elements = self.element_limit.unwrap_or(usize::MAX);
if self.value.len() >= accounted_elements || self.value.len() == self.value.capacity() {
return Err(StoreError::AllocationFailed {
needed: u64::try_from(self.value.len().saturating_add(1)).unwrap_or(u64::MAX),
component: "accounted vector capacity",
});
}
self.value.push(value);
Ok(())
}
pub(crate) fn extend_from_slice(&mut self, values: &[T]) -> Result<(), StoreError>
where
T: Clone,
{
let new_len =
self.value
.len()
.checked_add(values.len())
.ok_or(StoreError::AllocationFailed {
needed: u64::MAX,
component: "accounted vector capacity",
})?;
let accounted_elements = self.element_limit.unwrap_or(usize::MAX);
if new_len > accounted_elements || new_len > self.value.capacity() {
return Err(StoreError::AllocationFailed {
needed: u64::try_from(new_len).unwrap_or(u64::MAX),
component: "accounted vector capacity",
});
}
self.value.extend_from_slice(values);
Ok(())
}
pub(crate) fn replace_range(
&mut self,
range: std::ops::Range<usize>,
replacement: &[T],
) -> Result<(), StoreError>
where
T: Copy + Default,
{
let old_len = self.value.len();
if range.start > range.end || range.end > old_len {
return Err(StoreError::ActiveRowOverflow);
}
let removed = range.end.saturating_sub(range.start);
let new_len = old_len
.checked_sub(removed)
.and_then(|length| length.checked_add(replacement.len()))
.ok_or(StoreError::ActiveRowOverflow)?;
let accounted_elements = self.element_limit.unwrap_or(usize::MAX);
if new_len > accounted_elements || new_len > self.value.capacity() {
return Err(StoreError::AllocationFailed {
needed: u64::try_from(new_len).unwrap_or(u64::MAX),
component: "accounted vector capacity",
});
}
if new_len > old_len {
self.value.resize(new_len, T::default());
}
let replacement_end = range
.start
.checked_add(replacement.len())
.ok_or(StoreError::ActiveRowOverflow)?;
self.value.copy_within(range.end..old_len, replacement_end);
self.value
.get_mut(range.start..replacement_end)
.ok_or(StoreError::ActiveRowOverflow)?
.copy_from_slice(replacement);
self.value.truncate(new_len);
Ok(())
}
pub(crate) fn as_mut_slice(&mut self) -> &mut [T] {
self.value.as_mut_slice()
}
pub(crate) fn clear(&mut self) {
self.value.clear();
}
pub(crate) fn dedup(&mut self)
where
T: PartialEq,
{
self.value.dedup();
}
pub(crate) fn reserve_additional_bytes(&mut self, additional: u64) -> Result<(), StoreError> {
match self._reservation.as_mut() {
Some(reservation) => reservation.grow(additional),
None if additional == 0 => Ok(()),
None => Err(StoreError::AllocationFailed {
needed: additional,
component: "unaccounted vector",
}),
}
}
}
#[cfg(test)]
impl Accounted<Vec<u8>> {
fn try_zeroed(
accounting: &Arc<Accounting>,
bytes: u64,
component: AllocationComponent,
) -> Result<Self, StoreError> {
let capacity = usize::try_from(bytes).map_err(|_| StoreError::BudgetExceeded {
needed: bytes,
budget: usize::MAX as u64,
component: component.name(),
})?;
let mut accounted = Self::try_with_capacity(accounting, capacity, component)?;
accounted.value.resize(capacity, 0);
Ok(accounted)
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Stats {
pub resident_owned_bytes: u64,
pub snapshot_bytes: u64,
pub mapped_bytes: u64,
pub mapped_resident_bytes: u64,
pub segment_bytes: u64,
pub active_segment_bytes: u64,
pub retired_active_segment_bytes: u64,
pub active_row_count: u64,
pub tombstone_count: u64,
pub tombstone_bytes: u64,
pub wal_bytes: u64,
pub cache_bytes: u64,
pub temporary_bytes: u64,
pub query_pool_bytes: u64,
pub native_graph_bytes: u64,
pub open_files: u64,
pub active_queries: u64,
pub active_snapshot_leases: u64,
pub phys_footprint: Option<u64>,
pub scans_by_reason: [u64; crate::planner::SCAN_REASON_COUNTER_COUNT],
pub graph_segments_served: u64,
}
impl Store {
pub fn stats(&self) -> Result<Stats, StoreError> {
let state = self
.state
.lock()
.map_err(|_| StoreError::Synchronization { component: "state" })?;
match *state {
StoreState::Open => {}
StoreState::Closing => return Err(StoreError::Closing),
StoreState::Closed => return Err(StoreError::Closed),
}
let stats = self.stats_while_open();
drop(state);
stats
}
pub(crate) fn stats_while_open(&self) -> Result<Stats, StoreError> {
let writer_lock = self
.writer_lock
.lock()
.map_err(|_| StoreError::Synchronization {
component: "writer lock",
})?;
let wal_writer = self
.wal_writer
.lock()
.map_err(|_| StoreError::Synchronization {
component: "WAL writer",
})?;
let base_open_files =
u64::from(writer_lock.is_some()).saturating_add(u64::from(wal_writer.is_some()));
#[cfg(feature = "graph-cypher")]
let base_open_files = base_open_files.saturating_add(u64::from(
self.reader_store_lock
.lock()
.map_err(|_| StoreError::Synchronization {
component: "reader store lock",
})?
.is_some(),
));
let active_guard = self
.active
.lock()
.map_err(|_| StoreError::Synchronization {
component: "active segment",
})?;
#[cfg(feature = "graph-cypher")]
let native_only = active_guard.is_none() && self.native_graph.has_publication()?;
#[cfg(not(feature = "graph-cypher"))]
let native_only = false;
let (active_segment_bytes, active_row_count, tombstone_count, tombstone_bytes) =
if let Some(active_state) = active_guard.as_ref() {
let rows = u64::try_from(active_state.segment.row_count()).map_err(|_| {
StoreError::Statistics {
component: "active row count",
source: std::io::Error::other("active row count exceeds u64"),
}
})?;
(
active_state.segment.resident_bytes(),
rows,
active_state.segment.tombstone_count(),
active_state.segment.tombstone_bytes(),
)
} else if native_only {
(0, 0, 0, 0)
} else {
return Err(StoreError::Closed);
};
let active_queries = self
.active_queries
.load(std::sync::atomic::Ordering::Relaxed);
let accounting = self.accounting.audit()?;
if accounting.resident_owned_bytes != accounting.component_sum() {
return Err(StoreError::Statistics {
component: "resident accounting conservation",
source: std::io::Error::other(
"resident-owned total disagrees with component counters",
),
});
}
let retired_active_segment_bytes = accounting
.active_bytes
.checked_sub(active_segment_bytes)
.ok_or_else(|| StoreError::Statistics {
component: "active segment accounting",
source: std::io::Error::other(
"current active segment capacity exceeds the accounting component",
),
})?;
let snapshot_guard = self
.snapshot
.read()
.map_err(|_| StoreError::Synchronization {
component: "published snapshot",
})?;
let snapshot = snapshot_guard.as_ref();
if snapshot.is_none() != native_only {
return Err(StoreError::Closed);
}
let active_snapshot_leases = if let Some(snapshot) = snapshot {
u64::try_from(Arc::strong_count(snapshot).saturating_sub(1)).map_err(|_| {
StoreError::Statistics {
component: "active snapshot leases",
source: std::io::Error::other("snapshot lease count exceeds u64"),
}
})?
} else {
0
};
let mut mapped_resident_bytes = 0_u64;
let mut segment_bytes = 0_u64;
for segment in snapshot
.into_iter()
.flat_map(|snapshot| snapshot.all_segments())
{
let bytes =
u64::try_from(segment.mapped_bytes()).map_err(|_| StoreError::Statistics {
component: "segment mapped bytes",
source: std::io::Error::other("segment mapped byte count exceeds u64"),
})?;
segment_bytes =
segment_bytes
.checked_add(bytes)
.ok_or_else(|| StoreError::Statistics {
component: "segment mapped bytes",
source: std::io::Error::other("segment mapped byte count overflow"),
})?;
let resident =
segment
.mapped_resident_bytes()
.map_err(|source| StoreError::Statistics {
component: "mapped resident bytes",
source,
})?;
mapped_resident_bytes =
mapped_resident_bytes.checked_add(resident).ok_or_else(|| {
StoreError::Statistics {
component: "mapped resident bytes",
source: std::io::Error::other("mapped resident byte count overflow"),
}
})?;
}
#[cfg(feature = "graph-cypher")]
let (native_mapping_count, native_resident_bytes) = self.native_graph.mapping_stats()?;
#[cfg(not(feature = "graph-cypher"))]
let (native_mapping_count, native_resident_bytes) = (0, 0);
mapped_resident_bytes = mapped_resident_bytes
.checked_add(native_resident_bytes)
.ok_or_else(|| StoreError::Statistics {
component: "mapped resident bytes",
source: std::io::Error::other("mapped resident byte count overflow"),
})?;
let open_files = base_open_files
.checked_add(native_mapping_count)
.ok_or_else(|| StoreError::Statistics {
component: "open file count",
source: std::io::Error::other("open file count overflow"),
})?;
#[cfg(any(target_os = "macos", target_os = "ios"))]
let phys_footprint =
Some(
crate::sys::darwin::phys_footprint().map_err(|error| StoreError::Statistics {
component: "phys_footprint",
source: std::io::Error::other(error),
})?,
);
#[cfg(not(any(target_os = "macos", target_os = "ios")))]
let phys_footprint = None;
drop(snapshot_guard);
drop(active_guard);
drop(wal_writer);
drop(writer_lock);
Ok(Stats {
resident_owned_bytes: accounting.resident_owned_bytes,
snapshot_bytes: accounting.snapshot_bytes,
mapped_bytes: accounting.mapped_bytes,
mapped_resident_bytes,
segment_bytes,
active_segment_bytes,
retired_active_segment_bytes,
active_row_count,
tombstone_count,
tombstone_bytes,
wal_bytes: accounting.wal_bytes,
cache_bytes: accounting.cache_bytes,
temporary_bytes: accounting.temporary_bytes,
query_pool_bytes: accounting.query_pool_bytes,
native_graph_bytes: accounting.native_graph_bytes,
open_files,
active_queries,
active_snapshot_leases,
phys_footprint,
scans_by_reason: accounting.scans_by_reason,
graph_segments_served: accounting.graph_segments_served,
})
}
#[cfg(test)]
fn allocate_bytes(
&self,
bytes: u64,
component: AllocationComponent,
) -> Result<Accounted<Vec<u8>>, StoreError> {
let state = self
.state
.lock()
.map_err(|_| StoreError::Synchronization { component: "state" })?;
match *state {
StoreState::Open => {}
StoreState::Closing => return Err(StoreError::Closing),
StoreState::Closed => return Err(StoreError::Closed),
}
let allocation = Accounted::try_zeroed(&self.accounting, bytes, component)?;
drop(state);
Ok(allocation)
}
}
#[cfg(feature = "graph-cypher")]
pub(crate) struct ApplicationReservation {
accounting: Arc<Accounting>,
bytes: u64,
}
#[cfg(feature = "graph-cypher")]
impl Drop for ApplicationReservation {
fn drop(&mut self) {
let mut state = match self.accounting.state.lock() {
Ok(state) => state,
Err(poisoned) => poisoned.into_inner(),
};
state.application_bytes = state.application_bytes.saturating_sub(self.bytes);
}
}
#[cfg(test)]
#[allow(clippy::expect_used)]
mod tests {
use tempfile::{TempDir, tempdir};
use super::super::{CancelToken, OpenOptions, QueryControl, Store, StoreError};
use super::{Accounted, Accounting, AllocationComponent};
use crate::lifecycle::durability::{CommitTier, DurabilityMode, DurabilityPolicy};
use crate::manifest::Manifest;
use crate::manifest::io::commit_manifest;
use crate::meta::{AliveSet, ColumnStoreBuilder, Schema};
use crate::quant::Bit4Factors;
use crate::scan::{F32Rows, ScanOptions, ScanQuery, ScanRequest, ScanRows};
use crate::segment::SegmentId;
use crate::segment::layout::RegionEntry;
use crate::segment::reader::SegmentReader;
use crate::segment::writer::{SegmentBuild, SegmentFactors, write_segment};
use crate::vfs::StdVfs;
#[cfg(feature = "graph-cypher")]
#[test]
fn native_graph_bytes_is_zero_for_a_graph_free_store() {
let directory = tempdir().expect("directory");
let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
assert_eq!(store.stats().expect("stats").native_graph_bytes, 0);
store.enable_graph().expect("enable graph");
let enabled_bytes = store.stats().expect("enabled stats").native_graph_bytes;
assert!(enabled_bytes > 0);
store.enable_graph().expect("enable again");
assert_eq!(
store.stats().expect("stats").native_graph_bytes,
enabled_bytes
);
store.close().expect("close");
let reopened = Store::open(directory.path(), OpenOptions::default()).expect("reopen");
assert_eq!(
reopened.stats().expect("stats").native_graph_bytes,
enabled_bytes
);
reopened.close().expect("close");
}
#[cfg(feature = "graph-cypher")]
#[test]
fn ze204_native_only_stats_succeed_until_close() {
let directory = tempdir().expect("directory");
let store = Store::create_native_graph(
directory.path().join("native"),
OpenOptions::new()
.with_durability(DurabilityMode::Durable, CommitTier::Durable)
.with_max_resident_bytes(256 * 1024 * 1024),
None,
)
.expect("native store");
assert!(store.active.lock().expect("active").is_some());
assert!(store.snapshot.read().expect("snapshot").is_some());
let stats = store.stats().expect("open native stats");
assert_eq!(stats.segment_bytes, 0);
assert_eq!(stats.active_segment_bytes, 0);
assert_eq!(stats.retired_active_segment_bytes, 0);
assert_eq!(stats.active_row_count, 0);
assert_eq!(stats.tombstone_count, 0);
assert_eq!(stats.tombstone_bytes, 0);
assert!(stats.snapshot_bytes > 0);
assert_eq!(stats.active_snapshot_leases, 0);
assert!(stats.native_graph_bytes > 0);
assert_eq!(
stats.resident_owned_bytes,
store.accounting.audit().expect("audit").component_sum()
);
store.close().expect("close");
assert!(matches!(store.stats(), Err(StoreError::Closed)));
}
#[test]
fn stats_bytes_are_conserved() {
let directory = tempdir().expect("store directory");
let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
let baseline = store.stats().expect("baseline stats");
let wal = store
.allocate_bytes(37, AllocationComponent::Wal)
.expect("accounted WAL bytes");
let cache = store
.allocate_bytes(19, AllocationComponent::Cache)
.expect("accounted cache bytes");
let temporary = store
.allocate_bytes(11, AllocationComponent::Temporary)
.expect("accounted temporary bytes");
let stats = store.stats().expect("stats");
let audit = store.accounting.audit().expect("accounting audit");
assert_eq!(stats.wal_bytes, 37);
assert_eq!(stats.cache_bytes, 19);
assert_eq!(stats.temporary_bytes, 11);
assert_eq!(stats.retired_active_segment_bytes, 0);
assert_eq!(
stats.resident_owned_bytes,
baseline.resident_owned_bytes + 67
);
assert_eq!(stats.resident_owned_bytes, audit.component_sum());
assert_eq!(stats.resident_owned_bytes, audit.resident_owned_bytes);
drop((wal, cache, temporary));
assert_eq!(
store.stats().expect("released stats").resident_owned_bytes,
baseline.resident_owned_bytes
);
let published = published_store();
let mapped = Store::open(published.path(), OpenOptions::default()).expect("mapped open");
let snapshot = mapped.snapshot().expect("mapped snapshot");
let segment = snapshot.segments().first().expect("fixture segment");
let expected_snapshot_bytes = std::mem::size_of::<SegmentReader>()
+ std::mem::size_of_val(segment.directory())
+ segment.retained_validation_bytes()
+ usize::try_from(baseline.snapshot_bytes).expect("baseline snapshot bytes fit usize")
+ usize::try_from(baseline.native_graph_bytes)
.expect("baseline native graph bytes fit usize");
drop(snapshot);
let mapped_stats = mapped.stats().expect("mapped stats");
assert_eq!(
mapped_stats.resident_owned_bytes,
expected_snapshot_bytes as u64
);
assert_eq!(
mapped_stats.resident_owned_bytes,
mapped
.accounting
.audit()
.expect("mapped accounting audit")
.component_sum()
);
assert_eq!(mapped_stats.tombstone_count, 0);
assert_eq!(mapped_stats.tombstone_bytes, 0);
assert_eq!(mapped_stats.wal_bytes, 0);
assert_eq!(mapped_stats.cache_bytes, 0);
assert_eq!(mapped_stats.temporary_bytes, 0);
assert_eq!(mapped_stats.open_files, 2);
#[cfg(any(target_os = "macos", target_os = "ios"))]
assert!(mapped_stats.phys_footprint.is_some_and(|bytes| bytes > 0));
mapped.close().expect("close mapped store");
assert_eq!(
mapped
.accounting
.audit()
.expect("closed mapping accounting")
.mapped_bytes,
0,
"close must release the exact mapped-byte accounting"
);
assert!(matches!(mapped.stats(), Err(StoreError::Closed)));
}
#[test]
fn repeated_close_returns_exact_stats_counters_to_pre_open_baseline() {
const ITERATIONS: usize = 20;
const PRE_OPEN_MAPPED_BYTES: u64 = 0;
const PRE_OPEN_RESIDENT_OWNED_BYTES: u64 = 0;
let published = published_store();
for iteration in 0..ITERATIONS {
let store = Store::open(published.path(), OpenOptions::default()).expect("mapped open");
let query = [1.0_f32];
let rows = F32Rows::new(vec![1.0_f32; 128]);
store
.top_k_with_options(
ScanRequest {
query: ScanQuery::F32(&query),
rows: ScanRows::F32RowMajor(&rows),
row_mask: None,
},
1,
ScanOptions { thread_budget: 1 },
QueryControl::Cancel(CancelToken::new()),
)
.expect("start the persistent query pool");
let live = store.stats().expect("live stats");
assert!(live.mapped_bytes > PRE_OPEN_MAPPED_BYTES);
assert!(
live.resident_owned_bytes > PRE_OPEN_RESIDENT_OWNED_BYTES,
"iteration {iteration} must own mapped-snapshot bookkeeping"
);
assert!(
live.query_pool_bytes > 0,
"iteration {iteration} must own an accounted query pool"
);
store.close().expect("close mapped store");
let closed = store
.accounting
.audit()
.expect("post-close accounting source for Stats");
assert_eq!(
closed.mapped_bytes, PRE_OPEN_MAPPED_BYTES,
"Stats::mapped_bytes source leaked at iteration {iteration}"
);
assert_eq!(
closed.resident_owned_bytes, PRE_OPEN_RESIDENT_OWNED_BYTES,
"Stats::resident_owned_bytes source leaked at iteration {iteration}"
);
assert_eq!(
closed.query_pool_bytes, 0,
"Stats::query_pool_bytes source leaked at iteration {iteration}"
);
}
}
#[test]
fn budget_exceeded_is_typed_and_pre_allocation() {
let directory = tempdir().expect("store directory");
let baseline_store =
Store::open(directory.path(), OpenOptions::default()).expect("baseline open");
let baseline = baseline_store
.stats()
.expect("baseline stats")
.resident_owned_bytes;
baseline_store.close().expect("close baseline");
let options = OpenOptions::new()
.with_max_resident_bytes(baseline + 64)
.with_max_temp_bytes(32);
let store = Store::open(directory.path(), options).expect("open");
let rejected = store.allocate_bytes(40, AllocationComponent::Temporary);
assert!(matches!(
rejected,
Err(StoreError::BudgetExceeded {
needed: 40,
budget: 32,
component: "temporary",
})
));
let after_rejection = store.stats().expect("stats after rejection");
assert_eq!(after_rejection.resident_owned_bytes, baseline);
assert_eq!(after_rejection.temporary_bytes, 0);
let accepted = store
.allocate_bytes(24, AllocationComponent::Temporary)
.expect("smaller allocation on same handle");
assert_eq!(accepted.len(), 24);
assert_eq!(store.stats().expect("usable stats").temporary_bytes, 24);
assert!(matches!(
store.allocate_bytes(50, AllocationComponent::Wal),
Err(StoreError::BudgetExceeded {
needed,
budget,
component: "wal",
}) if needed == baseline + 74 && budget == baseline + 64
));
assert_eq!(
store.stats().expect("resident rejection stats").wal_bytes,
0
);
store.snapshot().expect("same handle remains usable");
let published = published_store();
let snapshot_bytes =
std::mem::size_of::<SegmentReader>() + 6 * std::mem::size_of::<RegionEntry>();
let budget =
u64::try_from(snapshot_bytes.saturating_sub(1)).expect("snapshot budget fits u64");
let needed = u64::try_from(snapshot_bytes).expect("snapshot bytes fit u64");
let rejected_open = Store::open(
published.path(),
OpenOptions::new().with_max_resident_bytes(budget),
);
assert!(matches!(
rejected_open,
Err(StoreError::BudgetExceeded {
needed: actual_needed,
budget: actual_budget,
component: "snapshot",
}) if actual_needed == needed && actual_budget == budget
));
}
#[test]
fn accounted_vector_rejects_growth_without_a_reservation() {
let mut unaccounted = Accounted::<Vec<u8>>::unaccounted_empty();
unaccounted
.reserve_additional_bytes(0)
.expect("zero bytes require no reservation");
assert!(matches!(
unaccounted.reserve_additional_bytes(1),
Err(StoreError::AllocationFailed {
needed: 1,
component: "unaccounted vector",
})
));
let accounting = std::sync::Arc::new(Accounting::new(u64::MAX, u64::MAX));
let mut fixed = Accounted::try_with_capacity(&accounting, 1, AllocationComponent::Wal)
.expect("fixed-capacity vector");
fixed.push(7_u8).expect("reserved element");
assert!(matches!(
fixed.push(8_u8),
Err(StoreError::AllocationFailed {
needed: 2,
component: "accounted vector capacity",
})
));
}
#[test]
fn accounted_allocation_after_close_is_typed_closed() {
let directory = tempdir().expect("store directory");
let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
store.close().expect("close");
assert!(matches!(
store.allocate_bytes(1, AllocationComponent::Wal),
Err(StoreError::Closed)
));
}
#[cfg(feature = "allocation-audit")]
#[test]
fn nothing_allocates_outside_accounting() {
let directory = tempdir().expect("store directory");
let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
let (unaccounted, negative_control) =
crate::allocation_audit::audit_engine_path(|| Vec::<u8>::with_capacity(13));
assert_eq!(negative_control.attributed_bytes, 0);
assert!(
negative_control.unattributed_bytes >= 13,
"the audit wrapper failed to detect an unattributed engine-path allocation"
);
drop(unaccounted);
store.stats().expect("warm accounting synchronization");
let (accounted, report) = crate::allocation_audit::audit_engine_path(|| {
store.allocate_bytes(23, AllocationComponent::Wal)
});
let accounted = accounted.expect("accounted allocation");
assert_eq!(report.unattributed_bytes, 0);
assert_eq!(report.attributed_bytes, 23);
assert_eq!(store.stats().expect("stats").wal_bytes, 23);
drop(accounted);
}
fn published_store() -> TempDir {
let directory = tempdir().expect("published store directory");
let schema = Schema::new(Vec::new()).expect("timestamp-only schema");
let mut columns = ColumnStoreBuilder::new(schema.clone());
columns.push_row(17, &[]).expect("fixture row");
let columns = columns.finish().expect("fixture columns");
let alive = AliveSet::new(1);
let id = SegmentId::new(0x0102_0304_0506, [0x09; 10]);
let codes = [0x88_u8];
let factors = [Bit4Factors::from_persisted(1.0, 1.0, 1.0)];
let rescore = [0.0_f32, 0.0_f32];
let policy = DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::Ordered)
.expect("derived policy");
let segment = write_segment(
&StdVfs,
directory.path(),
SegmentBuild {
id,
scheme: 4,
dims: 2,
codes: &codes,
factors: SegmentFactors::Bit4(&factors),
rescore: &rescore,
columns: &columns,
alive: &alive,
},
policy,
)
.expect("fixture segment");
commit_manifest(
&StdVfs,
directory.path(),
&Manifest {
#[cfg(feature = "graph-cypher")]
graph: None,
generation: 1,
log_seq: 0,
segments: vec![segment],
epochs: Vec::new(),
epoch_alias: None,
schema,
},
policy,
)
.expect("fixture manifest");
directory
}
}