risingwave_storage/hummock/sstable/
forward_sstable_iterator.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::ops::Bound::*;
16use std::sync::Arc;
17
18use await_tree::{InstrumentAwait, SpanExt};
19use risingwave_hummock_sdk::key::FullKey;
20use risingwave_hummock_sdk::sstable_info::SstableInfo;
21use sync_point::sync_point;
22use thiserror_ext::AsReport;
23
24use super::super::{HummockResult, HummockValue};
25use crate::hummock::block_stream::{PrefetchBlockStream, PrefetchLookup};
26use crate::hummock::iterator::{Forward, HummockIterator, ValueMeta};
27use crate::hummock::sstable::SstableIteratorReadOptions;
28use crate::hummock::{BlockIterator, SstableStoreRef, TableHolder};
29use crate::monitor::StoreLocalStatistic;
30
31pub trait SstableIteratorType: HummockIterator + 'static {
32    fn create(
33        sstable: TableHolder,
34        sstable_store: SstableStoreRef,
35        read_options: Arc<SstableIteratorReadOptions>,
36        sstable_info_ref: &SstableInfo,
37    ) -> Self;
38}
39
40/// Iterates on a sstable.
41pub struct SstableIterator {
42    /// The iterator of the current block.
43    block_iter: Option<BlockIterator>,
44
45    /// Current block index.
46    cur_idx: usize,
47
48    preload_stream: Option<Box<PrefetchBlockStream>>,
49    /// Reference to the sst
50    pub sst: TableHolder,
51    preload_end_block_idx: usize,
52
53    sstable_store: SstableStoreRef,
54    stats: StoreLocalStatistic,
55    options: Arc<SstableIteratorReadOptions>,
56
57    // The readable block window is first restricted by table IDs and then by `scan_end_user_key`.
58    // Included/Excluded bounds set its exclusive block end; outer iterators still filter keys
59    // precisely inside the boundary block.
60    block_start_idx_inclusive: usize,
61    block_end_idx_exclusive: usize,
62}
63
64impl SstableIterator {
65    pub fn new(
66        sstable: TableHolder,
67        sstable_store: SstableStoreRef,
68        options: Arc<SstableIteratorReadOptions>,
69        sstable_info_ref: &SstableInfo,
70    ) -> Self {
71        let mut block_start_idx_inclusive = 0;
72        let mut block_end_idx_exclusive = sstable.meta.block_metas.len();
73        assert!(
74            !sstable_info_ref.table_ids.is_empty(),
75            "SstableIterator: SST {} (object {}) has empty table_ids",
76            sstable_info_ref.sst_id,
77            sstable_info_ref.object_id,
78        );
79        let read_table_id_range = if let Some(read_table_id) = options.read_table_id
80            && sstable_info_ref
81                .table_ids
82                .binary_search(&read_table_id)
83                .is_ok()
84        {
85            (read_table_id, read_table_id)
86        } else {
87            (
88                *sstable_info_ref.table_ids.first().unwrap(),
89                *sstable_info_ref.table_ids.last().unwrap(),
90            )
91        };
92        assert!(
93            read_table_id_range.0 <= read_table_id_range.1,
94            "invalid table id range {} - {}",
95            read_table_id_range.0,
96            read_table_id_range.1
97        );
98        let block_meta_count = sstable.meta.block_metas.len();
99        assert!(block_meta_count > 0);
100        assert!(
101            sstable.meta.block_metas[0].table_id() <= read_table_id_range.0,
102            "table id {} not found table_ids in block_meta {:?}",
103            read_table_id_range.0,
104            sstable
105                .meta
106                .block_metas
107                .iter()
108                .map(|meta| meta.table_id())
109                .collect::<Vec<_>>()
110        );
111        assert!(
112            sstable.meta.block_metas[block_meta_count - 1].table_id() >= read_table_id_range.1,
113            "table id {} not found table_ids in block_meta {:?}",
114            read_table_id_range.1,
115            sstable
116                .meta
117                .block_metas
118                .iter()
119                .map(|meta| meta.table_id())
120                .collect::<Vec<_>>()
121        );
122
123        while block_start_idx_inclusive < block_meta_count
124            && sstable.meta.block_metas[block_start_idx_inclusive].table_id()
125                < read_table_id_range.0
126        {
127            block_start_idx_inclusive += 1;
128        }
129        // We assume that the table id read must exist in the sstable, otherwise it is a fatal error.
130        assert!(
131            block_start_idx_inclusive < block_meta_count,
132            "table id {} not found table_ids in block_meta {:?}",
133            read_table_id_range.0,
134            sstable
135                .meta
136                .block_metas
137                .iter()
138                .map(|meta| meta.table_id())
139                .collect::<Vec<_>>()
140        );
141
142        while block_end_idx_exclusive > block_start_idx_inclusive
143            && sstable.meta.block_metas[block_end_idx_exclusive - 1].table_id()
144                > read_table_id_range.1
145        {
146            block_end_idx_exclusive -= 1;
147        }
148        assert!(
149            block_end_idx_exclusive > block_start_idx_inclusive,
150            "block_end_idx_exclusive {} <= block_start_idx_inclusive {} block_meta_count {}",
151            block_end_idx_exclusive,
152            block_start_idx_inclusive,
153            block_meta_count
154        );
155
156        if let Some(end_bound) = options.scan_end_user_key.as_ref() {
157            let block_metas =
158                &sstable.meta.block_metas[block_start_idx_inclusive..block_end_idx_exclusive];
159            let range_end_idx_exclusive = match end_bound {
160                Unbounded => block_metas.len(),
161                Included(end_key) => block_metas.partition_point(|block_meta| {
162                    FullKey::decode(&block_meta.smallest_key).user_key <= end_key.as_ref()
163                }),
164                Excluded(end_key) => block_metas.partition_point(|block_meta| {
165                    FullKey::decode(&block_meta.smallest_key).user_key < end_key.as_ref()
166                }),
167            };
168            block_end_idx_exclusive = block_start_idx_inclusive + range_end_idx_exclusive;
169        }
170
171        Self {
172            block_iter: None,
173            cur_idx: 0,
174            preload_stream: None,
175            sst: sstable,
176            sstable_store,
177            stats: StoreLocalStatistic::default(),
178            options,
179            preload_end_block_idx: 0,
180            block_start_idx_inclusive,
181            block_end_idx_exclusive,
182        }
183    }
184
185    fn init_block_prefetch_range(&mut self, start_idx: usize) {
186        assert!(
187            start_idx >= self.block_start_idx_inclusive && start_idx < self.block_end_idx_exclusive
188        );
189
190        self.preload_end_block_idx = 0;
191        if !self.options.prefetch {
192            return;
193        }
194
195        // Prefetch never extends beyond the block window already clipped by the scan hard bound.
196        if start_idx + 1 < self.block_end_idx_exclusive {
197            self.preload_end_block_idx = self.block_end_idx_exclusive;
198        }
199    }
200
201    /// Seeks to a block, and then seeks to the key if `seek_key` is given.
202    async fn seek_idx(
203        &mut self,
204        idx: usize,
205        seek_key: Option<FullKey<&[u8]>>,
206    ) -> HummockResult<()> {
207        tracing::debug!(
208            target: "events::storage::sstable::block_seek",
209            "table iterator seek: sstable_object_id = {}, block_id = {}",
210            self.sst.id,
211            idx,
212        );
213
214        // When all data are in block cache, it is highly possible that this iterator will stay on a
215        // worker thread for a full time. Therefore, we use tokio's unstable API consume_budget to
216        // do cooperative scheduling.
217        tokio::task::consume_budget().await;
218
219        if idx >= self.block_end_idx_exclusive {
220            self.block_iter = None;
221            return Ok(());
222        }
223        let mut prefetched_block = None;
224        let should_prefetch = match self
225            .preload_stream
226            .as_mut()
227            .map(|stream| stream.take_block(idx))
228        {
229            Some(PrefetchLookup::Hit(block)) => {
230                prefetched_block = Some(block);
231                false
232            }
233            // Preserve the queue for subsequent forward reads after a backward seek.
234            Some(PrefetchLookup::BeforeStart) => false,
235            Some(PrefetchLookup::Exhausted) | None => {
236                // Release the old tracker before checking the prefetch memory budget again.
237                self.preload_stream = None;
238                true
239            }
240        };
241
242        // Try at most once per seek. Even an empty prefetch result must reach the block-get
243        // fallback instead of repeatedly recreating a stream without making progress.
244        if should_prefetch && idx + 1 < self.preload_end_block_idx {
245            match self
246                .sstable_store
247                .prefetch_blocks(
248                    &self.sst,
249                    idx,
250                    self.preload_end_block_idx,
251                    self.options.cache_policy,
252                    &mut self.stats,
253                )
254                .instrument_await("prefetch_blocks".verbose())
255                .await
256            {
257                Ok(mut stream) => {
258                    if let PrefetchLookup::Hit(block) = stream.take_block(idx) {
259                        prefetched_block = Some(block);
260                        self.preload_stream = Some(stream);
261                    }
262                }
263                Err(e) => {
264                    tracing::warn!(error = %e.as_report(), "failed to prefetch data, fall back to block get");
265                }
266            }
267        }
268        let block = match prefetched_block {
269            Some(block) => block,
270            None => {
271                self.sstable_store
272                    .get(&self.sst, idx, self.options.cache_policy, &mut self.stats)
273                    .await?
274            }
275        };
276        self.block_iter = Some(BlockIterator::new(block));
277        let block_iter = self.block_iter.as_mut().unwrap();
278        if let Some(key) = seek_key {
279            block_iter.seek(key);
280        } else {
281            block_iter.seek_to_first();
282        }
283
284        self.cur_idx = idx;
285
286        Ok(())
287    }
288
289    fn calculate_block_idx_by_key(&self, key: FullKey<&[u8]>) -> usize {
290        self.block_start_idx_inclusive
291            + self.sst.meta.block_metas
292                [self.block_start_idx_inclusive..self.block_end_idx_exclusive]
293                .partition_point(|block_meta| {
294                    // compare by version comparator
295                    // Note: we are comparing against the `smallest_key` of the `block`, thus the
296                    // partition point should be `prev(<=)` instead of `<`.
297                    FullKey::decode(&block_meta.smallest_key).le(&key)
298                })
299                .saturating_sub(1) // considering the boundary of 0
300    }
301}
302
303impl HummockIterator for SstableIterator {
304    type Direction = Forward;
305
306    async fn next(&mut self) -> HummockResult<()> {
307        self.stats.total_key_count += 1;
308        let block_iter = self.block_iter.as_mut().expect("no block iter");
309        if !block_iter.try_next() {
310            // seek to next block
311            self.seek_idx(self.cur_idx + 1, None).await?;
312        }
313
314        Ok(())
315    }
316
317    fn key(&self) -> FullKey<&[u8]> {
318        self.block_iter.as_ref().expect("no block iter").key()
319    }
320
321    fn value(&self) -> HummockValue<&[u8]> {
322        let raw_value = self.block_iter.as_ref().expect("no block iter").value();
323
324        HummockValue::from_slice(raw_value).expect("decode error")
325    }
326
327    fn is_valid(&self) -> bool {
328        self.block_iter.as_ref().is_some_and(|i| i.is_valid())
329    }
330
331    async fn rewind(&mut self) -> HummockResult<()> {
332        if self.block_start_idx_inclusive >= self.block_end_idx_exclusive {
333            self.block_iter = None;
334            return Ok(());
335        }
336        self.init_block_prefetch_range(self.block_start_idx_inclusive);
337        // seek_idx will update the current block iter state
338        self.seek_idx(self.block_start_idx_inclusive, None).await?;
339        Ok(())
340    }
341
342    async fn seek<'a>(&'a mut self, key: FullKey<&'a [u8]>) -> HummockResult<()> {
343        if self.block_start_idx_inclusive >= self.block_end_idx_exclusive {
344            self.block_iter = None;
345            return Ok(());
346        }
347        let block_idx = self.calculate_block_idx_by_key(key);
348        self.init_block_prefetch_range(block_idx);
349
350        self.seek_idx(block_idx, Some(key)).await?;
351        if !self.is_valid() {
352            // seek to next block
353            sync_point!("SSTABLE_ITERATOR::SEEK::BEFORE_NEXT_BLOCK");
354            self.seek_idx(block_idx + 1, None).await?;
355        }
356        Ok(())
357    }
358
359    fn collect_local_statistic(&self, stats: &mut StoreLocalStatistic) {
360        stats.add(&self.stats);
361    }
362
363    fn value_meta(&self) -> ValueMeta {
364        ValueMeta {
365            object_id: Some(self.sst.id),
366            block_id: Some(self.cur_idx as _),
367        }
368    }
369}
370
371impl SstableIteratorType for SstableIterator {
372    fn create(
373        sstable: TableHolder,
374        sstable_store: SstableStoreRef,
375        options: Arc<SstableIteratorReadOptions>,
376        sstable_info_ref: &SstableInfo,
377    ) -> Self {
378        SstableIterator::new(sstable, sstable_store, options, sstable_info_ref)
379    }
380}
381
382#[cfg(test)]
383mod tests {
384    use std::collections::Bound;
385
386    use bytes::Bytes;
387    use foyer::Hint;
388    use itertools::Itertools;
389    use rand::prelude::*;
390    use rand::rng as thread_rng;
391    use risingwave_common::catalog::TableId;
392    use risingwave_common::hash::VirtualNode;
393    use risingwave_common::util::epoch::test_epoch;
394    use risingwave_hummock_sdk::EpochWithGap;
395    use risingwave_hummock_sdk::key::{TableKey, UserKey};
396    use risingwave_hummock_sdk::sstable_info::{SstableInfo, SstableInfoInner};
397
398    use super::*;
399    use crate::assert_bytes_eq;
400    use crate::hummock::CachePolicy;
401    use crate::hummock::iterator::test_utils::mock_sstable_store;
402    use crate::hummock::test_utils::{
403        TEST_KEYS_COUNT, default_builder_opt_for_test, gen_default_test_sstable,
404        gen_test_sstable_info, gen_test_sstable_with_table_ids, test_key_of, test_value_of,
405    };
406
407    async fn inner_test_forward_iterator(
408        sstable_store: SstableStoreRef,
409        handle: TableHolder,
410        sstable_info: SstableInfo,
411    ) {
412        // We should have at least 10 blocks, so that sstable iterator test could cover more code
413        // path.
414        let mut sstable_iter = SstableIterator::create(
415            handle,
416            sstable_store,
417            Arc::new(SstableIteratorReadOptions::default()),
418            &sstable_info,
419        );
420        let mut cnt = 0;
421        sstable_iter.rewind().await.unwrap();
422
423        while sstable_iter.is_valid() {
424            let key = sstable_iter.key();
425            let value = sstable_iter.value();
426            assert_eq!(key, test_key_of(cnt).to_ref());
427            assert_bytes_eq!(value.into_user_value().unwrap(), test_value_of(cnt));
428            cnt += 1;
429            sstable_iter.next().await.unwrap();
430        }
431
432        assert_eq!(cnt, TEST_KEYS_COUNT);
433    }
434
435    #[tokio::test]
436    async fn test_table_iterator() {
437        // Build remote sstable
438        let sstable_store = mock_sstable_store().await;
439        let (sstable, sstable_info) =
440            gen_default_test_sstable(default_builder_opt_for_test(), 0, sstable_store.clone())
441                .await;
442        // We should have at least 10 blocks, so that sstable iterator test could cover more code
443        // path.
444        assert!(sstable.meta.block_metas.len() > 10);
445
446        inner_test_forward_iterator(sstable_store.clone(), sstable, sstable_info).await;
447    }
448
449    #[tokio::test]
450    async fn test_table_seek() {
451        let sstable_store = mock_sstable_store().await;
452        let (sstable, sstable_info) =
453            gen_default_test_sstable(default_builder_opt_for_test(), 0, sstable_store.clone())
454                .await;
455        // We should have at least 10 blocks, so that sstable iterator test could cover more code
456        // path.
457        assert!(sstable.meta.block_metas.len() > 10);
458        let mut sstable_iter = SstableIterator::create(
459            sstable,
460            sstable_store,
461            Arc::new(SstableIteratorReadOptions::default()),
462            &sstable_info,
463        );
464        let mut all_key_to_test = (0..TEST_KEYS_COUNT).collect_vec();
465        let mut rng = thread_rng();
466        all_key_to_test.shuffle(&mut rng);
467
468        // We seek and access all the keys in random order
469        for i in all_key_to_test {
470            sstable_iter.seek(test_key_of(i).to_ref()).await.unwrap();
471            // sstable_iter.next().await.unwrap();
472            let key = sstable_iter.key();
473            assert_eq!(key, test_key_of(i).to_ref());
474        }
475
476        // Seek to key #500 and start iterating.
477        sstable_iter.seek(test_key_of(500).to_ref()).await.unwrap();
478        for i in 500..TEST_KEYS_COUNT {
479            let key = sstable_iter.key();
480            assert_eq!(key, test_key_of(i).to_ref());
481            sstable_iter.next().await.unwrap();
482        }
483        assert!(!sstable_iter.is_valid());
484
485        // Seek to < first key
486        let smallest_key = FullKey::for_test(
487            TableId::default(),
488            [
489                VirtualNode::ZERO.to_be_bytes().as_slice(),
490                format!("key_aaaa_{:05}", 0).as_bytes(),
491            ]
492            .concat(),
493            test_epoch(233),
494        );
495        sstable_iter.seek(smallest_key.to_ref()).await.unwrap();
496        let key = sstable_iter.key();
497        assert_eq!(key, test_key_of(0).to_ref());
498
499        // Seek to > last key
500        let largest_key = FullKey::for_test(
501            TableId::default(),
502            [
503                VirtualNode::ZERO.to_be_bytes().as_slice(),
504                format!("key_zzzz_{:05}", 0).as_bytes(),
505            ]
506            .concat(),
507            test_epoch(233),
508        );
509        sstable_iter.seek(largest_key.to_ref()).await.unwrap();
510        assert!(!sstable_iter.is_valid());
511
512        // Seek to non-existing key
513        for idx in 1..TEST_KEYS_COUNT {
514            // Seek to the previous key of each existing key. e.g.,
515            // Our key space is `key_test_00000`, `key_test_00002`, `key_test_00004`, ...
516            // And we seek to `key_test_00001` (will produce `key_test_00002`), `key_test_00003`
517            // (will produce `key_test_00004`).
518            sstable_iter
519                .seek(
520                    FullKey::for_test(
521                        TableId::default(),
522                        [
523                            VirtualNode::ZERO.to_be_bytes().as_slice(),
524                            format!("key_test_{:05}", idx * 2 - 1).as_bytes(),
525                        ]
526                        .concat(),
527                        0,
528                    )
529                    .to_ref(),
530                )
531                .await
532                .unwrap();
533
534            let key = sstable_iter.key();
535            assert_eq!(key, test_key_of(idx).to_ref());
536            sstable_iter.next().await.unwrap();
537        }
538        assert!(!sstable_iter.is_valid());
539    }
540
541    #[tokio::test]
542    async fn test_prefetch_table_read() {
543        let sstable_store = mock_sstable_store().await;
544        // when upload data is successful, but upload meta is fail and delete is fail
545        let kv_iter =
546            (0..TEST_KEYS_COUNT).map(|i| (test_key_of(i), HummockValue::put(test_value_of(i))));
547        let sst_info = gen_test_sstable_info(
548            default_builder_opt_for_test(),
549            0,
550            kv_iter,
551            sstable_store.clone(),
552        )
553        .await;
554
555        let end_key = test_key_of(TEST_KEYS_COUNT);
556        let uk = UserKey::new(
557            end_key.user_key.table_id,
558            TableKey(Bytes::from(end_key.user_key.table_key.0)),
559        );
560        let options = Arc::new(SstableIteratorReadOptions {
561            cache_policy: CachePolicy::Fill(Hint::Normal),
562            read_table_id: None,
563            scan_end_user_key: Some(Bound::Included(uk.clone())),
564            prefetch: true,
565        });
566        let mut stats = StoreLocalStatistic::default();
567        let mut sstable_iter = SstableIterator::create(
568            sstable_store.sstable(&sst_info, &mut stats).await.unwrap(),
569            sstable_store.clone(),
570            options.clone(),
571            &sst_info,
572        );
573        let mut cnt = 1000;
574        sstable_iter.seek(test_key_of(cnt).to_ref()).await.unwrap();
575        while sstable_iter.is_valid() {
576            let key = sstable_iter.key();
577            let value = sstable_iter.value();
578            assert_eq!(
579                key,
580                test_key_of(cnt).to_ref(),
581                "fail at {}, get key :{:?}",
582                cnt,
583                String::from_utf8(key.user_key.table_key.key_part().to_vec()).unwrap()
584            );
585            assert_bytes_eq!(value.into_user_value().unwrap(), test_value_of(cnt));
586            cnt += 1;
587            sstable_iter.next().await.unwrap();
588        }
589        assert_eq!(cnt, TEST_KEYS_COUNT);
590        let mut sstable_iter = SstableIterator::create(
591            sstable_store.sstable(&sst_info, &mut stats).await.unwrap(),
592            sstable_store,
593            options.clone(),
594            &sst_info,
595        );
596
597        // Cache hits produce a singleton prefetch stream. Seeking past its end must
598        // fetch the requested block instead of repeatedly polling the exhausted stream.
599        assert!(
600            sstable_iter.calculate_block_idx_by_key(test_key_of(TEST_KEYS_COUNT / 2).to_ref())
601                > sstable_iter.calculate_block_idx_by_key(test_key_of(1000).to_ref()) + 1
602        );
603        for idx in [1000, TEST_KEYS_COUNT / 2, TEST_KEYS_COUNT - 1] {
604            sstable_iter.seek(test_key_of(idx).to_ref()).await.unwrap();
605            assert!(sstable_iter.is_valid());
606            assert_eq!(sstable_iter.key(), test_key_of(idx).to_ref());
607            assert_bytes_eq!(
608                sstable_iter.value().into_user_value().unwrap(),
609                test_value_of(idx)
610            );
611        }
612
613        let mut cnt = 1000;
614        sstable_iter.seek(test_key_of(cnt).to_ref()).await.unwrap();
615        while sstable_iter.is_valid() {
616            let key = sstable_iter.key();
617            let value = sstable_iter.value();
618            assert_eq!(key, test_key_of(cnt).to_ref());
619            assert_bytes_eq!(value.into_user_value().unwrap(), test_value_of(cnt));
620            cnt += 1;
621            sstable_iter.next().await.unwrap();
622        }
623        assert_eq!(cnt, TEST_KEYS_COUNT);
624    }
625
626    #[tokio::test]
627    async fn test_prefetch_seek_preserves_remaining_blocks() {
628        let sstable_store = mock_sstable_store().await;
629        let mut builder_options = default_builder_opt_for_test();
630        builder_options.block_capacity = 128;
631        let info = gen_test_sstable_info(
632            builder_options,
633            0,
634            (0..64).map(|i| (test_key_of(i), HummockValue::put(test_value_of(i)))),
635            sstable_store.clone(),
636        )
637        .await;
638        let sstable = sstable_store
639            .sstable(&info, &mut StoreLocalStatistic::default())
640            .await
641            .unwrap();
642        assert!(sstable.meta.block_metas.len() > 16);
643        let mut iter = SstableIterator::new(
644            sstable.clone(),
645            sstable_store.clone(),
646            Arc::new(SstableIteratorReadOptions {
647                prefetch: true,
648                cache_policy: CachePolicy::Disable,
649                ..Default::default()
650            }),
651            &info,
652        );
653        let mut reference = SstableIterator::new(
654            sstable.clone(),
655            sstable_store.clone(),
656            Arc::new(SstableIteratorReadOptions {
657                prefetch: false,
658                cache_policy: CachePolicy::Disable,
659                ..Default::default()
660            }),
661            &info,
662        );
663        let last = sstable.meta.block_metas.len() - 1;
664        // Initial prefetch, forward skip, backward seek, reuse, final-block fallback,
665        // rewind after exhaustion, and reread of an already consumed block.
666        for (idx, prefetches, gets) in [
667            (0, 1, 0),
668            (2, 1, 0),
669            (1, 1, 1),
670            (3, 1, 1),
671            (last, 1, 2),
672            (0, 2, 2),
673            (0, 2, 3),
674        ] {
675            let key = FullKey::decode(&sstable.meta.block_metas[idx].smallest_key);
676            iter.seek(key).await.unwrap();
677            reference.seek(key).await.unwrap();
678            assert_eq!(iter.key(), reference.key(), "block {idx}");
679            assert_bytes_eq!(
680                iter.value().into_user_value().unwrap(),
681                reference.value().into_user_value().unwrap()
682            );
683            assert_eq!(
684                iter.stats.cache_data_prefetch_count, prefetches,
685                "block {idx}"
686            );
687            assert_eq!(iter.stats.cache_data_block_total, gets, "block {idx}");
688            if idx == last {
689                assert_eq!(sstable_store.get_prefetch_memory_usage(), 0);
690            } else {
691                assert!(sstable_store.get_prefetch_memory_usage() > 0);
692            }
693        }
694        drop(iter);
695        assert_eq!(sstable_store.get_prefetch_memory_usage(), 0);
696    }
697
698    #[tokio::test]
699    async fn test_scan_end_with_prefetch_on_or_off() {
700        let sstable_store = mock_sstable_store().await;
701        let mut builder_options = default_builder_opt_for_test();
702        builder_options.block_capacity = 128;
703
704        let test_user_key = |table_id, idx| {
705            UserKey::new(
706                TableId::new(table_id),
707                TableKey(Bytes::from(
708                    [
709                        VirtualNode::ZERO.to_be_bytes().as_slice(),
710                        format!("key_{idx:05}").as_bytes(),
711                    ]
712                    .concat(),
713                )),
714            )
715        };
716        let test_key = |table_id, idx| FullKey {
717            user_key: test_user_key(table_id, idx),
718            epoch_with_gap: EpochWithGap::new_from_epoch(test_epoch(1)),
719        };
720        let kv_pairs = (1..=3).flat_map(|table_id| {
721            (0..8).map(move |idx| {
722                (
723                    test_key(table_id, idx),
724                    HummockValue::put(Bytes::from(test_value_of(idx))),
725                )
726            })
727        });
728        let (sstable, sstable_info) = gen_test_sstable_with_table_ids(
729            builder_options,
730            10,
731            kv_pairs,
732            sstable_store.clone(),
733            vec![1, 2, 3],
734        )
735        .await;
736
737        let table_2_block_start = sstable
738            .meta
739            .block_metas
740            .partition_point(|block_meta| block_meta.table_id() < TableId::new(2));
741        let table_3_block_start = sstable
742            .meta
743            .block_metas
744            .partition_point(|block_meta| block_meta.table_id() < TableId::new(3));
745        assert!(table_2_block_start + 2 < table_3_block_start);
746
747        let table_3_start = UserKey::new(TableId::new(3), TableKey(Bytes::from_static(b"")));
748        for (case, prefetch) in [("prefetch off", false), ("prefetch on", true)] {
749            let options = Arc::new(SstableIteratorReadOptions {
750                cache_policy: CachePolicy::Disable,
751                read_table_id: None,
752                scan_end_user_key: Some(Bound::Excluded(table_3_start.clone())),
753                prefetch,
754            });
755            let mut sstable_iter = SstableIterator::create(
756                sstable.clone(),
757                sstable_store.clone(),
758                options,
759                &sstable_info,
760            );
761
762            assert_eq!(sstable_iter.block_end_idx_exclusive, table_3_block_start);
763            sstable_iter.seek(test_key(2, 0).to_ref()).await.unwrap();
764            if prefetch {
765                assert_eq!(
766                    sstable_iter.preload_end_block_idx, sstable_iter.block_end_idx_exclusive,
767                    "{case}"
768                );
769            } else {
770                assert_eq!(sstable_iter.preload_end_block_idx, 0, "{case}");
771            }
772            let mut key_count = 0;
773            while sstable_iter.is_valid() {
774                assert_eq!(
775                    sstable_iter.key().user_key.table_id,
776                    TableId::new(2),
777                    "{case}"
778                );
779                key_count += 1;
780                sstable_iter.next().await.unwrap();
781            }
782            assert_eq!(key_count, 8, "{case}");
783
784            let mut stats = StoreLocalStatistic::default();
785            sstable_iter.collect_local_statistic(&mut stats);
786            assert_eq!(
787                stats.cache_data_block_total + stats.cache_data_prefetch_block_count,
788                (table_3_block_start - table_2_block_start) as u64,
789                "{case}"
790            );
791            if prefetch {
792                assert!(stats.cache_data_prefetch_block_count > 0, "{case}");
793            } else {
794                assert_eq!(stats.cache_data_prefetch_block_count, 0, "{case}");
795            }
796        }
797
798        let options = Arc::new(SstableIteratorReadOptions {
799            cache_policy: CachePolicy::Disable,
800            read_table_id: Some(TableId::new(2)),
801            scan_end_user_key: None,
802            prefetch: false,
803        });
804        let mut sstable_iter = SstableIterator::create(
805            sstable.clone(),
806            sstable_store.clone(),
807            options,
808            &sstable_info,
809        );
810        assert_eq!(sstable_iter.block_start_idx_inclusive, table_2_block_start);
811        assert_eq!(sstable_iter.block_end_idx_exclusive, table_3_block_start);
812        sstable_iter.rewind().await.unwrap();
813        assert!(sstable_iter.is_valid());
814        while sstable_iter.is_valid() {
815            assert_eq!(sstable_iter.key().user_key.table_id, TableId::new(2));
816            sstable_iter.next().await.unwrap();
817        }
818
819        for missing_table_id in [0, 4] {
820            let options = Arc::new(SstableIteratorReadOptions {
821                read_table_id: Some(TableId::new(missing_table_id)),
822                ..Default::default()
823            });
824            let mut sstable_iter = SstableIterator::create(
825                sstable.clone(),
826                sstable_store.clone(),
827                options,
828                &sstable_info,
829            );
830            sstable_iter.rewind().await.unwrap();
831            for table_id in 1..=3 {
832                for idx in 0..8 {
833                    assert!(sstable_iter.is_valid());
834                    assert_eq!(sstable_iter.key(), test_key(table_id, idx).to_ref());
835                    sstable_iter.next().await.unwrap();
836                }
837            }
838            assert!(!sstable_iter.is_valid());
839        }
840
841        let mut table_2_sstable_info = sstable_info.get_inner();
842        table_2_sstable_info.table_ids = vec![TableId::new(2)];
843        let table_2_sstable_info = SstableInfo::from(table_2_sstable_info);
844        let table_2_start: UserKey<Bytes> =
845            FullKey::decode(&sstable.meta.block_metas[table_2_block_start].smallest_key)
846                .user_key
847                .copy_into();
848        let options = Arc::new(SstableIteratorReadOptions {
849            cache_policy: CachePolicy::Disable,
850            read_table_id: None,
851            scan_end_user_key: Some(Bound::Excluded(table_2_start)),
852            prefetch: false,
853        });
854        let mut sstable_iter =
855            SstableIterator::create(sstable, sstable_store, options, &table_2_sstable_info);
856
857        assert_eq!(
858            sstable_iter.block_start_idx_inclusive,
859            sstable_iter.block_end_idx_exclusive
860        );
861        sstable_iter.seek(test_key(2, 0).to_ref()).await.unwrap();
862        assert!(!sstable_iter.is_valid());
863        let mut stats = StoreLocalStatistic::default();
864        sstable_iter.collect_local_statistic(&mut stats);
865        assert_eq!(stats.cache_data_block_total, 0);
866        assert_eq!(stats.cache_data_prefetch_block_count, 0);
867    }
868
869    #[tokio::test]
870    async fn test_read_table_id_range() {
871        {
872            let sstable_store = mock_sstable_store().await;
873            let (sstable, sstable_info) =
874                gen_default_test_sstable(default_builder_opt_for_test(), 0, sstable_store.clone())
875                    .await;
876            let mut sstable_iter = SstableIterator::create(
877                sstable,
878                sstable_store.clone(),
879                Arc::new(SstableIteratorReadOptions::default()),
880                &sstable_info,
881            );
882            sstable_iter.rewind().await.unwrap();
883            assert!(sstable_iter.is_valid());
884            assert_eq!(sstable_iter.key(), test_key_of(0).to_ref());
885        }
886
887        {
888            let sstable_store = mock_sstable_store().await;
889            // test key_range right
890            let k1 = {
891                let mut table_key = VirtualNode::ZERO.to_be_bytes().to_vec();
892                table_key.extend_from_slice(format!("key_test_{:05}", 1).as_bytes());
893                let uk = UserKey::for_test(TableId::from(1), table_key);
894                FullKey {
895                    user_key: uk,
896                    epoch_with_gap: EpochWithGap::new_from_epoch(test_epoch(1)),
897                }
898            };
899
900            let k2 = {
901                let mut table_key = VirtualNode::ZERO.to_be_bytes().to_vec();
902                table_key.extend_from_slice(format!("key_test_{:05}", 2).as_bytes());
903                let uk = UserKey::for_test(TableId::from(2), table_key);
904                FullKey {
905                    user_key: uk,
906                    epoch_with_gap: EpochWithGap::new_from_epoch(test_epoch(1)),
907                }
908            };
909
910            let k3 = {
911                let mut table_key = VirtualNode::ZERO.to_be_bytes().to_vec();
912                table_key.extend_from_slice(format!("key_test_{:05}", 3).as_bytes());
913                let uk = UserKey::for_test(TableId::from(3), table_key);
914                FullKey {
915                    user_key: uk,
916                    epoch_with_gap: EpochWithGap::new_from_epoch(test_epoch(1)),
917                }
918            };
919
920            {
921                let kv_pairs = vec![
922                    (k1.clone(), HummockValue::put(test_value_of(1))),
923                    (k2.clone(), HummockValue::put(test_value_of(2))),
924                    (k3.clone(), HummockValue::put(test_value_of(3))),
925                ];
926
927                let (sstable, _sstable_info) = gen_test_sstable_with_table_ids(
928                    default_builder_opt_for_test(),
929                    10,
930                    kv_pairs.into_iter(),
931                    sstable_store.clone(),
932                    vec![1, 2, 3],
933                )
934                .await;
935                let mut sstable_iter = SstableIterator::create(
936                    sstable,
937                    sstable_store.clone(),
938                    Arc::new(SstableIteratorReadOptions::default()),
939                    &SstableInfo::from(SstableInfoInner {
940                        table_ids: vec![1.into(), 2.into(), 3.into()],
941                        ..Default::default()
942                    }),
943                );
944                sstable_iter.rewind().await.unwrap();
945                assert!(sstable_iter.is_valid());
946                assert!(sstable_iter.key().eq(&k1.to_ref()));
947
948                let mut cnt = 0;
949                let mut last_key = k1.clone();
950                while sstable_iter.is_valid() {
951                    last_key = sstable_iter.key().to_vec();
952                    cnt += 1;
953                    sstable_iter.next().await.unwrap();
954                }
955
956                assert_eq!(3, cnt);
957                assert_eq!(last_key, k3.clone());
958            }
959
960            {
961                let kv_pairs = vec![
962                    (k1.clone(), HummockValue::put(test_value_of(1))),
963                    (k2.clone(), HummockValue::put(test_value_of(2))),
964                    (k3.clone(), HummockValue::put(test_value_of(3))),
965                ];
966
967                let (sstable, _sstable_info) = gen_test_sstable_with_table_ids(
968                    default_builder_opt_for_test(),
969                    10,
970                    kv_pairs.into_iter(),
971                    sstable_store.clone(),
972                    vec![1, 2, 3],
973                )
974                .await;
975
976                let mut sstable_iter = SstableIterator::create(
977                    sstable,
978                    sstable_store.clone(),
979                    Arc::new(SstableIteratorReadOptions::default()),
980                    &SstableInfo::from(SstableInfoInner {
981                        table_ids: vec![1.into(), 2.into()],
982                        ..Default::default()
983                    }),
984                );
985                sstable_iter.rewind().await.unwrap();
986                assert!(sstable_iter.is_valid());
987                assert!(sstable_iter.key().eq(&k1.to_ref()));
988
989                let mut cnt = 0;
990                let mut last_key = k1.clone();
991                while sstable_iter.is_valid() {
992                    last_key = sstable_iter.key().to_vec();
993                    cnt += 1;
994                    sstable_iter.next().await.unwrap();
995                }
996
997                assert_eq!(2, cnt);
998                assert_eq!(last_key, k2.clone());
999            }
1000
1001            {
1002                let kv_pairs = vec![
1003                    (k1.clone(), HummockValue::put(test_value_of(1))),
1004                    (k2.clone(), HummockValue::put(test_value_of(2))),
1005                    (k3.clone(), HummockValue::put(test_value_of(3))),
1006                ];
1007
1008                let (sstable, _sstable_info) = gen_test_sstable_with_table_ids(
1009                    default_builder_opt_for_test(),
1010                    10,
1011                    kv_pairs.into_iter(),
1012                    sstable_store.clone(),
1013                    vec![1, 2, 3],
1014                )
1015                .await;
1016
1017                let mut sstable_iter = SstableIterator::create(
1018                    sstable,
1019                    sstable_store.clone(),
1020                    Arc::new(SstableIteratorReadOptions::default()),
1021                    &SstableInfo::from(SstableInfoInner {
1022                        table_ids: vec![2.into(), 3.into()],
1023                        ..Default::default()
1024                    }),
1025                );
1026                sstable_iter.rewind().await.unwrap();
1027                assert!(sstable_iter.is_valid());
1028                assert!(sstable_iter.key().eq(&k2.to_ref()));
1029
1030                let mut cnt = 0;
1031                let mut last_key = k1.clone();
1032                while sstable_iter.is_valid() {
1033                    last_key = sstable_iter.key().to_vec();
1034                    cnt += 1;
1035                    sstable_iter.next().await.unwrap();
1036                }
1037
1038                assert_eq!(2, cnt);
1039                assert_eq!(last_key, k3.clone());
1040            }
1041
1042            {
1043                let kv_pairs = vec![
1044                    (k1.clone(), HummockValue::put(test_value_of(1))),
1045                    (k2.clone(), HummockValue::put(test_value_of(2))),
1046                    (k3.clone(), HummockValue::put(test_value_of(3))),
1047                ];
1048
1049                let (sstable, _sstable_info) = gen_test_sstable_with_table_ids(
1050                    default_builder_opt_for_test(),
1051                    10,
1052                    kv_pairs.into_iter(),
1053                    sstable_store.clone(),
1054                    vec![1, 2, 3],
1055                )
1056                .await;
1057
1058                let mut sstable_iter = SstableIterator::create(
1059                    sstable,
1060                    sstable_store.clone(),
1061                    Arc::new(SstableIteratorReadOptions::default()),
1062                    &SstableInfo::from(SstableInfoInner {
1063                        table_ids: vec![2.into()],
1064                        ..Default::default()
1065                    }),
1066                );
1067                sstable_iter.rewind().await.unwrap();
1068                assert!(sstable_iter.is_valid());
1069                assert!(sstable_iter.key().eq(&k2.to_ref()));
1070
1071                let mut cnt = 0;
1072                let mut last_key = k1.clone();
1073                while sstable_iter.is_valid() {
1074                    last_key = sstable_iter.key().to_vec();
1075                    cnt += 1;
1076                    sstable_iter.next().await.unwrap();
1077                }
1078
1079                assert_eq!(1, cnt);
1080                assert_eq!(last_key, k2.clone());
1081            }
1082        }
1083    }
1084}
risingwave_storage/hummock/sstable/forward_sstable_iterator.rs

risingwave_storage/hummock/sstable/
forward_sstable_iterator.rs
