tinymist_query/adt/
revision.rs

1use std::{
2    collections::HashMap,
3    num::NonZeroUsize,
4    sync::{Arc, OnceLock},
5};
6
7pub struct RevisionLock {
8    estimated: usize,
9    used: OnceLock<usize>,
10}
11
12impl RevisionLock {
13    pub fn access(&self, revision: NonZeroUsize) {
14        self.used
15            .set(revision.get())
16            .unwrap_or_else(|_| panic!("revision {revision} is determined"))
17    }
18}
19
20pub struct RevisionSlot<T> {
21    pub revision: usize,
22    pub data: T,
23}
24
25impl<T> std::ops::Deref for RevisionSlot<T> {
26    type Target = T;
27
28    fn deref(&self) -> &Self::Target {
29        &self.data
30    }
31}
32
33impl<T> std::ops::DerefMut for RevisionSlot<T> {
34    fn deref_mut(&mut self) -> &mut Self::Target {
35        &mut self.data
36    }
37}
38
39pub struct RevisionManager<T> {
40    estimated: usize,
41    locked: HashMap<usize, usize>,
42    slots: Vec<Arc<RevisionSlot<T>>>,
43}
44
45impl<T> Default for RevisionManager<T> {
46    fn default() -> Self {
47        Self {
48            estimated: 0,
49            locked: Default::default(),
50            slots: Default::default(),
51        }
52    }
53}
54
55impl<T> RevisionManager<T> {
56    pub fn clear(&mut self) {
57        self.slots.clear();
58    }
59
60    /// Lock the revision in *main thread*.
61    #[must_use]
62    pub fn lock(&mut self, used: NonZeroUsize) -> RevisionLock {
63        let lock = self.lock_estimated();
64        lock.access(used);
65        lock
66    }
67
68    /// Lock the revision in *main thread*.
69    #[must_use]
70    pub fn lock_estimated(&mut self) -> RevisionLock {
71        let estimated = self.estimated;
72        *self.locked.entry(estimated).or_default() += 1;
73        RevisionLock {
74            estimated,
75            used: OnceLock::new(),
76        }
77    }
78
79    /// Find the last revision slot by revision number.
80    pub fn find_revision(
81        &mut self,
82        revision: NonZeroUsize,
83        f: impl FnOnce(Option<&Arc<RevisionSlot<T>>>) -> T,
84    ) -> Arc<RevisionSlot<T>> {
85        let slot_base = self
86            .slots
87            .iter()
88            .filter(|slot| slot.revision <= revision.get())
89            .reduce(|x, y| if x.revision > y.revision { x } else { y });
90
91        if let Some(slot) = slot_base
92            && slot.revision == revision.get()
93        {
94            return slot.clone();
95        }
96
97        let slot = Arc::new(RevisionSlot {
98            revision: revision.get(),
99            data: f(slot_base),
100        });
101        self.slots.push(slot.clone());
102        self.estimated = revision.get().max(self.estimated);
103        slot
104    }
105
106    pub fn unlock(&mut self, rev: &mut RevisionLock) -> Option<usize> {
107        let rev = rev.estimated;
108        let revision_cnt = self
109            .locked
110            .entry(rev)
111            .or_insert_with(|| panic!("revision {rev} is not locked"));
112        *revision_cnt -= 1;
113        if *revision_cnt != 0 {
114            return None;
115        }
116
117        self.locked.remove(&rev);
118        let existing = self.locked.keys().min().copied();
119        existing.or_else(||
120            // if there is no locked revision, we only keep the latest revision
121            self.slots
122                .iter()
123                .map(|slot| slot.revision)
124                .max())
125    }
126}
127
128pub trait RevisionManagerLike {
129    fn gc(&mut self, min_rev: usize);
130}
131
132impl<T> RevisionManagerLike for RevisionManager<T> {
133    fn gc(&mut self, min_rev: usize) {
134        self.slots.retain(|r| r.revision >= min_rev);
135    }
136}
tinymist_query/adt/revision.rs

tinymist_query/adt/
revision.rs
