risingwave_meta/hummock/
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
15#![cfg(any(test, feature = "test"))]
16
17use std::collections::{BTreeSet, HashMap};
18use std::sync::{Arc, LazyLock};
19use std::time::Duration;
20
21use bytes::Bytes;
22use itertools::Itertools;
23use risingwave_common::catalog::{TableId, TableOption};
24use risingwave_common::util::epoch::test_epoch;
25use risingwave_hummock_sdk::key::key_with_epoch;
26use risingwave_hummock_sdk::key_range::KeyRange;
27use risingwave_hummock_sdk::level::Levels;
28use risingwave_hummock_sdk::sstable_info::{SstableInfo, SstableInfoInner};
29use risingwave_hummock_sdk::table_watermark::TableWatermarks;
30use risingwave_hummock_sdk::version::{HummockVersion, HummockVersionStateTableInfo};
31use risingwave_hummock_sdk::{
32    CompactionGroupId, HummockEpoch, HummockSstableObjectId, LocalSstableInfo, SyncResult,
33};
34use risingwave_meta_model::WorkerId;
35use risingwave_pb::common::worker_node::Property;
36use risingwave_pb::common::{HostAddress, WorkerType};
37use risingwave_pb::hummock::CompactionConfig;
38use risingwave_pb::hummock::compact_task::TaskStatus;
39use risingwave_rpc_client::HummockMetaClient;
40
41use crate::controller::catalog::CatalogController;
42use crate::controller::cluster::{ClusterController, ClusterControllerRef};
43use crate::hummock::compaction::compaction_config::CompactionConfigBuilder;
44use crate::hummock::compaction::in_progress_compaction::InProgressCompactionView;
45use crate::hummock::compaction::selector::{LocalSelectorStatistic, default_compaction_selector};
46use crate::hummock::compaction::{CompactionDeveloperConfig, CompactionSelectorContext};
47use crate::hummock::level_handler::LevelHandler;
48pub use crate::hummock::manager::CommitEpochInfo;
49use crate::hummock::model::CompactionGroup;
50use crate::hummock::{CompactorManager, HummockManager, HummockManagerRef};
51use crate::manager::MetaSrvEnv;
52use crate::rpc::metrics::MetaMetrics;
53
54static EMPTY_IN_PROGRESS_COMPACTION_VIEW: LazyLock<InProgressCompactionView> =
55    LazyLock::new(InProgressCompactionView::default);
56
57/// Advance virtual time without sleeping; Tokio tests must start with time paused.
58#[cfg(test)]
59pub(crate) async fn advance_time(duration: Duration) {
60    #[cfg(not(madsim))]
61    tokio::time::advance(duration).await;
62    #[cfg(madsim)]
63    tokio::time::advance(duration);
64}
65
66pub fn to_local_sstable_info(ssts: &[SstableInfo]) -> Vec<LocalSstableInfo> {
67    ssts.iter()
68        .map(|sst| LocalSstableInfo::for_test(sst.clone()))
69        .collect_vec()
70}
71
72// This function has 3 phases:
73// 1. add 3 ssts to
74// 2. trigger a compaction and replace the input from phase 1 with the 1 new sst
75// 3. add 1 new sst
76// Please make sure the function do what you want before using it.
77pub async fn add_test_tables(
78    hummock_manager: &HummockManager,
79    hummock_meta_client: Arc<dyn HummockMetaClient>,
80    compaction_group_id: CompactionGroupId,
81) -> Vec<Vec<SstableInfo>> {
82    // Increase version by 2.
83
84    use risingwave_common::util::epoch::EpochExt;
85
86    let mut epoch = test_epoch(1);
87    let sstable_ids = get_sst_ids(hummock_manager, 3).await;
88    let test_tables = generate_test_sstables_with_table_id(epoch, 1, sstable_ids);
89    register_sstable_infos_to_compaction_group(hummock_manager, &test_tables, compaction_group_id)
90        .await;
91    let test_local_tables = to_local_sstable_info(&test_tables);
92    hummock_meta_client
93        .commit_epoch(
94            epoch,
95            SyncResult {
96                uncommitted_ssts: test_local_tables,
97                ..Default::default()
98            },
99        )
100        .await
101        .unwrap();
102
103    // Simulate a compaction and increase version by 1.
104    let test_tables_2 = generate_test_tables(epoch, get_sst_ids(hummock_manager, 1).await);
105    register_sstable_infos_to_compaction_group(
106        hummock_manager,
107        &test_tables_2,
108        compaction_group_id,
109    )
110    .await;
111    let mut compact_task = hummock_manager
112        .get_compact_task(
113            1.into(),
114            compaction_group_id,
115            &mut *default_compaction_selector(),
116        )
117        .await
118        .unwrap()
119        .unwrap();
120    assert_eq!(
121        compact_task
122            .input_ssts
123            .iter()
124            .map(|i| i.table_infos.len())
125            .sum::<usize>(),
126        3
127    );
128
129    compact_task.target_level = 6;
130    hummock_manager
131        .report_compact_task_for_test(
132            compact_task.task_id,
133            Some(compact_task),
134            TaskStatus::Success,
135            test_tables_2.clone(),
136            None,
137        )
138        .await
139        .unwrap();
140    // Increase version by 1.
141    epoch.inc_epoch();
142    let test_tables_3 = generate_test_tables(epoch, get_sst_ids(hummock_manager, 1).await);
143    register_sstable_infos_to_compaction_group(
144        hummock_manager,
145        &test_tables_3,
146        compaction_group_id,
147    )
148    .await;
149    let test_local_tables_3 = to_local_sstable_info(&test_tables_3);
150    hummock_meta_client
151        .commit_epoch(
152            epoch,
153            SyncResult {
154                uncommitted_ssts: test_local_tables_3,
155                ..Default::default()
156            },
157        )
158        .await
159        .unwrap();
160    vec![test_tables, test_tables_2, test_tables_3]
161}
162
163pub fn generate_test_sstables_with_table_id(
164    epoch: u64,
165    table_id: impl Into<TableId>,
166    sst_ids: Vec<u64>,
167) -> Vec<SstableInfo> {
168    let table_id = table_id.into();
169    let mut sst_info = vec![];
170    for (i, sst_id) in sst_ids.into_iter().enumerate() {
171        let object_size = 2;
172        sst_info.push(
173            SstableInfoInner {
174                object_id: sst_id.into(),
175                sst_id: sst_id.into(),
176                key_range: KeyRange {
177                    left: Bytes::from(key_with_epoch(
178                        format!("{:03}\0\0_key_test_{:05}", table_id, i + 1)
179                            .as_bytes()
180                            .to_vec(),
181                        epoch,
182                    )),
183                    right: Bytes::from(key_with_epoch(
184                        format!("{:03}\0\0_key_test_{:05}", table_id, (i + 1) * 10)
185                            .as_bytes()
186                            .to_vec(),
187                        epoch,
188                    )),
189                    right_exclusive: false,
190                },
191                file_size: object_size,
192                table_ids: vec![table_id],
193                uncompressed_file_size: object_size,
194                max_epoch: epoch,
195                sst_size: object_size,
196                ..Default::default()
197            }
198            .into(),
199        );
200    }
201    sst_info
202}
203
204pub fn generate_test_tables(epoch: u64, sst_ids: Vec<u64>) -> Vec<SstableInfo> {
205    let mut sst_info = vec![];
206    for (i, sst_id) in sst_ids.into_iter().enumerate() {
207        let object_size = 2;
208        sst_info.push(
209            SstableInfoInner {
210                object_id: sst_id.into(),
211                sst_id: sst_id.into(),
212                key_range: KeyRange {
213                    left: Bytes::from(iterator_test_key_of_epoch(sst_id, i + 1, epoch)),
214                    right: Bytes::from(iterator_test_key_of_epoch(sst_id, (i + 1) * 10, epoch)),
215                    right_exclusive: false,
216                },
217                file_size: object_size,
218                table_ids: vec![(sst_id as u32).into(), (sst_id as u32 * 10000).into()],
219                uncompressed_file_size: object_size,
220                max_epoch: epoch,
221                sst_size: object_size,
222                ..Default::default()
223            }
224            .into(),
225        );
226    }
227    sst_info
228}
229
230pub async fn register_sstable_infos_to_compaction_group(
231    compaction_group_manager_ref: &HummockManager,
232    sstable_infos: &[SstableInfo],
233    compaction_group_id: CompactionGroupId,
234) {
235    let table_ids = sstable_infos
236        .iter()
237        .flat_map(|sstable_info| &sstable_info.table_ids)
238        .sorted()
239        .dedup()
240        .copied()
241        .collect_vec();
242    register_table_ids_to_compaction_group(
243        compaction_group_manager_ref,
244        &table_ids,
245        compaction_group_id,
246    )
247    .await;
248}
249
250pub async fn register_table_ids_to_compaction_group(
251    hummock_manager_ref: &HummockManager,
252    table_ids: &[impl Into<TableId> + Copy],
253    compaction_group_id: CompactionGroupId,
254) {
255    hummock_manager_ref
256        .register_table_ids_for_test(
257            &table_ids
258                .iter()
259                .map(|table_id| (*table_id, compaction_group_id))
260                .collect_vec(),
261        )
262        .await
263        .unwrap();
264}
265
266pub async fn unregister_table_ids_from_compaction_group(
267    hummock_manager_ref: &HummockManager,
268    table_ids: &[impl Into<TableId> + Copy],
269) {
270    hummock_manager_ref
271        .unregister_table_ids(table_ids.iter().map(|table_id| (*table_id).into()))
272        .await
273        .unwrap();
274}
275
276/// Generate keys like `001_key_test_00002` with timestamp `epoch`.
277pub fn iterator_test_key_of_epoch(table: u64, idx: usize, ts: HummockEpoch) -> Vec<u8> {
278    // key format: {prefix_index}_version
279    key_with_epoch(
280        format!("{:03}\0\0_key_test_{:05}", table, idx)
281            .as_bytes()
282            .to_vec(),
283        ts,
284    )
285}
286
287pub fn get_sorted_object_ids(sstables: &[SstableInfo]) -> Vec<HummockSstableObjectId> {
288    sstables
289        .iter()
290        .map(|table| table.object_id)
291        .sorted()
292        .collect_vec()
293}
294
295pub fn get_sorted_committed_object_ids(
296    hummock_version: &HummockVersion,
297    compaction_group_id: CompactionGroupId,
298) -> Vec<HummockSstableObjectId> {
299    let levels = match hummock_version.levels.get(&compaction_group_id) {
300        Some(levels) => levels,
301        None => return vec![],
302    };
303    levels
304        .levels
305        .iter()
306        .chain(levels.l0.sub_levels.iter())
307        .flat_map(|levels| levels.table_infos.iter().map(|info| info.object_id))
308        .sorted()
309        .collect_vec()
310}
311
312pub async fn setup_compute_env_with_config(
313    port: i32,
314    config: CompactionConfig,
315) -> (
316    MetaSrvEnv,
317    HummockManagerRef,
318    ClusterControllerRef,
319    WorkerId,
320) {
321    setup_compute_env_with_metric(port, config, None).await
322}
323
324pub async fn setup_compute_env_with_metric(
325    port: i32,
326    config: CompactionConfig,
327    meta_metric: Option<MetaMetrics>,
328) -> (
329    MetaSrvEnv,
330    HummockManagerRef,
331    ClusterControllerRef,
332    WorkerId,
333) {
334    let env = MetaSrvEnv::for_test().await;
335    let cluster_ctl = Arc::new(
336        ClusterController::new(env.clone(), Duration::from_secs(1))
337            .await
338            .unwrap(),
339    );
340    let catalog_ctl = Arc::new(CatalogController::new(env.clone()).await.unwrap());
341
342    let compactor_manager = Arc::new(CompactorManager::for_test());
343
344    let (compactor_streams_change_tx, _compactor_streams_change_rx) =
345        tokio::sync::mpsc::unbounded_channel();
346
347    let hummock_manager = HummockManager::with_config(
348        env.clone(),
349        cluster_ctl.clone(),
350        catalog_ctl,
351        Arc::new(meta_metric.unwrap_or_default()),
352        compactor_manager,
353        config,
354        compactor_streams_change_tx,
355    )
356    .await;
357
358    let fake_host_address = HostAddress {
359        host: "127.0.0.1".to_owned(),
360        port,
361    };
362    let fake_parallelism = 4;
363    let worker_id = cluster_ctl
364        .add_worker(
365            WorkerType::ComputeNode,
366            fake_host_address,
367            Property {
368                is_streaming: true,
369                is_serving: true,
370                parallelism: fake_parallelism as _,
371                ..Default::default()
372            },
373            Default::default(),
374        )
375        .await
376        .unwrap();
377    (env, hummock_manager, cluster_ctl, worker_id)
378}
379
380pub async fn setup_compute_env(
381    port: i32,
382) -> (
383    MetaSrvEnv,
384    HummockManagerRef,
385    ClusterControllerRef,
386    WorkerId,
387) {
388    let config = CompactionConfigBuilder::new()
389        .level0_tier_compact_file_number(1)
390        .level0_max_compact_file_number(130)
391        .level0_sub_level_compact_level_count(1)
392        .level0_overlapping_sub_level_compact_level_count(1)
393        .build();
394    setup_compute_env_with_config(port, config).await
395}
396
397pub async fn get_sst_ids(hummock_manager: &HummockManager, number: u32) -> Vec<u64> {
398    let range = hummock_manager.get_new_object_ids(number).await.unwrap();
399    (range.start_id..range.end_id)
400        .map(|id| id.as_raw_id())
401        .collect_vec()
402}
403
404pub async fn add_ssts(
405    epoch: HummockEpoch,
406    hummock_manager: &HummockManager,
407    hummock_meta_client: Arc<dyn HummockMetaClient>,
408) -> Vec<SstableInfo> {
409    let table_ids = get_sst_ids(hummock_manager, 3).await;
410    let test_tables = generate_test_sstables_with_table_id(test_epoch(epoch), 1, table_ids);
411    let ssts = to_local_sstable_info(&test_tables);
412    hummock_meta_client
413        .commit_epoch(
414            epoch,
415            SyncResult {
416                uncommitted_ssts: ssts,
417                ..Default::default()
418            },
419        )
420        .await
421        .unwrap();
422    test_tables
423}
424
425pub fn compaction_selector_context<'a>(
426    group: &'a CompactionGroup,
427    levels: &'a Levels,
428    member_table_ids: &'a BTreeSet<TableId>,
429    level_handlers: &'a mut [LevelHandler],
430    selector_stats: &'a mut LocalSelectorStatistic,
431    table_id_to_options: &'a HashMap<TableId, TableOption>,
432    developer_config: Arc<CompactionDeveloperConfig>,
433    table_watermarks: &'a HashMap<TableId, Arc<TableWatermarks>>,
434    state_table_info: &'a HummockVersionStateTableInfo,
435) -> CompactionSelectorContext<'a> {
436    CompactionSelectorContext {
437        group,
438        levels,
439        member_table_ids,
440        level_handlers,
441        selector_stats,
442        table_id_to_options,
443        developer_config,
444        table_watermarks,
445        state_table_info,
446        in_progress_compactions: &EMPTY_IN_PROGRESS_COMPACTION_VIEW,
447    }
448}
449
450pub async fn get_compaction_group_id_by_table_id(
451    hummock_manager_ref: HummockManagerRef,
452    table_id: impl Into<TableId>,
453) -> CompactionGroupId {
454    let version = hummock_manager_ref.get_current_version().await;
455    let mapping = version.state_table_info.build_table_compaction_group_id();
456    *mapping.get(&table_id.into()).unwrap()
457}
risingwave_meta/hummock/test_utils.rs

risingwave_meta/hummock/
test_utils.rs
