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

risingwave_storage/hummock/sstable/
forward_sstable_iterator.rs
