1use std::ops::Deref;
4use std::sync::{Arc, Mutex};
5
6use super::budget::Budgets;
7use super::{Store, StoreError, StoreState};
8
9#[allow(dead_code)]
10#[derive(Clone, Copy)]
11pub(crate) enum AllocationComponent {
12 Snapshot,
13 Active,
14 Wal,
15 Cache,
16 Temporary,
17 QueryPool,
18}
19
20impl AllocationComponent {
21 const fn name(self) -> &'static str {
22 match self {
23 Self::Snapshot => "snapshot",
24 Self::Active => "active segment",
25 Self::Wal => "wal",
26 Self::Cache => "cache",
27 Self::Temporary => "temporary",
28 Self::QueryPool => "query pool",
29 }
30 }
31
32 const fn is_temporary(self) -> bool {
33 matches!(self, Self::Temporary)
34 }
35}
36
37#[derive(Default)]
38struct AccountingState {
39 resident_owned_bytes: u64,
40 wal_bytes: u64,
41 cache_bytes: u64,
42 temporary_bytes: u64,
43 snapshot_bytes: u64,
44 active_bytes: u64,
45 query_pool_bytes: u64,
46 mapped_bytes: u64,
47 scans_by_reason: [u64; crate::planner::SCAN_REASON_COUNTER_COUNT],
48 graph_segments_served: u64,
49}
50
51pub(crate) struct Accounting {
52 budgets: Budgets,
53 state: Mutex<AccountingState>,
54}
55
56impl Accounting {
57 pub(crate) const fn new(max_resident_bytes: u64, max_temp_bytes: u64) -> Self {
58 Self {
59 budgets: Budgets::new(max_resident_bytes, max_temp_bytes),
60 state: Mutex::new(AccountingState {
61 resident_owned_bytes: 0,
62 wal_bytes: 0,
63 cache_bytes: 0,
64 temporary_bytes: 0,
65 snapshot_bytes: 0,
66 active_bytes: 0,
67 query_pool_bytes: 0,
68 mapped_bytes: 0,
69 scans_by_reason: [0; crate::planner::SCAN_REASON_COUNTER_COUNT],
70 graph_segments_served: 0,
71 }),
72 }
73 }
74
75 fn reserve(
76 self: &Arc<Self>,
77 bytes: u64,
78 component: AllocationComponent,
79 ) -> Result<Reservation, StoreError> {
80 self.add(bytes, component)?;
81 Ok(Reservation {
82 accounting: Arc::clone(self),
83 bytes,
84 component,
85 })
86 }
87
88 fn add(&self, bytes: u64, component: AllocationComponent) -> Result<(), StoreError> {
89 let mut state = self.state.lock().map_err(|_| StoreError::Synchronization {
90 component: "memory accounting",
91 })?;
92 let (resident, temporary) = self.budgets.check(
93 state.resident_owned_bytes,
94 state.temporary_bytes,
95 bytes,
96 component.name(),
97 component.is_temporary(),
98 )?;
99 state.resident_owned_bytes = resident;
100 match component {
101 AllocationComponent::Snapshot => {
102 state.snapshot_bytes = state.snapshot_bytes.saturating_add(bytes);
103 }
104 AllocationComponent::Active => {
105 state.active_bytes = state.active_bytes.saturating_add(bytes);
106 }
107 AllocationComponent::Wal => state.wal_bytes = state.wal_bytes.saturating_add(bytes),
108 AllocationComponent::Cache => {
109 state.cache_bytes = state.cache_bytes.saturating_add(bytes);
110 }
111 AllocationComponent::Temporary => state.temporary_bytes = temporary,
112 AllocationComponent::QueryPool => {
113 state.query_pool_bytes = state.query_pool_bytes.saturating_add(bytes);
114 }
115 }
116 drop(state);
117 Ok(())
118 }
119
120 pub(crate) fn audit(&self) -> Result<AccountingAudit, StoreError> {
121 let state = self.state.lock().map_err(|_| StoreError::Synchronization {
122 component: "memory accounting",
123 })?;
124 Ok(AccountingAudit {
125 resident_owned_bytes: state.resident_owned_bytes,
126 wal_bytes: state.wal_bytes,
127 cache_bytes: state.cache_bytes,
128 temporary_bytes: state.temporary_bytes,
129 snapshot_bytes: state.snapshot_bytes,
130 active_bytes: state.active_bytes,
131 query_pool_bytes: state.query_pool_bytes,
132 mapped_bytes: state.mapped_bytes,
133 scans_by_reason: state.scans_by_reason,
134 graph_segments_served: state.graph_segments_served,
135 })
136 }
137
138 pub(crate) fn record_plans(
139 &self,
140 plans: &[crate::planner::SegmentPlan],
141 ) -> Result<(), StoreError> {
142 let mut state = self.state.lock().map_err(|_| StoreError::Synchronization {
143 component: "query plan accounting",
144 })?;
145 let mut scans_by_reason = state.scans_by_reason;
146 let mut graph_segments_served = state.graph_segments_served;
147 for plan in plans {
148 if let Some(reason) = plan.scan_reason {
149 let slot = reason.counter().index();
150 let counter =
151 scans_by_reason
152 .get_mut(slot)
153 .ok_or_else(|| StoreError::Statistics {
154 component: "scan reason counter",
155 source: std::io::Error::other(
156 "scan reason slot is outside the registry",
157 ),
158 })?;
159 *counter = counter
160 .checked_add(1)
161 .ok_or_else(|| StoreError::Statistics {
162 component: "scan reason counter",
163 source: std::io::Error::other("scan reason counter overflow"),
164 })?;
165 }
166 if matches!(
167 plan.branch,
168 crate::planner::SegmentBranch::Graph | crate::planner::SegmentBranch::FilteredGraph
169 ) {
170 graph_segments_served =
171 graph_segments_served
172 .checked_add(1)
173 .ok_or_else(|| StoreError::Statistics {
174 component: "graph segment counter",
175 source: std::io::Error::other("graph segment counter overflow"),
176 })?;
177 }
178 }
179 state.scans_by_reason = scans_by_reason;
180 state.graph_segments_served = graph_segments_served;
181 drop(state);
182 Ok(())
183 }
184
185 pub(crate) fn track_mapping(
186 self: &Arc<Self>,
187 bytes: u64,
188 ) -> Result<MappingReservation, StoreError> {
189 let mut state = self.state.lock().map_err(|_| StoreError::Synchronization {
190 component: "memory accounting",
191 })?;
192 state.mapped_bytes =
193 state
194 .mapped_bytes
195 .checked_add(bytes)
196 .ok_or_else(|| StoreError::Statistics {
197 component: "mapped bytes",
198 source: std::io::Error::other("mapped byte count overflow"),
199 })?;
200 drop(state);
201 Ok(MappingReservation {
202 accounting: Arc::clone(self),
203 bytes,
204 })
205 }
206}
207
208#[derive(Clone, Copy)]
209pub(crate) struct AccountingAudit {
210 pub(crate) resident_owned_bytes: u64,
211 pub(crate) wal_bytes: u64,
212 pub(crate) cache_bytes: u64,
213 pub(crate) temporary_bytes: u64,
214 pub(crate) snapshot_bytes: u64,
215 pub(crate) active_bytes: u64,
216 pub(crate) query_pool_bytes: u64,
217 pub(crate) mapped_bytes: u64,
218 pub(crate) scans_by_reason: [u64; crate::planner::SCAN_REASON_COUNTER_COUNT],
219 pub(crate) graph_segments_served: u64,
220}
221
222impl AccountingAudit {
223 pub(crate) const fn component_sum(self) -> u64 {
224 self.wal_bytes
225 .saturating_add(self.cache_bytes)
226 .saturating_add(self.temporary_bytes)
227 .saturating_add(self.snapshot_bytes)
228 .saturating_add(self.active_bytes)
229 .saturating_add(self.query_pool_bytes)
230 }
231}
232
233pub(crate) struct MappingReservation {
234 accounting: Arc<Accounting>,
235 bytes: u64,
236}
237
238impl Drop for MappingReservation {
239 fn drop(&mut self) {
240 let mut state = match self.accounting.state.lock() {
241 Ok(state) => state,
242 Err(poisoned) => poisoned.into_inner(),
243 };
244 state.mapped_bytes = state.mapped_bytes.saturating_sub(self.bytes);
245 }
246}
247
248struct Reservation {
249 accounting: Arc<Accounting>,
250 bytes: u64,
251 component: AllocationComponent,
252}
253
254impl Reservation {
255 fn grow(&mut self, additional_bytes: u64) -> Result<(), StoreError> {
256 self.accounting.add(additional_bytes, self.component)?;
257 self.bytes =
258 self.bytes
259 .checked_add(additional_bytes)
260 .ok_or(StoreError::BudgetExceeded {
261 needed: u64::MAX,
262 budget: u64::MAX,
263 component: self.component.name(),
264 })?;
265 Ok(())
266 }
267
268 fn shrink(&mut self, released_bytes: u64) {
269 let released = released_bytes.min(self.bytes);
270 self.bytes = self.bytes.saturating_sub(released);
271 let mut state = match self.accounting.state.lock() {
272 Ok(state) => state,
273 Err(poisoned) => poisoned.into_inner(),
274 };
275 state.resident_owned_bytes = state.resident_owned_bytes.saturating_sub(released);
276 match self.component {
277 AllocationComponent::Snapshot => {
278 state.snapshot_bytes = state.snapshot_bytes.saturating_sub(released);
279 }
280 AllocationComponent::Active => {
281 state.active_bytes = state.active_bytes.saturating_sub(released);
282 }
283 AllocationComponent::Wal => {
284 state.wal_bytes = state.wal_bytes.saturating_sub(released);
285 }
286 AllocationComponent::Cache => {
287 state.cache_bytes = state.cache_bytes.saturating_sub(released);
288 }
289 AllocationComponent::Temporary => {
290 state.temporary_bytes = state.temporary_bytes.saturating_sub(released);
291 }
292 AllocationComponent::QueryPool => {
293 state.query_pool_bytes = state.query_pool_bytes.saturating_sub(released);
294 }
295 }
296 }
297}
298
299impl Drop for Reservation {
300 fn drop(&mut self) {
301 let mut state = match self.accounting.state.lock() {
302 Ok(state) => state,
303 Err(poisoned) => poisoned.into_inner(),
304 };
305 state.resident_owned_bytes = state.resident_owned_bytes.saturating_sub(self.bytes);
306 match self.component {
307 AllocationComponent::Snapshot => {
308 state.snapshot_bytes = state.snapshot_bytes.saturating_sub(self.bytes);
309 }
310 AllocationComponent::Active => {
311 state.active_bytes = state.active_bytes.saturating_sub(self.bytes);
312 }
313 AllocationComponent::Wal => {
314 state.wal_bytes = state.wal_bytes.saturating_sub(self.bytes);
315 }
316 AllocationComponent::Cache => {
317 state.cache_bytes = state.cache_bytes.saturating_sub(self.bytes);
318 }
319 AllocationComponent::Temporary => {
320 state.temporary_bytes = state.temporary_bytes.saturating_sub(self.bytes);
321 }
322 AllocationComponent::QueryPool => {
323 state.query_pool_bytes = state.query_pool_bytes.saturating_sub(self.bytes);
324 }
325 }
326 }
327}
328
329pub(crate) struct AccountedCounter {
332 reservation: Reservation,
333}
334
335impl AccountedCounter {
336 pub(crate) fn new(
337 accounting: &Arc<Accounting>,
338 component: AllocationComponent,
339 ) -> Result<Self, StoreError> {
340 Ok(Self {
341 reservation: accounting.reserve(0, component)?,
342 })
343 }
344
345 pub(crate) fn set(&mut self, bytes: usize) -> Result<(), StoreError> {
346 let bytes = u64::try_from(bytes).map_err(|_| StoreError::BudgetExceeded {
347 needed: u64::MAX,
348 budget: u64::MAX,
349 component: self.reservation.component.name(),
350 })?;
351 if bytes >= self.reservation.bytes {
352 return self
353 .reservation
354 .grow(bytes.saturating_sub(self.reservation.bytes));
355 }
356 self.reservation
357 .shrink(self.reservation.bytes.saturating_sub(bytes));
358 Ok(())
359 }
360
361 pub(crate) const fn bytes(&self) -> u64 {
362 self.reservation.bytes
363 }
364}
365
366pub(crate) struct Accounted<T> {
368 value: T,
369 _reservation: Option<Reservation>,
370 element_limit: Option<usize>,
371}
372
373impl<T> Deref for Accounted<T> {
374 type Target = T;
375
376 fn deref(&self) -> &Self::Target {
377 &self.value
378 }
379}
380
381impl<T> Accounted<Vec<T>> {
382 pub(crate) fn unaccounted_empty() -> Self {
383 Self {
384 value: Vec::new(),
385 _reservation: None,
386 element_limit: None,
387 }
388 }
389
390 pub(crate) fn try_with_capacity(
391 accounting: &Arc<Accounting>,
392 capacity: usize,
393 component: AllocationComponent,
394 ) -> Result<Self, StoreError> {
395 let requested =
396 capacity
397 .checked_mul(std::mem::size_of::<T>())
398 .ok_or(StoreError::BudgetExceeded {
399 needed: u64::MAX,
400 budget: u64::MAX,
401 component: component.name(),
402 })?;
403 let bytes = u64::try_from(requested).map_err(|_| StoreError::BudgetExceeded {
404 needed: u64::MAX,
405 budget: u64::MAX,
406 component: component.name(),
407 })?;
408 let reservation = accounting.reserve(bytes, component)?;
409 let mut value = Vec::new();
410 #[cfg(feature = "allocation-audit")]
411 let allocation = crate::allocation_audit::attributed(|| value.try_reserve_exact(capacity));
412 #[cfg(not(feature = "allocation-audit"))]
413 let allocation = value.try_reserve_exact(capacity);
414 allocation.map_err(|_| StoreError::AllocationFailed {
415 needed: bytes,
416 component: component.name(),
417 })?;
418 Ok(Self {
419 value,
420 _reservation: Some(reservation),
421 element_limit: Some(capacity),
422 })
423 }
424
425 pub(crate) fn try_from_vec(
426 accounting: &Arc<Accounting>,
427 value: Vec<T>,
428 component: AllocationComponent,
429 ) -> Result<Self, StoreError> {
430 let capacity = value.capacity();
431 let requested =
432 capacity
433 .checked_mul(std::mem::size_of::<T>())
434 .ok_or(StoreError::BudgetExceeded {
435 needed: u64::MAX,
436 budget: u64::MAX,
437 component: component.name(),
438 })?;
439 let bytes = u64::try_from(requested).map_err(|_| StoreError::BudgetExceeded {
440 needed: u64::MAX,
441 budget: u64::MAX,
442 component: component.name(),
443 })?;
444 let reservation = accounting.reserve(bytes, component)?;
445 Ok(Self {
446 value,
447 _reservation: Some(reservation),
448 element_limit: Some(capacity),
449 })
450 }
451
452 pub(crate) fn resident_bytes(&self) -> u64 {
453 self._reservation
454 .as_ref()
455 .map_or(0, |reservation| reservation.bytes)
456 }
457
458 pub(crate) fn try_reserve_total(
461 &mut self,
462 accounting: &Arc<Accounting>,
463 capacity: usize,
464 component: AllocationComponent,
465 ) -> Result<(), StoreError> {
466 if capacity <= self.element_limit.unwrap_or(0) {
467 return Ok(());
468 }
469 let mut replacement = Self::try_with_capacity(accounting, capacity, component)?;
470 if self.value.len() > capacity {
471 return Err(StoreError::AllocationFailed {
472 needed: self.value.len() as u64,
473 component: "accounted vector capacity",
474 });
475 }
476 replacement.value.append(&mut self.value);
479 *self = replacement;
480 Ok(())
481 }
482
483 pub(crate) fn push(&mut self, value: T) -> Result<(), StoreError> {
484 let accounted_elements = self.element_limit.unwrap_or(usize::MAX);
485 if self.value.len() >= accounted_elements || self.value.len() == self.value.capacity() {
486 return Err(StoreError::AllocationFailed {
487 needed: u64::try_from(self.value.len().saturating_add(1)).unwrap_or(u64::MAX),
488 component: "accounted vector capacity",
489 });
490 }
491 self.value.push(value);
492 Ok(())
493 }
494
495 pub(crate) fn extend_from_slice(&mut self, values: &[T]) -> Result<(), StoreError>
496 where
497 T: Clone,
498 {
499 let new_len =
500 self.value
501 .len()
502 .checked_add(values.len())
503 .ok_or(StoreError::AllocationFailed {
504 needed: u64::MAX,
505 component: "accounted vector capacity",
506 })?;
507 let accounted_elements = self.element_limit.unwrap_or(usize::MAX);
508 if new_len > accounted_elements || new_len > self.value.capacity() {
509 return Err(StoreError::AllocationFailed {
510 needed: u64::try_from(new_len).unwrap_or(u64::MAX),
511 component: "accounted vector capacity",
512 });
513 }
514 self.value.extend_from_slice(values);
515 Ok(())
516 }
517
518 pub(crate) fn replace_range(
519 &mut self,
520 range: std::ops::Range<usize>,
521 replacement: &[T],
522 ) -> Result<(), StoreError>
523 where
524 T: Copy + Default,
525 {
526 let old_len = self.value.len();
527 if range.start > range.end || range.end > old_len {
528 return Err(StoreError::ActiveRowOverflow);
529 }
530 let removed = range.end.saturating_sub(range.start);
531 let new_len = old_len
532 .checked_sub(removed)
533 .and_then(|length| length.checked_add(replacement.len()))
534 .ok_or(StoreError::ActiveRowOverflow)?;
535 let accounted_elements = self.element_limit.unwrap_or(usize::MAX);
536 if new_len > accounted_elements || new_len > self.value.capacity() {
537 return Err(StoreError::AllocationFailed {
538 needed: u64::try_from(new_len).unwrap_or(u64::MAX),
539 component: "accounted vector capacity",
540 });
541 }
542 if new_len > old_len {
543 self.value.resize(new_len, T::default());
544 }
545 let replacement_end = range
546 .start
547 .checked_add(replacement.len())
548 .ok_or(StoreError::ActiveRowOverflow)?;
549 self.value.copy_within(range.end..old_len, replacement_end);
550 self.value
551 .get_mut(range.start..replacement_end)
552 .ok_or(StoreError::ActiveRowOverflow)?
553 .copy_from_slice(replacement);
554 self.value.truncate(new_len);
555 Ok(())
556 }
557
558 pub(crate) fn as_mut_slice(&mut self) -> &mut [T] {
559 self.value.as_mut_slice()
560 }
561
562 pub(crate) fn clear(&mut self) {
564 self.value.clear();
565 }
566
567 pub(crate) fn dedup(&mut self)
568 where
569 T: PartialEq,
570 {
571 self.value.dedup();
572 }
573
574 pub(crate) fn reserve_additional_bytes(&mut self, additional: u64) -> Result<(), StoreError> {
575 match self._reservation.as_mut() {
576 Some(reservation) => reservation.grow(additional),
577 None if additional == 0 => Ok(()),
578 None => Err(StoreError::AllocationFailed {
579 needed: additional,
580 component: "unaccounted vector",
581 }),
582 }
583 }
584}
585
586#[cfg(test)]
587impl Accounted<Vec<u8>> {
588 fn try_zeroed(
589 accounting: &Arc<Accounting>,
590 bytes: u64,
591 component: AllocationComponent,
592 ) -> Result<Self, StoreError> {
593 let capacity = usize::try_from(bytes).map_err(|_| StoreError::BudgetExceeded {
594 needed: bytes,
595 budget: usize::MAX as u64,
596 component: component.name(),
597 })?;
598 let mut accounted = Self::try_with_capacity(accounting, capacity, component)?;
599 accounted.value.resize(capacity, 0);
600 Ok(accounted)
601 }
602}
603
604#[derive(Clone, Copy, Debug, Eq, PartialEq)]
606pub struct Stats {
607 pub resident_owned_bytes: u64,
609 pub snapshot_bytes: u64,
611 pub mapped_bytes: u64,
613 pub mapped_resident_bytes: u64,
615 pub segment_bytes: u64,
617 pub active_segment_bytes: u64,
620 pub retired_active_segment_bytes: u64,
623 pub active_row_count: u64,
625 pub tombstone_count: u64,
627 pub tombstone_bytes: u64,
629 pub wal_bytes: u64,
632 pub cache_bytes: u64,
634 pub temporary_bytes: u64,
636 pub query_pool_bytes: u64,
638 pub open_files: u64,
640 pub active_queries: u64,
642 pub active_snapshot_leases: u64,
644 pub phys_footprint: Option<u64>,
647 pub scans_by_reason: [u64; crate::planner::SCAN_REASON_COUNTER_COUNT],
649 pub graph_segments_served: u64,
651}
652
653impl Store {
654 pub fn stats(&self) -> Result<Stats, StoreError> {
660 let state = self
661 .state
662 .lock()
663 .map_err(|_| StoreError::Synchronization { component: "state" })?;
664 match *state {
665 StoreState::Open => {}
666 StoreState::Closing => return Err(StoreError::Closing),
667 StoreState::Closed => return Err(StoreError::Closed),
668 }
669 let stats = self.stats_while_open();
670 drop(state);
671 stats
672 }
673
674 pub(crate) fn stats_while_open(&self) -> Result<Stats, StoreError> {
675 let writer_lock = self
676 .writer_lock
677 .lock()
678 .map_err(|_| StoreError::Synchronization {
679 component: "writer lock",
680 })?;
681 let wal_writer = self
682 .wal_writer
683 .lock()
684 .map_err(|_| StoreError::Synchronization {
685 component: "WAL writer",
686 })?;
687 let open_files =
688 u64::from(writer_lock.is_some()).saturating_add(u64::from(wal_writer.is_some()));
689 let active_guard = self
690 .active
691 .lock()
692 .map_err(|_| StoreError::Synchronization {
693 component: "active segment",
694 })?;
695 let active_state = active_guard.as_ref().ok_or(StoreError::Closed)?;
696 let active_segment_bytes = active_state.segment.resident_bytes();
697 let active_row_count = u64::try_from(active_state.segment.row_count()).map_err(|_| {
698 StoreError::Statistics {
699 component: "active row count",
700 source: std::io::Error::other("active row count exceeds u64"),
701 }
702 })?;
703 let tombstone_count = active_state.segment.tombstone_count();
704 let tombstone_bytes = active_state.segment.tombstone_bytes();
705 let active_queries = self
706 .active_queries
707 .load(std::sync::atomic::Ordering::Relaxed);
708 let accounting = self.accounting.audit()?;
709 if accounting.resident_owned_bytes != accounting.component_sum() {
710 return Err(StoreError::Statistics {
711 component: "resident accounting conservation",
712 source: std::io::Error::other(
713 "resident-owned total disagrees with component counters",
714 ),
715 });
716 }
717 let retired_active_segment_bytes = accounting
718 .active_bytes
719 .checked_sub(active_segment_bytes)
720 .ok_or_else(|| StoreError::Statistics {
721 component: "active segment accounting",
722 source: std::io::Error::other(
723 "current active segment capacity exceeds the accounting component",
724 ),
725 })?;
726 let snapshot_guard = self
727 .snapshot
728 .read()
729 .map_err(|_| StoreError::Synchronization {
730 component: "published snapshot",
731 })?;
732 let snapshot = snapshot_guard.as_ref().ok_or(StoreError::Closed)?;
733 let active_snapshot_leases = u64::try_from(Arc::strong_count(snapshot).saturating_sub(1))
734 .map_err(|_| StoreError::Statistics {
735 component: "active snapshot leases",
736 source: std::io::Error::other("snapshot lease count exceeds u64"),
737 })?;
738 let mut mapped_resident_bytes = 0_u64;
739 for segment in snapshot.all_segments() {
740 let resident =
741 segment
742 .mapped_resident_bytes()
743 .map_err(|source| StoreError::Statistics {
744 component: "mapped resident bytes",
745 source,
746 })?;
747 mapped_resident_bytes =
748 mapped_resident_bytes.checked_add(resident).ok_or_else(|| {
749 StoreError::Statistics {
750 component: "mapped resident bytes",
751 source: std::io::Error::other("mapped resident byte count overflow"),
752 }
753 })?;
754 }
755 #[cfg(any(target_os = "macos", target_os = "ios"))]
756 let phys_footprint =
757 Some(
758 crate::sys::darwin::phys_footprint().map_err(|error| StoreError::Statistics {
759 component: "phys_footprint",
760 source: std::io::Error::other(error),
761 })?,
762 );
763 #[cfg(not(any(target_os = "macos", target_os = "ios")))]
764 let phys_footprint = None;
765 drop(snapshot_guard);
766 drop(active_guard);
767 drop(wal_writer);
768 drop(writer_lock);
769 Ok(Stats {
770 resident_owned_bytes: accounting.resident_owned_bytes,
771 snapshot_bytes: accounting.snapshot_bytes,
772 mapped_bytes: accounting.mapped_bytes,
773 mapped_resident_bytes,
774 segment_bytes: accounting.mapped_bytes,
775 active_segment_bytes,
776 retired_active_segment_bytes,
777 active_row_count,
778 tombstone_count,
779 tombstone_bytes,
780 wal_bytes: accounting.wal_bytes,
781 cache_bytes: accounting.cache_bytes,
782 temporary_bytes: accounting.temporary_bytes,
783 query_pool_bytes: accounting.query_pool_bytes,
784 open_files,
785 active_queries,
786 active_snapshot_leases,
787 phys_footprint,
788 scans_by_reason: accounting.scans_by_reason,
789 graph_segments_served: accounting.graph_segments_served,
790 })
791 }
792
793 #[cfg(test)]
794 fn allocate_bytes(
795 &self,
796 bytes: u64,
797 component: AllocationComponent,
798 ) -> Result<Accounted<Vec<u8>>, StoreError> {
799 let state = self
800 .state
801 .lock()
802 .map_err(|_| StoreError::Synchronization { component: "state" })?;
803 match *state {
804 StoreState::Open => {}
805 StoreState::Closing => return Err(StoreError::Closing),
806 StoreState::Closed => return Err(StoreError::Closed),
807 }
808 let allocation = Accounted::try_zeroed(&self.accounting, bytes, component)?;
809 drop(state);
810 Ok(allocation)
811 }
812}
813
814#[cfg(test)]
815#[allow(clippy::expect_used)]
816mod tests {
817 use tempfile::{TempDir, tempdir};
818
819 use super::super::{CancelToken, OpenOptions, QueryControl, Store, StoreError};
820 use super::{Accounted, Accounting, AllocationComponent};
821 use crate::lifecycle::durability::{CommitTier, DurabilityMode, DurabilityPolicy};
822 use crate::manifest::Manifest;
823 use crate::manifest::io::commit_manifest;
824 use crate::meta::{AliveSet, ColumnStoreBuilder, Schema};
825 use crate::quant::Bit4Factors;
826 use crate::scan::{F32Rows, ScanOptions, ScanQuery, ScanRequest, ScanRows};
827 use crate::segment::SegmentId;
828 use crate::segment::layout::RegionEntry;
829 use crate::segment::reader::SegmentReader;
830 use crate::segment::writer::{SegmentBuild, SegmentFactors, write_segment};
831 use crate::vfs::StdVfs;
832
833 #[test]
834 fn stats_bytes_are_conserved() {
835 let directory = tempdir().expect("store directory");
836 let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
837 let baseline = store.stats().expect("baseline stats");
838 let wal = store
839 .allocate_bytes(37, AllocationComponent::Wal)
840 .expect("accounted WAL bytes");
841 let cache = store
842 .allocate_bytes(19, AllocationComponent::Cache)
843 .expect("accounted cache bytes");
844 let temporary = store
845 .allocate_bytes(11, AllocationComponent::Temporary)
846 .expect("accounted temporary bytes");
847
848 let stats = store.stats().expect("stats");
849 let audit = store.accounting.audit().expect("accounting audit");
850 assert_eq!(stats.wal_bytes, 37);
851 assert_eq!(stats.cache_bytes, 19);
852 assert_eq!(stats.temporary_bytes, 11);
853 assert_eq!(stats.retired_active_segment_bytes, 0);
854 assert_eq!(
855 stats.resident_owned_bytes,
856 baseline.resident_owned_bytes + 67
857 );
858 assert_eq!(stats.resident_owned_bytes, audit.component_sum());
859 assert_eq!(stats.resident_owned_bytes, audit.resident_owned_bytes);
860
861 drop((wal, cache, temporary));
862 assert_eq!(
863 store.stats().expect("released stats").resident_owned_bytes,
864 baseline.resident_owned_bytes
865 );
866
867 let published = published_store();
868 let mapped = Store::open(published.path(), OpenOptions::default()).expect("mapped open");
869 let snapshot = mapped.snapshot().expect("mapped snapshot");
870 let segment = snapshot.segments().first().expect("fixture segment");
871 let expected_snapshot_bytes = std::mem::size_of::<SegmentReader>()
872 + std::mem::size_of_val(segment.directory())
873 + segment.retained_validation_bytes()
874 + usize::try_from(baseline.snapshot_bytes).expect("baseline snapshot bytes fit usize");
875 drop(snapshot);
876 let mapped_stats = mapped.stats().expect("mapped stats");
877 assert_eq!(
878 mapped_stats.resident_owned_bytes,
879 expected_snapshot_bytes as u64
880 );
881 assert_eq!(
882 mapped_stats.resident_owned_bytes,
883 mapped
884 .accounting
885 .audit()
886 .expect("mapped accounting audit")
887 .component_sum()
888 );
889 assert_eq!(mapped_stats.tombstone_count, 0);
890 assert_eq!(mapped_stats.tombstone_bytes, 0);
891 assert_eq!(mapped_stats.wal_bytes, 0);
892 assert_eq!(mapped_stats.cache_bytes, 0);
893 assert_eq!(mapped_stats.temporary_bytes, 0);
894 assert_eq!(mapped_stats.open_files, 2);
895 #[cfg(any(target_os = "macos", target_os = "ios"))]
896 assert!(mapped_stats.phys_footprint.is_some_and(|bytes| bytes > 0));
897
898 mapped.close().expect("close mapped store");
899 assert_eq!(
900 mapped
901 .accounting
902 .audit()
903 .expect("closed mapping accounting")
904 .mapped_bytes,
905 0,
906 "close must release the exact mapped-byte accounting"
907 );
908 assert!(matches!(mapped.stats(), Err(StoreError::Closed)));
909 }
910
911 #[test]
912 fn repeated_close_returns_exact_stats_counters_to_pre_open_baseline() {
913 const ITERATIONS: usize = 20;
914 const PRE_OPEN_MAPPED_BYTES: u64 = 0;
915 const PRE_OPEN_RESIDENT_OWNED_BYTES: u64 = 0;
916
917 let published = published_store();
918 for iteration in 0..ITERATIONS {
919 let store = Store::open(published.path(), OpenOptions::default()).expect("mapped open");
920 let query = [1.0_f32];
921 let rows = F32Rows::new(vec![1.0_f32; 128]);
922 store
923 .top_k_with_options(
924 ScanRequest {
925 query: ScanQuery::F32(&query),
926 rows: ScanRows::F32RowMajor(&rows),
927 row_mask: None,
928 },
929 1,
930 ScanOptions { thread_budget: 1 },
931 QueryControl::Cancel(CancelToken::new()),
932 )
933 .expect("start the persistent query pool");
934 let live = store.stats().expect("live stats");
935 assert!(live.mapped_bytes > PRE_OPEN_MAPPED_BYTES);
936 assert!(
937 live.resident_owned_bytes > PRE_OPEN_RESIDENT_OWNED_BYTES,
938 "iteration {iteration} must own mapped-snapshot bookkeeping"
939 );
940 assert!(
941 live.query_pool_bytes > 0,
942 "iteration {iteration} must own an accounted query pool"
943 );
944
945 store.close().expect("close mapped store");
946
947 let closed = store
948 .accounting
949 .audit()
950 .expect("post-close accounting source for Stats");
951 assert_eq!(
952 closed.mapped_bytes, PRE_OPEN_MAPPED_BYTES,
953 "Stats::mapped_bytes source leaked at iteration {iteration}"
954 );
955 assert_eq!(
956 closed.resident_owned_bytes, PRE_OPEN_RESIDENT_OWNED_BYTES,
957 "Stats::resident_owned_bytes source leaked at iteration {iteration}"
958 );
959 assert_eq!(
960 closed.query_pool_bytes, 0,
961 "Stats::query_pool_bytes source leaked at iteration {iteration}"
962 );
963 }
964 }
965
966 #[test]
967 fn budget_exceeded_is_typed_and_pre_allocation() {
968 let directory = tempdir().expect("store directory");
969 let baseline_store =
970 Store::open(directory.path(), OpenOptions::default()).expect("baseline open");
971 let baseline = baseline_store
972 .stats()
973 .expect("baseline stats")
974 .snapshot_bytes;
975 baseline_store.close().expect("close baseline");
976 let options = OpenOptions::new()
977 .with_max_resident_bytes(baseline + 64)
978 .with_max_temp_bytes(32);
979 let store = Store::open(directory.path(), options).expect("open");
980
981 let rejected = store.allocate_bytes(40, AllocationComponent::Temporary);
982 assert!(matches!(
983 rejected,
984 Err(StoreError::BudgetExceeded {
985 needed: 40,
986 budget: 32,
987 component: "temporary",
988 })
989 ));
990 let after_rejection = store.stats().expect("stats after rejection");
991 assert_eq!(after_rejection.resident_owned_bytes, baseline);
992 assert_eq!(after_rejection.temporary_bytes, 0);
993
994 let accepted = store
995 .allocate_bytes(24, AllocationComponent::Temporary)
996 .expect("smaller allocation on same handle");
997 assert_eq!(accepted.len(), 24);
998 assert_eq!(store.stats().expect("usable stats").temporary_bytes, 24);
999 assert!(matches!(
1000 store.allocate_bytes(50, AllocationComponent::Wal),
1001 Err(StoreError::BudgetExceeded {
1002 needed,
1003 budget,
1004 component: "wal",
1005 }) if needed == baseline + 74 && budget == baseline + 64
1006 ));
1007 assert_eq!(
1008 store.stats().expect("resident rejection stats").wal_bytes,
1009 0
1010 );
1011 store.snapshot().expect("same handle remains usable");
1012
1013 let published = published_store();
1014 let snapshot_bytes =
1015 std::mem::size_of::<SegmentReader>() + 6 * std::mem::size_of::<RegionEntry>();
1016 let budget =
1017 u64::try_from(snapshot_bytes.saturating_sub(1)).expect("snapshot budget fits u64");
1018 let needed = u64::try_from(snapshot_bytes).expect("snapshot bytes fit u64");
1019 let rejected_open = Store::open(
1020 published.path(),
1021 OpenOptions::new().with_max_resident_bytes(budget),
1022 );
1023 assert!(matches!(
1024 rejected_open,
1025 Err(StoreError::BudgetExceeded {
1026 needed: actual_needed,
1027 budget: actual_budget,
1028 component: "snapshot",
1029 }) if actual_needed == needed && actual_budget == budget
1030 ));
1031 }
1032
1033 #[test]
1034 fn accounted_vector_rejects_growth_without_a_reservation() {
1035 let mut unaccounted = Accounted::<Vec<u8>>::unaccounted_empty();
1036 unaccounted
1037 .reserve_additional_bytes(0)
1038 .expect("zero bytes require no reservation");
1039 assert!(matches!(
1040 unaccounted.reserve_additional_bytes(1),
1041 Err(StoreError::AllocationFailed {
1042 needed: 1,
1043 component: "unaccounted vector",
1044 })
1045 ));
1046
1047 let accounting = std::sync::Arc::new(Accounting::new(u64::MAX, u64::MAX));
1048 let mut fixed = Accounted::try_with_capacity(&accounting, 1, AllocationComponent::Wal)
1049 .expect("fixed-capacity vector");
1050 fixed.push(7_u8).expect("reserved element");
1051 assert!(matches!(
1052 fixed.push(8_u8),
1053 Err(StoreError::AllocationFailed {
1054 needed: 2,
1055 component: "accounted vector capacity",
1056 })
1057 ));
1058 }
1059
1060 #[test]
1061 fn accounted_allocation_after_close_is_typed_closed() {
1062 let directory = tempdir().expect("store directory");
1063 let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
1064 store.close().expect("close");
1065
1066 assert!(matches!(
1067 store.allocate_bytes(1, AllocationComponent::Wal),
1068 Err(StoreError::Closed)
1069 ));
1070 }
1071
1072 #[cfg(feature = "allocation-audit")]
1073 #[test]
1074 fn nothing_allocates_outside_accounting() {
1075 let directory = tempdir().expect("store directory");
1076 let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
1077
1078 let (unaccounted, negative_control) =
1079 crate::allocation_audit::audit_engine_path(|| Vec::<u8>::with_capacity(13));
1080 assert_eq!(negative_control.attributed_bytes, 0);
1081 assert!(
1082 negative_control.unattributed_bytes >= 13,
1083 "the audit wrapper failed to detect an unattributed engine-path allocation"
1084 );
1085 drop(unaccounted);
1086 store.stats().expect("warm accounting synchronization");
1087
1088 let (accounted, report) = crate::allocation_audit::audit_engine_path(|| {
1089 store.allocate_bytes(23, AllocationComponent::Wal)
1090 });
1091 let accounted = accounted.expect("accounted allocation");
1092 assert_eq!(report.unattributed_bytes, 0);
1093 assert_eq!(report.attributed_bytes, 23);
1094 assert_eq!(store.stats().expect("stats").wal_bytes, 23);
1095 drop(accounted);
1096 }
1097
1098 fn published_store() -> TempDir {
1099 let directory = tempdir().expect("published store directory");
1100 let schema = Schema::new(Vec::new()).expect("timestamp-only schema");
1101 let mut columns = ColumnStoreBuilder::new(schema.clone());
1102 columns.push_row(17, &[]).expect("fixture row");
1103 let columns = columns.finish().expect("fixture columns");
1104 let alive = AliveSet::new(1);
1105 let id = SegmentId::new(0x0102_0304_0506, [0x09; 10]);
1106 let codes = [0x88_u8];
1107 let factors = [Bit4Factors::from_persisted(1.0, 1.0, 1.0)];
1108 let rescore = [0.0_f32, 0.0_f32];
1109 let policy = DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::Ordered)
1110 .expect("derived policy");
1111 let segment = write_segment(
1112 &StdVfs,
1113 directory.path(),
1114 SegmentBuild {
1115 id,
1116 scheme: 4,
1117 dims: 2,
1118 codes: &codes,
1119 factors: SegmentFactors::Bit4(&factors),
1120 rescore: &rescore,
1121 columns: &columns,
1122 alive: &alive,
1123 },
1124 policy,
1125 )
1126 .expect("fixture segment");
1127 commit_manifest(
1128 &StdVfs,
1129 directory.path(),
1130 &Manifest {
1131 generation: 1,
1132 log_seq: 0,
1133 segments: vec![segment],
1134 epochs: Vec::new(),
1135 epoch_alias: None,
1136 schema,
1137 },
1138 policy,
1139 )
1140 .expect("fixture manifest");
1141 directory
1142 }
1143}