1use std::ops::Deref;
4use std::sync::{Arc, Mutex};
5
6use super::budget::Budgets;
7use super::{Store, StoreError, StoreState};
8
9#[cfg(feature = "graph-cypher")]
10mod allocation_interval;
11#[cfg(feature = "graph-cypher")]
12pub use allocation_interval::{
13 AllocationInterval, AllocationIntervalError, AllocationIntervalSnapshot,
14};
15
16#[cfg(feature = "graph-cypher")]
19#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
20pub struct GraphWorkLedger {
21 pub storage_lookups: u64,
23 pub storage_scans: u64,
25 pub storage_adjacency_entries: u64,
27 pub storage_copied_bytes: u64,
29 pub storage_pages_decoded: u64,
31 pub storage_pages_copied: u64,
33 pub storage_property_values: u64,
35 pub storage_property_bytes: u64,
37 pub storage_adjacency_physical_entries: u64,
39 pub storage_adjacency_merged_visits: u64,
41 pub storage_adjacency_merge_runs: u64,
43
44 pub canonical_comparison_bytes: u64,
46 pub canonical_encoding_bytes: u64,
48 pub wal_codec_units: u64,
50 pub encoded_wal_bytes: u64,
52 pub artifact_bytes_written: u64,
54 pub artifact_writes: u64,
56 pub wal_bytes_appended: u64,
58 pub wal_appends: u64,
60 pub full_sync_attempts: u64,
62 pub full_sync_successes: u64,
64 pub directory_sync_attempts: u64,
66 pub directory_sync_successes: u64,
68}
69#[cfg(feature = "graph-cypher")]
70#[derive(Clone, Copy)]
71pub(crate) enum GraphWorkKind {
72 StorageLookups,
73 StorageScans,
74 StorageAdjacencyEntries,
75 StorageCopiedBytes,
76 StoragePagesDecoded,
77 StoragePagesCopied,
78 StoragePropertyValues,
79 StoragePropertyBytes,
80 StorageAdjacencyPhysicalEntries,
81 StorageAdjacencyMergedVisits,
82 StorageAdjacencyMergeRuns,
83
84 CanonicalComparisonBytes,
85 CanonicalEncodingBytes,
86 WalCodecUnits,
87 EncodedWalBytes,
88 ArtifactBytesWritten,
89 ArtifactWrites,
90 WalBytesAppended,
91 WalAppends,
92 FullSyncAttempts,
93 FullSyncSuccesses,
94 DirectorySyncAttempts,
95 DirectorySyncSuccesses,
96}
97#[cfg(feature = "graph-cypher")]
98impl GraphWorkLedger {
99 pub(crate) fn add(&mut self, kind: GraphWorkKind, units: u64) {
100 let slot = match kind {
101 GraphWorkKind::StorageLookups => &mut self.storage_lookups,
102 GraphWorkKind::StorageScans => &mut self.storage_scans,
103 GraphWorkKind::StorageAdjacencyEntries => &mut self.storage_adjacency_entries,
104 GraphWorkKind::StorageCopiedBytes => &mut self.storage_copied_bytes,
105 GraphWorkKind::StoragePagesDecoded => &mut self.storage_pages_decoded,
106 GraphWorkKind::StoragePagesCopied => &mut self.storage_pages_copied,
107 GraphWorkKind::StoragePropertyValues => &mut self.storage_property_values,
108 GraphWorkKind::StoragePropertyBytes => &mut self.storage_property_bytes,
109 GraphWorkKind::StorageAdjacencyPhysicalEntries => {
110 &mut self.storage_adjacency_physical_entries
111 }
112 GraphWorkKind::StorageAdjacencyMergedVisits => {
113 &mut self.storage_adjacency_merged_visits
114 }
115 GraphWorkKind::StorageAdjacencyMergeRuns => &mut self.storage_adjacency_merge_runs,
116
117 GraphWorkKind::CanonicalComparisonBytes => &mut self.canonical_comparison_bytes,
118 GraphWorkKind::CanonicalEncodingBytes => &mut self.canonical_encoding_bytes,
119 GraphWorkKind::WalCodecUnits => &mut self.wal_codec_units,
120 GraphWorkKind::EncodedWalBytes => &mut self.encoded_wal_bytes,
121 GraphWorkKind::ArtifactBytesWritten => &mut self.artifact_bytes_written,
122 GraphWorkKind::ArtifactWrites => &mut self.artifact_writes,
123 GraphWorkKind::WalBytesAppended => &mut self.wal_bytes_appended,
124 GraphWorkKind::WalAppends => &mut self.wal_appends,
125 GraphWorkKind::FullSyncAttempts => &mut self.full_sync_attempts,
126 GraphWorkKind::FullSyncSuccesses => &mut self.full_sync_successes,
127 GraphWorkKind::DirectorySyncAttempts => &mut self.directory_sync_attempts,
128 GraphWorkKind::DirectorySyncSuccesses => &mut self.directory_sync_successes,
129 };
130 *slot = slot.saturating_add(units);
131 }
132}
133
134#[cfg(feature = "graph-cypher")]
135impl GraphWorkLedger {
136 pub(crate) fn merge(&mut self, other: Self) {
137 self.storage_lookups = self.storage_lookups.saturating_add(other.storage_lookups);
138 self.storage_scans = self.storage_scans.saturating_add(other.storage_scans);
139 self.storage_adjacency_entries = self
140 .storage_adjacency_entries
141 .saturating_add(other.storage_adjacency_entries);
142 self.storage_copied_bytes = self
143 .storage_copied_bytes
144 .saturating_add(other.storage_copied_bytes);
145 self.storage_pages_decoded = self
146 .storage_pages_decoded
147 .saturating_add(other.storage_pages_decoded);
148 self.storage_pages_copied = self
149 .storage_pages_copied
150 .saturating_add(other.storage_pages_copied);
151 self.storage_property_values = self
152 .storage_property_values
153 .saturating_add(other.storage_property_values);
154 self.storage_property_bytes = self
155 .storage_property_bytes
156 .saturating_add(other.storage_property_bytes);
157 self.storage_adjacency_physical_entries = self
158 .storage_adjacency_physical_entries
159 .saturating_add(other.storage_adjacency_physical_entries);
160 self.storage_adjacency_merged_visits = self
161 .storage_adjacency_merged_visits
162 .saturating_add(other.storage_adjacency_merged_visits);
163 self.storage_adjacency_merge_runs = self
164 .storage_adjacency_merge_runs
165 .saturating_add(other.storage_adjacency_merge_runs);
166 self.canonical_comparison_bytes = self
167 .canonical_comparison_bytes
168 .saturating_add(other.canonical_comparison_bytes);
169 self.canonical_encoding_bytes = self
170 .canonical_encoding_bytes
171 .saturating_add(other.canonical_encoding_bytes);
172 self.wal_codec_units = self.wal_codec_units.saturating_add(other.wal_codec_units);
173 self.encoded_wal_bytes = self
174 .encoded_wal_bytes
175 .saturating_add(other.encoded_wal_bytes);
176 self.artifact_bytes_written = self
177 .artifact_bytes_written
178 .saturating_add(other.artifact_bytes_written);
179 self.artifact_writes = self.artifact_writes.saturating_add(other.artifact_writes);
180 self.wal_bytes_appended = self
181 .wal_bytes_appended
182 .saturating_add(other.wal_bytes_appended);
183 self.wal_appends = self.wal_appends.saturating_add(other.wal_appends);
184 self.full_sync_attempts = self
185 .full_sync_attempts
186 .saturating_add(other.full_sync_attempts);
187 self.full_sync_successes = self
188 .full_sync_successes
189 .saturating_add(other.full_sync_successes);
190 self.directory_sync_attempts = self
191 .directory_sync_attempts
192 .saturating_add(other.directory_sync_attempts);
193 self.directory_sync_successes = self
194 .directory_sync_successes
195 .saturating_add(other.directory_sync_successes);
196 }
197}
198#[allow(dead_code)]
199#[derive(Clone, Copy)]
200pub(crate) enum AllocationComponent {
201 Snapshot,
202 Active,
203 Wal,
204 Cache,
205 Temporary,
206 QueryPool,
207 NativeGraph,
208}
209
210impl AllocationComponent {
211 const fn name(self) -> &'static str {
212 match self {
213 Self::Snapshot => "snapshot",
214 Self::Active => "active segment",
215 Self::Wal => "wal",
216 Self::Cache => "cache",
217 Self::Temporary => "temporary",
218 Self::QueryPool => "query pool",
219 Self::NativeGraph => "native graph registry",
220 }
221 }
222
223 const fn is_temporary(self) -> bool {
224 matches!(self, Self::Temporary)
225 }
226}
227
228#[derive(Default)]
229struct AccountingState {
230 resident_owned_bytes: u64,
231 resident_peak_bytes: u64,
232 wal_bytes: u64,
233 cache_bytes: u64,
234 temporary_bytes: u64,
235 snapshot_bytes: u64,
236 active_bytes: u64,
237 query_pool_bytes: u64,
238 native_graph_bytes: u64,
239 mapped_bytes: u64,
240 scans_by_reason: [u64; crate::planner::SCAN_REASON_COUNTER_COUNT],
241 graph_segments_served: u64,
242 #[cfg(feature = "graph-cypher")]
243 allocation_interval_owner: u64,
244 #[cfg(feature = "graph-cypher")]
245 allocation_interval: Option<allocation_interval::AllocationIntervalState>,
246}
247
248pub(crate) struct Accounting {
249 budgets: Budgets,
250 #[cfg(feature = "graph-cypher")]
251 resident_ceiling: std::sync::atomic::AtomicU64,
252 state: Mutex<AccountingState>,
253 #[cfg(feature = "graph-cypher")]
254 graph_work: Mutex<GraphWorkLedger>,
255 #[cfg(all(test, feature = "graph-cypher"))]
256 graph_work_merges: std::sync::atomic::AtomicU64,
257}
258
259impl Accounting {
260 pub(crate) const fn new(max_resident_bytes: u64, max_temp_bytes: u64) -> Self {
261 Self {
262 budgets: Budgets::new(max_resident_bytes, max_temp_bytes),
263 #[cfg(feature = "graph-cypher")]
264 resident_ceiling: std::sync::atomic::AtomicU64::new(max_resident_bytes),
265 #[cfg(all(test, feature = "graph-cypher"))]
266 graph_work_merges: std::sync::atomic::AtomicU64::new(0),
267 #[cfg(feature = "graph-cypher")]
268 graph_work: Mutex::new(GraphWorkLedger {
269 storage_lookups: 0,
270 storage_scans: 0,
271 storage_adjacency_entries: 0,
272 storage_copied_bytes: 0,
273 storage_pages_decoded: 0,
274 storage_pages_copied: 0,
275 storage_property_values: 0,
276 storage_property_bytes: 0,
277 storage_adjacency_physical_entries: 0,
278 storage_adjacency_merged_visits: 0,
279 storage_adjacency_merge_runs: 0,
280 canonical_comparison_bytes: 0,
281 canonical_encoding_bytes: 0,
282 wal_codec_units: 0,
283 encoded_wal_bytes: 0,
284 artifact_bytes_written: 0,
285 artifact_writes: 0,
286 wal_bytes_appended: 0,
287 wal_appends: 0,
288 full_sync_attempts: 0,
289 full_sync_successes: 0,
290 directory_sync_attempts: 0,
291 directory_sync_successes: 0,
292 }),
293 state: Mutex::new(AccountingState {
294 resident_owned_bytes: 0,
295 resident_peak_bytes: 0,
296 wal_bytes: 0,
297 cache_bytes: 0,
298 temporary_bytes: 0,
299 snapshot_bytes: 0,
300 active_bytes: 0,
301 query_pool_bytes: 0,
302 native_graph_bytes: 0,
303 mapped_bytes: 0,
304 scans_by_reason: [0; crate::planner::SCAN_REASON_COUNTER_COUNT],
305 graph_segments_served: 0,
306 #[cfg(feature = "graph-cypher")]
307 allocation_interval_owner: 0,
308 #[cfg(feature = "graph-cypher")]
309 allocation_interval: None,
310 }),
311 }
312 }
313
314 #[cfg(feature = "graph-cypher")]
315 pub(crate) fn resident_limit(&self) -> u64 {
316 self.resident_ceiling
317 .load(std::sync::atomic::Ordering::Relaxed)
318 }
319
320 #[cfg(feature = "graph-cypher")]
322 pub(crate) fn enable_graph_ceiling(&self) -> Result<(), StoreError> {
323 let state = self.state.lock().map_err(|_| StoreError::Synchronization {
324 component: "memory accounting",
325 })?;
326 let limit = self
327 .resident_limit()
328 .min(crate::property_graph::resources::MAX_GRAPH_RESIDENT_BYTES);
329 if state.resident_owned_bytes > limit {
330 return Err(StoreError::BudgetExceeded {
331 needed: state.resident_owned_bytes,
332 budget: limit,
333 component: "graph aggregate configuration",
334 });
335 }
336 self.resident_ceiling
337 .store(limit, std::sync::atomic::Ordering::Relaxed);
338 Ok(())
339 }
340
341 #[cfg(feature = "graph-cypher")]
342 pub(crate) fn resident_peak_bytes(&self) -> Result<u64, StoreError> {
343 Ok(self
344 .state
345 .lock()
346 .map_err(|_| StoreError::Synchronization {
347 component: "memory accounting",
348 })?
349 .resident_peak_bytes)
350 }
351
352 fn reserve(
353 self: &Arc<Self>,
354 bytes: u64,
355 component: AllocationComponent,
356 ) -> Result<Reservation, StoreError> {
357 self.add(bytes, component)?;
358 Ok(Reservation {
359 accounting: Arc::clone(self),
360 bytes,
361 component,
362 })
363 }
364
365 fn add(&self, bytes: u64, component: AllocationComponent) -> Result<(), StoreError> {
366 let mut state = self.state.lock().map_err(|_| StoreError::Synchronization {
367 component: "memory accounting",
368 })?;
369 let (resident, temporary) = self.budgets.check(
370 state.resident_owned_bytes,
371 state.temporary_bytes,
372 bytes,
373 component.name(),
374 component.is_temporary(),
375 )?;
376 #[cfg(feature = "graph-cypher")]
377 if resident > self.resident_limit() {
378 return Err(StoreError::BudgetExceeded {
379 needed: resident,
380 budget: self.resident_limit(),
381 component: component.name(),
382 });
383 }
384 state.resident_owned_bytes = resident;
385 state.resident_peak_bytes = state.resident_peak_bytes.max(resident);
386 #[cfg(feature = "graph-cypher")]
387 if let Some(interval) = state.allocation_interval.as_mut() {
388 interval.peak_reserved_bytes = interval.peak_reserved_bytes.max(resident);
389 }
390 match component {
391 AllocationComponent::Snapshot => {
392 state.snapshot_bytes = state.snapshot_bytes.saturating_add(bytes);
393 }
394 AllocationComponent::Active => {
395 state.active_bytes = state.active_bytes.saturating_add(bytes);
396 }
397 AllocationComponent::Wal => state.wal_bytes = state.wal_bytes.saturating_add(bytes),
398 AllocationComponent::Cache => {
399 state.cache_bytes = state.cache_bytes.saturating_add(bytes);
400 }
401 AllocationComponent::Temporary => state.temporary_bytes = temporary,
402 AllocationComponent::QueryPool => {
403 state.query_pool_bytes = state.query_pool_bytes.saturating_add(bytes);
404 }
405 AllocationComponent::NativeGraph => {
406 state.native_graph_bytes = state.native_graph_bytes.saturating_add(bytes);
407 }
408 }
409 drop(state);
410 Ok(())
411 }
412
413 pub(crate) fn audit(&self) -> Result<AccountingAudit, StoreError> {
414 let state = self.state.lock().map_err(|_| StoreError::Synchronization {
415 component: "memory accounting",
416 })?;
417 Ok(AccountingAudit {
418 resident_owned_bytes: state.resident_owned_bytes,
419 wal_bytes: state.wal_bytes,
420 cache_bytes: state.cache_bytes,
421 temporary_bytes: state.temporary_bytes,
422 snapshot_bytes: state.snapshot_bytes,
423 active_bytes: state.active_bytes,
424 query_pool_bytes: state.query_pool_bytes,
425 native_graph_bytes: state.native_graph_bytes,
426 mapped_bytes: state.mapped_bytes,
427 scans_by_reason: state.scans_by_reason,
428 graph_segments_served: state.graph_segments_served,
429 })
430 }
431
432 #[cfg(feature = "graph-cypher")]
433 pub(crate) fn graph_resource_bytes(&self) -> Result<(u64, u64), StoreError> {
434 let state = self.state.lock().map_err(|_| StoreError::Synchronization {
435 component: "graph resources",
436 })?;
437 Ok((state.resident_owned_bytes, state.resident_peak_bytes))
438 }
439 #[cfg(all(feature = "graph-cypher", any(test, feature = "test-seams")))]
440 pub(crate) fn graph_work(&self) -> GraphWorkLedger {
441 *self
442 .graph_work
443 .lock()
444 .unwrap_or_else(std::sync::PoisonError::into_inner)
445 }
446 #[cfg(feature = "graph-cypher")]
447 pub(crate) fn merge_graph_work(&self, delta: GraphWorkLedger) {
448 #[cfg(test)]
449 self.graph_work_merges
450 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
451 self.graph_work
452 .lock()
453 .unwrap_or_else(std::sync::PoisonError::into_inner)
454 .merge(delta);
455 }
456 #[cfg(all(test, feature = "graph-cypher"))]
457 pub(crate) fn ze76_work_merges(&self) -> u64 {
458 self.graph_work_merges
459 .load(std::sync::atomic::Ordering::Relaxed)
460 }
461 #[cfg(all(test, feature = "graph-cypher"))]
462 #[allow(clippy::unwrap_used)]
463 pub(crate) fn ze76_overflow_work(&self) {
464 self.graph_work.lock().unwrap().wal_appends = u64::MAX;
465 }
466 #[cfg(all(test, feature = "graph-cypher"))]
467 #[allow(clippy::unwrap_used, clippy::panic)]
468 pub(crate) fn ze76_poison_work(&self) {
469 let _ = std::panic::catch_unwind(|| {
470 let _guard = self.graph_work.lock().unwrap();
471 panic!("test-only diagnostic lock poison");
472 });
473 }
474 pub(crate) fn record_plans(
475 &self,
476 plans: &[crate::planner::SegmentPlan],
477 ) -> Result<(), StoreError> {
478 let mut state = self.state.lock().map_err(|_| StoreError::Synchronization {
479 component: "query plan accounting",
480 })?;
481 let mut scans_by_reason = state.scans_by_reason;
482 let mut graph_segments_served = state.graph_segments_served;
483 for plan in plans {
484 if let Some(reason) = plan.scan_reason {
485 let slot = reason.counter().index();
486 let counter =
487 scans_by_reason
488 .get_mut(slot)
489 .ok_or_else(|| StoreError::Statistics {
490 component: "scan reason counter",
491 source: std::io::Error::other(
492 "scan reason slot is outside the registry",
493 ),
494 })?;
495 *counter = counter
496 .checked_add(1)
497 .ok_or_else(|| StoreError::Statistics {
498 component: "scan reason counter",
499 source: std::io::Error::other("scan reason counter overflow"),
500 })?;
501 }
502 if matches!(
503 plan.branch,
504 crate::planner::SegmentBranch::Graph | crate::planner::SegmentBranch::FilteredGraph
505 ) {
506 graph_segments_served =
507 graph_segments_served
508 .checked_add(1)
509 .ok_or_else(|| StoreError::Statistics {
510 component: "graph segment counter",
511 source: std::io::Error::other("graph segment counter overflow"),
512 })?;
513 }
514 }
515 state.scans_by_reason = scans_by_reason;
516 state.graph_segments_served = graph_segments_served;
517 drop(state);
518 Ok(())
519 }
520
521 pub(crate) fn track_mapping(
522 self: &Arc<Self>,
523 bytes: u64,
524 ) -> Result<MappingReservation, StoreError> {
525 let mut state = self.state.lock().map_err(|_| StoreError::Synchronization {
526 component: "memory accounting",
527 })?;
528 state.mapped_bytes =
529 state
530 .mapped_bytes
531 .checked_add(bytes)
532 .ok_or_else(|| StoreError::Statistics {
533 component: "mapped bytes",
534 source: std::io::Error::other("mapped byte count overflow"),
535 })?;
536 drop(state);
537 Ok(MappingReservation {
538 accounting: Arc::clone(self),
539 bytes,
540 })
541 }
542}
543
544#[derive(Clone, Copy)]
545pub(crate) struct AccountingAudit {
546 pub(crate) resident_owned_bytes: u64,
547 pub(crate) wal_bytes: u64,
548 pub(crate) cache_bytes: u64,
549 pub(crate) temporary_bytes: u64,
550 pub(crate) snapshot_bytes: u64,
551 pub(crate) active_bytes: u64,
552 pub(crate) query_pool_bytes: u64,
553 pub(crate) native_graph_bytes: u64,
554 pub(crate) mapped_bytes: u64,
555 pub(crate) scans_by_reason: [u64; crate::planner::SCAN_REASON_COUNTER_COUNT],
556 pub(crate) graph_segments_served: u64,
557}
558
559impl AccountingAudit {
560 pub(crate) const fn component_sum(self) -> u64 {
561 self.wal_bytes
562 .saturating_add(self.cache_bytes)
563 .saturating_add(self.temporary_bytes)
564 .saturating_add(self.snapshot_bytes)
565 .saturating_add(self.active_bytes)
566 .saturating_add(self.query_pool_bytes)
567 .saturating_add(self.native_graph_bytes)
568 }
569}
570
571pub(crate) struct MappingReservation {
572 accounting: Arc<Accounting>,
573 bytes: u64,
574}
575
576impl Drop for MappingReservation {
577 fn drop(&mut self) {
578 let mut state = match self.accounting.state.lock() {
579 Ok(state) => state,
580 Err(poisoned) => poisoned.into_inner(),
581 };
582 state.mapped_bytes = state.mapped_bytes.saturating_sub(self.bytes);
583 }
584}
585
586struct Reservation {
587 accounting: Arc<Accounting>,
588 bytes: u64,
589 component: AllocationComponent,
590}
591
592impl Reservation {
593 fn grow(&mut self, additional_bytes: u64) -> Result<(), StoreError> {
594 self.accounting.add(additional_bytes, self.component)?;
595 self.bytes =
596 self.bytes
597 .checked_add(additional_bytes)
598 .ok_or(StoreError::BudgetExceeded {
599 needed: u64::MAX,
600 budget: u64::MAX,
601 component: self.component.name(),
602 })?;
603 Ok(())
604 }
605
606 fn shrink(&mut self, released_bytes: u64) {
607 let released = released_bytes.min(self.bytes);
608 self.bytes = self.bytes.saturating_sub(released);
609 let mut state = match self.accounting.state.lock() {
610 Ok(state) => state,
611 Err(poisoned) => poisoned.into_inner(),
612 };
613 state.resident_owned_bytes = state.resident_owned_bytes.saturating_sub(released);
614 match self.component {
615 AllocationComponent::Snapshot => {
616 state.snapshot_bytes = state.snapshot_bytes.saturating_sub(released);
617 }
618 AllocationComponent::Active => {
619 state.active_bytes = state.active_bytes.saturating_sub(released);
620 }
621 AllocationComponent::Wal => {
622 state.wal_bytes = state.wal_bytes.saturating_sub(released);
623 }
624 AllocationComponent::Cache => {
625 state.cache_bytes = state.cache_bytes.saturating_sub(released);
626 }
627 AllocationComponent::Temporary => {
628 state.temporary_bytes = state.temporary_bytes.saturating_sub(released);
629 }
630 AllocationComponent::QueryPool => {
631 state.query_pool_bytes = state.query_pool_bytes.saturating_sub(released);
632 }
633 AllocationComponent::NativeGraph => {
634 state.native_graph_bytes = state.native_graph_bytes.saturating_sub(released);
635 }
636 }
637 }
638}
639
640impl Drop for Reservation {
641 fn drop(&mut self) {
642 let mut state = match self.accounting.state.lock() {
643 Ok(state) => state,
644 Err(poisoned) => poisoned.into_inner(),
645 };
646 state.resident_owned_bytes = state.resident_owned_bytes.saturating_sub(self.bytes);
647 match self.component {
648 AllocationComponent::Snapshot => {
649 state.snapshot_bytes = state.snapshot_bytes.saturating_sub(self.bytes);
650 }
651 AllocationComponent::Active => {
652 state.active_bytes = state.active_bytes.saturating_sub(self.bytes);
653 }
654 AllocationComponent::Wal => {
655 state.wal_bytes = state.wal_bytes.saturating_sub(self.bytes);
656 }
657 AllocationComponent::Cache => {
658 state.cache_bytes = state.cache_bytes.saturating_sub(self.bytes);
659 }
660 AllocationComponent::Temporary => {
661 state.temporary_bytes = state.temporary_bytes.saturating_sub(self.bytes);
662 }
663 AllocationComponent::QueryPool => {
664 state.query_pool_bytes = state.query_pool_bytes.saturating_sub(self.bytes);
665 }
666 AllocationComponent::NativeGraph => {
667 state.native_graph_bytes = state.native_graph_bytes.saturating_sub(self.bytes);
668 }
669 }
670 }
671}
672
673pub(crate) struct AccountedCounter {
676 reservation: Reservation,
677}
678
679impl AccountedCounter {
680 pub(crate) fn new(
681 accounting: &Arc<Accounting>,
682 component: AllocationComponent,
683 ) -> Result<Self, StoreError> {
684 Ok(Self {
685 reservation: accounting.reserve(0, component)?,
686 })
687 }
688
689 pub(crate) fn set(&mut self, bytes: usize) -> Result<(), StoreError> {
690 let bytes = u64::try_from(bytes).map_err(|_| StoreError::BudgetExceeded {
691 needed: u64::MAX,
692 budget: u64::MAX,
693 component: self.reservation.component.name(),
694 })?;
695 if bytes >= self.reservation.bytes {
696 return self
697 .reservation
698 .grow(bytes.saturating_sub(self.reservation.bytes));
699 }
700 self.reservation
701 .shrink(self.reservation.bytes.saturating_sub(bytes));
702 Ok(())
703 }
704
705 pub(crate) const fn bytes(&self) -> u64 {
706 self.reservation.bytes
707 }
708
709 #[cfg(feature = "graph-cypher")]
710 pub(crate) fn belongs_to(&self, accounting: &Arc<Accounting>) -> bool {
711 Arc::ptr_eq(&self.reservation.accounting, accounting)
712 }
713}
714
715pub(crate) struct Accounted<T> {
717 value: T,
718 _reservation: Option<Reservation>,
719 element_limit: Option<usize>,
720}
721
722impl<T> Deref for Accounted<T> {
723 type Target = T;
724
725 fn deref(&self) -> &Self::Target {
726 &self.value
727 }
728}
729
730impl<T> Accounted<Vec<T>> {
731 pub(crate) fn unaccounted_empty() -> Self {
732 Self {
733 value: Vec::new(),
734 _reservation: None,
735 element_limit: None,
736 }
737 }
738
739 pub(crate) fn try_with_capacity(
740 accounting: &Arc<Accounting>,
741 capacity: usize,
742 component: AllocationComponent,
743 ) -> Result<Self, StoreError> {
744 let requested =
745 capacity
746 .checked_mul(std::mem::size_of::<T>())
747 .ok_or(StoreError::BudgetExceeded {
748 needed: u64::MAX,
749 budget: u64::MAX,
750 component: component.name(),
751 })?;
752 let bytes = u64::try_from(requested).map_err(|_| StoreError::BudgetExceeded {
753 needed: u64::MAX,
754 budget: u64::MAX,
755 component: component.name(),
756 })?;
757 let reservation = accounting.reserve(bytes, component)?;
758 let mut value = Vec::new();
759 #[cfg(feature = "allocation-audit")]
760 let allocation = crate::allocation_audit::attributed(|| value.try_reserve_exact(capacity));
761 #[cfg(not(feature = "allocation-audit"))]
762 let allocation = value.try_reserve_exact(capacity);
763 allocation.map_err(|_| StoreError::AllocationFailed {
764 needed: bytes,
765 component: component.name(),
766 })?;
767 Ok(Self {
768 value,
769 _reservation: Some(reservation),
770 element_limit: Some(capacity),
771 })
772 }
773
774 pub(crate) fn try_from_vec(
775 accounting: &Arc<Accounting>,
776 value: Vec<T>,
777 component: AllocationComponent,
778 ) -> Result<Self, StoreError> {
779 let capacity = value.capacity();
780 let requested =
781 capacity
782 .checked_mul(std::mem::size_of::<T>())
783 .ok_or(StoreError::BudgetExceeded {
784 needed: u64::MAX,
785 budget: u64::MAX,
786 component: component.name(),
787 })?;
788 let bytes = u64::try_from(requested).map_err(|_| StoreError::BudgetExceeded {
789 needed: u64::MAX,
790 budget: u64::MAX,
791 component: component.name(),
792 })?;
793 let reservation = accounting.reserve(bytes, component)?;
794 Ok(Self {
795 value,
796 _reservation: Some(reservation),
797 element_limit: Some(capacity),
798 })
799 }
800
801 pub(crate) fn resident_bytes(&self) -> u64 {
802 self._reservation
803 .as_ref()
804 .map_or(0, |reservation| reservation.bytes)
805 }
806
807 pub(crate) fn try_reserve_total(
810 &mut self,
811 accounting: &Arc<Accounting>,
812 capacity: usize,
813 component: AllocationComponent,
814 ) -> Result<(), StoreError> {
815 if capacity <= self.element_limit.unwrap_or(0) {
816 return Ok(());
817 }
818 let mut replacement = Self::try_with_capacity(accounting, capacity, component)?;
819 if self.value.len() > capacity {
820 return Err(StoreError::AllocationFailed {
821 needed: self.value.len() as u64,
822 component: "accounted vector capacity",
823 });
824 }
825 replacement.value.append(&mut self.value);
828 *self = replacement;
829 Ok(())
830 }
831
832 pub(crate) fn push(&mut self, value: T) -> Result<(), StoreError> {
833 let accounted_elements = self.element_limit.unwrap_or(usize::MAX);
834 if self.value.len() >= accounted_elements || self.value.len() == self.value.capacity() {
835 return Err(StoreError::AllocationFailed {
836 needed: u64::try_from(self.value.len().saturating_add(1)).unwrap_or(u64::MAX),
837 component: "accounted vector capacity",
838 });
839 }
840 self.value.push(value);
841 Ok(())
842 }
843
844 pub(crate) fn extend_from_slice(&mut self, values: &[T]) -> Result<(), StoreError>
845 where
846 T: Clone,
847 {
848 let new_len =
849 self.value
850 .len()
851 .checked_add(values.len())
852 .ok_or(StoreError::AllocationFailed {
853 needed: u64::MAX,
854 component: "accounted vector capacity",
855 })?;
856 let accounted_elements = self.element_limit.unwrap_or(usize::MAX);
857 if new_len > accounted_elements || new_len > self.value.capacity() {
858 return Err(StoreError::AllocationFailed {
859 needed: u64::try_from(new_len).unwrap_or(u64::MAX),
860 component: "accounted vector capacity",
861 });
862 }
863 self.value.extend_from_slice(values);
864 Ok(())
865 }
866
867 pub(crate) fn replace_range(
868 &mut self,
869 range: std::ops::Range<usize>,
870 replacement: &[T],
871 ) -> Result<(), StoreError>
872 where
873 T: Copy + Default,
874 {
875 let old_len = self.value.len();
876 if range.start > range.end || range.end > old_len {
877 return Err(StoreError::ActiveRowOverflow);
878 }
879 let removed = range.end.saturating_sub(range.start);
880 let new_len = old_len
881 .checked_sub(removed)
882 .and_then(|length| length.checked_add(replacement.len()))
883 .ok_or(StoreError::ActiveRowOverflow)?;
884 let accounted_elements = self.element_limit.unwrap_or(usize::MAX);
885 if new_len > accounted_elements || new_len > self.value.capacity() {
886 return Err(StoreError::AllocationFailed {
887 needed: u64::try_from(new_len).unwrap_or(u64::MAX),
888 component: "accounted vector capacity",
889 });
890 }
891 if new_len > old_len {
892 self.value.resize(new_len, T::default());
893 }
894 let replacement_end = range
895 .start
896 .checked_add(replacement.len())
897 .ok_or(StoreError::ActiveRowOverflow)?;
898 self.value.copy_within(range.end..old_len, replacement_end);
899 self.value
900 .get_mut(range.start..replacement_end)
901 .ok_or(StoreError::ActiveRowOverflow)?
902 .copy_from_slice(replacement);
903 self.value.truncate(new_len);
904 Ok(())
905 }
906
907 pub(crate) fn as_mut_slice(&mut self) -> &mut [T] {
908 self.value.as_mut_slice()
909 }
910
911 pub(crate) fn clear(&mut self) {
913 self.value.clear();
914 }
915
916 pub(crate) fn dedup(&mut self)
917 where
918 T: PartialEq,
919 {
920 self.value.dedup();
921 }
922
923 pub(crate) fn reserve_additional_bytes(&mut self, additional: u64) -> Result<(), StoreError> {
924 match self._reservation.as_mut() {
925 Some(reservation) => reservation.grow(additional),
926 None if additional == 0 => Ok(()),
927 None => Err(StoreError::AllocationFailed {
928 needed: additional,
929 component: "unaccounted vector",
930 }),
931 }
932 }
933}
934
935#[cfg(test)]
936impl Accounted<Vec<u8>> {
937 fn try_zeroed(
938 accounting: &Arc<Accounting>,
939 bytes: u64,
940 component: AllocationComponent,
941 ) -> Result<Self, StoreError> {
942 let capacity = usize::try_from(bytes).map_err(|_| StoreError::BudgetExceeded {
943 needed: bytes,
944 budget: usize::MAX as u64,
945 component: component.name(),
946 })?;
947 let mut accounted = Self::try_with_capacity(accounting, capacity, component)?;
948 accounted.value.resize(capacity, 0);
949 Ok(accounted)
950 }
951}
952
953#[derive(Clone, Copy, Debug, Eq, PartialEq)]
955pub struct Stats {
956 pub resident_owned_bytes: u64,
958 pub snapshot_bytes: u64,
960 pub mapped_bytes: u64,
962 pub mapped_resident_bytes: u64,
964 pub segment_bytes: u64,
966 pub active_segment_bytes: u64,
969 pub retired_active_segment_bytes: u64,
972 pub active_row_count: u64,
974 pub tombstone_count: u64,
976 pub tombstone_bytes: u64,
978 pub wal_bytes: u64,
981 pub cache_bytes: u64,
983 pub temporary_bytes: u64,
985 pub query_pool_bytes: u64,
987 pub native_graph_bytes: u64,
991 pub open_files: u64,
993 pub active_queries: u64,
995 pub active_snapshot_leases: u64,
997 pub phys_footprint: Option<u64>,
1000 pub scans_by_reason: [u64; crate::planner::SCAN_REASON_COUNTER_COUNT],
1002 pub graph_segments_served: u64,
1004}
1005
1006impl Store {
1007 pub fn stats(&self) -> Result<Stats, StoreError> {
1013 let state = self
1014 .state
1015 .lock()
1016 .map_err(|_| StoreError::Synchronization { component: "state" })?;
1017 match *state {
1018 StoreState::Open => {}
1019 StoreState::Closing => return Err(StoreError::Closing),
1020 StoreState::Closed => return Err(StoreError::Closed),
1021 }
1022 let stats = self.stats_while_open();
1023 drop(state);
1024 stats
1025 }
1026
1027 pub(crate) fn stats_while_open(&self) -> Result<Stats, StoreError> {
1028 let writer_lock = self
1029 .writer_lock
1030 .lock()
1031 .map_err(|_| StoreError::Synchronization {
1032 component: "writer lock",
1033 })?;
1034 let wal_writer = self
1035 .wal_writer
1036 .lock()
1037 .map_err(|_| StoreError::Synchronization {
1038 component: "WAL writer",
1039 })?;
1040 let base_open_files =
1041 u64::from(writer_lock.is_some()).saturating_add(u64::from(wal_writer.is_some()));
1042 #[cfg(feature = "graph-cypher")]
1043 let base_open_files = base_open_files.saturating_add(u64::from(
1044 self.reader_store_lock
1045 .lock()
1046 .map_err(|_| StoreError::Synchronization {
1047 component: "reader store lock",
1048 })?
1049 .is_some(),
1050 ));
1051 let active_guard = self
1052 .active
1053 .lock()
1054 .map_err(|_| StoreError::Synchronization {
1055 component: "active segment",
1056 })?;
1057 #[cfg(feature = "graph-cypher")]
1058 let native_only = active_guard.is_none() && self.native_graph.has_publication()?;
1059 #[cfg(not(feature = "graph-cypher"))]
1060 let native_only = false;
1061 let (active_segment_bytes, active_row_count, tombstone_count, tombstone_bytes) =
1062 if let Some(active_state) = active_guard.as_ref() {
1063 let rows = u64::try_from(active_state.segment.row_count()).map_err(|_| {
1064 StoreError::Statistics {
1065 component: "active row count",
1066 source: std::io::Error::other("active row count exceeds u64"),
1067 }
1068 })?;
1069 (
1070 active_state.segment.resident_bytes(),
1071 rows,
1072 active_state.segment.tombstone_count(),
1073 active_state.segment.tombstone_bytes(),
1074 )
1075 } else if native_only {
1076 (0, 0, 0, 0)
1077 } else {
1078 return Err(StoreError::Closed);
1079 };
1080 let active_queries = self
1081 .active_queries
1082 .load(std::sync::atomic::Ordering::Relaxed);
1083 let accounting = self.accounting.audit()?;
1084 if accounting.resident_owned_bytes != accounting.component_sum() {
1085 return Err(StoreError::Statistics {
1086 component: "resident accounting conservation",
1087 source: std::io::Error::other(
1088 "resident-owned total disagrees with component counters",
1089 ),
1090 });
1091 }
1092 let retired_active_segment_bytes = accounting
1093 .active_bytes
1094 .checked_sub(active_segment_bytes)
1095 .ok_or_else(|| StoreError::Statistics {
1096 component: "active segment accounting",
1097 source: std::io::Error::other(
1098 "current active segment capacity exceeds the accounting component",
1099 ),
1100 })?;
1101 let snapshot_guard = self
1102 .snapshot
1103 .read()
1104 .map_err(|_| StoreError::Synchronization {
1105 component: "published snapshot",
1106 })?;
1107 let snapshot = snapshot_guard.as_ref();
1108 if snapshot.is_none() != native_only {
1109 return Err(StoreError::Closed);
1110 }
1111 let active_snapshot_leases = if let Some(snapshot) = snapshot {
1112 u64::try_from(Arc::strong_count(snapshot).saturating_sub(1)).map_err(|_| {
1113 StoreError::Statistics {
1114 component: "active snapshot leases",
1115 source: std::io::Error::other("snapshot lease count exceeds u64"),
1116 }
1117 })?
1118 } else {
1119 0
1120 };
1121 let mut mapped_resident_bytes = 0_u64;
1122 let mut segment_bytes = 0_u64;
1123 for segment in snapshot
1124 .into_iter()
1125 .flat_map(|snapshot| snapshot.all_segments())
1126 {
1127 let bytes =
1128 u64::try_from(segment.mapped_bytes()).map_err(|_| StoreError::Statistics {
1129 component: "segment mapped bytes",
1130 source: std::io::Error::other("segment mapped byte count exceeds u64"),
1131 })?;
1132 segment_bytes =
1133 segment_bytes
1134 .checked_add(bytes)
1135 .ok_or_else(|| StoreError::Statistics {
1136 component: "segment mapped bytes",
1137 source: std::io::Error::other("segment mapped byte count overflow"),
1138 })?;
1139 let resident =
1140 segment
1141 .mapped_resident_bytes()
1142 .map_err(|source| StoreError::Statistics {
1143 component: "mapped resident bytes",
1144 source,
1145 })?;
1146 mapped_resident_bytes =
1147 mapped_resident_bytes.checked_add(resident).ok_or_else(|| {
1148 StoreError::Statistics {
1149 component: "mapped resident bytes",
1150 source: std::io::Error::other("mapped resident byte count overflow"),
1151 }
1152 })?;
1153 }
1154 #[cfg(feature = "graph-cypher")]
1155 let (native_mapping_count, native_resident_bytes) = self.native_graph.mapping_stats()?;
1156 #[cfg(not(feature = "graph-cypher"))]
1157 let (native_mapping_count, native_resident_bytes) = (0, 0);
1158 mapped_resident_bytes = mapped_resident_bytes
1159 .checked_add(native_resident_bytes)
1160 .ok_or_else(|| StoreError::Statistics {
1161 component: "mapped resident bytes",
1162 source: std::io::Error::other("mapped resident byte count overflow"),
1163 })?;
1164 let open_files = base_open_files
1165 .checked_add(native_mapping_count)
1166 .ok_or_else(|| StoreError::Statistics {
1167 component: "open file count",
1168 source: std::io::Error::other("open file count overflow"),
1169 })?;
1170 #[cfg(any(target_os = "macos", target_os = "ios"))]
1171 let phys_footprint =
1172 Some(
1173 crate::sys::darwin::phys_footprint().map_err(|error| StoreError::Statistics {
1174 component: "phys_footprint",
1175 source: std::io::Error::other(error),
1176 })?,
1177 );
1178 #[cfg(not(any(target_os = "macos", target_os = "ios")))]
1179 let phys_footprint = None;
1180 drop(snapshot_guard);
1181 drop(active_guard);
1182 drop(wal_writer);
1183 drop(writer_lock);
1184 Ok(Stats {
1185 resident_owned_bytes: accounting.resident_owned_bytes,
1186 snapshot_bytes: accounting.snapshot_bytes,
1187 mapped_bytes: accounting.mapped_bytes,
1188 mapped_resident_bytes,
1189 segment_bytes,
1190 active_segment_bytes,
1191 retired_active_segment_bytes,
1192 active_row_count,
1193 tombstone_count,
1194 tombstone_bytes,
1195 wal_bytes: accounting.wal_bytes,
1196 cache_bytes: accounting.cache_bytes,
1197 temporary_bytes: accounting.temporary_bytes,
1198 query_pool_bytes: accounting.query_pool_bytes,
1199 native_graph_bytes: accounting.native_graph_bytes,
1200 open_files,
1201 active_queries,
1202 active_snapshot_leases,
1203 phys_footprint,
1204 scans_by_reason: accounting.scans_by_reason,
1205 graph_segments_served: accounting.graph_segments_served,
1206 })
1207 }
1208
1209 #[cfg(test)]
1210 fn allocate_bytes(
1211 &self,
1212 bytes: u64,
1213 component: AllocationComponent,
1214 ) -> Result<Accounted<Vec<u8>>, StoreError> {
1215 let state = self
1216 .state
1217 .lock()
1218 .map_err(|_| StoreError::Synchronization { component: "state" })?;
1219 match *state {
1220 StoreState::Open => {}
1221 StoreState::Closing => return Err(StoreError::Closing),
1222 StoreState::Closed => return Err(StoreError::Closed),
1223 }
1224 let allocation = Accounted::try_zeroed(&self.accounting, bytes, component)?;
1225 drop(state);
1226 Ok(allocation)
1227 }
1228}
1229
1230#[cfg(test)]
1231#[allow(clippy::expect_used)]
1232mod tests {
1233 use tempfile::{TempDir, tempdir};
1234
1235 use super::super::{CancelToken, OpenOptions, QueryControl, Store, StoreError};
1236 use super::{Accounted, Accounting, AllocationComponent};
1237 use crate::lifecycle::durability::{CommitTier, DurabilityMode, DurabilityPolicy};
1238 use crate::manifest::Manifest;
1239 use crate::manifest::io::commit_manifest;
1240 use crate::meta::{AliveSet, ColumnStoreBuilder, Schema};
1241 use crate::quant::Bit4Factors;
1242 use crate::scan::{F32Rows, ScanOptions, ScanQuery, ScanRequest, ScanRows};
1243 use crate::segment::SegmentId;
1244 use crate::segment::layout::RegionEntry;
1245 use crate::segment::reader::SegmentReader;
1246 use crate::segment::writer::{SegmentBuild, SegmentFactors, write_segment};
1247 use crate::vfs::StdVfs;
1248
1249 #[cfg(feature = "graph-cypher")]
1250 #[test]
1251 fn native_graph_bytes_is_zero_for_a_graph_free_store() {
1252 let directory = tempdir().expect("directory");
1253 let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
1254 assert_eq!(store.stats().expect("stats").native_graph_bytes, 0);
1255 store.enable_graph().expect("enable graph");
1256 let enabled_bytes = store.stats().expect("enabled stats").native_graph_bytes;
1257 assert!(enabled_bytes > 0);
1258 store.enable_graph().expect("enable again");
1259 assert_eq!(
1260 store.stats().expect("stats").native_graph_bytes,
1261 enabled_bytes
1262 );
1263 store.close().expect("close");
1264 let reopened = Store::open(directory.path(), OpenOptions::default()).expect("reopen");
1265 assert_eq!(
1266 reopened.stats().expect("stats").native_graph_bytes,
1267 enabled_bytes
1268 );
1269 reopened.close().expect("close");
1270 }
1271
1272 #[cfg(feature = "graph-cypher")]
1273 #[test]
1274 fn ze204_native_only_stats_succeed_until_close() {
1275 let directory = tempdir().expect("directory");
1276 let store = Store::create_native_graph(
1277 directory.path().join("native"),
1278 OpenOptions::new()
1279 .with_durability(DurabilityMode::Durable, CommitTier::Durable)
1280 .with_max_resident_bytes(256 * 1024 * 1024),
1281 None,
1282 )
1283 .expect("native store");
1284 assert!(store.active.lock().expect("active").is_some());
1286 assert!(store.snapshot.read().expect("snapshot").is_some());
1287 let stats = store.stats().expect("open native stats");
1288 assert_eq!(stats.segment_bytes, 0);
1289 assert_eq!(stats.active_segment_bytes, 0);
1290 assert_eq!(stats.retired_active_segment_bytes, 0);
1291 assert_eq!(stats.active_row_count, 0);
1292 assert_eq!(stats.tombstone_count, 0);
1293 assert_eq!(stats.tombstone_bytes, 0);
1294 assert!(stats.snapshot_bytes > 0);
1295 assert_eq!(stats.active_snapshot_leases, 0);
1296 assert!(stats.native_graph_bytes > 0);
1297 assert_eq!(
1298 stats.resident_owned_bytes,
1299 store.accounting.audit().expect("audit").component_sum()
1300 );
1301 store.close().expect("close");
1302 assert!(matches!(store.stats(), Err(StoreError::Closed)));
1303 }
1304
1305 #[test]
1306 fn stats_bytes_are_conserved() {
1307 let directory = tempdir().expect("store directory");
1308 let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
1309 let baseline = store.stats().expect("baseline stats");
1310 let wal = store
1311 .allocate_bytes(37, AllocationComponent::Wal)
1312 .expect("accounted WAL bytes");
1313 let cache = store
1314 .allocate_bytes(19, AllocationComponent::Cache)
1315 .expect("accounted cache bytes");
1316 let temporary = store
1317 .allocate_bytes(11, AllocationComponent::Temporary)
1318 .expect("accounted temporary bytes");
1319
1320 let stats = store.stats().expect("stats");
1321 let audit = store.accounting.audit().expect("accounting audit");
1322 assert_eq!(stats.wal_bytes, 37);
1323 assert_eq!(stats.cache_bytes, 19);
1324 assert_eq!(stats.temporary_bytes, 11);
1325 assert_eq!(stats.retired_active_segment_bytes, 0);
1326 assert_eq!(
1327 stats.resident_owned_bytes,
1328 baseline.resident_owned_bytes + 67
1329 );
1330 assert_eq!(stats.resident_owned_bytes, audit.component_sum());
1331 assert_eq!(stats.resident_owned_bytes, audit.resident_owned_bytes);
1332
1333 drop((wal, cache, temporary));
1334 assert_eq!(
1335 store.stats().expect("released stats").resident_owned_bytes,
1336 baseline.resident_owned_bytes
1337 );
1338
1339 let published = published_store();
1340 let mapped = Store::open(published.path(), OpenOptions::default()).expect("mapped open");
1341 let snapshot = mapped.snapshot().expect("mapped snapshot");
1342 let segment = snapshot.segments().first().expect("fixture segment");
1343 let expected_snapshot_bytes = std::mem::size_of::<SegmentReader>()
1344 + std::mem::size_of_val(segment.directory())
1345 + segment.retained_validation_bytes()
1346 + usize::try_from(baseline.snapshot_bytes).expect("baseline snapshot bytes fit usize")
1347 + usize::try_from(baseline.native_graph_bytes)
1350 .expect("baseline native graph bytes fit usize");
1351 drop(snapshot);
1352 let mapped_stats = mapped.stats().expect("mapped stats");
1353 assert_eq!(
1354 mapped_stats.resident_owned_bytes,
1355 expected_snapshot_bytes as u64
1356 );
1357 assert_eq!(
1358 mapped_stats.resident_owned_bytes,
1359 mapped
1360 .accounting
1361 .audit()
1362 .expect("mapped accounting audit")
1363 .component_sum()
1364 );
1365 assert_eq!(mapped_stats.tombstone_count, 0);
1366 assert_eq!(mapped_stats.tombstone_bytes, 0);
1367 assert_eq!(mapped_stats.wal_bytes, 0);
1368 assert_eq!(mapped_stats.cache_bytes, 0);
1369 assert_eq!(mapped_stats.temporary_bytes, 0);
1370 assert_eq!(mapped_stats.open_files, 2);
1371 #[cfg(any(target_os = "macos", target_os = "ios"))]
1372 assert!(mapped_stats.phys_footprint.is_some_and(|bytes| bytes > 0));
1373
1374 mapped.close().expect("close mapped store");
1375 assert_eq!(
1376 mapped
1377 .accounting
1378 .audit()
1379 .expect("closed mapping accounting")
1380 .mapped_bytes,
1381 0,
1382 "close must release the exact mapped-byte accounting"
1383 );
1384 assert!(matches!(mapped.stats(), Err(StoreError::Closed)));
1385 }
1386
1387 #[test]
1388 fn repeated_close_returns_exact_stats_counters_to_pre_open_baseline() {
1389 const ITERATIONS: usize = 20;
1390 const PRE_OPEN_MAPPED_BYTES: u64 = 0;
1391 const PRE_OPEN_RESIDENT_OWNED_BYTES: u64 = 0;
1392
1393 let published = published_store();
1394 for iteration in 0..ITERATIONS {
1395 let store = Store::open(published.path(), OpenOptions::default()).expect("mapped open");
1396 let query = [1.0_f32];
1397 let rows = F32Rows::new(vec![1.0_f32; 128]);
1398 store
1399 .top_k_with_options(
1400 ScanRequest {
1401 query: ScanQuery::F32(&query),
1402 rows: ScanRows::F32RowMajor(&rows),
1403 row_mask: None,
1404 },
1405 1,
1406 ScanOptions { thread_budget: 1 },
1407 QueryControl::Cancel(CancelToken::new()),
1408 )
1409 .expect("start the persistent query pool");
1410 let live = store.stats().expect("live stats");
1411 assert!(live.mapped_bytes > PRE_OPEN_MAPPED_BYTES);
1412 assert!(
1413 live.resident_owned_bytes > PRE_OPEN_RESIDENT_OWNED_BYTES,
1414 "iteration {iteration} must own mapped-snapshot bookkeeping"
1415 );
1416 assert!(
1417 live.query_pool_bytes > 0,
1418 "iteration {iteration} must own an accounted query pool"
1419 );
1420
1421 store.close().expect("close mapped store");
1422
1423 let closed = store
1424 .accounting
1425 .audit()
1426 .expect("post-close accounting source for Stats");
1427 assert_eq!(
1428 closed.mapped_bytes, PRE_OPEN_MAPPED_BYTES,
1429 "Stats::mapped_bytes source leaked at iteration {iteration}"
1430 );
1431 assert_eq!(
1432 closed.resident_owned_bytes, PRE_OPEN_RESIDENT_OWNED_BYTES,
1433 "Stats::resident_owned_bytes source leaked at iteration {iteration}"
1434 );
1435 assert_eq!(
1436 closed.query_pool_bytes, 0,
1437 "Stats::query_pool_bytes source leaked at iteration {iteration}"
1438 );
1439 }
1440 }
1441
1442 #[test]
1443 fn budget_exceeded_is_typed_and_pre_allocation() {
1444 let directory = tempdir().expect("store directory");
1445 let baseline_store =
1446 Store::open(directory.path(), OpenOptions::default()).expect("baseline open");
1447 let baseline = baseline_store
1450 .stats()
1451 .expect("baseline stats")
1452 .resident_owned_bytes;
1453 baseline_store.close().expect("close baseline");
1454 let options = OpenOptions::new()
1455 .with_max_resident_bytes(baseline + 64)
1456 .with_max_temp_bytes(32);
1457 let store = Store::open(directory.path(), options).expect("open");
1458
1459 let rejected = store.allocate_bytes(40, AllocationComponent::Temporary);
1460 assert!(matches!(
1461 rejected,
1462 Err(StoreError::BudgetExceeded {
1463 needed: 40,
1464 budget: 32,
1465 component: "temporary",
1466 })
1467 ));
1468 let after_rejection = store.stats().expect("stats after rejection");
1469 assert_eq!(after_rejection.resident_owned_bytes, baseline);
1470 assert_eq!(after_rejection.temporary_bytes, 0);
1471
1472 let accepted = store
1473 .allocate_bytes(24, AllocationComponent::Temporary)
1474 .expect("smaller allocation on same handle");
1475 assert_eq!(accepted.len(), 24);
1476 assert_eq!(store.stats().expect("usable stats").temporary_bytes, 24);
1477 assert!(matches!(
1478 store.allocate_bytes(50, AllocationComponent::Wal),
1479 Err(StoreError::BudgetExceeded {
1480 needed,
1481 budget,
1482 component: "wal",
1483 }) if needed == baseline + 74 && budget == baseline + 64
1484 ));
1485 assert_eq!(
1486 store.stats().expect("resident rejection stats").wal_bytes,
1487 0
1488 );
1489 store.snapshot().expect("same handle remains usable");
1490
1491 let published = published_store();
1492 let snapshot_bytes =
1493 std::mem::size_of::<SegmentReader>() + 6 * std::mem::size_of::<RegionEntry>();
1494 let budget =
1495 u64::try_from(snapshot_bytes.saturating_sub(1)).expect("snapshot budget fits u64");
1496 let needed = u64::try_from(snapshot_bytes).expect("snapshot bytes fit u64");
1497 let rejected_open = Store::open(
1498 published.path(),
1499 OpenOptions::new().with_max_resident_bytes(budget),
1500 );
1501 assert!(matches!(
1502 rejected_open,
1503 Err(StoreError::BudgetExceeded {
1504 needed: actual_needed,
1505 budget: actual_budget,
1506 component: "snapshot",
1507 }) if actual_needed == needed && actual_budget == budget
1508 ));
1509 }
1510
1511 #[test]
1512 fn accounted_vector_rejects_growth_without_a_reservation() {
1513 let mut unaccounted = Accounted::<Vec<u8>>::unaccounted_empty();
1514 unaccounted
1515 .reserve_additional_bytes(0)
1516 .expect("zero bytes require no reservation");
1517 assert!(matches!(
1518 unaccounted.reserve_additional_bytes(1),
1519 Err(StoreError::AllocationFailed {
1520 needed: 1,
1521 component: "unaccounted vector",
1522 })
1523 ));
1524
1525 let accounting = std::sync::Arc::new(Accounting::new(u64::MAX, u64::MAX));
1526 let mut fixed = Accounted::try_with_capacity(&accounting, 1, AllocationComponent::Wal)
1527 .expect("fixed-capacity vector");
1528 fixed.push(7_u8).expect("reserved element");
1529 assert!(matches!(
1530 fixed.push(8_u8),
1531 Err(StoreError::AllocationFailed {
1532 needed: 2,
1533 component: "accounted vector capacity",
1534 })
1535 ));
1536 }
1537
1538 #[test]
1539 fn accounted_allocation_after_close_is_typed_closed() {
1540 let directory = tempdir().expect("store directory");
1541 let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
1542 store.close().expect("close");
1543
1544 assert!(matches!(
1545 store.allocate_bytes(1, AllocationComponent::Wal),
1546 Err(StoreError::Closed)
1547 ));
1548 }
1549
1550 #[cfg(feature = "allocation-audit")]
1551 #[test]
1552 fn nothing_allocates_outside_accounting() {
1553 let directory = tempdir().expect("store directory");
1554 let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
1555
1556 let (unaccounted, negative_control) =
1557 crate::allocation_audit::audit_engine_path(|| Vec::<u8>::with_capacity(13));
1558 assert_eq!(negative_control.attributed_bytes, 0);
1559 assert!(
1560 negative_control.unattributed_bytes >= 13,
1561 "the audit wrapper failed to detect an unattributed engine-path allocation"
1562 );
1563 drop(unaccounted);
1564 store.stats().expect("warm accounting synchronization");
1565
1566 let (accounted, report) = crate::allocation_audit::audit_engine_path(|| {
1567 store.allocate_bytes(23, AllocationComponent::Wal)
1568 });
1569 let accounted = accounted.expect("accounted allocation");
1570 assert_eq!(report.unattributed_bytes, 0);
1571 assert_eq!(report.attributed_bytes, 23);
1572 assert_eq!(store.stats().expect("stats").wal_bytes, 23);
1573 drop(accounted);
1574 }
1575
1576 fn published_store() -> TempDir {
1577 let directory = tempdir().expect("published store directory");
1578 let schema = Schema::new(Vec::new()).expect("timestamp-only schema");
1579 let mut columns = ColumnStoreBuilder::new(schema.clone());
1580 columns.push_row(17, &[]).expect("fixture row");
1581 let columns = columns.finish().expect("fixture columns");
1582 let alive = AliveSet::new(1);
1583 let id = SegmentId::new(0x0102_0304_0506, [0x09; 10]);
1584 let codes = [0x88_u8];
1585 let factors = [Bit4Factors::from_persisted(1.0, 1.0, 1.0)];
1586 let rescore = [0.0_f32, 0.0_f32];
1587 let policy = DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::Ordered)
1588 .expect("derived policy");
1589 let segment = write_segment(
1590 &StdVfs,
1591 directory.path(),
1592 SegmentBuild {
1593 id,
1594 scheme: 4,
1595 dims: 2,
1596 codes: &codes,
1597 factors: SegmentFactors::Bit4(&factors),
1598 rescore: &rescore,
1599 columns: &columns,
1600 alive: &alive,
1601 },
1602 policy,
1603 )
1604 .expect("fixture segment");
1605 commit_manifest(
1606 &StdVfs,
1607 directory.path(),
1608 &Manifest {
1609 #[cfg(feature = "graph-cypher")]
1610 graph: None,
1611 generation: 1,
1612 log_seq: 0,
1613 segments: vec![segment],
1614 epochs: Vec::new(),
1615 epoch_alias: None,
1616 schema,
1617 },
1618 policy,
1619 )
1620 .expect("fixture manifest");
1621 directory
1622 }
1623}