risingwave_storage/hummock/iterator/
test_utils.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::sync::Arc;
16
17use bytes::Bytes;
18use foyer::{CacheBuilder, HybridCacheBuilder};
19use itertools::Itertools;
20use risingwave_common::catalog::TableId;
21use risingwave_common::config::{MetricLevel, ObjectStoreConfig};
22use risingwave_common::hash::VirtualNode;
23use risingwave_common::util::epoch::test_epoch;
24use risingwave_hummock_sdk::key::{FullKey, TableKey, UserKey, prefix_slice_with_vnode};
25use risingwave_hummock_sdk::sstable_info::SstableInfo;
26use risingwave_hummock_sdk::{EpochWithGap, HummockEpoch, HummockSstableObjectId};
27use risingwave_object_store::object::{
28    InMemObjectStore, ObjectStore, ObjectStoreImpl, ObjectStoreRef,
29};
30
31use crate::hummock::none::NoneRecentFilter;
32use crate::hummock::shared_buffer::shared_buffer_batch::SharedBufferValue;
33use crate::hummock::sstable::SstableIteratorReadOptions;
34use crate::hummock::sstable_store::SstableStore;
35pub use crate::hummock::test_utils::default_builder_opt_for_test;
36use crate::hummock::test_utils::{
37    gen_test_sstable, gen_test_sstable_info, gen_test_sstable_with_range_tombstone,
38};
39use crate::hummock::{
40    HummockValue, RecentFilter, SstableBuilderOptions, SstableIterator, SstableIteratorType,
41    SstableStoreConfig, SstableStoreRef, TableHolder,
42};
43use crate::monitor::{ObjectStoreMetrics, global_hummock_state_store_metrics};
44
45/// `assert_eq` two `Vec<u8>` with human-readable format.
46#[macro_export]
47macro_rules! assert_bytes_eq {
48    ($left:expr, $right:expr) => {{
49        use bytes::Bytes;
50        assert_eq!(
51            Bytes::copy_from_slice(&$left),
52            Bytes::copy_from_slice(&$right)
53        )
54    }};
55}
56
57pub const TEST_KEYS_COUNT: usize = 10;
58
59pub async fn mock_sstable_store() -> SstableStoreRef {
60    mock_sstable_store_with_recent_filter(Arc::new(NoneRecentFilter::default().into())).await
61}
62
63pub async fn mock_sstable_store_with_recent_filter(
64    recent_filter: Arc<RecentFilter<(HummockSstableObjectId, usize)>>,
65) -> SstableStoreRef {
66    mock_sstable_store_with_object_store_and_recent_filter(
67        Arc::new(ObjectStoreImpl::InMem(
68            InMemObjectStore::for_test().monitored(
69                Arc::new(ObjectStoreMetrics::unused()),
70                Arc::new(ObjectStoreConfig::default()),
71            ),
72        )),
73        recent_filter,
74    )
75    .await
76}
77
78pub async fn mock_sstable_store_with_object_store(store: ObjectStoreRef) -> SstableStoreRef {
79    mock_sstable_store_with_object_store_and_recent_filter(
80        store,
81        Arc::new(NoneRecentFilter::default().into()),
82    )
83    .await
84}
85
86pub async fn mock_sstable_store_with_object_store_and_recent_filter(
87    store: ObjectStoreRef,
88    recent_filter: Arc<RecentFilter<(HummockSstableObjectId, usize)>>,
89) -> SstableStoreRef {
90    let path = "test".to_owned();
91    let meta_cache = HybridCacheBuilder::new()
92        .memory(64 << 20)
93        .with_shards(2)
94        .storage()
95        .build()
96        .await
97        .unwrap();
98    let block_cache = HybridCacheBuilder::new()
99        .memory(64 << 20)
100        .with_shards(2)
101        .storage()
102        .build()
103        .await
104        .unwrap();
105    Arc::new(SstableStore::new(SstableStoreConfig {
106        store,
107        path,
108
109        prefetch_buffer_capacity: 64 << 20,
110        max_prefetch_block_number: 16,
111
112        recent_filter,
113        state_store_metrics: Arc::new(global_hummock_state_store_metrics(MetricLevel::Disabled)),
114        use_new_object_prefix_strategy: true,
115        skip_bloom_filter_in_serde: false,
116
117        meta_cache,
118        block_cache,
119        vector_meta_cache: CacheBuilder::new(64 << 20).build(),
120        vector_block_cache: CacheBuilder::new(64 << 20).build(),
121    }))
122}
123
124// Generate test table key with vnode 0
125pub fn iterator_test_table_key_of(idx: usize) -> Vec<u8> {
126    prefix_slice_with_vnode(VirtualNode::ZERO, format!("key_test_{:05}", idx).as_bytes()).to_vec()
127}
128
129pub fn iterator_test_user_key_of(idx: usize) -> UserKey<Vec<u8>> {
130    UserKey::for_test(TableId::default(), iterator_test_table_key_of(idx))
131}
132
133pub fn iterator_test_bytes_user_key_of(idx: usize) -> UserKey<Bytes> {
134    UserKey::for_test(
135        TableId::default(),
136        Bytes::from(iterator_test_table_key_of(idx)),
137    )
138}
139
140/// Generates keys like `{table_id=0}key_test_00002` with epoch 233.
141pub fn iterator_test_key_of(idx: usize) -> FullKey<Vec<u8>> {
142    FullKey {
143        user_key: iterator_test_user_key_of(idx),
144        epoch_with_gap: EpochWithGap::new_from_epoch(test_epoch(233)),
145    }
146}
147
148/// Generates keys like `{table_id=0}key_test_00002` with epoch 233.
149pub fn iterator_test_bytes_key_of(idx: usize) -> FullKey<Bytes> {
150    iterator_test_key_of(idx).into_bytes()
151}
152
153/// Generates keys like `{table_id=0}key_test_00002` with epoch `epoch` .
154pub fn iterator_test_key_of_epoch(idx: usize, epoch: HummockEpoch) -> FullKey<Vec<u8>> {
155    FullKey {
156        user_key: iterator_test_user_key_of(idx),
157        epoch_with_gap: EpochWithGap::new_from_epoch(epoch),
158    }
159}
160
161/// Generates keys like `{table_id=0}key_test_00002` with epoch `epoch` .
162pub fn iterator_test_bytes_key_of_epoch(idx: usize, epoch: HummockEpoch) -> FullKey<Bytes> {
163    iterator_test_key_of_epoch(idx, test_epoch(epoch)).into_bytes()
164}
165
166/// The value of an index, like `value_test_00002` without value meta
167pub fn iterator_test_value_of(idx: usize) -> Vec<u8> {
168    format!("value_test_{:05}", idx).as_bytes().to_vec()
169}
170
171pub fn transform_shared_buffer(
172    batches: Vec<(Vec<u8>, SharedBufferValue<Bytes>)>,
173) -> Vec<(TableKey<Bytes>, SharedBufferValue<Bytes>)> {
174    batches
175        .into_iter()
176        .map(|(k, v)| (TableKey(k.into()), v))
177        .collect_vec()
178}
179
180/// Generates a test table used in almost all table-related tests. Developers may verify the
181/// correctness of their implementations by comparing the got value and the expected value
182/// generated by `test_key_of` and `test_value_of`.
183pub async fn gen_iterator_test_sstable_info(
184    object_id: u64,
185    opts: SstableBuilderOptions,
186    idx_mapping: impl Fn(usize) -> usize,
187    sstable_store: SstableStoreRef,
188    total: usize,
189) -> SstableInfo {
190    gen_test_sstable_info(
191        opts,
192        object_id,
193        (0..total).map(|i| {
194            (
195                iterator_test_key_of(idx_mapping(i)),
196                HummockValue::put(iterator_test_value_of(idx_mapping(i))),
197            )
198        }),
199        sstable_store,
200    )
201    .await
202}
203
204/// Generates a test table used in almost all table-related tests. Developers may verify the
205/// correctness of their implementations by comparing the got value and the expected value
206/// generated by `test_key_of` and `test_value_of`.
207pub async fn gen_iterator_test_sstable_base(
208    object_id: u64,
209    opts: SstableBuilderOptions,
210    idx_mapping: impl Fn(usize) -> usize,
211    sstable_store: SstableStoreRef,
212    total: usize,
213) -> (TableHolder, SstableInfo) {
214    gen_test_sstable(
215        opts,
216        object_id,
217        (0..total).map(|i| {
218            (
219                iterator_test_key_of(idx_mapping(i)),
220                HummockValue::put(iterator_test_value_of(idx_mapping(i))),
221            )
222        }),
223        sstable_store,
224    )
225    .await
226}
227
228// key=[idx, epoch], value
229pub async fn gen_iterator_test_sstable_from_kv_pair(
230    object_id: u64,
231    kv_pairs: Vec<(usize, u64, HummockValue<Vec<u8>>)>,
232    sstable_store: SstableStoreRef,
233) -> (TableHolder, SstableInfo) {
234    gen_test_sstable(
235        default_builder_opt_for_test(),
236        object_id,
237        kv_pairs
238            .into_iter()
239            .map(|kv| (iterator_test_key_of_epoch(kv.0, test_epoch(kv.1)), kv.2)),
240        sstable_store,
241    )
242    .await
243}
244
245// key=[idx, epoch], value
246pub async fn gen_iterator_test_sstable_with_range_tombstones(
247    object_id: u64,
248    kv_pairs: Vec<(usize, u64, HummockValue<Vec<u8>>)>,
249    sstable_store: SstableStoreRef,
250) -> SstableInfo {
251    gen_test_sstable_with_range_tombstone(
252        default_builder_opt_for_test(),
253        object_id,
254        kv_pairs
255            .into_iter()
256            .map(|kv| (iterator_test_key_of_epoch(kv.0, test_epoch(kv.1)), kv.2)),
257        sstable_store,
258    )
259    .await
260}
261
262pub async fn gen_merge_iterator_interleave_test_sstable_iters(
263    key_count: usize,
264    count: usize,
265) -> Vec<SstableIterator> {
266    let sstable_store = mock_sstable_store().await;
267    let mut result = vec![];
268    for i in 0..count {
269        let (table, sstable_info) = gen_iterator_test_sstable_base(
270            i as u64,
271            default_builder_opt_for_test(),
272            |x| x * count + i,
273            sstable_store.clone(),
274            key_count,
275        )
276        .await;
277        result.push(SstableIterator::create(
278            table,
279            sstable_store.clone(),
280            Arc::new(SstableIteratorReadOptions::default()),
281            &sstable_info,
282        ));
283    }
284    result
285}
286
287pub async fn gen_iterator_test_sstable_with_incr_epoch(
288    object_id: u64,
289    opts: SstableBuilderOptions,
290    idx_mapping: impl Fn(usize) -> usize,
291    sstable_store: SstableStoreRef,
292    total: usize,
293    epoch_base: u64,
294) -> (TableHolder, SstableInfo) {
295    gen_test_sstable(
296        opts,
297        object_id,
298        (0..total).map(|i| {
299            (
300                iterator_test_key_of_epoch(idx_mapping(i), test_epoch(epoch_base + i as u64)),
301                HummockValue::put(iterator_test_value_of(idx_mapping(i))),
302            )
303        }),
304        sstable_store,
305    )
306    .await
307}
risingwave_storage/hummock/iterator/test_utils.rs

risingwave_storage/hummock/iterator/
test_utils.rs
