TLA Line data Source code
1 : use crate::repository::{key_range_size, singleton_range, Key};
2 : use postgres_ffi::BLCKSZ;
3 : use std::ops::Range;
4 :
5 : ///
6 : /// Represents a set of Keys, in a compact form.
7 : ///
8 CBC 23863 : #[derive(Clone, Debug, Default)]
9 : pub struct KeySpace {
10 : /// Contiguous ranges of keys that belong to the key space. In key order,
11 : /// and with no overlap.
12 : pub ranges: Vec<Range<Key>>,
13 : }
14 :
15 : impl KeySpace {
16 : ///
17 : /// Partition a key space into roughly chunks of roughly 'target_size' bytes
18 : /// in each partition.
19 : ///
20 662 : pub fn partition(&self, target_size: u64) -> KeyPartitioning {
21 662 : // Assume that each value is 8k in size.
22 662 : let target_nblocks = (target_size / BLCKSZ as u64) as usize;
23 662 :
24 662 : let mut parts = Vec::new();
25 662 : let mut current_part = Vec::new();
26 662 : let mut current_part_size: usize = 0;
27 966151 : for range in &self.ranges {
28 : // If appending the next contiguous range in the keyspace to the current
29 : // partition would cause it to be too large, start a new partition.
30 965489 : let this_size = key_range_size(range) as usize;
31 965489 : if current_part_size + this_size > target_nblocks && !current_part.is_empty() {
32 13753 : parts.push(KeySpace {
33 13753 : ranges: current_part,
34 13753 : });
35 13753 : current_part = Vec::new();
36 13753 : current_part_size = 0;
37 951736 : }
38 :
39 : // If the next range is larger than 'target_size', split it into
40 : // 'target_size' chunks.
41 965489 : let mut remain_size = this_size;
42 965489 : let mut start = range.start;
43 970152 : while remain_size > target_nblocks {
44 4663 : let next = start.add(target_nblocks as u32);
45 4663 : parts.push(KeySpace {
46 4663 : ranges: vec![start..next],
47 4663 : });
48 4663 : start = next;
49 4663 : remain_size -= target_nblocks
50 : }
51 965489 : current_part.push(start..range.end);
52 965489 : current_part_size += remain_size;
53 : }
54 :
55 : // add last partition that wasn't full yet.
56 662 : if !current_part.is_empty() {
57 662 : parts.push(KeySpace {
58 662 : ranges: current_part,
59 662 : });
60 662 : }
61 :
62 662 : KeyPartitioning { parts }
63 662 : }
64 :
65 : ///
66 : /// Check if key space contains overlapping range
67 : ///
68 3666 : pub fn overlaps(&self, range: &Range<Key>) -> bool {
69 3666 : match self.ranges.binary_search_by_key(&range.end, |r| r.start) {
70 1 : Ok(0) => false,
71 3655 : Err(0) => false,
72 2 : Ok(index) => self.ranges[index - 1].end > range.start,
73 8 : Err(index) => self.ranges[index - 1].end > range.start,
74 : }
75 3666 : }
76 : }
77 :
78 : ///
79 : /// Represents a partitioning of the key space.
80 : ///
81 : /// The only kind of partitioning we do is to partition the key space into
82 : /// partitions that are roughly equal in physical size (see KeySpace::partition).
83 : /// But this data structure could represent any partitioning.
84 : ///
85 1301 : #[derive(Clone, Debug, Default)]
86 : pub struct KeyPartitioning {
87 : pub parts: Vec<KeySpace>,
88 : }
89 :
90 : impl KeyPartitioning {
91 1302 : pub fn new() -> Self {
92 1302 : KeyPartitioning { parts: Vec::new() }
93 1302 : }
94 : }
95 :
96 : ///
97 : /// A helper object, to collect a set of keys and key ranges into a KeySpace
98 : /// object. This takes care of merging adjacent keys and key ranges into
99 : /// contiguous ranges.
100 : ///
101 UBC 0 : #[derive(Clone, Debug, Default)]
102 : pub struct KeySpaceAccum {
103 : accum: Option<Range<Key>>,
104 :
105 : ranges: Vec<Range<Key>>,
106 : }
107 :
108 : impl KeySpaceAccum {
109 CBC 672 : pub fn new() -> Self {
110 672 : Self {
111 672 : accum: None,
112 672 : ranges: Vec::new(),
113 672 : }
114 672 : }
115 :
116 1067974 : pub fn add_key(&mut self, key: Key) {
117 1067974 : self.add_range(singleton_range(key))
118 1067974 : }
119 :
120 1623917 : pub fn add_range(&mut self, range: Range<Key>) {
121 1623917 : match self.accum.as_mut() {
122 1623245 : Some(accum) => {
123 1623245 : if range.start == accum.end {
124 648482 : accum.end = range.end;
125 648482 : } else {
126 974763 : assert!(range.start > accum.end);
127 974763 : self.ranges.push(accum.clone());
128 974763 : *accum = range;
129 : }
130 : }
131 672 : None => self.accum = Some(range),
132 : }
133 1623917 : }
134 :
135 : pub fn to_keyspace(mut self) -> KeySpace {
136 668 : if let Some(accum) = self.accum.take() {
137 668 : self.ranges.push(accum);
138 668 : }
139 668 : KeySpace {
140 668 : ranges: self.ranges,
141 668 : }
142 668 : }
143 : }
144 :
145 : ///
146 : /// A helper object, to collect a set of keys and key ranges into a KeySpace
147 : /// object. Key ranges may be inserted in any order and can overlap.
148 : ///
149 596 : #[derive(Clone, Debug, Default)]
150 : pub struct KeySpaceRandomAccum {
151 : ranges: Vec<Range<Key>>,
152 : }
153 :
154 : impl KeySpaceRandomAccum {
155 UBC 0 : pub fn new() -> Self {
156 0 : Self { ranges: Vec::new() }
157 0 : }
158 :
159 0 : pub fn add_key(&mut self, key: Key) {
160 0 : self.add_range(singleton_range(key))
161 0 : }
162 :
163 CBC 21 : pub fn add_range(&mut self, range: Range<Key>) {
164 21 : self.ranges.push(range);
165 21 : }
166 :
167 596 : pub fn to_keyspace(mut self) -> KeySpace {
168 596 : let mut ranges = Vec::new();
169 596 : if !self.ranges.is_empty() {
170 30 : self.ranges.sort_by_key(|r| r.start);
171 9 : let mut start = self.ranges.first().unwrap().start;
172 9 : let mut end = self.ranges.first().unwrap().end;
173 30 : for r in self.ranges {
174 21 : assert!(r.start >= start);
175 21 : if r.start > end {
176 2 : ranges.push(start..end);
177 2 : start = r.start;
178 2 : end = r.end;
179 19 : } else if r.end > end {
180 6 : end = r.end;
181 13 : }
182 : }
183 9 : ranges.push(start..end);
184 587 : }
185 596 : KeySpace { ranges }
186 596 : }
187 : }
188 :
189 : #[cfg(test)]
190 : mod tests {
191 : use super::*;
192 : use std::fmt::Write;
193 :
194 : // Helper function to create a key range.
195 : //
196 : // Make the tests below less verbose.
197 43 : fn kr(irange: Range<i128>) -> Range<Key> {
198 43 : Key::from_i128(irange.start)..Key::from_i128(irange.end)
199 43 : }
200 :
201 : #[allow(dead_code)]
202 UBC 0 : fn dump_keyspace(ks: &KeySpace) {
203 0 : for r in ks.ranges.iter() {
204 0 : println!(" {}..{}", r.start.to_i128(), r.end.to_i128());
205 0 : }
206 0 : }
207 :
208 CBC 8 : fn assert_ks_eq(actual: &KeySpace, expected: Vec<Range<Key>>) {
209 8 : if actual.ranges != expected {
210 UBC 0 : let mut msg = String::new();
211 0 :
212 0 : writeln!(msg, "expected:").unwrap();
213 0 : for r in &expected {
214 0 : writeln!(msg, " {}..{}", r.start.to_i128(), r.end.to_i128()).unwrap();
215 0 : }
216 0 : writeln!(msg, "got:").unwrap();
217 0 : for r in &actual.ranges {
218 0 : writeln!(msg, " {}..{}", r.start.to_i128(), r.end.to_i128()).unwrap();
219 0 : }
220 0 : panic!("{}", msg);
221 CBC 8 : }
222 8 : }
223 :
224 1 : #[test]
225 1 : fn keyspace_add_range() {
226 1 : // two separate ranges
227 1 : //
228 1 : // #####
229 1 : // #####
230 1 : let mut ks = KeySpaceRandomAccum::default();
231 1 : ks.add_range(kr(0..10));
232 1 : ks.add_range(kr(20..30));
233 1 : assert_ks_eq(&ks.to_keyspace(), vec![kr(0..10), kr(20..30)]);
234 1 :
235 1 : // two separate ranges, added in reverse order
236 1 : //
237 1 : // #####
238 1 : // #####
239 1 : let mut ks = KeySpaceRandomAccum::default();
240 1 : ks.add_range(kr(20..30));
241 1 : ks.add_range(kr(0..10));
242 1 :
243 1 : // add range that is adjacent to the end of an existing range
244 1 : //
245 1 : // #####
246 1 : // #####
247 1 : ks.add_range(kr(0..10));
248 1 : ks.add_range(kr(10..30));
249 1 : assert_ks_eq(&ks.to_keyspace(), vec![kr(0..30)]);
250 1 :
251 1 : // add range that is adjacent to the start of an existing range
252 1 : //
253 1 : // #####
254 1 : // #####
255 1 : let mut ks = KeySpaceRandomAccum::default();
256 1 : ks.add_range(kr(10..30));
257 1 : ks.add_range(kr(0..10));
258 1 : assert_ks_eq(&ks.to_keyspace(), vec![kr(0..30)]);
259 1 :
260 1 : // add range that overlaps with the end of an existing range
261 1 : //
262 1 : // #####
263 1 : // #####
264 1 : let mut ks = KeySpaceRandomAccum::default();
265 1 : ks.add_range(kr(0..10));
266 1 : ks.add_range(kr(5..30));
267 1 : assert_ks_eq(&ks.to_keyspace(), vec![kr(0..30)]);
268 1 :
269 1 : // add range that overlaps with the start of an existing range
270 1 : //
271 1 : // #####
272 1 : // #####
273 1 : let mut ks = KeySpaceRandomAccum::default();
274 1 : ks.add_range(kr(5..30));
275 1 : ks.add_range(kr(0..10));
276 1 : assert_ks_eq(&ks.to_keyspace(), vec![kr(0..30)]);
277 1 :
278 1 : // add range that is fully covered by an existing range
279 1 : //
280 1 : // #########
281 1 : // #####
282 1 : let mut ks = KeySpaceRandomAccum::default();
283 1 : ks.add_range(kr(0..30));
284 1 : ks.add_range(kr(10..20));
285 1 : assert_ks_eq(&ks.to_keyspace(), vec![kr(0..30)]);
286 1 :
287 1 : // add range that extends an existing range from both ends
288 1 : //
289 1 : // #####
290 1 : // #########
291 1 : let mut ks = KeySpaceRandomAccum::default();
292 1 : ks.add_range(kr(10..20));
293 1 : ks.add_range(kr(0..30));
294 1 : assert_ks_eq(&ks.to_keyspace(), vec![kr(0..30)]);
295 1 :
296 1 : // add a range that overlaps with two existing ranges, joining them
297 1 : //
298 1 : // ##### #####
299 1 : // #######
300 1 : let mut ks = KeySpaceRandomAccum::default();
301 1 : ks.add_range(kr(0..10));
302 1 : ks.add_range(kr(20..30));
303 1 : ks.add_range(kr(5..25));
304 1 : assert_ks_eq(&ks.to_keyspace(), vec![kr(0..30)]);
305 1 : }
306 :
307 1 : #[test]
308 1 : fn keyspace_overlaps() {
309 1 : let mut ks = KeySpaceRandomAccum::default();
310 1 : ks.add_range(kr(10..20));
311 1 : ks.add_range(kr(30..40));
312 1 : let ks = ks.to_keyspace();
313 1 :
314 1 : // ##### #####
315 1 : // xxxx
316 1 : assert!(!ks.overlaps(&kr(0..5)));
317 :
318 : // ##### #####
319 : // xxxx
320 1 : assert!(!ks.overlaps(&kr(5..9)));
321 :
322 : // ##### #####
323 : // xxxx
324 1 : assert!(!ks.overlaps(&kr(5..10)));
325 :
326 : // ##### #####
327 : // xxxx
328 1 : assert!(ks.overlaps(&kr(5..11)));
329 :
330 : // ##### #####
331 : // xxxx
332 1 : assert!(ks.overlaps(&kr(10..15)));
333 :
334 : // ##### #####
335 : // xxxx
336 1 : assert!(ks.overlaps(&kr(15..20)));
337 :
338 : // ##### #####
339 : // xxxx
340 1 : assert!(ks.overlaps(&kr(15..25)));
341 :
342 : // ##### #####
343 : // xxxx
344 1 : assert!(!ks.overlaps(&kr(22..28)));
345 :
346 : // ##### #####
347 : // xxxx
348 1 : assert!(!ks.overlaps(&kr(25..30)));
349 :
350 : // ##### #####
351 : // xxxx
352 1 : assert!(ks.overlaps(&kr(35..35)));
353 :
354 : // ##### #####
355 : // xxxx
356 1 : assert!(!ks.overlaps(&kr(40..45)));
357 :
358 : // ##### #####
359 : // xxxx
360 1 : assert!(!ks.overlaps(&kr(45..50)));
361 :
362 : // ##### #####
363 : // xxxxxxxxxxx
364 1 : assert!(ks.overlaps(&kr(0..30))); // XXXXX This fails currently!
365 1 : }
366 : }
|
