cranpose_core/
snapshot_pinning.rs

1use std::cell::RefCell;
2
3use crate::{
4    snapshot_double_index_heap::{SnapshotDoubleIndexHeap, SnapshotDoubleIndexHeapDebugStats},
5    snapshot_id_set::{SnapshotId, SnapshotIdSet},
6};
7
8/// A handle to a pinned snapshot. Dropping this handle releases the pin.
9///
10/// Internally stores a heap handle for O(log N) removal.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
12pub struct PinHandle(usize);
13
14impl PinHandle {
15    /// Invalid pin handle constant (0 is reserved as invalid).
16    pub const INVALID: PinHandle = PinHandle(0);
17
18    /// Check if this handle is valid (non-zero).
19    pub fn is_valid(&self) -> bool {
20        self.0 != 0
21    }
22}
23
24struct PinningTable {
25    heap: SnapshotDoubleIndexHeap,
26}
27
28impl PinningTable {
29    fn new() -> Self {
30        Self {
31            heap: SnapshotDoubleIndexHeap::new(),
32        }
33    }
34
35    fn add(&mut self, snapshot_id: SnapshotId) -> PinHandle {
36        let heap_handle = self.heap.add(snapshot_id);
37        PinHandle(heap_handle + 1)
38    }
39
40    fn remove(&mut self, handle: PinHandle) -> bool {
41        if !handle.is_valid() {
42            return false;
43        }
44
45        let heap_handle = handle.0 - 1;
46
47        if heap_handle < usize::MAX {
48            self.heap.remove(heap_handle);
49            true
50        } else {
51            false
52        }
53    }
54
55    fn lowest_pinned(&self) -> Option<SnapshotId> {
56        if self.heap.is_empty() {
57            None
58        } else {
59            Some(self.heap.lowest_or_default(0))
60        }
61    }
62
63    fn pin_count(&self) -> usize {
64        self.heap.len()
65    }
66
67    fn debug_stats(&self) -> SnapshotPinningDebugStats {
68        SnapshotPinningDebugStats {
69            pin_count: self.pin_count(),
70            lowest_pinned_snapshot: self.lowest_pinned(),
71            heap: self.heap.debug_stats(),
72        }
73    }
74}
75
76thread_local! {
77    static PINNING_TABLE: RefCell<PinningTable> = RefCell::new(PinningTable::new());
78}
79
80#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
81pub struct SnapshotPinningDebugStats {
82    pub pin_count: usize,
83    pub lowest_pinned_snapshot: Option<SnapshotId>,
84    pub heap: SnapshotDoubleIndexHeapDebugStats,
85}
86
87/// Pin a snapshot and its invalid set, returning a handle.
88///
89/// This should be called when a snapshot is created to ensure that state records
90/// from the pinned snapshot and all its dependencies remain valid.
91///
92/// # Arguments
93/// * `snapshot_id` - The ID of the snapshot being created
94/// * `invalid` - The set of invalid snapshot IDs for this snapshot
95///
96/// # Returns
97/// A pin handle that should be released when the snapshot is disposed.
98///
99/// # Time Complexity
100/// O(log N) where N is the number of pinned snapshots
101pub fn track_pinning(snapshot_id: SnapshotId, invalid: &SnapshotIdSet) -> PinHandle {
102    let pinned_id = invalid.lowest(snapshot_id);
103
104    PINNING_TABLE.with(|cell| cell.borrow_mut().add(pinned_id))
105}
106
107/// Release a pinned snapshot.
108///
109/// # Arguments
110/// * `handle` - The pin handle returned by `track_pinning`
111///
112/// This must be called while holding the appropriate lock (sync).
113///
114/// # Time Complexity
115/// O(log N) where N is the number of pinned snapshots
116pub fn release_pinning(handle: PinHandle) {
117    if !handle.is_valid() {
118        return;
119    }
120
121    PINNING_TABLE.with(|cell| {
122        cell.borrow_mut().remove(handle);
123    });
124}
125
126/// Get the lowest currently pinned snapshot ID.
127///
128/// This is used to determine which state records can be safely garbage collected.
129/// Any state records from snapshots older than this ID are still potentially in use.
130///
131/// # Time Complexity
132/// O(1)
133pub fn lowest_pinned_snapshot() -> Option<SnapshotId> {
134    PINNING_TABLE.with(|cell| cell.borrow().lowest_pinned())
135}
136
137/// Get the current count of pinned snapshots (for testing).
138/// Get the current count of pinned snapshots (for testing/debugging).
139pub fn pin_count() -> usize {
140    PINNING_TABLE.with(|cell| cell.borrow().pin_count())
141}
142
143pub fn debug_snapshot_pinning_stats() -> SnapshotPinningDebugStats {
144    PINNING_TABLE.with(|cell| cell.borrow().debug_stats())
145}
146
147#[cfg(test)]
148pub fn reset_pinning_table() {
149    PINNING_TABLE.with(|cell| {
150        let mut table = cell.borrow_mut();
151        table.heap = SnapshotDoubleIndexHeap::new();
152    });
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158
159    fn setup() {
160        reset_pinning_table();
161    }
162
163    #[test]
164    fn test_invalid_handle() {
165        let handle = PinHandle::INVALID;
166        assert!(!handle.is_valid());
167        assert_eq!(handle.0, 0);
168    }
169
170    #[test]
171    fn test_valid_handle() {
172        setup();
173        let invalid = SnapshotIdSet::new().set(10);
174        let handle = track_pinning(20, &invalid);
175        assert!(handle.is_valid());
176        assert!(handle.0 > 0);
177    }
178
179    #[test]
180    fn test_track_and_release() {
181        setup();
182
183        let invalid = SnapshotIdSet::new().set(10);
184        let handle = track_pinning(20, &invalid);
185
186        assert_eq!(pin_count(), 1);
187        assert_eq!(lowest_pinned_snapshot(), Some(10));
188
189        release_pinning(handle);
190        assert_eq!(pin_count(), 0);
191        assert_eq!(lowest_pinned_snapshot(), None);
192    }
193
194    #[test]
195    fn test_multiple_pins() {
196        setup();
197
198        let invalid1 = SnapshotIdSet::new().set(10);
199        let handle1 = track_pinning(20, &invalid1);
200
201        let invalid2 = SnapshotIdSet::new().set(5).set(15);
202        let handle2 = track_pinning(30, &invalid2);
203
204        assert_eq!(pin_count(), 2);
205        assert_eq!(lowest_pinned_snapshot(), Some(5));
206
207        release_pinning(handle1);
208        assert_eq!(pin_count(), 1);
209        assert_eq!(lowest_pinned_snapshot(), Some(5));
210
211        release_pinning(handle2);
212        assert_eq!(pin_count(), 0);
213        assert_eq!(lowest_pinned_snapshot(), None);
214    }
215
216    #[test]
217    fn test_duplicate_pins() {
218        setup();
219
220        let invalid = SnapshotIdSet::new().set(10);
221        let handle1 = track_pinning(20, &invalid);
222        let handle2 = track_pinning(25, &invalid);
223
224        assert_eq!(pin_count(), 2);
225        assert_eq!(lowest_pinned_snapshot(), Some(10));
226
227        release_pinning(handle1);
228        assert_eq!(pin_count(), 1);
229        assert_eq!(lowest_pinned_snapshot(), Some(10));
230
231        release_pinning(handle2);
232        assert_eq!(pin_count(), 0);
233        assert_eq!(lowest_pinned_snapshot(), None);
234    }
235
236    #[test]
237    fn test_pin_ordering() {
238        setup();
239
240        let invalid1 = SnapshotIdSet::new().set(30);
241        let _handle1 = track_pinning(40, &invalid1);
242
243        let invalid2 = SnapshotIdSet::new().set(10);
244        let _handle2 = track_pinning(20, &invalid2);
245
246        let invalid3 = SnapshotIdSet::new().set(20);
247        let _handle3 = track_pinning(30, &invalid3);
248
249        assert_eq!(lowest_pinned_snapshot(), Some(10));
250    }
251
252    #[test]
253    fn test_release_invalid_handle() {
254        setup();
255
256        release_pinning(PinHandle::INVALID);
257        assert_eq!(pin_count(), 0);
258    }
259
260    #[test]
261    fn test_empty_invalid_set() {
262        setup();
263
264        let invalid = SnapshotIdSet::new();
265        let handle = track_pinning(100, &invalid);
266
267        assert_eq!(pin_count(), 1);
268        assert_eq!(lowest_pinned_snapshot(), Some(100));
269
270        release_pinning(handle);
271    }
272
273    #[test]
274    fn test_lowest_from_invalid_set() {
275        setup();
276
277        let invalid = SnapshotIdSet::new().set(5).set(10).set(15).set(20);
278        let handle = track_pinning(25, &invalid);
279
280        assert_eq!(lowest_pinned_snapshot(), Some(5));
281
282        release_pinning(handle);
283    }
284
285    #[test]
286    fn test_concurrent_snapshots() {
287        setup();
288
289        let handles: Vec<_> = (0..10)
290            .map(|i| {
291                let invalid = SnapshotIdSet::new().set(i * 10);
292                track_pinning(i * 10 + 5, &invalid)
293            })
294            .collect();
295
296        assert_eq!(pin_count(), 10);
297        assert_eq!(lowest_pinned_snapshot(), Some(0));
298
299        for handle in handles {
300            release_pinning(handle);
301        }
302
303        assert_eq!(pin_count(), 0);
304        assert_eq!(lowest_pinned_snapshot(), None);
305    }
306
307    #[test]
308    fn test_heap_handle_based_removal() {
309        setup();
310
311        let invalid1 = SnapshotIdSet::new().set(42);
312        let invalid2 = SnapshotIdSet::new().set(17);
313        let invalid3 = SnapshotIdSet::new().set(99);
314
315        let h1 = track_pinning(50, &invalid1);
316        let h2 = track_pinning(25, &invalid2);
317        let h3 = track_pinning(100, &invalid3);
318
319        assert_eq!(pin_count(), 3);
320        assert_eq!(lowest_pinned_snapshot(), Some(17));
321
322        release_pinning(h1);
323        assert_eq!(pin_count(), 2);
324        assert_eq!(lowest_pinned_snapshot(), Some(17));
325
326        release_pinning(h2);
327        assert_eq!(pin_count(), 1);
328        assert_eq!(lowest_pinned_snapshot(), Some(99));
329
330        release_pinning(h3);
331        assert!(pin_count() == 0);
332    }
333}
cranpose_core/snapshot_pinning.rs

cranpose_core/
snapshot_pinning.rs
