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