Line data Source code
1 : use std::ops::Range;
2 :
3 : // NOTE the `im` crate has 20x more downloads and also has
4 : // persistent/immutable BTree. But it's bugged so rpds is a
5 : // better choice <https://github.com/neondatabase/neon/issues/3395>
6 : use rpds::RedBlackTreeMapSync;
7 :
8 : /// Data structure that can efficiently:
9 : /// - find the latest layer by lsn.end at a given key
10 : /// - iterate the latest layers in a key range
11 : /// - insert layers in non-decreasing lsn.start order
12 : ///
13 : /// For a detailed explanation and justification of this approach, see:
14 : /// <https://neon.tech/blog/persistent-structures-in-neons-wal-indexing>
15 : ///
16 : /// NOTE The struct is parameterized over Value for easier
17 : /// testing, but in practice it's some sort of layer.
18 : pub struct LayerCoverage<Value> {
19 : /// For every change in coverage (as we sweep the key space)
20 : /// we store (lsn.end, value).
21 : ///
22 : /// NOTE We use an immutable/persistent tree so that we can keep historic
23 : /// versions of this coverage without cloning the whole thing and
24 : /// incurring quadratic memory cost. See HistoricLayerCoverage.
25 : ///
26 : /// NOTE We use the Sync version of the map because we want Self to
27 : /// be Sync. Using nonsync might be faster, if we can work with
28 : /// that.
29 : nodes: RedBlackTreeMapSync<i128, Option<(u64, Value)>>,
30 : }
31 :
32 : impl<T: Clone> Default for LayerCoverage<T> {
33 11112 : fn default() -> Self {
34 11112 : Self::new()
35 11112 : }
36 : }
37 :
38 : impl<Value: Clone> LayerCoverage<Value> {
39 11112 : pub fn new() -> Self {
40 11112 : Self {
41 11112 : nodes: RedBlackTreeMapSync::default(),
42 11112 : }
43 11112 : }
44 :
45 : /// Helper function to subdivide the key range without changing any values
46 : ///
47 : /// This operation has no semantic effect by itself. It only helps us pin in
48 : /// place the part of the coverage we don't want to change when inserting.
49 : ///
50 : /// As an analogy, think of a polygon. If you add a vertex along one of the
51 : /// segments, the polygon is still the same, but it behaves differently when
52 : /// we move or delete one of the other points.
53 : ///
54 : /// Complexity: O(log N)
55 62760 : fn add_node(&mut self, key: i128) {
56 62760 : let value = match self.nodes.range(..=key).next_back() {
57 41388 : Some((_, Some(v))) => Some(v.clone()),
58 17304 : Some((_, None)) => None,
59 4068 : None => None,
60 : };
61 62760 : self.nodes.insert_mut(key, value);
62 62760 : }
63 :
64 : /// Insert a layer.
65 : ///
66 : /// Complexity: worst case O(N), in practice O(log N). See NOTE in implementation.
67 31380 : pub fn insert(&mut self, key: Range<i128>, lsn: Range<u64>, value: Value) {
68 31380 : // Add nodes at endpoints
69 31380 : //
70 31380 : // NOTE The order of lines is important. We add nodes at the start
71 31380 : // and end of the key range **before updating any nodes** in order
72 31380 : // to pin down the current coverage outside of the relevant key range.
73 31380 : // Only the coverage inside the layer's key range should change.
74 31380 : self.add_node(key.start);
75 31380 : self.add_node(key.end);
76 31380 :
77 31380 : // Raise the height where necessary
78 31380 : //
79 31380 : // NOTE This loop is worst case O(N), but amortized O(log N) in the special
80 31380 : // case when rectangles have no height. In practice I don't think we'll see
81 31380 : // the kind of layer intersections needed to trigger O(N) behavior. The worst
82 31380 : // case is N/2 horizontal layers overlapped with N/2 vertical layers in a
83 31380 : // grid pattern.
84 31380 : let mut to_update = Vec::new();
85 31380 : let mut to_remove = Vec::new();
86 31380 : let mut prev_covered = false;
87 41964 : for (k, node) in self.nodes.range(key) {
88 41964 : let needs_cover = match node {
89 7440 : None => true,
90 34524 : Some((h, _)) => h < &lsn.end,
91 : };
92 41964 : if needs_cover {
93 41880 : match prev_covered {
94 10464 : true => to_remove.push(*k),
95 31416 : false => to_update.push(*k),
96 : }
97 84 : }
98 41964 : prev_covered = needs_cover;
99 : }
100 : // TODO check if the nodes inserted at key.start and key.end are safe
101 : // to remove. It's fine to keep them but they could be redundant.
102 62796 : for k in to_update {
103 31416 : self.nodes.insert_mut(k, Some((lsn.end, value.clone())));
104 31416 : }
105 41844 : for k in to_remove {
106 10464 : self.nodes.remove_mut(&k);
107 10464 : }
108 31380 : }
109 :
110 : /// Get the latest (by lsn.end) layer at a given key
111 : ///
112 : /// Complexity: O(log N)
113 11495844 : pub fn query(&self, key: i128) -> Option<Value> {
114 11495844 : self.nodes
115 11495844 : .range(..=key)
116 11495844 : .next_back()?
117 : .1
118 8880394 : .as_ref()
119 8880394 : .map(|(_, v)| v.clone())
120 11495844 : }
121 :
122 : /// Iterate the changes in layer coverage in a given range. You will likely
123 : /// want to start with self.query(key.start), and then follow up with self.range
124 : ///
125 : /// Complexity: O(log N + result_size)
126 9767616 : pub fn range(&self, key: Range<i128>) -> impl '_ + Iterator<Item = (i128, Option<Value>)> {
127 9767616 : self.nodes
128 9767616 : .range(key)
129 9767616 : .map(|(k, v)| (*k, v.as_ref().map(|x| x.1.clone())))
130 9767616 : }
131 :
132 : /// Returns an iterator which includes all coverage changes for layers that intersect
133 : /// with the provided range.
134 9767436 : pub fn range_overlaps(
135 9767436 : &self,
136 9767436 : key_range: &Range<i128>,
137 9767436 : ) -> impl Iterator<Item = (i128, Option<Value>)> + '_
138 9767436 : where
139 9767436 : Value: Eq,
140 9767436 : {
141 9767436 : let first_change = self.query(key_range.start);
142 9767436 : match first_change {
143 7323286 : Some(change) => {
144 7323286 : // If the start of the range is covered, we have to deal with two cases:
145 7323286 : // 1. Start of the range is aligned with the start of a layer.
146 7323286 : // In this case the return of `self.range` will contain the layer which aligns with the start of the key range.
147 7323286 : // We advance said iterator to avoid duplicating the first change.
148 7323286 : // 2. Start of the range is not aligned with the start of a layer.
149 7323286 : let range = key_range.start..key_range.end;
150 7323286 : let mut range_coverage = self.range(range).peekable();
151 7323286 : if range_coverage
152 7323286 : .peek()
153 7323286 : .is_some_and(|c| c.1.as_ref() == Some(&change))
154 42192 : {
155 42192 : range_coverage.next();
156 7281094 : }
157 7323286 : itertools::Either::Left(
158 7323286 : std::iter::once((key_range.start, Some(change))).chain(range_coverage),
159 7323286 : )
160 : }
161 : None => {
162 2444150 : let range = key_range.start..key_range.end;
163 2444150 : let coverage = self.range(range);
164 2444150 : itertools::Either::Right(coverage)
165 : }
166 : }
167 9767436 : }
168 : /// O(1) clone
169 75024 : pub fn clone(&self) -> Self {
170 75024 : Self {
171 75024 : nodes: self.nodes.clone(),
172 75024 : }
173 75024 : }
174 : }
175 :
176 : /// Image and delta coverage at a specific LSN.
177 : pub struct LayerCoverageTuple<Value> {
178 : pub image_coverage: LayerCoverage<Value>,
179 : pub delta_coverage: LayerCoverage<Value>,
180 : }
181 :
182 : impl<T: Clone> Default for LayerCoverageTuple<T> {
183 5556 : fn default() -> Self {
184 5556 : Self {
185 5556 : image_coverage: LayerCoverage::default(),
186 5556 : delta_coverage: LayerCoverage::default(),
187 5556 : }
188 5556 : }
189 : }
190 :
191 : impl<Value: Clone> LayerCoverageTuple<Value> {
192 37512 : pub fn clone(&self) -> Self {
193 37512 : Self {
194 37512 : image_coverage: self.image_coverage.clone(),
195 37512 : delta_coverage: self.delta_coverage.clone(),
196 37512 : }
197 37512 : }
198 : }
|
