Skip to main content

zeppelin_embed/lifecycle/
stats.rs

1//! Exact engine-owned memory and kernel-backed process statistics.
2
3use 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
329/// Exact non-allocating component bytes whose backing allocation is owned by
330/// another subsystem, such as the WAL's shared encoded record buffers.
331pub(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
366/// An owned value whose exact requested allocation bytes have one component owner.
367pub(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    /// Grows retained capacity while charging both old and replacement
459    /// allocations until the old buffer has actually been released.
460    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        // The destination was reserved for every element. Moving them keeps
477        // nested allocation owners intact and cannot trigger another growth.
478        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    /// Drops elements while retaining the complete accounted allocation.
563    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/// One point-in-time snapshot of resources owned by an open [`Store`].
605#[derive(Clone, Copy, Debug, Eq, PartialEq)]
606pub struct Stats {
607    /// Exact bytes in engine-owned anonymous allocation arenas and vectors.
608    pub resident_owned_bytes: u64,
609    /// Exact bytes retained by published snapshot metadata and decoded views.
610    pub snapshot_bytes: u64,
611    /// Exact byte length of the store's live read-only mappings.
612    pub mapped_bytes: u64,
613    /// Bytes in resident mapped pages, counted by `mincore`.
614    pub mapped_resident_bytes: u64,
615    /// Exact bytes in immutable segment mappings.
616    pub segment_bytes: u64,
617    /// Exact allocation capacity held by the currently published in-RAM active
618    /// segment.
619    pub active_segment_bytes: u64,
620    /// Exact allocation capacity held by retired active generations that are
621    /// still pinned by admitted queries.
622    pub retired_active_segment_bytes: u64,
623    /// Exact number of rows currently represented by the in-RAM active segment.
624    pub active_row_count: u64,
625    /// Exact number of in-memory active-segment tombstones.
626    pub tombstone_count: u64,
627    /// Exact bytes owned by active-segment tombstones.
628    pub tombstone_bytes: u64,
629    /// Exact bytes retained by the store's in-memory WAL component. This is
630    /// zero while [`Store`] has no WAL writer component.
631    pub wal_bytes: u64,
632    /// Exact bytes owned by reusable per-segment graph-search scratch caches.
633    pub cache_bytes: u64,
634    /// Exact bytes in live, explicitly accounted temporary allocations.
635    pub temporary_bytes: u64,
636    /// Exact capacity in bytes of the persistent query pool's worker registry.
637    pub query_pool_bytes: u64,
638    /// Exact number of file descriptors retained by this store handle.
639    pub open_files: u64,
640    /// Exact number of store-admitted queries that have not yet returned.
641    pub active_queries: u64,
642    /// Exact number of snapshot leases held outside the published snapshot slot.
643    pub active_snapshot_leases: u64,
644    /// Darwin's `TASK_VM_INFO` physical-footprint kernel counter, or `None` on
645    /// platforms where that counter does not exist.
646    pub phys_footprint: Option<u64>,
647    /// Monotonic count of successfully executed scan plans by typed reason.
648    pub scans_by_reason: [u64; crate::planner::SCAN_REASON_COUNTER_COUNT],
649    /// Monotonic count of segment plans served by graph traversal.
650    pub graph_segments_served: u64,
651}
652
653impl Store {
654    /// Returns exact accounting and kernel counters for this open handle.
655    ///
656    /// This operation never estimates. Engine-owned counters are maintained at
657    /// allocation/mapping boundaries, mapped residency comes from `mincore`,
658    /// and `phys_footprint` comes from Darwin `TASK_VM_INFO`.
659    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}