1use 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>();
32pub 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
57pub fn vnode_range(range: &TableKeyRange) -> (usize, usize) {
67 let (left, right) = range;
68 let left = match left {
69 Included(key) | Excluded(key) => key.vnode_part().to_index(),
70 Unbounded => 0,
71 };
72 let right = match right {
73 Included(key) => key.vnode_part().to_index() + 1,
74 Excluded(key) => {
75 let (vnode, inner_key) = key.split_vnode();
76 if inner_key.is_empty() {
77 vnode.to_index()
80 } else {
81 vnode.to_index() + 1
82 }
83 }
84 Unbounded => VirtualNode::MAX_REPRESENTABLE.to_index() + 1,
85 };
86 (left, right)
87}
88
89pub fn vnode(range: &TableKeyRange) -> VirtualNode {
99 let (l, r_exclusive) = vnode_range(range);
100 assert_eq!(r_exclusive - l, 1);
101 VirtualNode::from_index(l)
102}
103
104pub fn key_with_epoch(mut user_key: Vec<u8>, epoch: HummockEpoch) -> Vec<u8> {
106 let res = epoch.to_be();
107 user_key.reserve(EPOCH_LEN);
108 let buf = user_key.chunk_mut();
109
110 unsafe {
112 ptr::copy_nonoverlapping(
113 &res as *const _ as *const u8,
114 buf.as_mut_ptr() as *mut _,
115 EPOCH_LEN,
116 );
117 user_key.advance_mut(EPOCH_LEN);
118 }
119
120 user_key
121}
122
123#[inline]
125pub fn split_key_epoch(full_key: &[u8]) -> (&[u8], &[u8]) {
126 let pos = full_key
127 .len()
128 .checked_sub(EPOCH_LEN)
129 .unwrap_or_else(|| panic!("bad full key format: {:?}", full_key));
130 full_key.split_at(pos)
131}
132
133pub fn user_key(full_key: &[u8]) -> &[u8] {
135 split_key_epoch(full_key).0
136}
137
138pub fn table_key(user_key: &[u8]) -> &[u8] {
140 &user_key[TABLE_PREFIX_LEN..]
141}
142
143#[inline(always)]
144pub fn get_user_key(full_key: &[u8]) -> Vec<u8> {
146 if full_key.is_empty() {
147 vec![]
148 } else {
149 user_key(full_key).to_vec()
150 }
151}
152
153#[inline(always)]
155pub fn get_table_id(full_key: &[u8]) -> u32 {
156 let mut buf = full_key;
157 buf.get_u32()
158}
159
160pub fn next_key(key: &[u8]) -> Vec<u8> {
179 if let Some((s, e)) = next_key_no_alloc(key) {
180 let mut res = Vec::with_capacity(s.len() + 1);
181 res.extend_from_slice(s);
182 res.push(e);
183 res
184 } else {
185 Vec::new()
186 }
187}
188
189pub fn prev_key(key: &[u8]) -> Vec<u8> {
206 let pos = key.iter().rposition(|b| *b != 0x00);
207 match pos {
208 Some(pos) => {
209 let mut res = Vec::with_capacity(key.len());
210 res.extend_from_slice(&key[0..pos]);
211 res.push(key[pos] - 1);
212 if pos + 1 < key.len() {
213 res.push(b"\xff".to_owned()[0]);
214 }
215 res
216 }
217 None => {
218 vec![0xff; key.len()]
219 }
220 }
221}
222
223fn next_key_no_alloc(key: &[u8]) -> Option<(&[u8], u8)> {
224 let pos = key.iter().rposition(|b| *b != 0xff)?;
225 Some((&key[..pos], key[pos] + 1))
226}
227
228pub fn next_epoch(epoch: &[u8]) -> Vec<u8> {
243 let pos = epoch.iter().rposition(|b| *b != 0xff);
244 match pos {
245 Some(mut pos) => {
246 let mut res = Vec::with_capacity(epoch.len());
247 res.extend_from_slice(&epoch[0..pos]);
248 res.push(epoch[pos] + 1);
249 while pos + 1 < epoch.len() {
250 res.push(0x00);
251 pos += 1;
252 }
253 res
254 }
255 None => {
256 vec![0x00; epoch.len()]
257 }
258 }
259}
260
261pub fn prev_epoch(epoch: &[u8]) -> Vec<u8> {
274 let pos = epoch.iter().rposition(|b| *b != 0x00);
275 match pos {
276 Some(mut pos) => {
277 let mut res = Vec::with_capacity(epoch.len());
278 res.extend_from_slice(&epoch[0..pos]);
279 res.push(epoch[pos] - 1);
280 while pos + 1 < epoch.len() {
281 res.push(0xff);
282 pos += 1;
283 }
284 res
285 }
286 None => {
287 vec![0xff; epoch.len()]
288 }
289 }
290}
291
292pub fn next_full_key(full_key: &[u8]) -> Vec<u8> {
296 let (user_key, epoch) = split_key_epoch(full_key);
297 let prev_epoch = prev_epoch(epoch);
298 let mut res = Vec::with_capacity(full_key.len());
299 if prev_epoch.cmp(&vec![0xff; prev_epoch.len()]) == Ordering::Equal {
300 let next_user_key = next_key(user_key);
301 if next_user_key.is_empty() {
302 return Vec::new();
303 }
304 res.extend_from_slice(next_user_key.as_slice());
305 res.extend_from_slice(prev_epoch.as_slice());
306 res
307 } else {
308 res.extend_from_slice(user_key);
309 res.extend_from_slice(prev_epoch.as_slice());
310 res
311 }
312}
313
314pub fn prev_full_key(full_key: &[u8]) -> Vec<u8> {
318 let (user_key, epoch) = split_key_epoch(full_key);
319 let next_epoch = next_epoch(epoch);
320 let mut res = Vec::with_capacity(full_key.len());
321 if next_epoch.cmp(&vec![0x00; next_epoch.len()]) == Ordering::Equal {
322 let prev_user_key = prev_key(user_key);
323 if prev_user_key.cmp(&vec![0xff; prev_user_key.len()]) == Ordering::Equal {
324 return Vec::new();
325 }
326 res.extend_from_slice(prev_user_key.as_slice());
327 res.extend_from_slice(next_epoch.as_slice());
328 res
329 } else {
330 res.extend_from_slice(user_key);
331 res.extend_from_slice(next_epoch.as_slice());
332 res
333 }
334}
335
336pub fn end_bound_of_vnode(vnode: VirtualNode) -> Bound<Bytes> {
344 if vnode == VirtualNode::MAX_REPRESENTABLE {
345 Unbounded
346 } else {
347 let end_bound_index = vnode.to_index() + 1;
348 Excluded(Bytes::copy_from_slice(
349 &VirtualNode::from_index(end_bound_index).to_be_bytes(),
350 ))
351 }
352}
353
354pub fn end_bound_of_prefix(prefix: &[u8]) -> Bound<Bytes> {
356 if let Some((s, e)) = next_key_no_alloc(prefix) {
357 let mut buf = BytesMut::with_capacity(s.len() + 1);
358 buf.extend_from_slice(s);
359 buf.put_u8(e);
360 Excluded(buf.freeze())
361 } else {
362 Unbounded
363 }
364}
365
366pub fn start_bound_of_excluded_prefix(prefix: &[u8]) -> Bound<Bytes> {
368 if let Some((s, e)) = next_key_no_alloc(prefix) {
369 let mut buf = BytesMut::with_capacity(s.len() + 1);
370 buf.extend_from_slice(s);
371 buf.put_u8(e);
372 Included(buf.freeze())
373 } else {
374 panic!("the prefix is the maximum value")
375 }
376}
377
378pub fn range_of_prefix(prefix: &[u8]) -> (Bound<Bytes>, Bound<Bytes>) {
380 if prefix.is_empty() {
381 (Unbounded, Unbounded)
382 } else {
383 (
384 Included(Bytes::copy_from_slice(prefix)),
385 end_bound_of_prefix(prefix),
386 )
387 }
388}
389
390pub fn prefix_slice_with_vnode(vnode: VirtualNode, slice: &[u8]) -> Bytes {
391 let prefix = vnode.to_be_bytes();
392 let mut buf = BytesMut::with_capacity(prefix.len() + slice.len());
393 buf.extend_from_slice(&prefix);
394 buf.extend_from_slice(slice);
395 buf.freeze()
396}
397
398pub fn prefixed_range_with_vnode<B: AsRef<[u8]>>(
400 range: impl RangeBounds<B>,
401 vnode: VirtualNode,
402) -> TableKeyRange {
403 let prefixed = |b: &B| -> Bytes { prefix_slice_with_vnode(vnode, b.as_ref()) };
404
405 let start: Bound<Bytes> = match range.start_bound() {
406 Included(b) => Included(prefixed(b)),
407 Excluded(b) => {
408 assert!(!b.as_ref().is_empty());
409 Excluded(prefixed(b))
410 }
411 Unbounded => Included(Bytes::copy_from_slice(&vnode.to_be_bytes())),
412 };
413
414 let end = match range.end_bound() {
415 Included(b) => Included(prefixed(b)),
416 Excluded(b) => {
417 assert!(!b.as_ref().is_empty());
418 Excluded(prefixed(b))
419 }
420 Unbounded => end_bound_of_vnode(vnode),
421 };
422
423 map_table_key_range((start, end))
424}
425
426pub trait SetSlice<S: AsRef<[u8]> + ?Sized> {
427 fn set(&mut self, value: &S);
428}
429
430impl<S: AsRef<[u8]> + ?Sized> SetSlice<S> for Vec<u8> {
431 fn set(&mut self, value: &S) {
432 self.clear();
433 self.extend_from_slice(value.as_ref());
434 }
435}
436
437impl SetSlice<Bytes> for Bytes {
438 fn set(&mut self, value: &Bytes) {
439 *self = value.clone()
440 }
441}
442
443pub trait CopyFromSlice: Send + 'static {
444 fn copy_from_slice(slice: &[u8]) -> Self;
445}
446
447impl CopyFromSlice for Vec<u8> {
448 fn copy_from_slice(slice: &[u8]) -> Self {
449 Vec::from(slice)
450 }
451}
452
453impl CopyFromSlice for Bytes {
454 fn copy_from_slice(slice: &[u8]) -> Self {
455 Bytes::copy_from_slice(slice)
456 }
457}
458
459impl CopyFromSlice for () {
460 fn copy_from_slice(_: &[u8]) -> Self {}
461}
462
463#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
469pub struct TableKey<T: AsRef<[u8]>>(pub T);
470
471impl<T: AsRef<[u8]>> Debug for TableKey<T> {
472 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
473 write!(f, "TableKey {{ {} }}", hex::encode(self.0.as_ref()))
474 }
475}
476
477impl<T: AsRef<[u8]>> Deref for TableKey<T> {
478 type Target = T;
479
480 fn deref(&self) -> &Self::Target {
481 &self.0
482 }
483}
484
485impl<T: AsRef<[u8]>> DerefMut for TableKey<T> {
486 fn deref_mut(&mut self) -> &mut Self::Target {
487 &mut self.0
488 }
489}
490
491impl<T: AsRef<[u8]>> AsRef<[u8]> for TableKey<T> {
492 fn as_ref(&self) -> &[u8] {
493 self.0.as_ref()
494 }
495}
496
497impl TableKey<Bytes> {
498 pub fn split_vnode_bytes(&self) -> (VirtualNode, Bytes) {
499 debug_assert!(
500 self.0.len() >= VirtualNode::SIZE,
501 "too short table key: {:?}",
502 self.0.as_ref()
503 );
504 let (vnode, _) = self.0.split_first_chunk::<{ VirtualNode::SIZE }>().unwrap();
505 (
506 VirtualNode::from_be_bytes(*vnode),
507 self.0.slice(VirtualNode::SIZE..),
508 )
509 }
510}
511
512impl<T: AsRef<[u8]>> TableKey<T> {
513 pub fn split_vnode(&self) -> (VirtualNode, &[u8]) {
514 debug_assert!(
515 self.0.as_ref().len() >= VirtualNode::SIZE,
516 "too short table key: {:?}",
517 self.0.as_ref()
518 );
519 let (vnode, inner_key) = self
520 .0
521 .as_ref()
522 .split_first_chunk::<{ VirtualNode::SIZE }>()
523 .unwrap();
524 (VirtualNode::from_be_bytes(*vnode), inner_key)
525 }
526
527 pub fn vnode_part(&self) -> VirtualNode {
528 self.split_vnode().0
529 }
530
531 pub fn key_part(&self) -> &[u8] {
532 self.split_vnode().1
533 }
534
535 pub fn to_ref(&self) -> TableKey<&[u8]> {
536 TableKey(self.0.as_ref())
537 }
538}
539
540impl<T: AsRef<[u8]>> Borrow<[u8]> for TableKey<T> {
541 fn borrow(&self) -> &[u8] {
542 self.0.as_ref()
543 }
544}
545
546impl EstimateSize for TableKey<Bytes> {
547 fn estimated_heap_size(&self) -> usize {
548 self.0.estimated_heap_size()
549 }
550}
551
552impl TableKey<&[u8]> {
553 pub fn copy_into<T: CopyFromSlice + AsRef<[u8]>>(&self) -> TableKey<T> {
554 TableKey(T::copy_from_slice(self.as_ref()))
555 }
556}
557
558#[inline]
559pub fn map_table_key_range(range: (Bound<KeyPayloadType>, Bound<KeyPayloadType>)) -> TableKeyRange {
560 (range.0.map(TableKey), range.1.map(TableKey))
561}
562
563pub fn gen_key_from_bytes(vnode: VirtualNode, payload: &[u8]) -> TableKey<Bytes> {
564 TableKey(Bytes::from(
565 [vnode.to_be_bytes().as_slice(), payload].concat(),
566 ))
567}
568
569pub fn gen_key_from_str(vnode: VirtualNode, payload: &str) -> TableKey<Bytes> {
570 gen_key_from_bytes(vnode, payload.as_bytes())
571}
572
573#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
579pub struct UserKey<T: AsRef<[u8]>> {
580 pub table_id: TableId,
583 pub table_key: TableKey<T>,
584}
585
586impl<T: AsRef<[u8]>> Debug for UserKey<T> {
587 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
588 write!(
589 f,
590 "UserKey {{ {}, {:?} }}",
591 self.table_id.table_id, self.table_key
592 )
593 }
594}
595
596impl<T: AsRef<[u8]>> UserKey<T> {
597 pub fn new(table_id: TableId, table_key: TableKey<T>) -> Self {
598 Self {
599 table_id,
600 table_key,
601 }
602 }
603
604 pub fn for_test(table_id: TableId, table_key: T) -> Self {
606 Self {
607 table_id,
608 table_key: TableKey(table_key),
609 }
610 }
611
612 pub fn encode_into(&self, buf: &mut impl BufMut) {
614 buf.put_u32(self.table_id.table_id());
615 buf.put_slice(self.table_key.as_ref());
616 }
617
618 pub fn encode_table_key_into(&self, buf: &mut impl BufMut) {
619 buf.put_slice(self.table_key.as_ref());
620 }
621
622 pub fn encode(&self) -> Vec<u8> {
623 let mut ret = Vec::with_capacity(TABLE_PREFIX_LEN + self.table_key.as_ref().len());
624 self.encode_into(&mut ret);
625 ret
626 }
627
628 pub fn is_empty(&self) -> bool {
629 self.table_key.as_ref().is_empty()
630 }
631
632 pub fn encoded_len(&self) -> usize {
634 self.table_key.as_ref().len() + TABLE_PREFIX_LEN
635 }
636
637 pub fn get_vnode_id(&self) -> usize {
638 self.table_key.vnode_part().to_index()
639 }
640}
641
642impl<'a> UserKey<&'a [u8]> {
643 pub fn decode(slice: &'a [u8]) -> Self {
646 let table_id: u32 = (&slice[..]).get_u32();
647
648 Self {
649 table_id: TableId::new(table_id),
650 table_key: TableKey(&slice[TABLE_PREFIX_LEN..]),
651 }
652 }
653
654 pub fn to_vec(self) -> UserKey<Vec<u8>> {
655 self.copy_into()
656 }
657
658 pub fn copy_into<T: CopyFromSlice + AsRef<[u8]>>(self) -> UserKey<T> {
659 UserKey {
660 table_id: self.table_id,
661 table_key: TableKey(T::copy_from_slice(self.table_key.0)),
662 }
663 }
664}
665
666impl<T: AsRef<[u8]> + Clone> UserKey<&T> {
667 pub fn cloned(self) -> UserKey<T> {
668 UserKey {
669 table_id: self.table_id,
670 table_key: TableKey(self.table_key.0.clone()),
671 }
672 }
673}
674
675impl<T: AsRef<[u8]>> UserKey<T> {
676 pub fn as_ref(&self) -> UserKey<&[u8]> {
677 UserKey::new(self.table_id, TableKey(self.table_key.as_ref()))
678 }
679}
680
681impl<T: AsRef<[u8]>> UserKey<T> {
682 pub fn set<F>(&mut self, other: UserKey<F>)
685 where
686 T: SetSlice<F>,
687 F: AsRef<[u8]>,
688 {
689 self.table_id = other.table_id;
690 self.table_key.0.set(&other.table_key.0);
691 }
692}
693
694impl UserKey<Vec<u8>> {
695 pub fn into_bytes(self) -> UserKey<Bytes> {
696 UserKey {
697 table_id: self.table_id,
698 table_key: TableKey(Bytes::from(self.table_key.0)),
699 }
700 }
701}
702
703#[derive(Clone, Copy, PartialEq, Eq, Hash, Default)]
707pub struct FullKey<T: AsRef<[u8]>> {
708 pub user_key: UserKey<T>,
709 pub epoch_with_gap: EpochWithGap,
710}
711
712impl<T: AsRef<[u8]>> Debug for FullKey<T> {
713 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
714 write!(
715 f,
716 "FullKey {{ {:?}, epoch: {}, epoch_with_gap: {}, spill_offset: {}}}",
717 self.user_key,
718 self.epoch_with_gap.pure_epoch(),
719 self.epoch_with_gap.as_u64(),
720 self.epoch_with_gap.as_u64() - self.epoch_with_gap.pure_epoch(),
721 )
722 }
723}
724
725impl<T: AsRef<[u8]>> FullKey<T> {
726 pub fn new(table_id: TableId, table_key: TableKey<T>, epoch: HummockEpoch) -> Self {
727 Self {
728 user_key: UserKey::new(table_id, table_key),
729 epoch_with_gap: EpochWithGap::new(epoch, 0),
730 }
731 }
732
733 pub fn new_with_gap_epoch(
734 table_id: TableId,
735 table_key: TableKey<T>,
736 epoch_with_gap: EpochWithGap,
737 ) -> Self {
738 Self {
739 user_key: UserKey::new(table_id, table_key),
740 epoch_with_gap,
741 }
742 }
743
744 pub fn from_user_key(user_key: UserKey<T>, epoch: HummockEpoch) -> Self {
745 Self {
746 user_key,
747 epoch_with_gap: EpochWithGap::new_from_epoch(epoch),
748 }
749 }
750
751 pub fn for_test(table_id: TableId, table_key: T, epoch: HummockEpoch) -> Self {
753 Self {
754 user_key: UserKey::for_test(table_id, table_key),
755 epoch_with_gap: EpochWithGap::new(epoch, 0),
756 }
757 }
758
759 pub fn encode_into(&self, buf: &mut impl BufMut) {
761 self.user_key.encode_into(buf);
762 buf.put_u64(self.epoch_with_gap.as_u64());
763 }
764
765 pub fn encode(&self) -> Vec<u8> {
766 let mut buf = Vec::with_capacity(
767 TABLE_PREFIX_LEN + self.user_key.table_key.as_ref().len() + EPOCH_LEN,
768 );
769 self.encode_into(&mut buf);
770 buf
771 }
772
773 pub fn encode_into_without_table_id(&self, buf: &mut impl BufMut) {
775 self.user_key.encode_table_key_into(buf);
776 buf.put_u64(self.epoch_with_gap.as_u64());
777 }
778
779 pub fn encode_reverse_epoch(&self) -> Vec<u8> {
780 let mut buf = Vec::with_capacity(
781 TABLE_PREFIX_LEN + self.user_key.table_key.as_ref().len() + EPOCH_LEN,
782 );
783 self.user_key.encode_into(&mut buf);
784 buf.put_u64(u64::MAX - self.epoch_with_gap.as_u64());
785 buf
786 }
787
788 pub fn is_empty(&self) -> bool {
789 self.user_key.is_empty()
790 }
791
792 pub fn encoded_len(&self) -> usize {
794 self.user_key.encoded_len() + EPOCH_LEN
795 }
796}
797
798impl<'a> FullKey<&'a [u8]> {
799 pub fn decode(slice: &'a [u8]) -> Self {
801 let epoch_pos = slice.len() - EPOCH_LEN;
802 let epoch = (&slice[epoch_pos..]).get_u64();
803
804 Self {
805 user_key: UserKey::decode(&slice[..epoch_pos]),
806 epoch_with_gap: EpochWithGap::from_u64(epoch),
807 }
808 }
809
810 pub fn from_slice_without_table_id(
812 table_id: TableId,
813 slice_without_table_id: &'a [u8],
814 ) -> Self {
815 let epoch_pos = slice_without_table_id.len() - EPOCH_LEN;
816 let epoch = (&slice_without_table_id[epoch_pos..]).get_u64();
817
818 Self {
819 user_key: UserKey::new(table_id, TableKey(&slice_without_table_id[..epoch_pos])),
820 epoch_with_gap: EpochWithGap::from_u64(epoch),
821 }
822 }
823
824 pub fn decode_reverse_epoch(slice: &'a [u8]) -> Self {
826 let epoch_pos = slice.len() - EPOCH_LEN;
827 let epoch = (&slice[epoch_pos..]).get_u64();
828
829 Self {
830 user_key: UserKey::decode(&slice[..epoch_pos]),
831 epoch_with_gap: EpochWithGap::from_u64(u64::MAX - epoch),
832 }
833 }
834
835 pub fn to_vec(self) -> FullKey<Vec<u8>> {
836 self.copy_into()
837 }
838
839 pub fn copy_into<T: CopyFromSlice + AsRef<[u8]>>(self) -> FullKey<T> {
840 FullKey {
841 user_key: self.user_key.copy_into(),
842 epoch_with_gap: self.epoch_with_gap,
843 }
844 }
845}
846
847impl FullKey<Vec<u8>> {
848 pub fn into_bytes(self) -> FullKey<Bytes> {
851 FullKey {
852 epoch_with_gap: self.epoch_with_gap,
853 user_key: self.user_key.into_bytes(),
854 }
855 }
856}
857
858impl<T: AsRef<[u8]>> FullKey<T> {
859 pub fn to_ref(&self) -> FullKey<&[u8]> {
860 FullKey {
861 user_key: self.user_key.as_ref(),
862 epoch_with_gap: self.epoch_with_gap,
863 }
864 }
865}
866
867impl<T: AsRef<[u8]>> FullKey<T> {
868 pub fn set<F>(&mut self, other: FullKey<F>)
871 where
872 T: SetSlice<F>,
873 F: AsRef<[u8]>,
874 {
875 self.user_key.set(other.user_key);
876 self.epoch_with_gap = other.epoch_with_gap;
877 }
878}
879
880impl<T: AsRef<[u8]> + Ord + Eq> Ord for FullKey<T> {
881 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
882 self.user_key
884 .cmp(&other.user_key)
885 .then_with(|| other.epoch_with_gap.cmp(&self.epoch_with_gap))
886 }
887}
888
889impl<T: AsRef<[u8]> + Ord + Eq> PartialOrd for FullKey<T> {
890 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
891 Some(self.cmp(other))
892 }
893}
894
895pub mod range_delete_backward_compatibility_serde_struct {
896 use bytes::{Buf, BufMut};
897 use risingwave_common::catalog::TableId;
898 use serde::{Deserialize, Serialize};
899
900 #[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
901 pub struct TableKey(Vec<u8>);
902
903 #[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
904 pub struct UserKey {
905 pub table_id: TableId,
908 pub table_key: TableKey,
909 }
910
911 impl UserKey {
912 pub fn decode_length_prefixed(buf: &mut &[u8]) -> Self {
913 let table_id = buf.get_u32();
914 let len = buf.get_u32() as usize;
915 let data = buf[..len].to_vec();
916 buf.advance(len);
917 UserKey {
918 table_id: TableId::new(table_id),
919 table_key: TableKey(data),
920 }
921 }
922
923 pub fn encode_length_prefixed(&self, mut buf: impl BufMut) {
924 buf.put_u32(self.table_id.table_id());
925 buf.put_u32(self.table_key.0.as_slice().len() as u32);
926 buf.put_slice(self.table_key.0.as_slice());
927 }
928 }
929
930 #[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
931 pub struct PointRange {
932 pub left_user_key: UserKey,
935 pub is_exclude_left_key: bool,
939 }
940}
941
942pub trait EmptySliceRef {
943 fn empty_slice_ref<'a>() -> &'a Self;
944}
945
946static EMPTY_BYTES: Bytes = Bytes::new();
947impl EmptySliceRef for Bytes {
948 fn empty_slice_ref<'a>() -> &'a Self {
949 &EMPTY_BYTES
950 }
951}
952
953static EMPTY_VEC: Vec<u8> = Vec::new();
954impl EmptySliceRef for Vec<u8> {
955 fn empty_slice_ref<'a>() -> &'a Self {
956 &EMPTY_VEC
957 }
958}
959
960const EMPTY_SLICE: &[u8] = b"";
961impl EmptySliceRef for &[u8] {
962 fn empty_slice_ref<'b>() -> &'b Self {
963 &EMPTY_SLICE
964 }
965}
966
967pub fn bound_table_key_range<T: AsRef<[u8]> + EmptySliceRef>(
969 table_id: TableId,
970 table_key_range: &impl RangeBounds<TableKey<T>>,
971) -> (Bound<UserKey<&T>>, Bound<UserKey<&T>>) {
972 let start = match table_key_range.start_bound() {
973 Included(b) => Included(UserKey::new(table_id, TableKey(&b.0))),
974 Excluded(b) => Excluded(UserKey::new(table_id, TableKey(&b.0))),
975 Unbounded => Included(UserKey::new(table_id, TableKey(T::empty_slice_ref()))),
976 };
977
978 let end = match table_key_range.end_bound() {
979 Included(b) => Included(UserKey::new(table_id, TableKey(&b.0))),
980 Excluded(b) => Excluded(UserKey::new(table_id, TableKey(&b.0))),
981 Unbounded => {
982 if let Some(next_table_id) = table_id.table_id().checked_add(1) {
983 Excluded(UserKey::new(
984 next_table_id.into(),
985 TableKey(T::empty_slice_ref()),
986 ))
987 } else {
988 Unbounded
989 }
990 }
991 };
992
993 (start, end)
994}
995
996pub struct FullKeyTracker<T: AsRef<[u8]> + Ord + Eq, const SKIP_DEDUP: bool = false> {
998 pub latest_full_key: FullKey<T>,
999 last_observed_epoch_with_gap: EpochWithGap,
1000}
1001
1002impl<T: AsRef<[u8]> + Ord + Eq, const SKIP_DEDUP: bool> FullKeyTracker<T, SKIP_DEDUP> {
1003 pub fn new(init_full_key: FullKey<T>) -> Self {
1004 let epoch_with_gap = init_full_key.epoch_with_gap;
1005 Self {
1006 latest_full_key: init_full_key,
1007 last_observed_epoch_with_gap: epoch_with_gap,
1008 }
1009 }
1010
1011 pub fn observe<F>(&mut self, key: FullKey<F>) -> bool
1046 where
1047 T: SetSlice<F>,
1048 F: AsRef<[u8]>,
1049 {
1050 self.observe_multi_version(key.user_key, once(key.epoch_with_gap))
1051 }
1052
1053 pub fn observe_multi_version<F>(
1055 &mut self,
1056 user_key: UserKey<F>,
1057 mut epochs: impl Iterator<Item = EpochWithGap>,
1058 ) -> bool
1059 where
1060 T: SetSlice<F>,
1061 F: AsRef<[u8]>,
1062 {
1063 let max_epoch_with_gap = epochs.next().expect("non-empty");
1064 let min_epoch_with_gap = epochs.fold(
1065 max_epoch_with_gap,
1066 |prev_epoch_with_gap, curr_epoch_with_gap| {
1067 assert!(
1068 prev_epoch_with_gap > curr_epoch_with_gap,
1069 "epoch list not sorted. prev: {:?}, curr: {:?}, user_key: {:?}",
1070 prev_epoch_with_gap,
1071 curr_epoch_with_gap,
1072 user_key
1073 );
1074 curr_epoch_with_gap
1075 },
1076 );
1077 match self
1078 .latest_full_key
1079 .user_key
1080 .as_ref()
1081 .cmp(&user_key.as_ref())
1082 {
1083 Ordering::Less => {
1084 self.last_observed_epoch_with_gap = min_epoch_with_gap;
1088
1089 self.latest_full_key.set(FullKey {
1091 user_key,
1092 epoch_with_gap: min_epoch_with_gap,
1093 });
1094 true
1095 }
1096 Ordering::Equal => {
1097 if max_epoch_with_gap > self.last_observed_epoch_with_gap
1098 || (!SKIP_DEDUP && max_epoch_with_gap == self.last_observed_epoch_with_gap)
1099 {
1100 panic!(
1102 "key {:?} epoch {:?} >= prev epoch {:?}",
1103 user_key, max_epoch_with_gap, self.last_observed_epoch_with_gap
1104 );
1105 }
1106 self.last_observed_epoch_with_gap = min_epoch_with_gap;
1107 false
1108 }
1109 Ordering::Greater => {
1110 panic!(
1112 "key {:?} <= prev key {:?}",
1113 user_key,
1114 FullKey {
1115 user_key: self.latest_full_key.user_key.as_ref(),
1116 epoch_with_gap: self.last_observed_epoch_with_gap
1117 }
1118 );
1119 }
1120 }
1121 }
1122
1123 pub fn latest_user_key(&self) -> &UserKey<T> {
1124 &self.latest_full_key.user_key
1125 }
1126}
1127
1128#[cfg(test)]
1129mod tests {
1130 use risingwave_common::util::epoch::test_epoch;
1131
1132 use super::*;
1133
1134 #[test]
1135 fn test_encode_decode() {
1136 let epoch = test_epoch(1);
1137 let table_key = b"abc".to_vec();
1138 let key = FullKey::for_test(TableId::new(0), &table_key[..], 0);
1139 let buf = key.encode();
1140 assert_eq!(FullKey::decode(&buf), key);
1141 let key = FullKey::for_test(TableId::new(1), &table_key[..], epoch);
1142 let buf = key.encode();
1143 assert_eq!(FullKey::decode(&buf), key);
1144 let mut table_key = vec![1];
1145 let a = FullKey::for_test(TableId::new(1), table_key.clone(), epoch);
1146 table_key[0] = 2;
1147 let b = FullKey::for_test(TableId::new(1), table_key.clone(), epoch);
1148 table_key[0] = 129;
1149 let c = FullKey::for_test(TableId::new(1), table_key, epoch);
1150 assert!(a.lt(&b));
1151 assert!(b.lt(&c));
1152 }
1153
1154 #[test]
1155 fn test_key_cmp() {
1156 let epoch = test_epoch(1);
1157 let epoch2 = test_epoch(2);
1158 let key1 = FullKey::for_test(TableId::new(0), b"0".to_vec(), epoch);
1160 let key2 = FullKey::for_test(TableId::new(1), b"0".to_vec(), epoch);
1161 let key3 = FullKey::for_test(TableId::new(1), b"1".to_vec(), epoch2);
1162 let key4 = FullKey::for_test(TableId::new(1), b"1".to_vec(), epoch);
1163
1164 assert_eq!(key1.cmp(&key1), Ordering::Equal);
1165 assert_eq!(key1.cmp(&key2), Ordering::Less);
1166 assert_eq!(key2.cmp(&key3), Ordering::Less);
1167 assert_eq!(key3.cmp(&key4), Ordering::Less);
1168 }
1169
1170 #[test]
1171 fn test_prev_key() {
1172 assert_eq!(prev_key(b"123"), b"122");
1173 assert_eq!(prev_key(b"12\x00"), b"11\xff");
1174 assert_eq!(prev_key(b"\x00\x00"), b"\xff\xff");
1175 assert_eq!(prev_key(b"\x00\x01"), b"\x00\x00");
1176 assert_eq!(prev_key(b"T"), b"S");
1177 assert_eq!(prev_key(b""), b"");
1178 }
1179
1180 #[test]
1181 fn test_bound_table_key_range() {
1182 assert_eq!(
1183 bound_table_key_range(
1184 TableId::default(),
1185 &(
1186 Included(TableKey(b"a".to_vec())),
1187 Included(TableKey(b"b".to_vec()))
1188 )
1189 ),
1190 (
1191 Included(UserKey::for_test(TableId::default(), &b"a".to_vec())),
1192 Included(UserKey::for_test(TableId::default(), &b"b".to_vec()),)
1193 )
1194 );
1195 assert_eq!(
1196 bound_table_key_range(
1197 TableId::from(1),
1198 &(Included(TableKey(b"a".to_vec())), Unbounded)
1199 ),
1200 (
1201 Included(UserKey::for_test(TableId::from(1), &b"a".to_vec())),
1202 Excluded(UserKey::for_test(TableId::from(2), &b"".to_vec()),)
1203 )
1204 );
1205 assert_eq!(
1206 bound_table_key_range(
1207 TableId::from(u32::MAX),
1208 &(Included(TableKey(b"a".to_vec())), Unbounded)
1209 ),
1210 (
1211 Included(UserKey::for_test(TableId::from(u32::MAX), &b"a".to_vec())),
1212 Unbounded,
1213 )
1214 );
1215 }
1216
1217 #[test]
1218 fn test_next_full_key() {
1219 let user_key = b"aaa".to_vec();
1220 let epoch: HummockEpoch = 3;
1221 let mut full_key = key_with_epoch(user_key, epoch);
1222 full_key = next_full_key(full_key.as_slice());
1223 assert_eq!(full_key, key_with_epoch(b"aaa".to_vec(), 2));
1224 full_key = next_full_key(full_key.as_slice());
1225 assert_eq!(full_key, key_with_epoch(b"aaa".to_vec(), 1));
1226 full_key = next_full_key(full_key.as_slice());
1227 assert_eq!(full_key, key_with_epoch(b"aaa".to_vec(), 0));
1228 full_key = next_full_key(full_key.as_slice());
1229 assert_eq!(
1230 full_key,
1231 key_with_epoch("aab".as_bytes().to_vec(), HummockEpoch::MAX)
1232 );
1233 assert_eq!(
1234 next_full_key(&key_with_epoch(b"\xff".to_vec(), 0)),
1235 Vec::<u8>::new()
1236 );
1237 }
1238
1239 #[test]
1240 fn test_prev_full_key() {
1241 let user_key = b"aab";
1242 let epoch: HummockEpoch = HummockEpoch::MAX - 3;
1243 let mut full_key = key_with_epoch(user_key.to_vec(), epoch);
1244 full_key = prev_full_key(full_key.as_slice());
1245 assert_eq!(
1246 full_key,
1247 key_with_epoch(b"aab".to_vec(), HummockEpoch::MAX - 2)
1248 );
1249 full_key = prev_full_key(full_key.as_slice());
1250 assert_eq!(
1251 full_key,
1252 key_with_epoch(b"aab".to_vec(), HummockEpoch::MAX - 1)
1253 );
1254 full_key = prev_full_key(full_key.as_slice());
1255 assert_eq!(full_key, key_with_epoch(b"aab".to_vec(), HummockEpoch::MAX));
1256 full_key = prev_full_key(full_key.as_slice());
1257 assert_eq!(full_key, key_with_epoch(b"aaa".to_vec(), 0));
1258
1259 assert_eq!(
1260 prev_full_key(&key_with_epoch(b"\x00".to_vec(), HummockEpoch::MAX)),
1261 Vec::<u8>::new()
1262 );
1263 }
1264
1265 #[test]
1266 fn test_uesr_key_order() {
1267 let a = UserKey::new(TableId::new(1), TableKey(b"aaa".to_vec()));
1268 let b = UserKey::new(TableId::new(2), TableKey(b"aaa".to_vec()));
1269 let c = UserKey::new(TableId::new(2), TableKey(b"bbb".to_vec()));
1270 assert!(a.lt(&b));
1271 assert!(b.lt(&c));
1272 let a = a.encode();
1273 let b = b.encode();
1274 let c = c.encode();
1275 assert!(a.lt(&b));
1276 assert!(b.lt(&c));
1277 }
1278
1279 #[test]
1280 fn test_prefixed_range_with_vnode() {
1281 let concat = |vnode: usize, b: &[u8]| -> Bytes {
1282 prefix_slice_with_vnode(VirtualNode::from_index(vnode), b)
1283 };
1284 assert_eq!(
1285 prefixed_range_with_vnode(
1286 (Included(Bytes::from("1")), Included(Bytes::from("2"))),
1287 VirtualNode::from_index(233),
1288 ),
1289 (
1290 Included(TableKey(concat(233, b"1"))),
1291 Included(TableKey(concat(233, b"2")))
1292 )
1293 );
1294 assert_eq!(
1295 prefixed_range_with_vnode(
1296 (Excluded(Bytes::from("1")), Excluded(Bytes::from("2"))),
1297 VirtualNode::from_index(233),
1298 ),
1299 (
1300 Excluded(TableKey(concat(233, b"1"))),
1301 Excluded(TableKey(concat(233, b"2")))
1302 )
1303 );
1304 assert_eq!(
1305 prefixed_range_with_vnode(
1306 (Bound::<Bytes>::Unbounded, Bound::<Bytes>::Unbounded),
1307 VirtualNode::from_index(233),
1308 ),
1309 (
1310 Included(TableKey(concat(233, b""))),
1311 Excluded(TableKey(concat(234, b"")))
1312 )
1313 );
1314 let max_vnode = VirtualNode::MAX_REPRESENTABLE.to_index();
1315 assert_eq!(
1316 prefixed_range_with_vnode(
1317 (Bound::<Bytes>::Unbounded, Bound::<Bytes>::Unbounded),
1318 VirtualNode::from_index(max_vnode),
1319 ),
1320 (Included(TableKey(concat(max_vnode, b""))), Unbounded)
1321 );
1322 let second_max_vnode = max_vnode - 1;
1323 assert_eq!(
1324 prefixed_range_with_vnode(
1325 (Bound::<Bytes>::Unbounded, Bound::<Bytes>::Unbounded),
1326 VirtualNode::from_index(second_max_vnode),
1327 ),
1328 (
1329 Included(TableKey(concat(second_max_vnode, b""))),
1330 Excluded(TableKey(concat(max_vnode, b"")))
1331 )
1332 );
1333 }
1334
1335 #[test]
1336 fn test_single_vnode_range() {
1337 let left_bound = vec![
1338 Included(b"0".as_slice()),
1339 Excluded(b"0".as_slice()),
1340 Unbounded,
1341 ];
1342 let right_bound = vec![
1343 Included(b"1".as_slice()),
1344 Excluded(b"1".as_slice()),
1345 Unbounded,
1346 ];
1347 for vnode in 0..VirtualNode::MAX_COUNT {
1348 for left in &left_bound {
1349 for right in &right_bound {
1350 assert_eq!(
1351 (vnode, vnode + 1),
1352 vnode_range(&prefixed_range_with_vnode::<&[u8]>(
1353 (*left, *right),
1354 VirtualNode::from_index(vnode)
1355 ))
1356 )
1357 }
1358 }
1359 }
1360 }
1361}