clt_database/alloc/
allocation_site.rs

1use std::cell::Cell;
2
3#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
4pub enum AllocationSite {
5    MvStore(MvStoreAllocationSite),
6    MvccCheckpoint(MvccCheckpointAllocationSite),
7    Schema(SchemaAllocationSite),
8    NoFaultInjection,
9}
10
11#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
12pub enum MvStoreAllocationSite {
13    RootpageMappingInsert,
14    TxInsert,
15    FinalizedTxStateInsert,
16    TableRowsEntry,
17    IndexRowsEntry,
18    IndexKeyEntry,
19    RowVersionReserve,
20    RowPayload,
21    SchemaRowPayload,
22}
23
24#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
25pub enum SchemaAllocationSite {
26    MakeMut,
27}
28
29#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
30pub enum MvccCheckpointAllocationSite {
31    CheckpointWriteSet,
32    CheckpointIndexWriteSet,
33    CheckpointMetadataPayload,
34    CheckpointSequenceCompactions,
35}
36
37impl From<MvStoreAllocationSite> for AllocationSite {
38    fn from(site: MvStoreAllocationSite) -> Self {
39        Self::MvStore(site)
40    }
41}
42
43impl From<MvccCheckpointAllocationSite> for AllocationSite {
44    fn from(site: MvccCheckpointAllocationSite) -> Self {
45        Self::MvccCheckpoint(site)
46    }
47}
48
49impl From<SchemaAllocationSite> for AllocationSite {
50    fn from(site: SchemaAllocationSite) -> Self {
51        Self::Schema(site)
52    }
53}
54
55thread_local! {
56    static CURRENT_ALLOCATION_SITE: Cell<Option<AllocationSite>> = const { Cell::new(None) };
57}
58
59pub struct AllocationSiteGuard {
60    previous: Option<AllocationSite>,
61}
62
63impl Drop for AllocationSiteGuard {
64    fn drop(&mut self) {
65        CURRENT_ALLOCATION_SITE.with(|slot| slot.set(self.previous));
66    }
67}
68
69pub fn enter_allocation_site(site: impl Into<AllocationSite>) -> AllocationSiteGuard {
70    let site = site.into();
71    let previous = CURRENT_ALLOCATION_SITE.with(|slot| {
72        let previous = slot.get();
73        let site = if matches!(previous, Some(AllocationSite::NoFaultInjection)) {
74            AllocationSite::NoFaultInjection
75        } else {
76            site
77        };
78        slot.set(Some(site));
79        previous
80    });
81    AllocationSiteGuard { previous }
82}
83
84pub fn current_allocation_site() -> Option<AllocationSite> {
85    CURRENT_ALLOCATION_SITE.with(Cell::get)
86}
87
88#[macro_export]
89macro_rules! without_allocation_faults {
90    ($expr:expr) => {{
91        #[cfg(clt_turso_feature = "allocation_metric")]
92        let _turso_allocation_site_guard =
93            $crate::alloc::enter_allocation_site($crate::alloc::AllocationSite::NoFaultInjection);
94        $expr
95    }};
96}
97
98#[macro_export]
99macro_rules! with_mv_store_allocation_site {
100    ($site:ident, $expr:expr) => {{
101        #[cfg(clt_turso_feature = "allocation_metric")]
102        let _turso_allocation_site_guard =
103            $crate::alloc::enter_allocation_site($crate::alloc::MvStoreAllocationSite::$site);
104        $expr
105    }};
106}
107
108#[cfg(clt_turso_tests)]
109mod tests {
110    use super::{
111        current_allocation_site, enter_allocation_site, AllocationSite, MvStoreAllocationSite,
112    };
113
114    #[test]
115    fn allocation_site_guard_restores_previous_site() {
116        assert_eq!(current_allocation_site(), None);
117        {
118            let _outer = enter_allocation_site(MvStoreAllocationSite::RootpageMappingInsert);
119            assert_eq!(
120                current_allocation_site(),
121                Some(AllocationSite::MvStore(
122                    MvStoreAllocationSite::RootpageMappingInsert
123                ))
124            );
125
126            {
127                let _inner = enter_allocation_site(AllocationSite::NoFaultInjection);
128                assert_eq!(
129                    current_allocation_site(),
130                    Some(AllocationSite::NoFaultInjection)
131                );
132            }
133
134            assert_eq!(
135                current_allocation_site(),
136                Some(AllocationSite::MvStore(
137                    MvStoreAllocationSite::RootpageMappingInsert
138                ))
139            );
140        }
141        assert_eq!(current_allocation_site(), None);
142    }
143
144    #[test]
145    fn no_fault_injection_site_dominates_nested_sites() {
146        let _outer = enter_allocation_site(AllocationSite::NoFaultInjection);
147        assert_eq!(
148            current_allocation_site(),
149            Some(AllocationSite::NoFaultInjection)
150        );
151        {
152            let _inner = enter_allocation_site(MvStoreAllocationSite::RowVersionReserve);
153            assert_eq!(
154                current_allocation_site(),
155                Some(AllocationSite::NoFaultInjection)
156            );
157        }
158        assert_eq!(
159            current_allocation_site(),
160            Some(AllocationSite::NoFaultInjection)
161        );
162    }
163}
clt_database/alloc/allocation_site.rs

clt_database/alloc/
allocation_site.rs
