risingwave_storage/hummock/iterator/
concat_inner.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::{Equal, Greater, Less};
16use std::sync::Arc;
17
18use fail::fail_point;
19use risingwave_hummock_sdk::key::FullKey;
20use risingwave_hummock_sdk::sstable_info::SstableInfo;
21
22use crate::hummock::iterator::{
23    DirectionEnum, HummockIterator, HummockIteratorDirection, ValueMeta,
24};
25use crate::hummock::sstable::SstableIteratorReadOptions;
26use crate::hummock::value::HummockValue;
27use crate::hummock::{HummockResult, IteratorStatsGuard, SstableIteratorType, SstableStoreRef};
28use crate::monitor::StoreLocalStatistic;
29
30fn smallest_key(sstable_info: &SstableInfo) -> &[u8] {
31    &sstable_info.key_range.left
32}
33
34fn largest_key(sstable_info: &SstableInfo) -> &[u8] {
35    &sstable_info.key_range.right
36}
37
38/// Served as the concrete implementation of `ConcatIterator` and `BackwardConcatIterator`.
39pub struct ConcatIteratorInner<TI: SstableIteratorType> {
40    /// The iterator of the current table.
41    sstable_iter: Option<TI>,
42
43    /// Current table index.
44    cur_idx: usize,
45
46    /// All non-overlapping `sstable_infos`.
47    sstable_infos: Vec<SstableInfo>,
48
49    sstable_store: SstableStoreRef,
50
51    stats: StoreLocalStatistic,
52    read_options: Arc<SstableIteratorReadOptions>,
53}
54
55impl<TI: SstableIteratorType> ConcatIteratorInner<TI> {
56    /// Caller should make sure that `sstable_infos` are non-overlapping,
57    /// arranged in ascending order when it serves as a forward iterator,
58    /// and arranged in descending order when it serves as a backward iterator.
59    pub fn new(
60        sstable_infos: Vec<SstableInfo>,
61        sstable_store: SstableStoreRef,
62        read_options: Arc<SstableIteratorReadOptions>,
63    ) -> Self {
64        Self {
65            sstable_iter: None,
66            cur_idx: 0,
67            sstable_infos,
68            sstable_store,
69            stats: StoreLocalStatistic::default(),
70            read_options,
71        }
72    }
73
74    /// Seeks to a table, and then seeks to the key if `seek_key` is given.
75    async fn seek_idx(
76        &mut self,
77        idx: usize,
78        seek_key: Option<FullKey<&[u8]>>,
79    ) -> HummockResult<()> {
80        if idx >= self.sstable_infos.len() {
81            if let Some(old_iter) = self.sstable_iter.take() {
82                old_iter.collect_local_statistic(&mut self.stats);
83            }
84            self.cur_idx = self.sstable_infos.len();
85        } else {
86            let table = self
87                .sstable_store
88                .sstable(&self.sstable_infos[idx], &mut self.stats)
89                .await?;
90            let sstable_iter = TI::create(
91                table,
92                self.sstable_store.clone(),
93                self.read_options.clone(),
94                &self.sstable_infos[idx],
95            );
96
97            let mut pending = IteratorStatsGuard::new(sstable_iter, &mut self.stats);
98            if let Some(key) = seek_key {
99                pending.iter_mut().seek(key).await?;
100            } else {
101                pending.iter_mut().rewind().await?;
102            }
103            fail_point!("concat_iter_after_seek_before_init_complete", |_| Err(
104                crate::hummock::HummockError::meta_error(
105                    "test error after concat iterator seek before initialization completes"
106                )
107            ));
108
109            let sstable_iter = pending.handoff();
110            if let Some(old_iter) = self.sstable_iter.take() {
111                old_iter.collect_local_statistic(&mut self.stats);
112            }
113            self.sstable_iter = Some(sstable_iter);
114            self.cur_idx = idx;
115        }
116        Ok(())
117    }
118}
119
120impl<TI: SstableIteratorType> HummockIterator for ConcatIteratorInner<TI> {
121    type Direction = TI::Direction;
122
123    async fn next(&mut self) -> HummockResult<()> {
124        let sstable_iter = self.sstable_iter.as_mut().expect("no table iter");
125        sstable_iter.next().await?;
126
127        if sstable_iter.is_valid() {
128            Ok(())
129        } else {
130            // seek to next table
131            let mut table_idx = self.cur_idx + 1;
132            while !self.is_valid() && table_idx < self.sstable_infos.len() {
133                self.seek_idx(table_idx, None).await?;
134                table_idx += 1;
135            }
136            Ok(())
137        }
138    }
139
140    fn key(&self) -> FullKey<&[u8]> {
141        self.sstable_iter.as_ref().expect("no table iter").key()
142    }
143
144    fn value(&self) -> HummockValue<&[u8]> {
145        self.sstable_iter.as_ref().expect("no table iter").value()
146    }
147
148    fn is_valid(&self) -> bool {
149        self.sstable_iter.as_ref().is_some_and(|i| i.is_valid())
150    }
151
152    async fn rewind(&mut self) -> HummockResult<()> {
153        self.seek_idx(0, None).await?;
154        let mut table_idx = 1;
155        while !self.is_valid() && table_idx < self.sstable_infos.len() {
156            // Seek to next table
157            self.seek_idx(table_idx, None).await?;
158            table_idx += 1;
159        }
160        Ok(())
161    }
162
163    async fn seek<'a>(&'a mut self, key: FullKey<&'a [u8]>) -> HummockResult<()> {
164        let mut table_idx = self
165            .sstable_infos
166            .partition_point(|table| match Self::Direction::direction() {
167                DirectionEnum::Forward => {
168                    let ord = FullKey::decode(smallest_key(table)).cmp(&key);
169
170                    ord == Less || ord == Equal
171                }
172                DirectionEnum::Backward => {
173                    let ord = FullKey::decode(largest_key(table)).cmp(&key);
174                    ord == Greater || (ord == Equal && !table.key_range.right_exclusive)
175                }
176            })
177            .saturating_sub(1); // considering the boundary of 0
178
179        self.seek_idx(table_idx, Some(key)).await?;
180        table_idx += 1;
181        while !self.is_valid() && table_idx < self.sstable_infos.len() {
182            // Seek to next table
183            self.seek_idx(table_idx, None).await?;
184            table_idx += 1;
185        }
186        Ok(())
187    }
188
189    fn collect_local_statistic(&self, stats: &mut StoreLocalStatistic) {
190        stats.add(&self.stats);
191        if let Some(iter) = &self.sstable_iter {
192            iter.collect_local_statistic(stats);
193        }
194    }
195
196    fn value_meta(&self) -> ValueMeta {
197        self.sstable_iter
198            .as_ref()
199            .expect("no table iter")
200            .value_meta()
201    }
202}
risingwave_storage/hummock/iterator/concat_inner.rs

risingwave_storage/hummock/iterator/
concat_inner.rs
