risingwave_storage/hummock/shared_buffer/
shared_buffer_batch.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::cmp::Ordering;
16use std::fmt::Debug;
17use std::marker::PhantomData;
18use std::mem::size_of_val;
19use std::ops::Bound::Included;
20use std::ops::RangeBounds;
21use std::sync::atomic::AtomicU64;
22use std::sync::atomic::Ordering::Relaxed;
23use std::sync::{Arc, LazyLock};
24
25use bytes::Bytes;
26use risingwave_common::catalog::TableId;
27use risingwave_common::metrics::LabelGuardedIntGauge;
28use risingwave_hummock_sdk::EpochWithGap;
29use risingwave_hummock_sdk::key::{FullKey, TableKey, TableKeyRange, UserKey};
30
31use crate::hummock::iterator::{
32    Backward, DirectionEnum, Forward, HummockIterator, HummockIteratorDirection, ValueMeta,
33};
34use crate::hummock::shared_buffer::TableMemoryMetrics;
35use crate::hummock::utils::range_overlap;
36use crate::hummock::value::HummockValue;
37use crate::hummock::{HummockEpoch, HummockResult};
38use crate::store::ReadOptions;
39
40#[derive(Clone, Copy, Debug, PartialEq)]
41pub enum SharedBufferValue<T> {
42    Insert(T),
43    Update(T),
44    Delete,
45}
46
47impl<T> SharedBufferValue<T> {
48    fn to_ref(&self) -> SharedBufferValue<&T> {
49        match self {
50            SharedBufferValue::Insert(val) => SharedBufferValue::Insert(val),
51            SharedBufferValue::Update(val) => SharedBufferValue::Update(val),
52            SharedBufferValue::Delete => SharedBufferValue::Delete,
53        }
54    }
55}
56
57impl<T> From<SharedBufferValue<T>> for HummockValue<T> {
58    fn from(val: SharedBufferValue<T>) -> HummockValue<T> {
59        match val {
60            SharedBufferValue::Insert(val) | SharedBufferValue::Update(val) => {
61                HummockValue::Put(val)
62            }
63            SharedBufferValue::Delete => HummockValue::Delete,
64        }
65    }
66}
67
68impl<'a, T: AsRef<[u8]>> SharedBufferValue<&'a T> {
69    pub(crate) fn to_slice(self) -> SharedBufferValue<&'a [u8]> {
70        match self {
71            SharedBufferValue::Insert(val) => SharedBufferValue::Insert(val.as_ref()),
72            SharedBufferValue::Update(val) => SharedBufferValue::Update(val.as_ref()),
73            SharedBufferValue::Delete => SharedBufferValue::Delete,
74        }
75    }
76}
77
78/// The key is `table_key`, which does not contain table id or epoch.
79pub(crate) type SharedBufferItem = (TableKey<Bytes>, SharedBufferValue<Bytes>);
80pub type SharedBufferBatchId = u64;
81
82#[derive(Debug, PartialEq)]
83pub(crate) struct SharedBufferEntry {
84    pub(crate) key: TableKey<Bytes>,
85    pub(crate) value: SharedBufferValue<Bytes>,
86}
87
88#[derive(Debug)]
89pub(crate) struct SharedBufferBatchOldValues {
90    /// Store the old values. If some, the length should be the same as `entries`. It contains empty `Bytes` when the
91    /// corresponding `value` is `Insert`, and contains the old values of `Update` and `Delete`.
92    values: Vec<Bytes>,
93    pub size: usize,
94    pub global_old_value_size: LabelGuardedIntGauge,
95}
96
97impl Drop for SharedBufferBatchOldValues {
98    fn drop(&mut self) {
99        self.global_old_value_size.sub(self.size as _);
100    }
101}
102
103impl SharedBufferBatchOldValues {
104    pub(crate) fn new(
105        values: Vec<Bytes>,
106        size: usize,
107        global_old_value_size: LabelGuardedIntGauge,
108    ) -> Self {
109        global_old_value_size.add(size as _);
110        Self {
111            values,
112            size,
113            global_old_value_size,
114        }
115    }
116
117    pub(crate) fn for_test(values: Vec<Bytes>, size: usize) -> Self {
118        Self::new(values, size, LabelGuardedIntGauge::test_int_gauge::<1>())
119    }
120}
121
122#[derive(Debug)]
123pub(crate) struct SharedBufferBatchInner {
124    entries: Vec<SharedBufferEntry>,
125    old_values: Option<SharedBufferBatchOldValues>,
126    epoch_with_gap: EpochWithGap,
127    /// Total size of all key-value items (excluding the `epoch` of value versions)
128    size: usize,
129    per_table_tracker: Arc<TableMemoryMetrics>,
130    /// For a batch created from multiple batches, this will be
131    /// the largest batch id among input batches
132    batch_id: SharedBufferBatchId,
133}
134
135impl SharedBufferBatchInner {
136    pub(crate) fn new(
137        epoch: HummockEpoch,
138        spill_offset: u16,
139        payload: Vec<SharedBufferItem>,
140        old_values: Option<SharedBufferBatchOldValues>,
141        size: usize,
142        table_metrics: Arc<TableMemoryMetrics>,
143    ) -> Self {
144        assert!(!payload.is_empty());
145        debug_assert!(payload.iter().is_sorted_by_key(|(key, _)| key));
146        if let Some(old_values) = &old_values {
147            assert_eq!(old_values.values.len(), payload.len());
148        }
149
150        let epoch_with_gap = EpochWithGap::new(epoch, spill_offset);
151        let entries = payload
152            .into_iter()
153            .map(|(key, value)| SharedBufferEntry { key, value })
154            .collect();
155
156        let batch_id = SHARED_BUFFER_BATCH_ID_GENERATOR.fetch_add(1, Relaxed);
157
158        table_metrics.inc_imm(size);
159
160        SharedBufferBatchInner {
161            entries,
162            old_values,
163            epoch_with_gap,
164            size,
165            per_table_tracker: table_metrics,
166            batch_id,
167        }
168    }
169
170    /// Return `None` if cannot find a visible version
171    /// Return `HummockValue::Delete` if the key has been deleted by some epoch <= `read_epoch`
172    fn get_value<'a>(
173        &'a self,
174        table_key: TableKey<&[u8]>,
175        read_epoch: HummockEpoch,
176    ) -> Option<(HummockValue<&'a Bytes>, EpochWithGap)> {
177        // Perform binary search on table key to find the corresponding entry
178        if let Ok(i) = self
179            .entries
180            .binary_search_by(|m| (m.key.as_ref()).cmp(*table_key))
181        {
182            let SharedBufferEntry { key, value } = &self.entries[i];
183            debug_assert_eq!(key.as_ref(), *table_key);
184            if read_epoch >= self.epoch_with_gap.pure_epoch() {
185                return Some((value.to_ref().into(), self.epoch_with_gap));
186            }
187        }
188
189        None
190    }
191}
192
193impl Drop for SharedBufferBatchInner {
194    fn drop(&mut self) {
195        self.per_table_tracker.dec_imm(self.size);
196    }
197}
198
199pub static SHARED_BUFFER_BATCH_ID_GENERATOR: LazyLock<AtomicU64> =
200    LazyLock::new(|| AtomicU64::new(0));
201
202/// A write batch stored in the shared buffer.
203#[derive(Clone, Debug)]
204pub struct SharedBufferBatch {
205    pub(crate) inner: Arc<SharedBufferBatchInner>,
206    pub table_id: TableId,
207}
208
209impl SharedBufferBatch {
210    pub fn for_test(
211        sorted_items: Vec<SharedBufferItem>,
212        epoch: HummockEpoch,
213        table_id: TableId,
214    ) -> Self {
215        Self::for_test_inner(sorted_items, None, epoch, table_id)
216    }
217
218    pub fn for_test_with_old_values(
219        sorted_items: Vec<SharedBufferItem>,
220        old_values: Vec<Bytes>,
221        epoch: HummockEpoch,
222        table_id: TableId,
223    ) -> Self {
224        Self::for_test_inner(sorted_items, Some(old_values), epoch, table_id)
225    }
226
227    fn for_test_inner(
228        sorted_items: Vec<SharedBufferItem>,
229        old_values: Option<Vec<Bytes>>,
230        epoch: HummockEpoch,
231        table_id: TableId,
232    ) -> Self {
233        let (size, old_value_size) = Self::measure_batch_size(&sorted_items, old_values.as_deref());
234
235        let old_values = old_values
236            .map(|old_values| SharedBufferBatchOldValues::for_test(old_values, old_value_size));
237
238        Self {
239            inner: Arc::new(SharedBufferBatchInner::new(
240                epoch,
241                0,
242                sorted_items,
243                old_values,
244                size,
245                TableMemoryMetrics::for_test(),
246            )),
247            table_id,
248        }
249    }
250
251    /// Return (total size, old value size or 0)
252    pub fn measure_batch_size(
253        batch_items: &[SharedBufferItem],
254        old_values: Option<&[Bytes]>,
255    ) -> (usize, usize) {
256        let old_value_size = old_values
257            .iter()
258            .flat_map(|slice| slice.iter().map(|value| size_of_val(value) + value.len()))
259            .sum::<usize>();
260        // size = Sum(length of full key + length of user value)
261        let kv_size = batch_items
262            .iter()
263            .map(|(k, v)| {
264                k.len() + {
265                    match v {
266                        SharedBufferValue::Insert(val) | SharedBufferValue::Update(val) => {
267                            val.len()
268                        }
269                        SharedBufferValue::Delete => 0,
270                    }
271                }
272            })
273            .sum::<usize>();
274        (kv_size + old_value_size, old_value_size)
275    }
276
277    pub fn filter<R, B>(&self, table_id: TableId, table_key_range: &R) -> bool
278    where
279        R: RangeBounds<TableKey<B>>,
280        B: AsRef<[u8]>,
281    {
282        let left = table_key_range
283            .start_bound()
284            .as_ref()
285            .map(|key| TableKey(key.0.as_ref()));
286        let right = table_key_range
287            .end_bound()
288            .as_ref()
289            .map(|key| TableKey(key.0.as_ref()));
290        self.table_id == table_id
291            && range_overlap(
292                &(left, right),
293                &self.start_table_key(),
294                Included(&self.end_table_key()),
295            )
296    }
297
298    pub fn table_id(&self) -> TableId {
299        self.table_id
300    }
301
302    pub fn key_count(&self) -> usize {
303        self.inner.entries.len()
304    }
305
306    pub fn value_count(&self) -> usize {
307        self.inner.entries.len()
308    }
309
310    pub fn has_old_value(&self) -> bool {
311        self.inner.old_values.is_some()
312    }
313
314    pub fn get<'a>(
315        &'a self,
316        table_key: TableKey<&[u8]>,
317        read_epoch: HummockEpoch,
318        _read_options: &ReadOptions,
319    ) -> Option<(HummockValue<&'a Bytes>, EpochWithGap)> {
320        self.inner.get_value(table_key, read_epoch)
321    }
322
323    pub fn range_exists(&self, table_key_range: &TableKeyRange) -> bool {
324        self.inner
325            .entries
326            .binary_search_by(|m| {
327                let key = &m.key;
328                let too_left = match &table_key_range.0 {
329                    std::ops::Bound::Included(range_start) => range_start.as_ref() > key.as_ref(),
330                    std::ops::Bound::Excluded(range_start) => range_start.as_ref() >= key.as_ref(),
331                    std::ops::Bound::Unbounded => false,
332                };
333                if too_left {
334                    return Ordering::Less;
335                }
336
337                let too_right = match &table_key_range.1 {
338                    std::ops::Bound::Included(range_end) => range_end.as_ref() < key.as_ref(),
339                    std::ops::Bound::Excluded(range_end) => range_end.as_ref() <= key.as_ref(),
340                    std::ops::Bound::Unbounded => false,
341                };
342                if too_right {
343                    return Ordering::Greater;
344                }
345
346                Ordering::Equal
347            })
348            .is_ok()
349    }
350
351    pub fn into_directed_iter<D: HummockIteratorDirection, const IS_NEW_VALUE: bool>(
352        self,
353    ) -> SharedBufferBatchIterator<D, IS_NEW_VALUE> {
354        SharedBufferBatchIterator::<D, IS_NEW_VALUE>::new(self.inner, self.table_id)
355    }
356
357    pub fn into_old_value_iter(self) -> SharedBufferBatchIterator<Forward, false> {
358        self.into_directed_iter()
359    }
360
361    pub fn into_forward_iter(self) -> SharedBufferBatchIterator<Forward> {
362        self.into_directed_iter()
363    }
364
365    pub fn into_backward_iter(self) -> SharedBufferBatchIterator<Backward> {
366        self.into_directed_iter()
367    }
368
369    #[inline(always)]
370    pub fn start_table_key(&self) -> TableKey<&[u8]> {
371        TableKey(self.inner.entries.first().expect("non-empty").key.as_ref())
372    }
373
374    #[inline(always)]
375    pub fn end_table_key(&self) -> TableKey<&[u8]> {
376        TableKey(self.inner.entries.last().expect("non-empty").key.as_ref())
377    }
378
379    /// return inclusive left endpoint, which means that all data in this batch should be larger or
380    /// equal than this key.
381    pub fn start_user_key(&self) -> UserKey<&[u8]> {
382        UserKey::new(self.table_id, self.start_table_key())
383    }
384
385    pub fn size(&self) -> usize {
386        self.inner.size
387    }
388
389    pub fn batch_id(&self) -> SharedBufferBatchId {
390        self.inner.batch_id
391    }
392
393    pub fn epoch(&self) -> HummockEpoch {
394        self.inner.epoch_with_gap.pure_epoch()
395    }
396
397    pub(crate) fn build_shared_buffer_batch(
398        epoch: HummockEpoch,
399        spill_offset: u16,
400        sorted_items: Vec<SharedBufferItem>,
401        old_values: Option<SharedBufferBatchOldValues>,
402        size: usize,
403        table_id: TableId,
404        table_metrics: Arc<TableMemoryMetrics>,
405    ) -> Self {
406        let inner = SharedBufferBatchInner::new(
407            epoch,
408            spill_offset,
409            sorted_items,
410            old_values,
411            size,
412            table_metrics,
413        );
414        SharedBufferBatch {
415            inner: Arc::new(inner),
416            table_id,
417        }
418    }
419
420    #[cfg(any(test, feature = "test"))]
421    pub fn build_shared_buffer_batch_for_test(
422        epoch: HummockEpoch,
423        spill_offset: u16,
424        sorted_items: Vec<SharedBufferItem>,
425        size: usize,
426        table_id: TableId,
427    ) -> Self {
428        let inner = SharedBufferBatchInner::new(
429            epoch,
430            spill_offset,
431            sorted_items,
432            None,
433            size,
434            TableMemoryMetrics::for_test(),
435        );
436        SharedBufferBatch {
437            inner: Arc::new(inner),
438            table_id,
439        }
440    }
441}
442
443/// Iterate all the items in the shared buffer batch
444/// If there are multiple versions of a key, the iterator will return all versions
445pub struct SharedBufferBatchIterator<D: HummockIteratorDirection, const IS_NEW_VALUE: bool = true> {
446    inner: Arc<SharedBufferBatchInner>,
447    /// The index of the current entry in the payload
448    current_entry_idx: usize,
449    table_id: TableId,
450    _phantom: PhantomData<D>,
451}
452
453impl<D: HummockIteratorDirection, const IS_NEW_VALUE: bool>
454    SharedBufferBatchIterator<D, IS_NEW_VALUE>
455{
456    pub(crate) fn new(inner: Arc<SharedBufferBatchInner>, table_id: TableId) -> Self {
457        if !IS_NEW_VALUE {
458            assert!(
459                inner.old_values.is_some(),
460                "create old value iter with no old value: {:?}",
461                table_id
462            );
463        }
464        Self {
465            inner,
466            current_entry_idx: 0,
467            table_id,
468            _phantom: Default::default(),
469        }
470    }
471
472    fn is_valid_entry_idx(&self) -> bool {
473        self.current_entry_idx < self.inner.entries.len()
474    }
475
476    fn invalidate(&mut self) {
477        self.current_entry_idx = self.inner.entries.len();
478    }
479
480    fn advance_to_next_entry(&mut self) {
481        debug_assert!(self.is_valid_entry_idx());
482        match D::direction() {
483            DirectionEnum::Forward => {
484                self.current_entry_idx += 1;
485            }
486            DirectionEnum::Backward => {
487                if self.current_entry_idx == 0 {
488                    self.invalidate();
489                } else {
490                    self.current_entry_idx -= 1;
491                }
492            }
493        }
494    }
495
496    fn assert_valid_idx(&self) {
497        debug_assert!(self.is_valid_entry_idx());
498        if !IS_NEW_VALUE {
499            debug_assert!(!matches!(
500                self.inner.entries[self.current_entry_idx].value,
501                SharedBufferValue::Insert(_)
502            ));
503        }
504    }
505
506    fn advance_until_valid_old_value(&mut self) {
507        debug_assert!(!IS_NEW_VALUE);
508        while self.is_valid_entry_idx()
509            && matches!(
510                self.inner.entries[self.current_entry_idx].value,
511                SharedBufferValue::Insert(_)
512            )
513        {
514            self.advance_to_next_entry();
515        }
516    }
517}
518
519impl<D: HummockIteratorDirection, const IS_NEW_VALUE: bool> HummockIterator
520    for SharedBufferBatchIterator<D, IS_NEW_VALUE>
521{
522    type Direction = D;
523
524    async fn next(&mut self) -> HummockResult<()> {
525        self.advance_to_next_entry();
526        if !IS_NEW_VALUE {
527            self.advance_until_valid_old_value();
528        }
529        Ok(())
530    }
531
532    fn key(&self) -> FullKey<&[u8]> {
533        self.assert_valid_idx();
534        let entry = &self.inner.entries[self.current_entry_idx];
535        FullKey::new_with_gap_epoch(
536            self.table_id,
537            TableKey(entry.key.as_ref()),
538            self.inner.epoch_with_gap,
539        )
540    }
541
542    fn value(&self) -> HummockValue<&[u8]> {
543        self.assert_valid_idx();
544        if IS_NEW_VALUE {
545            self.inner.entries[self.current_entry_idx]
546                .value
547                .to_ref()
548                .to_slice()
549                .into()
550        } else {
551            HummockValue::put(
552                self.inner.old_values.as_ref().unwrap().values[self.current_entry_idx].as_ref(),
553            )
554        }
555    }
556
557    fn is_valid(&self) -> bool {
558        self.is_valid_entry_idx()
559    }
560
561    async fn rewind(&mut self) -> HummockResult<()> {
562        match D::direction() {
563            DirectionEnum::Forward => {
564                self.current_entry_idx = 0;
565            }
566            DirectionEnum::Backward => {
567                self.current_entry_idx = self.inner.entries.len() - 1;
568            }
569        };
570        if !IS_NEW_VALUE {
571            self.advance_until_valid_old_value();
572        }
573        Ok(())
574    }
575
576    async fn seek<'a>(&'a mut self, key: FullKey<&'a [u8]>) -> HummockResult<()> {
577        match key.user_key.table_id.cmp(&self.table_id) {
578            Ordering::Less => {
579                match D::direction() {
580                    DirectionEnum::Forward => {
581                        // seek key table id < batch table id, so seek to beginning
582                        self.rewind().await?;
583                        return Ok(());
584                    }
585                    DirectionEnum::Backward => {
586                        self.invalidate();
587                        return Ok(());
588                    }
589                };
590            }
591            Ordering::Greater => {
592                match D::direction() {
593                    DirectionEnum::Forward => {
594                        self.invalidate();
595                        return Ok(());
596                    }
597                    DirectionEnum::Backward => {
598                        // seek key table id > batch table id, so seek to end
599                        self.rewind().await?;
600                        return Ok(());
601                    }
602                };
603            }
604            Ordering::Equal => (),
605        }
606        // Perform binary search on table key because the items in SharedBufferBatch is ordered
607        // by table key.
608        let partition_point = self
609            .inner
610            .entries
611            .binary_search_by(|probe| probe.key.as_ref().cmp(*key.user_key.table_key));
612        match partition_point {
613            Ok(i) => {
614                self.current_entry_idx = i;
615                // Epoch order --------->
616                // Epochs:  epoch300 epoch200 epoch100
617                // Forward: ------------------------->
618                // Backward:<-------------------------
619                // Assume self.inner.epoch_with_gap is epoch200
620                let skip_on_epoch = match D::direction() {
621                    DirectionEnum::Forward => {
622                        // should advance when key.epoch_with_gap is epoch100
623                        key.epoch_with_gap < self.inner.epoch_with_gap
624                    }
625                    DirectionEnum::Backward => {
626                        // should advance when key.epoch_with_gap is epoch300
627                        key.epoch_with_gap > self.inner.epoch_with_gap
628                    }
629                };
630                if skip_on_epoch {
631                    self.advance_to_next_entry()
632                }
633            }
634            Err(i) => match D::direction() {
635                DirectionEnum::Forward => {
636                    self.current_entry_idx = i;
637                }
638                DirectionEnum::Backward => {
639                    if i == 0 {
640                        self.invalidate();
641                    } else {
642                        self.current_entry_idx = i - 1;
643                    }
644                }
645            },
646        };
647        if !IS_NEW_VALUE {
648            self.advance_until_valid_old_value();
649        }
650        Ok(())
651    }
652
653    fn collect_local_statistic(&self, _stats: &mut crate::monitor::StoreLocalStatistic) {}
654
655    fn value_meta(&self) -> ValueMeta {
656        ValueMeta::default()
657    }
658}
659
660#[cfg(test)]
661mod tests {
662    use std::ops::Bound::Excluded;
663
664    use itertools::{Itertools, zip_eq};
665    use risingwave_common::util::epoch::{EpochExt, test_epoch};
666    use risingwave_hummock_sdk::key::map_table_key_range;
667
668    use super::*;
669    use crate::hummock::iterator::test_utils::{
670        iterator_test_key_of_epoch, iterator_test_table_key_of, transform_shared_buffer,
671    };
672
673    fn to_hummock_value_batch(
674        items: Vec<(Vec<u8>, SharedBufferValue<Bytes>)>,
675    ) -> Vec<(Vec<u8>, HummockValue<Bytes>)> {
676        items.into_iter().map(|(k, v)| (k, v.into())).collect()
677    }
678
679    #[tokio::test]
680    async fn test_shared_buffer_batch_basic() {
681        let epoch = test_epoch(1);
682        let shared_buffer_items: Vec<(Vec<u8>, SharedBufferValue<Bytes>)> = vec![
683            (
684                iterator_test_table_key_of(0),
685                SharedBufferValue::Insert(Bytes::from("value1")),
686            ),
687            (
688                iterator_test_table_key_of(1),
689                SharedBufferValue::Insert(Bytes::from("value1")),
690            ),
691            (
692                iterator_test_table_key_of(2),
693                SharedBufferValue::Insert(Bytes::from("value1")),
694            ),
695        ];
696        let shared_buffer_batch = SharedBufferBatch::for_test(
697            transform_shared_buffer(shared_buffer_items.clone()),
698            epoch,
699            Default::default(),
700        );
701        let shared_buffer_items = to_hummock_value_batch(shared_buffer_items);
702
703        // Sketch
704        assert_eq!(
705            *shared_buffer_batch.start_table_key(),
706            shared_buffer_items[0].0
707        );
708        assert_eq!(
709            *shared_buffer_batch.end_table_key(),
710            shared_buffer_items[2].0
711        );
712
713        // Point lookup
714        for (k, v) in &shared_buffer_items {
715            assert_eq!(
716                shared_buffer_batch
717                    .get(TableKey(k.as_slice()), epoch, &ReadOptions::default())
718                    .unwrap()
719                    .0
720                    .as_slice(),
721                v.as_slice()
722            );
723        }
724        assert_eq!(
725            shared_buffer_batch.get(
726                TableKey(iterator_test_table_key_of(3).as_slice()),
727                epoch,
728                &ReadOptions::default()
729            ),
730            None
731        );
732        assert_eq!(
733            shared_buffer_batch.get(
734                TableKey(iterator_test_table_key_of(4).as_slice()),
735                epoch,
736                &ReadOptions::default()
737            ),
738            None
739        );
740
741        // Forward iterator
742        let mut iter = shared_buffer_batch.clone().into_forward_iter();
743        iter.rewind().await.unwrap();
744        let mut output = vec![];
745        while iter.is_valid() {
746            output.push((
747                iter.key().user_key.table_key.to_vec(),
748                iter.value().to_bytes(),
749            ));
750            iter.next().await.unwrap();
751        }
752        assert_eq!(output, shared_buffer_items);
753
754        // Backward iterator
755        let mut backward_iter = shared_buffer_batch.clone().into_backward_iter();
756        backward_iter.rewind().await.unwrap();
757        let mut output = vec![];
758        while backward_iter.is_valid() {
759            output.push((
760                backward_iter.key().user_key.table_key.to_vec(),
761                backward_iter.value().to_bytes(),
762            ));
763            backward_iter.next().await.unwrap();
764        }
765        output.reverse();
766        assert_eq!(output, shared_buffer_items);
767    }
768
769    #[tokio::test]
770    async fn test_shared_buffer_batch_seek() {
771        let epoch = test_epoch(1);
772        let shared_buffer_items = vec![
773            (
774                iterator_test_table_key_of(1),
775                SharedBufferValue::Insert(Bytes::from("value1")),
776            ),
777            (
778                iterator_test_table_key_of(2),
779                SharedBufferValue::Insert(Bytes::from("value2")),
780            ),
781            (
782                iterator_test_table_key_of(3),
783                SharedBufferValue::Insert(Bytes::from("value3")),
784            ),
785        ];
786        let shared_buffer_batch = SharedBufferBatch::for_test(
787            transform_shared_buffer(shared_buffer_items.clone()),
788            epoch,
789            Default::default(),
790        );
791        let shared_buffer_items = to_hummock_value_batch(shared_buffer_items);
792
793        // FORWARD: Seek to a key < 1st key, expect all three items to return
794        let mut iter = shared_buffer_batch.clone().into_forward_iter();
795        iter.seek(iterator_test_key_of_epoch(0, epoch).to_ref())
796            .await
797            .unwrap();
798        for item in &shared_buffer_items {
799            assert!(iter.is_valid());
800            assert_eq!(*iter.key().user_key.table_key, item.0);
801            assert_eq!(iter.value(), item.1.as_slice());
802            iter.next().await.unwrap();
803        }
804        assert!(!iter.is_valid());
805
806        // FORWARD: Seek to a key > the last key, expect no items to return
807        let mut iter = shared_buffer_batch.clone().into_forward_iter();
808        iter.seek(iterator_test_key_of_epoch(4, epoch).to_ref())
809            .await
810            .unwrap();
811        assert!(!iter.is_valid());
812
813        // FORWARD: Seek to 2nd key with current epoch, expect last two items to return
814        let mut iter = shared_buffer_batch.clone().into_forward_iter();
815        iter.seek(iterator_test_key_of_epoch(2, epoch).to_ref())
816            .await
817            .unwrap();
818        for item in &shared_buffer_items[1..] {
819            assert!(iter.is_valid());
820            assert_eq!(*iter.key().user_key.table_key, item.0);
821            assert_eq!(iter.value(), item.1.as_slice());
822            iter.next().await.unwrap();
823        }
824        assert!(!iter.is_valid());
825
826        // FORWARD: Seek to 2nd key with future epoch, expect last two items to return
827        let mut iter = shared_buffer_batch.clone().into_forward_iter();
828        iter.seek(iterator_test_key_of_epoch(2, test_epoch(2)).to_ref())
829            .await
830            .unwrap();
831        for item in &shared_buffer_items[1..] {
832            assert!(iter.is_valid());
833            assert_eq!(*iter.key().user_key.table_key, item.0.as_slice());
834            assert_eq!(iter.value(), item.1.as_slice());
835            iter.next().await.unwrap();
836        }
837        assert!(!iter.is_valid());
838
839        // FORWARD: Seek to 2nd key with old epoch, expect last item to return
840        let mut iter = shared_buffer_batch.clone().into_forward_iter();
841        iter.seek(iterator_test_key_of_epoch(2, test_epoch(0)).to_ref())
842            .await
843            .unwrap();
844        let item = shared_buffer_items.last().unwrap();
845        assert!(iter.is_valid());
846        assert_eq!(*iter.key().user_key.table_key, item.0.as_slice());
847        assert_eq!(iter.value(), item.1.as_slice());
848        iter.next().await.unwrap();
849        assert!(!iter.is_valid());
850
851        // BACKWARD: Seek to a key < 1st key, expect no items to return
852        let mut iter = shared_buffer_batch.clone().into_backward_iter();
853        iter.seek(iterator_test_key_of_epoch(0, epoch).to_ref())
854            .await
855            .unwrap();
856        assert!(!iter.is_valid());
857
858        // BACKWARD: Seek to a key > the last key, expect all items to return
859        let mut iter = shared_buffer_batch.clone().into_backward_iter();
860        iter.seek(iterator_test_key_of_epoch(4, epoch).to_ref())
861            .await
862            .unwrap();
863        for item in shared_buffer_items.iter().rev() {
864            assert!(iter.is_valid());
865            assert_eq!(*iter.key().user_key.table_key, item.0.as_slice());
866            assert_eq!(iter.value(), item.1.as_slice());
867            iter.next().await.unwrap();
868        }
869        assert!(!iter.is_valid());
870
871        // BACKWARD: Seek to 2nd key with current epoch, expect first two items to return
872        let mut iter = shared_buffer_batch.clone().into_backward_iter();
873        iter.seek(iterator_test_key_of_epoch(2, epoch).to_ref())
874            .await
875            .unwrap();
876        for item in shared_buffer_items[0..=1].iter().rev() {
877            assert!(iter.is_valid());
878            assert_eq!(*iter.key().user_key.table_key, item.0.as_slice());
879            assert_eq!(iter.value(), item.1.as_slice());
880            iter.next().await.unwrap();
881        }
882        assert!(!iter.is_valid());
883
884        // BACKWARD: Seek to 2nd key with old epoch, expect first two item to return
885        let mut iter = shared_buffer_batch.clone().into_backward_iter();
886        iter.seek(iterator_test_key_of_epoch(2, epoch.prev_epoch()).to_ref())
887            .await
888            .unwrap();
889        for item in shared_buffer_items[0..=1].iter().rev() {
890            assert!(iter.is_valid());
891            assert_eq!(*iter.key().user_key.table_key, item.0.as_slice());
892            assert_eq!(iter.value(), item.1.as_slice());
893            iter.next().await.unwrap();
894        }
895        assert!(!iter.is_valid());
896
897        // BACKWARD: Seek to 2nd key with future epoch, expect first item to return
898        let mut iter = shared_buffer_batch.clone().into_backward_iter();
899        iter.seek(iterator_test_key_of_epoch(2, epoch.next_epoch()).to_ref())
900            .await
901            .unwrap();
902        assert!(iter.is_valid());
903        let item = shared_buffer_items.first().unwrap();
904        assert_eq!(*iter.key().user_key.table_key, item.0.as_slice());
905        assert_eq!(iter.value(), item.1.as_slice());
906        iter.next().await.unwrap();
907        assert!(!iter.is_valid());
908    }
909
910    #[tokio::test]
911    async fn test_shared_buffer_batch_old_value_iter() {
912        let epoch = test_epoch(1);
913        let key_values = vec![
914            (
915                iterator_test_table_key_of(1),
916                SharedBufferValue::Insert(Bytes::from("value1")),
917            ),
918            (
919                iterator_test_table_key_of(2),
920                SharedBufferValue::Update(Bytes::from("value2")),
921            ),
922            (
923                iterator_test_table_key_of(3),
924                SharedBufferValue::Insert(Bytes::from("value3")),
925            ),
926            (iterator_test_table_key_of(4), SharedBufferValue::Delete),
927        ];
928        let old_values = vec![
929            Bytes::new(),
930            Bytes::from("old_value2"),
931            Bytes::new(),
932            Bytes::from("old_value4"),
933        ];
934        let shared_buffer_batch = SharedBufferBatch::for_test_with_old_values(
935            transform_shared_buffer(key_values.clone()),
936            old_values.clone(),
937            epoch,
938            Default::default(),
939        );
940        let shared_buffer_items = to_hummock_value_batch(key_values.clone());
941        let expected_old_value_iter_items = zip_eq(&key_values, &old_values)
942            .filter(|((_, new_value), _)| !matches!(new_value, SharedBufferValue::Insert(_)))
943            .map(|((key, _), old_value)| (key.clone(), HummockValue::Put(old_value)))
944            .collect_vec();
945
946        let mut iter = shared_buffer_batch.clone().into_old_value_iter();
947        iter.rewind().await.unwrap();
948        for item in &expected_old_value_iter_items {
949            assert!(iter.is_valid());
950            assert_eq!(*iter.key().user_key.table_key, item.0);
951            assert_eq!(iter.value(), item.1.as_slice());
952            iter.next().await.unwrap();
953        }
954        assert!(!iter.is_valid());
955
956        // FORWARD: Seek to a key < 1st key, expect all three items to return
957        let mut iter = shared_buffer_batch.clone().into_forward_iter();
958        iter.seek(iterator_test_key_of_epoch(0, epoch).to_ref())
959            .await
960            .unwrap();
961        for item in &shared_buffer_items {
962            assert!(iter.is_valid());
963            assert_eq!(*iter.key().user_key.table_key, item.0);
964            assert_eq!(iter.value(), item.1.as_slice());
965            iter.next().await.unwrap();
966        }
967        assert!(!iter.is_valid());
968
969        let mut iter = shared_buffer_batch.clone().into_old_value_iter();
970        iter.seek(iterator_test_key_of_epoch(0, epoch).to_ref())
971            .await
972            .unwrap();
973        for item in &expected_old_value_iter_items {
974            assert!(iter.is_valid());
975            assert_eq!(*iter.key().user_key.table_key, item.0);
976            assert_eq!(iter.value(), item.1.as_slice());
977            iter.next().await.unwrap();
978        }
979        assert!(!iter.is_valid());
980
981        // FORWARD: Seek to a key > the last key, expect no items to return
982        let mut iter = shared_buffer_batch.clone().into_forward_iter();
983        iter.seek(iterator_test_key_of_epoch(5, epoch).to_ref())
984            .await
985            .unwrap();
986        assert!(!iter.is_valid());
987
988        let mut iter = shared_buffer_batch.clone().into_old_value_iter();
989        iter.seek(iterator_test_key_of_epoch(5, epoch).to_ref())
990            .await
991            .unwrap();
992        assert!(!iter.is_valid());
993
994        // FORWARD: Seek to 2nd key with current epoch, expect last two items to return
995        let mut iter = shared_buffer_batch.clone().into_forward_iter();
996        iter.seek(iterator_test_key_of_epoch(2, epoch).to_ref())
997            .await
998            .unwrap();
999        for item in &shared_buffer_items[1..] {
1000            assert!(iter.is_valid());
1001            assert_eq!(*iter.key().user_key.table_key, item.0);
1002            assert_eq!(iter.value(), item.1.as_slice());
1003            iter.next().await.unwrap();
1004        }
1005        assert!(!iter.is_valid());
1006
1007        let mut iter = shared_buffer_batch.clone().into_old_value_iter();
1008        iter.seek(iterator_test_key_of_epoch(2, epoch).to_ref())
1009            .await
1010            .unwrap();
1011        for item in &expected_old_value_iter_items {
1012            assert!(iter.is_valid());
1013            assert_eq!(*iter.key().user_key.table_key, item.0);
1014            assert_eq!(iter.value(), item.1.as_slice());
1015            iter.next().await.unwrap();
1016        }
1017        assert!(!iter.is_valid());
1018
1019        // FORWARD: Seek to 2nd key with future epoch, expect last two items to return
1020        let mut iter = shared_buffer_batch.clone().into_forward_iter();
1021        iter.seek(iterator_test_key_of_epoch(2, epoch.next_epoch()).to_ref())
1022            .await
1023            .unwrap();
1024        for item in &shared_buffer_items[1..] {
1025            assert!(iter.is_valid());
1026            assert_eq!(*iter.key().user_key.table_key, item.0.as_slice());
1027            assert_eq!(iter.value(), item.1.as_slice());
1028            iter.next().await.unwrap();
1029        }
1030        assert!(!iter.is_valid());
1031
1032        let mut iter = shared_buffer_batch.clone().into_old_value_iter();
1033        iter.seek(iterator_test_key_of_epoch(2, epoch.next_epoch()).to_ref())
1034            .await
1035            .unwrap();
1036        for item in &expected_old_value_iter_items {
1037            assert!(iter.is_valid());
1038            assert_eq!(*iter.key().user_key.table_key, item.0);
1039            assert_eq!(iter.value(), item.1.as_slice());
1040            iter.next().await.unwrap();
1041        }
1042        assert!(!iter.is_valid());
1043
1044        let mut iter = shared_buffer_batch.clone().into_old_value_iter();
1045        iter.seek(iterator_test_key_of_epoch(2, epoch.prev_epoch()).to_ref())
1046            .await
1047            .unwrap();
1048        for item in &expected_old_value_iter_items[1..] {
1049            assert!(iter.is_valid());
1050            assert_eq!(*iter.key().user_key.table_key, item.0);
1051            assert_eq!(iter.value(), item.1.as_slice());
1052            iter.next().await.unwrap();
1053        }
1054        assert!(!iter.is_valid());
1055
1056        // FORWARD: Seek to 2nd key with old epoch, expect last item to return
1057        let mut iter = shared_buffer_batch.clone().into_forward_iter();
1058        iter.seek(iterator_test_key_of_epoch(3, epoch.prev_epoch()).to_ref())
1059            .await
1060            .unwrap();
1061        let item = shared_buffer_items.last().unwrap();
1062        assert!(iter.is_valid());
1063        assert_eq!(*iter.key().user_key.table_key, item.0.as_slice());
1064        assert_eq!(iter.value(), item.1.as_slice());
1065        iter.next().await.unwrap();
1066        assert!(!iter.is_valid());
1067
1068        // Seek to an insert key
1069        let mut iter = shared_buffer_batch.clone().into_old_value_iter();
1070        iter.seek(iterator_test_key_of_epoch(3, epoch).to_ref())
1071            .await
1072            .unwrap();
1073        for item in &expected_old_value_iter_items[1..] {
1074            assert!(iter.is_valid());
1075            assert_eq!(*iter.key().user_key.table_key, item.0);
1076            assert_eq!(iter.value(), item.1.as_slice());
1077            iter.next().await.unwrap();
1078        }
1079        assert!(!iter.is_valid());
1080    }
1081
1082    #[tokio::test]
1083    #[should_panic]
1084    async fn test_invalid_table_id() {
1085        let epoch = test_epoch(1);
1086        let shared_buffer_batch = SharedBufferBatch::for_test(vec![], epoch, Default::default());
1087        // Seeking to non-current epoch should panic
1088        let mut iter = shared_buffer_batch.into_forward_iter();
1089        iter.seek(FullKey::for_test(TableId::new(1), vec![], epoch).to_ref())
1090            .await
1091            .unwrap();
1092    }
1093
1094    #[tokio::test]
1095    async fn test_shared_buffer_batch_range_existx() {
1096        let epoch = test_epoch(1);
1097        let shared_buffer_items = vec![
1098            (
1099                Vec::from("a_1"),
1100                SharedBufferValue::Insert(Bytes::from("value1")),
1101            ),
1102            (
1103                Vec::from("a_3"),
1104                SharedBufferValue::Insert(Bytes::from("value2")),
1105            ),
1106            (
1107                Vec::from("a_5"),
1108                SharedBufferValue::Insert(Bytes::from("value3")),
1109            ),
1110            (
1111                Vec::from("b_2"),
1112                SharedBufferValue::Insert(Bytes::from("value3")),
1113            ),
1114        ];
1115        let shared_buffer_batch = SharedBufferBatch::for_test(
1116            transform_shared_buffer(shared_buffer_items),
1117            epoch,
1118            Default::default(),
1119        );
1120
1121        let range = (Included(Bytes::from("a")), Excluded(Bytes::from("b")));
1122        assert!(shared_buffer_batch.range_exists(&map_table_key_range(range)));
1123        let range = (Included(Bytes::from("a_")), Excluded(Bytes::from("b_")));
1124        assert!(shared_buffer_batch.range_exists(&map_table_key_range(range)));
1125        let range = (Included(Bytes::from("a_1")), Included(Bytes::from("a_1")));
1126        assert!(shared_buffer_batch.range_exists(&map_table_key_range(range)));
1127        let range = (Included(Bytes::from("a_1")), Included(Bytes::from("a_2")));
1128        assert!(shared_buffer_batch.range_exists(&map_table_key_range(range)));
1129        let range = (Included(Bytes::from("a_0x")), Included(Bytes::from("a_2x")));
1130        assert!(shared_buffer_batch.range_exists(&map_table_key_range(range)));
1131        let range = (Included(Bytes::from("a_")), Excluded(Bytes::from("c_")));
1132        assert!(shared_buffer_batch.range_exists(&map_table_key_range(range)));
1133        let range = (Included(Bytes::from("b_0x")), Included(Bytes::from("b_2x")));
1134        assert!(shared_buffer_batch.range_exists(&map_table_key_range(range)));
1135        let range = (Included(Bytes::from("b_2")), Excluded(Bytes::from("c_1x")));
1136        assert!(shared_buffer_batch.range_exists(&map_table_key_range(range)));
1137
1138        let range = (Included(Bytes::from("a_0")), Excluded(Bytes::from("a_1")));
1139        assert!(!shared_buffer_batch.range_exists(&map_table_key_range(range)));
1140        let range = (Included(Bytes::from("a__0")), Excluded(Bytes::from("a__5")));
1141        assert!(!shared_buffer_batch.range_exists(&map_table_key_range(range)));
1142        let range = (Included(Bytes::from("b_1")), Excluded(Bytes::from("b_2")));
1143        assert!(!shared_buffer_batch.range_exists(&map_table_key_range(range)));
1144        let range = (Included(Bytes::from("b_3")), Excluded(Bytes::from("c_1")));
1145        assert!(!shared_buffer_batch.range_exists(&map_table_key_range(range)));
1146        let range = (Included(Bytes::from("b__x")), Excluded(Bytes::from("c__x")));
1147        assert!(!shared_buffer_batch.range_exists(&map_table_key_range(range)));
1148    }
1149
1150    #[tokio::test]
1151    async fn test_shared_buffer_batch_seek_bug() {
1152        // Reproduce the bug where seek falls through to binary_search when table_id mismatch
1153        let epoch = test_epoch(1);
1154        let table_id = TableId::new(100);
1155        let shared_buffer_items: Vec<(Vec<u8>, SharedBufferValue<Bytes>)> = vec![(
1156            iterator_test_table_key_of(1), // "key_test_000...001"
1157            SharedBufferValue::Insert(Bytes::from("value1")),
1158        )];
1159        let shared_buffer_batch = SharedBufferBatch::for_test(
1160            transform_shared_buffer(shared_buffer_items.clone()),
1161            epoch,
1162            table_id,
1163        );
1164
1165        // Case 1: Seek with smaller TableId (99), but larger TableKey ("key_test_000...002").
1166        // "key_test...002" > "key_test...001".
1167        // Forward Iterator.
1168        // Expected: Should land on the first item (TableId 100 > TableId 99).
1169        // Bug description: rewinds (index 0), then binary_search("key_2") in batch (only "key_1").
1170        // "key_2" > "key_1", so index becomes 1 (end). Iterator invalid.
1171        let mut iter = shared_buffer_batch.clone().into_forward_iter();
1172        let seek_key = FullKey::for_test(
1173            TableId::new(99),
1174            iterator_test_table_key_of(2), // larger key
1175            epoch,
1176        );
1177        iter.seek(seek_key.to_ref()).await.unwrap();
1178
1179        assert!(
1180            iter.is_valid(),
1181            "Iterator should be valid when seeking with smaller table_id, even if the key part is larger"
1182        );
1183        assert_eq!(iter.key().user_key.table_id, table_id);
1184
1185        // Case 2: Seek with larger TableId (101), but smaller TableKey ("key_test_000...000").
1186        // "key_test...000" < "key_test...001".
1187        // Backward Iterator.
1188        // Expected: Should land on the last item (TableId 100 < TableId 101).
1189        // Bug description: rewinds (index valid), then binary_search("key_0") which returns Err(0) (insertion point at start).
1190        // Backward generic logic for Err(0) is `invalidate()`.
1191        let mut iter = shared_buffer_batch.clone().into_backward_iter();
1192        let seek_key = FullKey::for_test(
1193            TableId::new(101),
1194            iterator_test_table_key_of(0), // smaller key
1195            epoch,
1196        );
1197        iter.seek(seek_key.to_ref()).await.unwrap();
1198
1199        assert!(
1200            iter.is_valid(),
1201            "Iterator should be valid when seeking with larger table_id, even if the key part is smaller"
1202        );
1203        assert_eq!(iter.key().user_key.table_id, table_id);
1204    }
1205}
risingwave_storage/hummock/shared_buffer/shared_buffer_batch.rs

risingwave_storage/hummock/shared_buffer/
shared_buffer_batch.rs
