risingwave_hummock_sdk/
key.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::borrow::Borrow;
16use std::cmp::Ordering;
17use std::fmt::Debug;
18use std::iter::once;
19use std::ops::Bound::*;
20use std::ops::{Bound, Deref, DerefMut, RangeBounds};
21use std::ptr;
22
23use bytes::{Buf, BufMut, Bytes, BytesMut};
24use risingwave_common::catalog::TableId;
25use risingwave_common::hash::VirtualNode;
26use risingwave_common_estimate_size::EstimateSize;
27
28use crate::{EpochWithGap, HummockEpoch};
29
30pub const EPOCH_LEN: usize = std::mem::size_of::<HummockEpoch>();
31pub const TABLE_PREFIX_LEN: usize = std::mem::size_of::<u32>();
32// Max length for key overlap and diff length. See KeyPrefix::encode.
33pub const MAX_KEY_LEN: usize = u16::MAX as usize;
34
35pub type KeyPayloadType = Bytes;
36pub type TableKeyRange = (
37    Bound<TableKey<KeyPayloadType>>,
38    Bound<TableKey<KeyPayloadType>>,
39);
40pub type UserKeyRange = (
41    Bound<UserKey<KeyPayloadType>>,
42    Bound<UserKey<KeyPayloadType>>,
43);
44pub type UserKeyRangeRef<'a> = (Bound<UserKey<&'a [u8]>>, Bound<UserKey<&'a [u8]>>);
45pub type FullKeyRange = (
46    Bound<FullKey<KeyPayloadType>>,
47    Bound<FullKey<KeyPayloadType>>,
48);
49
50pub fn is_empty_key_range(key_range: &TableKeyRange) -> bool {
51    match key_range {
52        (Included(start), Excluded(end)) => start == end,
53        _ => false,
54    }
55}
56
57/// Returns left inclusive and right exclusive vnode index of the given range.
58///
59/// # Vnode count unawareness
60///
61/// Note that this function is not aware of the vnode count that is actually used in this table.
62/// For example, if the total vnode count is 256, `Unbounded` can be a correct end bound for vnode 255,
63/// but this function will still return `Excluded(256)`.
64///
65/// See also [`vnode`] and [`end_bound_of_vnode`] which hold such invariant.
66pub fn vnode_range<T: AsRef<[u8]>>(range: &impl RangeBounds<TableKey<T>>) -> (usize, usize) {
67    let left = match range.start_bound() {
68        Included(key) | Excluded(key) => key.vnode_part().to_index(),
69        Unbounded => 0,
70    };
71    let right = match range.end_bound() {
72        Included(key) => key.vnode_part().to_index() + 1,
73        Excluded(key) => {
74            let (vnode, inner_key) = key.split_vnode();
75            if inner_key.is_empty() {
76                // When the exclusive end key range contains only a vnode,
77                // the whole vnode is excluded.
78                vnode.to_index()
79            } else {
80                vnode.to_index() + 1
81            }
82        }
83        Unbounded => VirtualNode::MAX_REPRESENTABLE.to_index() + 1,
84    };
85    (left, right)
86}
87
88/// Ensure there is only one vnode involved in table key range and return the vnode.
89///
90/// # Vnode count unawareness
91///
92/// Note that this function is not aware of the vnode count that is actually used in this table.
93/// For example, if the total vnode count is 256, `Unbounded` can be a correct end bound for vnode 255,
94/// but this function will still require `Excluded(256)`.
95///
96/// See also [`vnode_range`] and [`end_bound_of_vnode`] which hold such invariant.
97pub fn vnode(range: &TableKeyRange) -> VirtualNode {
98    let (l, r_exclusive) = vnode_range(range);
99    assert_eq!(r_exclusive - l, 1);
100    VirtualNode::from_index(l)
101}
102
103/// Converts user key to full key by appending `epoch` to the user key.
104pub fn key_with_epoch(mut user_key: Vec<u8>, epoch: HummockEpoch) -> Vec<u8> {
105    let res = epoch.to_be();
106    user_key.reserve(EPOCH_LEN);
107    let buf = user_key.chunk_mut();
108
109    // TODO: check whether this hack improves performance
110    unsafe {
111        ptr::copy_nonoverlapping(
112            &res as *const _ as *const u8,
113            buf.as_mut_ptr() as *mut _,
114            EPOCH_LEN,
115        );
116        user_key.advance_mut(EPOCH_LEN);
117    }
118
119    user_key
120}
121
122/// Splits a full key into its user key part and epoch part.
123#[inline]
124pub fn split_key_epoch(full_key: &[u8]) -> (&[u8], &[u8]) {
125    let pos = full_key
126        .len()
127        .checked_sub(EPOCH_LEN)
128        .unwrap_or_else(|| panic!("bad full key format: {:?}", full_key));
129    full_key.split_at(pos)
130}
131
132/// Extract encoded [`UserKey`] from encoded [`FullKey`] without epoch part
133pub fn user_key(full_key: &[u8]) -> &[u8] {
134    split_key_epoch(full_key).0
135}
136
137/// Extract table key from encoded [`UserKey`] without table id part
138pub fn table_key(user_key: &[u8]) -> &[u8] {
139    &user_key[TABLE_PREFIX_LEN..]
140}
141
142#[inline(always)]
143/// Extract encoded [`UserKey`] from encoded [`FullKey`] but allow empty slice
144pub fn get_user_key(full_key: &[u8]) -> Vec<u8> {
145    if full_key.is_empty() {
146        vec![]
147    } else {
148        user_key(full_key).to_vec()
149    }
150}
151
152/// Extract table id from encoded [`FullKey`]
153#[inline(always)]
154pub fn get_table_id(full_key: &[u8]) -> u32 {
155    let mut buf = full_key;
156    buf.get_u32()
157}
158
159// Copyright 2016 TiKV Project Authors. Licensed under Apache-2.0.
160
161/// Computes the next key of the given key.
162///
163/// If the key has no successor key (e.g. the input is "\xff\xff"), the result
164/// would be an empty vector.
165///
166/// # Examples
167///
168/// ```rust
169/// use risingwave_hummock_sdk::key::next_key;
170/// assert_eq!(next_key(b"123"), b"124");
171/// assert_eq!(next_key(b"12\xff"), b"13");
172/// assert_eq!(next_key(b"\xff\xff"), b"");
173/// assert_eq!(next_key(b"\xff\xfe"), b"\xff\xff");
174/// assert_eq!(next_key(b"T"), b"U");
175/// assert_eq!(next_key(b""), b"");
176/// ```
177pub fn next_key(key: &[u8]) -> Vec<u8> {
178    if let Some((s, e)) = next_key_no_alloc(key) {
179        let mut res = Vec::with_capacity(s.len() + 1);
180        res.extend_from_slice(s);
181        res.push(e);
182        res
183    } else {
184        Vec::new()
185    }
186}
187
188/// Computes the previous key of the given key.
189///
190/// If the key has no predecessor key (e.g. the input is "\x00\x00"), the result
191/// would be a "\xff\xff" vector.
192///
193/// # Examples
194///
195/// ```rust
196/// use risingwave_hummock_sdk::key::prev_key;
197/// assert_eq!(prev_key(b"123"), b"122");
198/// assert_eq!(prev_key(b"12\x00"), b"11\xff");
199/// assert_eq!(prev_key(b"\x00\x00"), b"\xff\xff");
200/// assert_eq!(prev_key(b"\x00\x01"), b"\x00\x00");
201/// assert_eq!(prev_key(b"T"), b"S");
202/// assert_eq!(prev_key(b""), b"");
203/// ```
204pub fn prev_key(key: &[u8]) -> Vec<u8> {
205    let pos = key.iter().rposition(|b| *b != 0x00);
206    match pos {
207        Some(pos) => {
208            let mut res = Vec::with_capacity(key.len());
209            res.extend_from_slice(&key[0..pos]);
210            res.push(key[pos] - 1);
211            if pos + 1 < key.len() {
212                res.push(b"\xff".to_owned()[0]);
213            }
214            res
215        }
216        None => {
217            vec![0xff; key.len()]
218        }
219    }
220}
221
222fn next_key_no_alloc(key: &[u8]) -> Option<(&[u8], u8)> {
223    let pos = key.iter().rposition(|b| *b != 0xff)?;
224    Some((&key[..pos], key[pos] + 1))
225}
226
227// End Copyright 2016 TiKV Project Authors. Licensed under Apache-2.0.
228
229/// compute the next epoch, and don't change the bytes of the u8 slice.
230/// # Examples
231///
232/// ```rust
233/// use risingwave_hummock_sdk::key::next_epoch;
234/// assert_eq!(next_epoch(b"123"), b"124");
235/// assert_eq!(next_epoch(b"\xff\x00\xff"), b"\xff\x01\x00");
236/// assert_eq!(next_epoch(b"\xff\xff"), b"\x00\x00");
237/// assert_eq!(next_epoch(b"\x00\x00"), b"\x00\x01");
238/// assert_eq!(next_epoch(b"S"), b"T");
239/// assert_eq!(next_epoch(b""), b"");
240/// ```
241pub fn next_epoch(epoch: &[u8]) -> Vec<u8> {
242    let pos = epoch.iter().rposition(|b| *b != 0xff);
243    match pos {
244        Some(mut pos) => {
245            let mut res = Vec::with_capacity(epoch.len());
246            res.extend_from_slice(&epoch[0..pos]);
247            res.push(epoch[pos] + 1);
248            while pos + 1 < epoch.len() {
249                res.push(0x00);
250                pos += 1;
251            }
252            res
253        }
254        None => {
255            vec![0x00; epoch.len()]
256        }
257    }
258}
259
260/// compute the prev epoch, and don't change the bytes of the u8 slice.
261/// # Examples
262///
263/// ```rust
264/// use risingwave_hummock_sdk::key::prev_epoch;
265/// assert_eq!(prev_epoch(b"124"), b"123");
266/// assert_eq!(prev_epoch(b"\xff\x01\x00"), b"\xff\x00\xff");
267/// assert_eq!(prev_epoch(b"\x00\x00"), b"\xff\xff");
268/// assert_eq!(prev_epoch(b"\x00\x01"), b"\x00\x00");
269/// assert_eq!(prev_epoch(b"T"), b"S");
270/// assert_eq!(prev_epoch(b""), b"");
271/// ```
272pub fn prev_epoch(epoch: &[u8]) -> Vec<u8> {
273    let pos = epoch.iter().rposition(|b| *b != 0x00);
274    match pos {
275        Some(mut pos) => {
276            let mut res = Vec::with_capacity(epoch.len());
277            res.extend_from_slice(&epoch[0..pos]);
278            res.push(epoch[pos] - 1);
279            while pos + 1 < epoch.len() {
280                res.push(0xff);
281                pos += 1;
282            }
283            res
284        }
285        None => {
286            vec![0xff; epoch.len()]
287        }
288    }
289}
290
291/// compute the next full key of the given full key
292///
293/// if the `user_key` has no successor key, the result will be a empty vec
294pub fn next_full_key(full_key: &[u8]) -> Vec<u8> {
295    let (user_key, epoch) = split_key_epoch(full_key);
296    let prev_epoch = prev_epoch(epoch);
297    let mut res = Vec::with_capacity(full_key.len());
298    if prev_epoch.cmp(&vec![0xff; prev_epoch.len()]) == Ordering::Equal {
299        let next_user_key = next_key(user_key);
300        if next_user_key.is_empty() {
301            return Vec::new();
302        }
303        res.extend_from_slice(next_user_key.as_slice());
304        res.extend_from_slice(prev_epoch.as_slice());
305        res
306    } else {
307        res.extend_from_slice(user_key);
308        res.extend_from_slice(prev_epoch.as_slice());
309        res
310    }
311}
312
313/// compute the prev full key of the given full key
314///
315/// if the `user_key` has no predecessor key, the result will be a empty vec
316pub fn prev_full_key(full_key: &[u8]) -> Vec<u8> {
317    let (user_key, epoch) = split_key_epoch(full_key);
318    let next_epoch = next_epoch(epoch);
319    let mut res = Vec::with_capacity(full_key.len());
320    if next_epoch.cmp(&vec![0x00; next_epoch.len()]) == Ordering::Equal {
321        let prev_user_key = prev_key(user_key);
322        if prev_user_key.cmp(&vec![0xff; prev_user_key.len()]) == Ordering::Equal {
323            return Vec::new();
324        }
325        res.extend_from_slice(prev_user_key.as_slice());
326        res.extend_from_slice(next_epoch.as_slice());
327        res
328    } else {
329        res.extend_from_slice(user_key);
330        res.extend_from_slice(next_epoch.as_slice());
331        res
332    }
333}
334
335/// [`Unbounded`] if the vnode is the maximum representable value (i.e. [`VirtualNode::MAX_REPRESENTABLE`]),
336/// otherwise [`Excluded`] the next vnode.
337///
338/// Note that this function is not aware of the vnode count that is actually used in this table.
339/// For example, if the total vnode count is 256, `Unbounded` can be a correct end bound for vnode 255,
340/// but this function will still return `Excluded(256)`. See also [`vnode`] and [`vnode_range`] which
341/// rely on such invariant.
342pub fn end_bound_of_vnode(vnode: VirtualNode) -> Bound<Bytes> {
343    if vnode == VirtualNode::MAX_REPRESENTABLE {
344        Unbounded
345    } else {
346        let end_bound_index = vnode.to_index() + 1;
347        Excluded(Bytes::copy_from_slice(
348            &VirtualNode::from_index(end_bound_index).to_be_bytes(),
349        ))
350    }
351}
352
353/// Get the end bound of the given `prefix` when transforming it to a key range.
354pub fn end_bound_of_prefix(prefix: &[u8]) -> Bound<Bytes> {
355    if let Some((s, e)) = next_key_no_alloc(prefix) {
356        let mut buf = BytesMut::with_capacity(s.len() + 1);
357        buf.extend_from_slice(s);
358        buf.put_u8(e);
359        Excluded(buf.freeze())
360    } else {
361        Unbounded
362    }
363}
364
365/// Get the start bound of the given `prefix` when it is excluded from the range.
366pub fn start_bound_of_excluded_prefix(prefix: &[u8]) -> Bound<Bytes> {
367    if let Some((s, e)) = next_key_no_alloc(prefix) {
368        let mut buf = BytesMut::with_capacity(s.len() + 1);
369        buf.extend_from_slice(s);
370        buf.put_u8(e);
371        Included(buf.freeze())
372    } else {
373        panic!("the prefix is the maximum value")
374    }
375}
376
377/// Transform the given `prefix` to a key range.
378pub fn range_of_prefix(prefix: &[u8]) -> (Bound<Bytes>, Bound<Bytes>) {
379    if prefix.is_empty() {
380        (Unbounded, Unbounded)
381    } else {
382        (
383            Included(Bytes::copy_from_slice(prefix)),
384            end_bound_of_prefix(prefix),
385        )
386    }
387}
388
389pub fn prefix_slice_with_vnode(vnode: VirtualNode, slice: &[u8]) -> Bytes {
390    let prefix = vnode.to_be_bytes();
391    let mut buf = BytesMut::with_capacity(prefix.len() + slice.len());
392    buf.extend_from_slice(&prefix);
393    buf.extend_from_slice(slice);
394    buf.freeze()
395}
396
397/// Prepend the `prefix` to the given key `range`.
398pub fn prefixed_range_with_vnode<B: AsRef<[u8]>>(
399    range: impl RangeBounds<B>,
400    vnode: VirtualNode,
401) -> TableKeyRange {
402    let prefixed = |b: &B| -> Bytes { prefix_slice_with_vnode(vnode, b.as_ref()) };
403
404    let start: Bound<Bytes> = match range.start_bound() {
405        Included(b) => Included(prefixed(b)),
406        Excluded(b) => {
407            assert!(!b.as_ref().is_empty());
408            Excluded(prefixed(b))
409        }
410        Unbounded => Included(Bytes::copy_from_slice(&vnode.to_be_bytes())),
411    };
412
413    let end = match range.end_bound() {
414        Included(b) => Included(prefixed(b)),
415        Excluded(b) => {
416            assert!(!b.as_ref().is_empty());
417            Excluded(prefixed(b))
418        }
419        Unbounded => end_bound_of_vnode(vnode),
420    };
421
422    map_table_key_range((start, end))
423}
424
425pub trait SetSlice<S: AsRef<[u8]> + ?Sized> {
426    fn set(&mut self, value: &S);
427}
428
429impl<S: AsRef<[u8]> + ?Sized> SetSlice<S> for Vec<u8> {
430    fn set(&mut self, value: &S) {
431        self.clear();
432        self.extend_from_slice(value.as_ref());
433    }
434}
435
436impl SetSlice<Bytes> for Bytes {
437    fn set(&mut self, value: &Bytes) {
438        *self = value.clone()
439    }
440}
441
442pub trait CopyFromSlice: Send + 'static {
443    fn copy_from_slice(slice: &[u8]) -> Self;
444}
445
446impl CopyFromSlice for Vec<u8> {
447    fn copy_from_slice(slice: &[u8]) -> Self {
448        Vec::from(slice)
449    }
450}
451
452impl CopyFromSlice for Bytes {
453    fn copy_from_slice(slice: &[u8]) -> Self {
454        Bytes::copy_from_slice(slice)
455    }
456}
457
458impl CopyFromSlice for () {
459    fn copy_from_slice(_: &[u8]) -> Self {}
460}
461
462/// [`TableKey`] is an internal concept in storage. It's a wrapper around the key directly from the
463/// user, to make the code clearer and avoid confusion with encoded [`UserKey`] and [`FullKey`].
464///
465/// Its name come from the assumption that Hummock is always accessed by a table-like structure
466/// identified by a [`TableId`].
467#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
468pub struct TableKey<T: AsRef<[u8]>>(pub T);
469
470impl<T: AsRef<[u8]>> Debug for TableKey<T> {
471    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
472        write!(f, "TableKey {{ {} }}", hex::encode(self.0.as_ref()))
473    }
474}
475
476impl<T: AsRef<[u8]>> Deref for TableKey<T> {
477    type Target = T;
478
479    fn deref(&self) -> &Self::Target {
480        &self.0
481    }
482}
483
484impl<T: AsRef<[u8]>> DerefMut for TableKey<T> {
485    fn deref_mut(&mut self) -> &mut Self::Target {
486        &mut self.0
487    }
488}
489
490impl<T: AsRef<[u8]>> AsRef<[u8]> for TableKey<T> {
491    fn as_ref(&self) -> &[u8] {
492        self.0.as_ref()
493    }
494}
495
496impl TableKey<Bytes> {
497    pub fn split_vnode_bytes(&self) -> (VirtualNode, Bytes) {
498        debug_assert!(
499            self.0.len() >= VirtualNode::SIZE,
500            "too short table key: {:?}",
501            self.0.as_ref()
502        );
503        let (vnode, _) = self.0.split_first_chunk::<{ VirtualNode::SIZE }>().unwrap();
504        (
505            VirtualNode::from_be_bytes(*vnode),
506            self.0.slice(VirtualNode::SIZE..),
507        )
508    }
509}
510
511impl<T: AsRef<[u8]>> TableKey<T> {
512    pub fn split_vnode(&self) -> (VirtualNode, &[u8]) {
513        debug_assert!(
514            self.0.as_ref().len() >= VirtualNode::SIZE,
515            "too short table key: {:?}",
516            self.0.as_ref()
517        );
518        let (vnode, inner_key) = self
519            .0
520            .as_ref()
521            .split_first_chunk::<{ VirtualNode::SIZE }>()
522            .unwrap();
523        (VirtualNode::from_be_bytes(*vnode), inner_key)
524    }
525
526    pub fn vnode_part(&self) -> VirtualNode {
527        self.split_vnode().0
528    }
529
530    pub fn key_part(&self) -> &[u8] {
531        self.split_vnode().1
532    }
533
534    pub fn to_ref(&self) -> TableKey<&[u8]> {
535        TableKey(self.0.as_ref())
536    }
537}
538
539impl<T: AsRef<[u8]>> Borrow<[u8]> for TableKey<T> {
540    fn borrow(&self) -> &[u8] {
541        self.0.as_ref()
542    }
543}
544
545impl EstimateSize for TableKey<Bytes> {
546    fn estimated_heap_size(&self) -> usize {
547        self.0.estimated_heap_size()
548    }
549}
550
551impl TableKey<&[u8]> {
552    pub fn copy_into<T: CopyFromSlice + AsRef<[u8]>>(&self) -> TableKey<T> {
553        TableKey(T::copy_from_slice(self.as_ref()))
554    }
555}
556
557#[inline]
558pub fn map_table_key_range(range: (Bound<KeyPayloadType>, Bound<KeyPayloadType>)) -> TableKeyRange {
559    (range.0.map(TableKey), range.1.map(TableKey))
560}
561
562pub fn gen_key_from_bytes(vnode: VirtualNode, payload: &[u8]) -> TableKey<Bytes> {
563    TableKey(Bytes::from(
564        [vnode.to_be_bytes().as_slice(), payload].concat(),
565    ))
566}
567
568pub fn gen_key_from_str(vnode: VirtualNode, payload: &str) -> TableKey<Bytes> {
569    gen_key_from_bytes(vnode, payload.as_bytes())
570}
571
572/// [`UserKey`] is is an internal concept in storage. In the storage interface, user specifies
573/// `table_key` and `table_id` (in `ReadOptions` or `WriteOptions`) as the input. The storage
574/// will group these two values into one struct for convenient filtering.
575///
576/// The encoded format is | `table_id` | `table_key` |.
577#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
578pub struct UserKey<T: AsRef<[u8]>> {
579    // When comparing `UserKey`, we first compare `table_id`, then `table_key`. So the order of
580    // declaration matters.
581    pub table_id: TableId,
582    pub table_key: TableKey<T>,
583}
584
585impl<T: AsRef<[u8]>> Debug for UserKey<T> {
586    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
587        write!(f, "UserKey {{ {}, {:?} }}", self.table_id, self.table_key)
588    }
589}
590
591impl<T: AsRef<[u8]>> UserKey<T> {
592    pub fn new(table_id: TableId, table_key: TableKey<T>) -> Self {
593        Self {
594            table_id,
595            table_key,
596        }
597    }
598
599    /// Pass the inner type of `table_key` to make the code less verbose.
600    pub fn for_test(table_id: TableId, table_key: T) -> Self {
601        Self {
602            table_id,
603            table_key: TableKey(table_key),
604        }
605    }
606
607    /// Encode in to a buffer.
608    pub fn encode_into(&self, buf: &mut impl BufMut) {
609        buf.put_u32(self.table_id.as_raw_id());
610        buf.put_slice(self.table_key.as_ref());
611    }
612
613    pub fn encode_table_key_into(&self, buf: &mut impl BufMut) {
614        buf.put_slice(self.table_key.as_ref());
615    }
616
617    pub fn encode(&self) -> Vec<u8> {
618        let mut ret = Vec::with_capacity(TABLE_PREFIX_LEN + self.table_key.as_ref().len());
619        self.encode_into(&mut ret);
620        ret
621    }
622
623    pub fn is_empty(&self) -> bool {
624        self.table_key.as_ref().is_empty()
625    }
626
627    /// Get the length of the encoded format.
628    pub fn encoded_len(&self) -> usize {
629        self.table_key.as_ref().len() + TABLE_PREFIX_LEN
630    }
631
632    pub fn get_vnode_id(&self) -> usize {
633        self.table_key.vnode_part().to_index()
634    }
635}
636
637impl<'a> UserKey<&'a [u8]> {
638    /// Construct a [`UserKey`] from a byte slice. Its `table_key` will be a part of the input
639    /// `slice`.
640    pub fn decode(slice: &'a [u8]) -> Self {
641        let table_id: u32 = (&slice[..]).get_u32();
642
643        Self {
644            table_id: TableId::new(table_id),
645            table_key: TableKey(&slice[TABLE_PREFIX_LEN..]),
646        }
647    }
648
649    pub fn to_vec(self) -> UserKey<Vec<u8>> {
650        self.copy_into()
651    }
652
653    pub fn copy_into<T: CopyFromSlice + AsRef<[u8]>>(self) -> UserKey<T> {
654        UserKey {
655            table_id: self.table_id,
656            table_key: TableKey(T::copy_from_slice(self.table_key.0)),
657        }
658    }
659}
660
661impl<T: AsRef<[u8]> + Clone> UserKey<&T> {
662    pub fn cloned(self) -> UserKey<T> {
663        UserKey {
664            table_id: self.table_id,
665            table_key: TableKey(self.table_key.0.clone()),
666        }
667    }
668}
669
670impl<T: AsRef<[u8]>> UserKey<T> {
671    pub fn as_ref(&self) -> UserKey<&[u8]> {
672        UserKey::new(self.table_id, TableKey(self.table_key.as_ref()))
673    }
674}
675
676impl<T: AsRef<[u8]>> UserKey<T> {
677    /// Use this method to override an old `UserKey<Vec<u8>>` with a `UserKey<&[u8]>` to own the
678    /// table key without reallocating a new `UserKey` object.
679    pub fn set<F>(&mut self, other: UserKey<F>)
680    where
681        T: SetSlice<F>,
682        F: AsRef<[u8]>,
683    {
684        self.table_id = other.table_id;
685        self.table_key.0.set(&other.table_key.0);
686    }
687}
688
689impl UserKey<Vec<u8>> {
690    pub fn into_bytes(self) -> UserKey<Bytes> {
691        UserKey {
692            table_id: self.table_id,
693            table_key: TableKey(Bytes::from(self.table_key.0)),
694        }
695    }
696}
697
698/// [`FullKey`] is an internal concept in storage. It associates [`UserKey`] with an epoch.
699///
700/// The encoded format is | `user_key` | `epoch` |.
701#[derive(Clone, Copy, PartialEq, Eq, Hash, Default)]
702pub struct FullKey<T: AsRef<[u8]>> {
703    pub user_key: UserKey<T>,
704    pub epoch_with_gap: EpochWithGap,
705}
706
707impl<T: AsRef<[u8]>> Debug for FullKey<T> {
708    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
709        write!(
710            f,
711            "FullKey {{ {:?}, epoch: {}, epoch_with_gap: {}, spill_offset: {}}}",
712            self.user_key,
713            self.epoch_with_gap.pure_epoch(),
714            self.epoch_with_gap.as_u64(),
715            self.epoch_with_gap.as_u64() - self.epoch_with_gap.pure_epoch(),
716        )
717    }
718}
719
720impl<T: AsRef<[u8]>> FullKey<T> {
721    pub fn new(table_id: TableId, table_key: TableKey<T>, epoch: HummockEpoch) -> Self {
722        Self {
723            user_key: UserKey::new(table_id, table_key),
724            epoch_with_gap: EpochWithGap::new(epoch, 0),
725        }
726    }
727
728    pub fn new_with_gap_epoch(
729        table_id: TableId,
730        table_key: TableKey<T>,
731        epoch_with_gap: EpochWithGap,
732    ) -> Self {
733        Self {
734            user_key: UserKey::new(table_id, table_key),
735            epoch_with_gap,
736        }
737    }
738
739    pub fn from_user_key(user_key: UserKey<T>, epoch: HummockEpoch) -> Self {
740        Self {
741            user_key,
742            epoch_with_gap: EpochWithGap::new_from_epoch(epoch),
743        }
744    }
745
746    /// Pass the inner type of `table_key` to make the code less verbose.
747    pub fn for_test(table_id: TableId, table_key: T, epoch: HummockEpoch) -> Self {
748        Self {
749            user_key: UserKey::for_test(table_id, table_key),
750            epoch_with_gap: EpochWithGap::new(epoch, 0),
751        }
752    }
753
754    /// Encode in to a buffer.
755    pub fn encode_into(&self, buf: &mut impl BufMut) {
756        self.user_key.encode_into(buf);
757        buf.put_u64(self.epoch_with_gap.as_u64());
758    }
759
760    pub fn encode(&self) -> Vec<u8> {
761        let mut buf = Vec::with_capacity(
762            TABLE_PREFIX_LEN + self.user_key.table_key.as_ref().len() + EPOCH_LEN,
763        );
764        self.encode_into(&mut buf);
765        buf
766    }
767
768    pub fn encode_reverse_epoch(&self) -> Vec<u8> {
769        let mut buf = Vec::with_capacity(
770            TABLE_PREFIX_LEN + self.user_key.table_key.as_ref().len() + EPOCH_LEN,
771        );
772        self.user_key.encode_into(&mut buf);
773        buf.put_u64(u64::MAX - self.epoch_with_gap.as_u64());
774        buf
775    }
776
777    pub fn is_empty(&self) -> bool {
778        self.user_key.is_empty()
779    }
780
781    /// Get the length of the encoded format.
782    pub fn encoded_len(&self) -> usize {
783        self.user_key.encoded_len() + EPOCH_LEN
784    }
785}
786
787impl<'a> FullKey<&'a [u8]> {
788    /// Construct a [`FullKey`] from a byte slice.
789    pub fn decode(slice: &'a [u8]) -> Self {
790        let epoch_pos = slice.len() - EPOCH_LEN;
791        let epoch = (&slice[epoch_pos..]).get_u64();
792
793        Self {
794            user_key: UserKey::decode(&slice[..epoch_pos]),
795            epoch_with_gap: EpochWithGap::from_u64(epoch),
796        }
797    }
798
799    /// Construct a [`FullKey`] from a byte slice without `table_id` encoded.
800    pub fn from_slice_without_table_id(
801        table_id: TableId,
802        slice_without_table_id: &'a [u8],
803    ) -> Self {
804        let epoch_pos = slice_without_table_id.len() - EPOCH_LEN;
805        let epoch = (&slice_without_table_id[epoch_pos..]).get_u64();
806
807        Self {
808            user_key: UserKey::new(table_id, TableKey(&slice_without_table_id[..epoch_pos])),
809            epoch_with_gap: EpochWithGap::from_u64(epoch),
810        }
811    }
812
813    /// Construct a [`FullKey`] from a byte slice.
814    pub fn decode_reverse_epoch(slice: &'a [u8]) -> Self {
815        let epoch_pos = slice.len() - EPOCH_LEN;
816        let epoch = (&slice[epoch_pos..]).get_u64();
817
818        Self {
819            user_key: UserKey::decode(&slice[..epoch_pos]),
820            epoch_with_gap: EpochWithGap::from_u64(u64::MAX - epoch),
821        }
822    }
823
824    pub fn to_vec(self) -> FullKey<Vec<u8>> {
825        self.copy_into()
826    }
827
828    pub fn copy_into<T: CopyFromSlice + AsRef<[u8]>>(self) -> FullKey<T> {
829        FullKey {
830            user_key: self.user_key.copy_into(),
831            epoch_with_gap: self.epoch_with_gap,
832        }
833    }
834}
835
836impl FullKey<Vec<u8>> {
837    /// Calling this method may accidentally cause memory allocation when converting `Vec` into
838    /// `Bytes`
839    pub fn into_bytes(self) -> FullKey<Bytes> {
840        FullKey {
841            epoch_with_gap: self.epoch_with_gap,
842            user_key: self.user_key.into_bytes(),
843        }
844    }
845}
846
847impl<T: AsRef<[u8]>> FullKey<T> {
848    pub fn to_ref(&self) -> FullKey<&[u8]> {
849        FullKey {
850            user_key: self.user_key.as_ref(),
851            epoch_with_gap: self.epoch_with_gap,
852        }
853    }
854}
855
856impl<T: AsRef<[u8]>> FullKey<T> {
857    /// Use this method to override an old `FullKey<Vec<u8>>` with a `FullKey<&[u8]>` to own the
858    /// table key without reallocating a new `FullKey` object.
859    pub fn set<F>(&mut self, other: FullKey<F>)
860    where
861        T: SetSlice<F>,
862        F: AsRef<[u8]>,
863    {
864        self.user_key.set(other.user_key);
865        self.epoch_with_gap = other.epoch_with_gap;
866    }
867}
868
869impl<T: AsRef<[u8]> + Ord + Eq> Ord for FullKey<T> {
870    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
871        // When `user_key` is the same, greater epoch comes first.
872        self.user_key
873            .cmp(&other.user_key)
874            .then_with(|| other.epoch_with_gap.cmp(&self.epoch_with_gap))
875    }
876}
877
878impl<T: AsRef<[u8]> + Ord + Eq> PartialOrd for FullKey<T> {
879    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
880        Some(self.cmp(other))
881    }
882}
883
884pub mod range_delete_backward_compatibility_serde_struct {
885    use bytes::{Buf, BufMut};
886    use risingwave_common::catalog::TableId;
887    use serde::{Deserialize, Serialize};
888
889    #[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
890    pub struct TableKey(Vec<u8>);
891
892    #[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
893    #[serde(from = "UserKeySerde", into = "UserKeySerde")]
894    pub struct UserKey {
895        // When comparing `UserKey`, we first compare `table_id`, then `table_key`. So the order of
896        // declaration matters.
897        pub table_id: TableId,
898        pub table_key: TableKey,
899    }
900
901    #[derive(Deserialize, Serialize)]
902    pub struct TableIdSerde {
903        table_id: u32,
904    }
905
906    #[derive(Deserialize, Serialize)]
907    struct UserKeySerde {
908        table_id: TableIdSerde,
909        table_key: TableKey,
910    }
911
912    impl From<UserKeySerde> for UserKey {
913        fn from(value: UserKeySerde) -> Self {
914            Self {
915                table_id: TableId::new(value.table_id.table_id),
916                table_key: value.table_key,
917            }
918        }
919    }
920
921    impl From<UserKey> for UserKeySerde {
922        fn from(value: UserKey) -> Self {
923            Self {
924                table_id: TableIdSerde {
925                    table_id: value.table_id.as_raw_id(),
926                },
927                table_key: value.table_key,
928            }
929        }
930    }
931
932    impl UserKey {
933        pub fn decode_length_prefixed(buf: &mut &[u8]) -> Self {
934            let table_id = buf.get_u32();
935            let len = buf.get_u32() as usize;
936            let data = buf[..len].to_vec();
937            buf.advance(len);
938            UserKey {
939                table_id: TableId::new(table_id),
940                table_key: TableKey(data),
941            }
942        }
943
944        pub fn encode_length_prefixed(&self, mut buf: impl BufMut) {
945            buf.put_u32(self.table_id.as_raw_id());
946            buf.put_u32(self.table_key.0.as_slice().len() as u32);
947            buf.put_slice(self.table_key.0.as_slice());
948        }
949    }
950
951    #[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
952    pub struct PointRange {
953        // When comparing `PointRange`, we first compare `left_user_key`, then
954        // `is_exclude_left_key`. Therefore the order of declaration matters.
955        pub left_user_key: UserKey,
956        /// `PointRange` represents the left user key itself if `is_exclude_left_key==false`
957        /// while represents the right δ Neighborhood of the left user key if
958        /// `is_exclude_left_key==true`.
959        pub is_exclude_left_key: bool,
960    }
961}
962
963pub trait EmptySliceRef {
964    fn empty_slice_ref<'a>() -> &'a Self;
965}
966
967static EMPTY_BYTES: Bytes = Bytes::new();
968impl EmptySliceRef for Bytes {
969    fn empty_slice_ref<'a>() -> &'a Self {
970        &EMPTY_BYTES
971    }
972}
973
974static EMPTY_VEC: Vec<u8> = Vec::new();
975impl EmptySliceRef for Vec<u8> {
976    fn empty_slice_ref<'a>() -> &'a Self {
977        &EMPTY_VEC
978    }
979}
980
981const EMPTY_SLICE: &[u8] = b"";
982impl EmptySliceRef for &[u8] {
983    fn empty_slice_ref<'b>() -> &'b Self {
984        &EMPTY_SLICE
985    }
986}
987
988/// Bound table key range with table id to generate a new user key range.
989pub fn bound_table_key_range<T: AsRef<[u8]> + EmptySliceRef>(
990    table_id: TableId,
991    table_key_range: &impl RangeBounds<TableKey<T>>,
992) -> (Bound<UserKey<&T>>, Bound<UserKey<&T>>) {
993    let start = match table_key_range.start_bound() {
994        Included(b) => Included(UserKey::new(table_id, TableKey(&b.0))),
995        Excluded(b) => Excluded(UserKey::new(table_id, TableKey(&b.0))),
996        Unbounded => Included(UserKey::new(table_id, TableKey(T::empty_slice_ref()))),
997    };
998
999    let end = match table_key_range.end_bound() {
1000        Included(b) => Included(UserKey::new(table_id, TableKey(&b.0))),
1001        Excluded(b) => Excluded(UserKey::new(table_id, TableKey(&b.0))),
1002        Unbounded => {
1003            if let Some(next_table_id) = table_id.as_raw_id().checked_add(1) {
1004                Excluded(UserKey::new(
1005                    next_table_id.into(),
1006                    TableKey(T::empty_slice_ref()),
1007                ))
1008            } else {
1009                Unbounded
1010            }
1011        }
1012    };
1013
1014    (start, end)
1015}
1016
1017/// TODO: Temporary bypass full key check. Remove this field after #15099 is resolved.
1018pub struct FullKeyTracker<T: AsRef<[u8]> + Ord + Eq, const SKIP_DEDUP: bool = false> {
1019    pub latest_full_key: FullKey<T>,
1020    last_observed_epoch_with_gap: EpochWithGap,
1021}
1022
1023impl<T: AsRef<[u8]> + Ord + Eq, const SKIP_DEDUP: bool> FullKeyTracker<T, SKIP_DEDUP> {
1024    pub fn new(init_full_key: FullKey<T>) -> Self {
1025        let epoch_with_gap = init_full_key.epoch_with_gap;
1026        Self {
1027            latest_full_key: init_full_key,
1028            last_observed_epoch_with_gap: epoch_with_gap,
1029        }
1030    }
1031
1032    /// Check and observe a new full key during iteration
1033    ///
1034    /// # Examples:
1035    /// ```
1036    /// use bytes::Bytes;
1037    /// use risingwave_common::catalog::TableId;
1038    /// use risingwave_common::util::epoch::EPOCH_AVAILABLE_BITS;
1039    /// use risingwave_hummock_sdk::EpochWithGap;
1040    /// use risingwave_hummock_sdk::key::{FullKey, FullKeyTracker, TableKey};
1041    ///
1042    /// let table_id = TableId::new(1);
1043    /// let full_key1 = FullKey::new(table_id, TableKey(Bytes::from("c")), 5 << EPOCH_AVAILABLE_BITS);
1044    /// let mut a: FullKeyTracker<_> = FullKeyTracker::<Bytes>::new(full_key1.clone());
1045    ///
1046    /// // Panic on non-decreasing epoch observed for the same user key.
1047    /// // let full_key_with_larger_epoch = FullKey::new(table_id, TableKey(Bytes::from("c")), 6 << EPOCH_AVAILABLE_BITS);
1048    /// // a.observe(full_key_with_larger_epoch);
1049    ///
1050    /// // Panic on non-increasing user key observed.
1051    /// // let full_key_with_smaller_user_key = FullKey::new(table_id, TableKey(Bytes::from("b")), 3 << EPOCH_AVAILABLE_BITS);
1052    /// // a.observe(full_key_with_smaller_user_key);
1053    ///
1054    /// let full_key2 = FullKey::new(table_id, TableKey(Bytes::from("c")), 3 << EPOCH_AVAILABLE_BITS);
1055    /// assert_eq!(a.observe(full_key2.clone()), false);
1056    /// assert_eq!(a.latest_user_key(), &full_key2.user_key);
1057    ///
1058    /// let full_key3 = FullKey::new(table_id, TableKey(Bytes::from("f")), 4 << EPOCH_AVAILABLE_BITS);
1059    /// assert_eq!(a.observe(full_key3.clone()), true);
1060    /// assert_eq!(a.latest_user_key(), &full_key3.user_key);
1061    /// ```
1062    ///
1063    /// Return:
1064    /// - If the provided `key` contains a new user key, return true.
1065    /// - Otherwise: return false
1066    pub fn observe<F>(&mut self, key: FullKey<F>) -> bool
1067    where
1068        T: SetSlice<F>,
1069        F: AsRef<[u8]>,
1070    {
1071        self.observe_multi_version(key.user_key, once(key.epoch_with_gap))
1072    }
1073
1074    /// `epochs` comes from greater to smaller
1075    pub fn observe_multi_version<F>(
1076        &mut self,
1077        user_key: UserKey<F>,
1078        mut epochs: impl Iterator<Item = EpochWithGap>,
1079    ) -> bool
1080    where
1081        T: SetSlice<F>,
1082        F: AsRef<[u8]>,
1083    {
1084        let max_epoch_with_gap = epochs.next().expect("non-empty");
1085        let min_epoch_with_gap = epochs.fold(
1086            max_epoch_with_gap,
1087            |prev_epoch_with_gap, curr_epoch_with_gap| {
1088                assert!(
1089                    prev_epoch_with_gap > curr_epoch_with_gap,
1090                    "epoch list not sorted. prev: {:?}, curr: {:?}, user_key: {:?}",
1091                    prev_epoch_with_gap,
1092                    curr_epoch_with_gap,
1093                    user_key
1094                );
1095                curr_epoch_with_gap
1096            },
1097        );
1098        match self
1099            .latest_full_key
1100            .user_key
1101            .as_ref()
1102            .cmp(&user_key.as_ref())
1103        {
1104            Ordering::Less => {
1105                // Observe a new user key
1106
1107                // Reset epochs
1108                self.last_observed_epoch_with_gap = min_epoch_with_gap;
1109
1110                // Take the previous key and set latest key
1111                self.latest_full_key.set(FullKey {
1112                    user_key,
1113                    epoch_with_gap: min_epoch_with_gap,
1114                });
1115                true
1116            }
1117            Ordering::Equal => {
1118                if max_epoch_with_gap > self.last_observed_epoch_with_gap
1119                    || (!SKIP_DEDUP && max_epoch_with_gap == self.last_observed_epoch_with_gap)
1120                {
1121                    // Epoch from the same user key should be monotonically decreasing
1122                    panic!(
1123                        "key {:?} epoch {:?} >= prev epoch {:?}",
1124                        user_key, max_epoch_with_gap, self.last_observed_epoch_with_gap
1125                    );
1126                }
1127                self.last_observed_epoch_with_gap = min_epoch_with_gap;
1128                false
1129            }
1130            Ordering::Greater => {
1131                // User key should be monotonically increasing
1132                panic!(
1133                    "key {:?} <= prev key {:?}",
1134                    user_key,
1135                    FullKey {
1136                        user_key: self.latest_full_key.user_key.as_ref(),
1137                        epoch_with_gap: self.last_observed_epoch_with_gap
1138                    }
1139                );
1140            }
1141        }
1142    }
1143
1144    pub fn latest_user_key(&self) -> &UserKey<T> {
1145        &self.latest_full_key.user_key
1146    }
1147}
1148
1149#[cfg(test)]
1150mod tests {
1151    use risingwave_common::util::epoch::test_epoch;
1152
1153    use super::*;
1154
1155    #[test]
1156    fn test_encode_decode() {
1157        let epoch = test_epoch(1);
1158        let table_key = b"abc".to_vec();
1159        let key = FullKey::for_test(TableId::new(0), &table_key[..], 0);
1160        let buf = key.encode();
1161        assert_eq!(FullKey::decode(&buf), key);
1162        let key = FullKey::for_test(TableId::new(1), &table_key[..], epoch);
1163        let buf = key.encode();
1164        assert_eq!(FullKey::decode(&buf), key);
1165        let mut table_key = vec![1];
1166        let a = FullKey::for_test(TableId::new(1), table_key.clone(), epoch);
1167        table_key[0] = 2;
1168        let b = FullKey::for_test(TableId::new(1), table_key.clone(), epoch);
1169        table_key[0] = 129;
1170        let c = FullKey::for_test(TableId::new(1), table_key, epoch);
1171        assert!(a.lt(&b));
1172        assert!(b.lt(&c));
1173    }
1174
1175    #[test]
1176    fn test_key_cmp() {
1177        let epoch = test_epoch(1);
1178        let epoch2 = test_epoch(2);
1179        // 1 compared with 256 under little-endian encoding would return wrong result.
1180        let key1 = FullKey::for_test(TableId::new(0), b"0".to_vec(), epoch);
1181        let key2 = FullKey::for_test(TableId::new(1), b"0".to_vec(), epoch);
1182        let key3 = FullKey::for_test(TableId::new(1), b"1".to_vec(), epoch2);
1183        let key4 = FullKey::for_test(TableId::new(1), b"1".to_vec(), epoch);
1184
1185        assert_eq!(key1.cmp(&key1), Ordering::Equal);
1186        assert_eq!(key1.cmp(&key2), Ordering::Less);
1187        assert_eq!(key2.cmp(&key3), Ordering::Less);
1188        assert_eq!(key3.cmp(&key4), Ordering::Less);
1189    }
1190
1191    #[test]
1192    fn test_prev_key() {
1193        assert_eq!(prev_key(b"123"), b"122");
1194        assert_eq!(prev_key(b"12\x00"), b"11\xff");
1195        assert_eq!(prev_key(b"\x00\x00"), b"\xff\xff");
1196        assert_eq!(prev_key(b"\x00\x01"), b"\x00\x00");
1197        assert_eq!(prev_key(b"T"), b"S");
1198        assert_eq!(prev_key(b""), b"");
1199    }
1200
1201    #[test]
1202    fn test_bound_table_key_range() {
1203        assert_eq!(
1204            bound_table_key_range(
1205                TableId::default(),
1206                &(
1207                    Included(TableKey(b"a".to_vec())),
1208                    Included(TableKey(b"b".to_vec()))
1209                )
1210            ),
1211            (
1212                Included(UserKey::for_test(TableId::default(), &b"a".to_vec())),
1213                Included(UserKey::for_test(TableId::default(), &b"b".to_vec()),)
1214            )
1215        );
1216        assert_eq!(
1217            bound_table_key_range(
1218                TableId::from(1),
1219                &(Included(TableKey(b"a".to_vec())), Unbounded)
1220            ),
1221            (
1222                Included(UserKey::for_test(TableId::from(1), &b"a".to_vec())),
1223                Excluded(UserKey::for_test(TableId::from(2), &b"".to_vec()),)
1224            )
1225        );
1226        assert_eq!(
1227            bound_table_key_range(
1228                TableId::from(u32::MAX),
1229                &(Included(TableKey(b"a".to_vec())), Unbounded)
1230            ),
1231            (
1232                Included(UserKey::for_test(TableId::from(u32::MAX), &b"a".to_vec())),
1233                Unbounded,
1234            )
1235        );
1236    }
1237
1238    #[test]
1239    fn test_next_full_key() {
1240        let user_key = b"aaa".to_vec();
1241        let epoch: HummockEpoch = 3;
1242        let mut full_key = key_with_epoch(user_key, epoch);
1243        full_key = next_full_key(full_key.as_slice());
1244        assert_eq!(full_key, key_with_epoch(b"aaa".to_vec(), 2));
1245        full_key = next_full_key(full_key.as_slice());
1246        assert_eq!(full_key, key_with_epoch(b"aaa".to_vec(), 1));
1247        full_key = next_full_key(full_key.as_slice());
1248        assert_eq!(full_key, key_with_epoch(b"aaa".to_vec(), 0));
1249        full_key = next_full_key(full_key.as_slice());
1250        assert_eq!(
1251            full_key,
1252            key_with_epoch("aab".as_bytes().to_vec(), HummockEpoch::MAX)
1253        );
1254        assert_eq!(
1255            next_full_key(&key_with_epoch(b"\xff".to_vec(), 0)),
1256            Vec::<u8>::new()
1257        );
1258    }
1259
1260    #[test]
1261    fn test_prev_full_key() {
1262        let user_key = b"aab";
1263        let epoch: HummockEpoch = HummockEpoch::MAX - 3;
1264        let mut full_key = key_with_epoch(user_key.to_vec(), epoch);
1265        full_key = prev_full_key(full_key.as_slice());
1266        assert_eq!(
1267            full_key,
1268            key_with_epoch(b"aab".to_vec(), HummockEpoch::MAX - 2)
1269        );
1270        full_key = prev_full_key(full_key.as_slice());
1271        assert_eq!(
1272            full_key,
1273            key_with_epoch(b"aab".to_vec(), HummockEpoch::MAX - 1)
1274        );
1275        full_key = prev_full_key(full_key.as_slice());
1276        assert_eq!(full_key, key_with_epoch(b"aab".to_vec(), HummockEpoch::MAX));
1277        full_key = prev_full_key(full_key.as_slice());
1278        assert_eq!(full_key, key_with_epoch(b"aaa".to_vec(), 0));
1279
1280        assert_eq!(
1281            prev_full_key(&key_with_epoch(b"\x00".to_vec(), HummockEpoch::MAX)),
1282            Vec::<u8>::new()
1283        );
1284    }
1285
1286    #[test]
1287    fn test_user_key_order() {
1288        let a = UserKey::new(TableId::new(1), TableKey(b"aaa".to_vec()));
1289        let b = UserKey::new(TableId::new(2), TableKey(b"aaa".to_vec()));
1290        let c = UserKey::new(TableId::new(2), TableKey(b"bbb".to_vec()));
1291        assert!(a.lt(&b));
1292        assert!(b.lt(&c));
1293        let a = a.encode();
1294        let b = b.encode();
1295        let c = c.encode();
1296        assert!(a.lt(&b));
1297        assert!(b.lt(&c));
1298    }
1299
1300    #[test]
1301    fn test_prefixed_range_with_vnode() {
1302        let concat = |vnode: usize, b: &[u8]| -> Bytes {
1303            prefix_slice_with_vnode(VirtualNode::from_index(vnode), b)
1304        };
1305        assert_eq!(
1306            prefixed_range_with_vnode(
1307                (Included(Bytes::from("1")), Included(Bytes::from("2"))),
1308                VirtualNode::from_index(233),
1309            ),
1310            (
1311                Included(TableKey(concat(233, b"1"))),
1312                Included(TableKey(concat(233, b"2")))
1313            )
1314        );
1315        assert_eq!(
1316            prefixed_range_with_vnode(
1317                (Excluded(Bytes::from("1")), Excluded(Bytes::from("2"))),
1318                VirtualNode::from_index(233),
1319            ),
1320            (
1321                Excluded(TableKey(concat(233, b"1"))),
1322                Excluded(TableKey(concat(233, b"2")))
1323            )
1324        );
1325        assert_eq!(
1326            prefixed_range_with_vnode(
1327                (Bound::<Bytes>::Unbounded, Bound::<Bytes>::Unbounded),
1328                VirtualNode::from_index(233),
1329            ),
1330            (
1331                Included(TableKey(concat(233, b""))),
1332                Excluded(TableKey(concat(234, b"")))
1333            )
1334        );
1335        let max_vnode = VirtualNode::MAX_REPRESENTABLE.to_index();
1336        assert_eq!(
1337            prefixed_range_with_vnode(
1338                (Bound::<Bytes>::Unbounded, Bound::<Bytes>::Unbounded),
1339                VirtualNode::from_index(max_vnode),
1340            ),
1341            (Included(TableKey(concat(max_vnode, b""))), Unbounded)
1342        );
1343        let second_max_vnode = max_vnode - 1;
1344        assert_eq!(
1345            prefixed_range_with_vnode(
1346                (Bound::<Bytes>::Unbounded, Bound::<Bytes>::Unbounded),
1347                VirtualNode::from_index(second_max_vnode),
1348            ),
1349            (
1350                Included(TableKey(concat(second_max_vnode, b""))),
1351                Excluded(TableKey(concat(max_vnode, b"")))
1352            )
1353        );
1354    }
1355
1356    #[test]
1357    fn test_single_vnode_range() {
1358        let left_bound = vec![
1359            Included(b"0".as_slice()),
1360            Excluded(b"0".as_slice()),
1361            Unbounded,
1362        ];
1363        let right_bound = vec![
1364            Included(b"1".as_slice()),
1365            Excluded(b"1".as_slice()),
1366            Unbounded,
1367        ];
1368        for vnode in 0..VirtualNode::MAX_COUNT {
1369            for left in &left_bound {
1370                for right in &right_bound {
1371                    assert_eq!(
1372                        (vnode, vnode + 1),
1373                        vnode_range(&prefixed_range_with_vnode::<&[u8]>(
1374                            (*left, *right),
1375                            VirtualNode::from_index(vnode)
1376                        ))
1377                    )
1378                }
1379            }
1380        }
1381    }
1382}
risingwave_hummock_sdk/key.rs

risingwave_hummock_sdk/
key.rs
