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#[cfg(feature = "graph-cypher")]
10mod allocation_interval;
11#[cfg(feature = "graph-cypher")]
12pub use allocation_interval::{
13    AllocationInterval, AllocationIntervalError, AllocationIntervalSnapshot,
14};
15
16/// Store-cumulative native write work, including work before rejected writes.
17/// Byte metrics are logical formats, not allocation capacities or DRAM traffic.
18#[cfg(feature = "graph-cypher")]
19#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
20pub struct GraphWorkLedger {
21    /// Non-query storage work, including preparation and maintenance.
22    pub storage_lookups: u64,
23    /// Non-query storage work, including preparation and maintenance.
24    pub storage_scans: u64,
25    /// Non-query storage work, including preparation and maintenance.
26    pub storage_adjacency_entries: u64,
27    /// Non-query storage work, including preparation and maintenance.
28    pub storage_copied_bytes: u64,
29    /// Non-query storage work, including preparation and maintenance.
30    pub storage_pages_decoded: u64,
31    /// Non-query storage work, including preparation and maintenance.
32    pub storage_pages_copied: u64,
33    /// Non-query storage work, including preparation and maintenance.
34    pub storage_property_values: u64,
35    /// Non-query storage work, including preparation and maintenance.
36    pub storage_property_bytes: u64,
37    /// Non-query storage work, including preparation and maintenance.
38    pub storage_adjacency_physical_entries: u64,
39    /// Non-query storage work, including preparation and maintenance.
40    pub storage_adjacency_merged_visits: u64,
41    /// Non-query storage work, including preparation and maintenance.
42    pub storage_adjacency_merge_runs: u64,
43
44    /// Canonical byte pairs actually compared.
45    pub canonical_comparison_bytes: u64,
46    /// Canonical bytes emitted into real staging backing, excluding sizing passes.
47    pub canonical_encoding_bytes: u64,
48    /// Admitted WAL codec byte/descriptor units, including repeated passes.
49    pub wal_codec_units: u64,
50    /// Complete native commit envelope bytes encoded before append.
51    pub encoded_wal_bytes: u64,
52    /// Bytes supplied to successful commit-artifact writes.
53    pub artifact_bytes_written: u64,
54    /// Successful commit-artifact writes.
55    pub artifact_writes: u64,
56    /// Bytes supplied to successful native commit WAL appends.
57    pub wal_bytes_appended: u64,
58    /// Successful native commit WAL appends.
59    pub wal_appends: u64,
60    /// Attempted full file syncs at native commit publication.
61    pub full_sync_attempts: u64,
62    /// Successful full file syncs at native commit publication.
63    pub full_sync_successes: u64,
64    /// Attempted commit directory syncs, separate from file syncs.
65    pub directory_sync_attempts: u64,
66    /// Successful commit directory syncs.
67    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    /// Enabling graph shares its hard ceiling with every store participant.
329    #[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    /// Checks batch-derived temporary headroom without reserving or allocating.
375    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
743/// Exact non-allocating component bytes whose backing allocation is owned by
744/// another subsystem, such as the WAL's shared encoded record buffers.
745pub(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
790/// An owned value whose exact requested allocation bytes have one component owner.
791pub(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    /// Grows retained capacity while charging both old and replacement
883    /// allocations until the old buffer has actually been released.
884    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        // The destination was reserved for every element. Moving them keeps
901        // nested allocation owners intact and cannot trigger another growth.
902        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    /// Drops elements while retaining the complete accounted allocation.
987    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/// One point-in-time snapshot of resources owned by an open [`Store`].
1029#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1030pub struct Stats {
1031    /// Exact bytes in engine-owned anonymous allocation arenas and vectors.
1032    pub resident_owned_bytes: u64,
1033    /// Exact bytes retained by published snapshot metadata and decoded views.
1034    pub snapshot_bytes: u64,
1035    /// Exact byte length of the store's live read-only mappings.
1036    pub mapped_bytes: u64,
1037    /// Bytes in resident mapped pages, counted by `mincore`.
1038    pub mapped_resident_bytes: u64,
1039    /// Exact bytes in immutable segment mappings.
1040    pub segment_bytes: u64,
1041    /// Exact allocation capacity held by the currently published in-RAM active
1042    /// segment.
1043    pub active_segment_bytes: u64,
1044    /// Exact allocation capacity held by retired active generations that are
1045    /// still pinned by admitted queries.
1046    pub retired_active_segment_bytes: u64,
1047    /// Exact number of rows currently represented by the in-RAM active segment.
1048    pub active_row_count: u64,
1049    /// Exact number of in-memory active-segment tombstones.
1050    pub tombstone_count: u64,
1051    /// Exact bytes owned by active-segment tombstones.
1052    pub tombstone_bytes: u64,
1053    /// Exact bytes retained by the store's in-memory WAL component. This is
1054    /// zero while [`Store`] has no WAL writer component.
1055    pub wal_bytes: u64,
1056    /// Exact bytes owned by reusable per-segment graph-search scratch caches.
1057    pub cache_bytes: u64,
1058    /// Exact bytes in live, explicitly accounted temporary allocations.
1059    pub temporary_bytes: u64,
1060    /// Exact capacity in bytes of the persistent query pool's worker registry.
1061    pub query_pool_bytes: u64,
1062    /// Exact capacity in bytes of the native graph publication's store-lifetime
1063    /// lease, mapping, preparation, and spill registries. This is zero while
1064    /// [`Store`] has no native graph publication.
1065    pub native_graph_bytes: u64,
1066    /// Exact number of file descriptors retained by this store handle.
1067    pub open_files: u64,
1068    /// Exact number of store-admitted queries that have not yet returned.
1069    pub active_queries: u64,
1070    /// Exact number of snapshot leases held outside the published snapshot slot.
1071    pub active_snapshot_leases: u64,
1072    /// Darwin's `TASK_VM_INFO` physical-footprint kernel counter, or `None` on
1073    /// platforms where that counter does not exist.
1074    pub phys_footprint: Option<u64>,
1075    /// Monotonic count of successfully executed scan plans by typed reason.
1076    pub scans_by_reason: [u64; crate::planner::SCAN_REASON_COUNTER_COUNT],
1077    /// Monotonic count of segment plans served by graph traversal.
1078    pub graph_segments_served: u64,
1079}
1080
1081impl Store {
1082    /// Returns exact accounting and kernel counters for this open handle.
1083    ///
1084    /// This operation never estimates. Engine-owned counters are maintained at
1085    /// allocation/mapping boundaries, mapped residency comes from `mincore`,
1086    /// and `phys_footprint` comes from Darwin `TASK_VM_INFO`.
1087    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        // Unified graph creation owns an empty document state and a v3 snapshot.
1376        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            // Store-lifetime components are resident from open, not just the
1439            // mapped snapshot. Zero unless `graph-cypher` is enabled.
1440            + 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        // The resident total, not just the snapshot component: a store-lifetime
1539        // component such as the native graph registry is resident from open.
1540        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}