risingwave_meta/hummock/compaction/picker/

base_level_compaction_picker.rs

// Copyright 2023 RisingWave Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::cell::RefCell;
use std::sync::Arc;

use itertools::Itertools;
use risingwave_common::config::meta::default::compaction_config;
use risingwave_hummock_sdk::level::{InputLevel, Level, Levels, OverlappingLevel};
use risingwave_pb::hummock::{CompactionConfig, LevelType};

use super::non_overlap_sub_level_picker::NonOverlapSubLevelPicker;
use super::{
    CompactionInput, CompactionPicker, CompactionTaskValidator, LocalPickerStatistic,
    ValidationRuleType,
};
use crate::hummock::compaction::picker::TrivialMovePicker;
use crate::hummock::compaction::{CompactionDeveloperConfig, create_overlap_strategy};
use crate::hummock::level_handler::LevelHandler;

std::thread_local! {
    static LOG_COUNTER: RefCell<usize> = const { RefCell::new(0) };
}

pub struct LevelCompactionPicker {
    target_level: usize,
    config: Arc<CompactionConfig>,
    compaction_task_validator: Arc<CompactionTaskValidator>,
    developer_config: Arc<CompactionDeveloperConfig>,
}

impl CompactionPicker for LevelCompactionPicker {
    fn pick_compaction(
        &mut self,
        levels: &Levels,
        level_handlers: &[LevelHandler],
        stats: &mut LocalPickerStatistic,
    ) -> Option<CompactionInput> {
        let l0 = &levels.l0;
        if l0.sub_levels.is_empty() {
            return None;
        }
        if l0.sub_levels[0].level_type != LevelType::Nonoverlapping
            && l0.sub_levels[0].table_infos.len() > 1
        {
            stats.skip_by_overlapping += 1;
            return None;
        }

        let is_l0_pending_compact =
            level_handlers[0].is_level_all_pending_compact(&l0.sub_levels[0]);

        if is_l0_pending_compact {
            stats.skip_by_pending_files += 1;
            return None;
        }

        if let Some(mut ret) = self.pick_base_trivial_move(
            l0,
            levels.get_level(self.target_level),
            level_handlers,
            stats,
        ) {
            ret.vnode_partition_count = self.config.split_weight_by_vnode;
            return Some(ret);
        }

        debug_assert!(self.target_level == levels.get_level(self.target_level).level_idx as usize);
        if let Some(ret) = self.pick_multi_level_to_base(
            l0,
            levels.get_level(self.target_level),
            self.config.split_weight_by_vnode,
            level_handlers,
            stats,
        ) {
            return Some(ret);
        }

        None
    }
}

impl LevelCompactionPicker {
    #[cfg(test)]
    pub fn new(
        target_level: usize,
        config: Arc<CompactionConfig>,
        developer_config: Arc<CompactionDeveloperConfig>,
    ) -> LevelCompactionPicker {
        LevelCompactionPicker {
            target_level,
            compaction_task_validator: Arc::new(CompactionTaskValidator::new(config.clone())),
            config,
            developer_config,
        }
    }

    pub fn new_with_validator(
        target_level: usize,
        config: Arc<CompactionConfig>,
        compaction_task_validator: Arc<CompactionTaskValidator>,
        developer_config: Arc<CompactionDeveloperConfig>,
    ) -> LevelCompactionPicker {
        LevelCompactionPicker {
            target_level,
            config,
            compaction_task_validator,
            developer_config,
        }
    }

    fn pick_base_trivial_move(
        &self,
        l0: &OverlappingLevel,
        target_level: &Level,
        level_handlers: &[LevelHandler],
        stats: &mut LocalPickerStatistic,
    ) -> Option<CompactionInput> {
        if !self.developer_config.enable_trivial_move {
            return None;
        }

        let overlap_strategy = create_overlap_strategy(self.config.compaction_mode());
        let trivial_move_picker = TrivialMovePicker::new(
            0,
            self.target_level,
            overlap_strategy.clone(),
            if self.compaction_task_validator.is_enable() {
                // tips: Older versions of the compaction group will be upgraded without this configuration, we leave it with its default behaviour and enable it manually when needed.
                self.config.sst_allowed_trivial_move_min_size.unwrap_or(0)
            } else {
                0
            },
            self.config
                .sst_allowed_trivial_move_max_count
                .unwrap_or(compaction_config::sst_allowed_trivial_move_max_count())
                as usize,
        );

        trivial_move_picker.pick_trivial_move_task(
            &l0.sub_levels[0].table_infos,
            &target_level.table_infos,
            level_handlers,
            stats,
        )
    }

    fn pick_multi_level_to_base(
        &self,
        l0: &OverlappingLevel,
        target_level: &Level,
        vnode_partition_count: u32,
        level_handlers: &[LevelHandler],
        stats: &mut LocalPickerStatistic,
    ) -> Option<CompactionInput> {
        let overlap_strategy = create_overlap_strategy(self.config.compaction_mode());
        let min_compaction_bytes = self.config.sub_level_max_compaction_bytes;
        let non_overlap_sub_level_picker = NonOverlapSubLevelPicker::new(
            min_compaction_bytes,
            // divide by 2 because we need to select files of base level and it need use the other
            // half quota.
            std::cmp::max(
                self.config.max_bytes_for_level_base,
                self.config.max_compaction_bytes / 2,
            ),
            1,
            // The maximum number of sub_level compact level per task
            self.config.level0_max_compact_file_number,
            overlap_strategy.clone(),
            self.developer_config.enable_check_task_level_overlap,
            self.config
                .max_l0_compact_level_count
                .unwrap_or(compaction_config::max_l0_compact_level_count()) as usize,
            self.config
                .enable_optimize_l0_interval_selection
                .unwrap_or(compaction_config::enable_optimize_l0_interval_selection()),
        );

        let mut max_vnode_partition_idx = 0;
        for (idx, level) in l0.sub_levels.iter().enumerate() {
            if level.vnode_partition_count < vnode_partition_count {
                break;
            }
            max_vnode_partition_idx = idx;
        }

        let candidate_l0_plans = non_overlap_sub_level_picker.pick_l0_multi_non_overlap_level(
            &l0.sub_levels[..=max_vnode_partition_idx],
            &level_handlers[0],
        );
        if candidate_l0_plans.is_empty() {
            stats.skip_by_pending_files += 1;
            return None;
        }

        let mut all_pending = true;
        for input in candidate_l0_plans {
            let mut overlap_info = overlap_strategy.create_overlap_info();
            for (_, tables) in &input.sstable_infos {
                for sst in tables {
                    overlap_info.update(&sst.key_range);
                }
            }
            let target_range = overlap_info.check_multiple_overlap(&target_level.table_infos);
            let target_level_files = if target_range.is_empty() {
                &[][..]
            } else {
                &target_level.table_infos[target_range]
            };

            let mut target_file_size = 0;
            let mut pending_compact = false;
            for sst in target_level_files {
                if level_handlers[target_level.level_idx as usize].is_pending_compact(&sst.sst_id) {
                    pending_compact = true;
                    break;
                }
                target_file_size += sst.sst_size;
            }
            if pending_compact {
                continue;
            }
            all_pending = false;

            let mut select_level_inputs = input
                .sstable_infos
                .into_iter()
                .map(|(_, table_infos)| InputLevel {
                    level_idx: 0,
                    level_type: LevelType::Nonoverlapping,
                    table_infos,
                })
                .collect_vec();
            select_level_inputs.reverse();
            let target_file_count = target_level_files.len();
            select_level_inputs.push(InputLevel {
                level_idx: target_level.level_idx,
                level_type: target_level.level_type,
                table_infos: target_level_files.to_vec(),
            });

            let result = CompactionInput {
                input_levels: select_level_inputs,
                target_level: self.target_level,
                select_input_size: input.total_file_size,
                target_input_size: target_file_size,
                total_file_count: (input.total_file_count + target_file_count) as u64,
                vnode_partition_count,
                ..Default::default()
            };

            if !self.compaction_task_validator.valid_compact_task(
                &result,
                ValidationRuleType::ToBase,
                stats,
            ) {
                if l0.total_file_size > target_level.total_file_size * 8 {
                    let log_counter = LOG_COUNTER.with_borrow_mut(|counter| {
                        *counter += 1;
                        *counter
                    });

                    // reduce log
                    if log_counter.is_multiple_of(100) {
                        tracing::warn!(
                            "skip task with level count: {}, file count: {}, select size: {}, target size: {}, target level size: {}",
                            result.input_levels.len(),
                            result.total_file_count,
                            result.select_input_size,
                            result.target_input_size,
                            target_level.total_file_size,
                        );
                    }
                }
                continue;
            }

            return Some(result);
        }
        // Preserve the pending statistic: validator rejection is a separate reason to skip.
        if all_pending {
            stats.skip_by_pending_files += 1;
        }
        None
    }
}

#[cfg(test)]
pub mod tests {

    use super::*;
    use crate::hummock::compaction::compaction_config::CompactionConfigBuilder;
    use crate::hummock::compaction::selector::tests::*;
    use crate::hummock::compaction::{CompactionMode, TierCompactionPicker};

    fn candidate_fixture(count: usize, target_size: u64) -> (OverlappingLevel, Level) {
        let sub_levels = (0..2)
            .map(|level_idx| {
                let table_infos = (0..count)
                    .map(|idx| {
                        generate_table(
                            (level_idx * count + idx) as u64,
                            1,
                            idx * 10,
                            idx * 10 + 9,
                            1,
                        )
                    })
                    .collect();
                Level {
                    sub_level_id: level_idx as u64,
                    level_type: LevelType::Nonoverlapping,
                    table_infos,
                    total_file_size: count as u64 * 10,
                    ..Default::default()
                }
            })
            .collect();
        let table_infos = (0..count)
            .map(|idx| {
                let mut sst =
                    generate_table_impl((2 * count + idx) as u64, 1, idx * 10, idx * 10 + 9, 1);
                sst.sst_size = target_size;
                sst.into()
            })
            .collect();
        (
            OverlappingLevel {
                sub_levels,
                total_file_size: count as u64 * 20,
                ..Default::default()
            },
            Level {
                level_idx: 1,
                level_type: LevelType::Nonoverlapping,
                table_infos,
                total_file_size: count as u64 * target_size,
                ..Default::default()
            },
        )
    }

    #[test]
    fn test_candidate_pending_and_validation_statistics() {
        let picker = create_compaction_picker_for_test();
        for (target_size, pending_count, succeeds, write_amp_skips, pending_skips) in [
            (10, 0, true, 0, 0),
            (10, 7, true, 0, 0),
            (10, 8, false, 0, 1),
            (100, 0, false, 8, 0),
            (100, 7, false, 1, 0),
            (100, 8, false, 0, 1),
        ] {
            let (l0, target) = candidate_fixture(8, target_size);
            let mut handlers = vec![LevelHandler::new(0), LevelHandler::new(1)];
            for sst in target.table_infos.iter().take(pending_count) {
                handlers[1].test_add_pending_sst(sst.sst_id, 1);
            }
            let mut stats = LocalPickerStatistic::default();
            let result = picker.pick_multi_level_to_base(&l0, &target, 0, &handlers, &mut stats);
            assert_eq!(
                result.is_some(),
                succeeds,
                "target_size={target_size}, pending={pending_count}"
            );
            assert_eq!(stats.skip_by_write_amp_limit, write_amp_skips);
            assert_eq!(stats.skip_by_pending_files, pending_skips);
            assert_eq!(stats.skip_by_count_limit, 0);
            assert_eq!(stats.skip_by_overlapping, 0);
            if let Some(result) = result {
                assert_eq!(
                    result.input_levels[0].table_infos[0].sst_id,
                    (8 + pending_count) as u64
                );
                assert_eq!(
                    result.input_levels[1].table_infos[0].sst_id,
                    pending_count as u64
                );
                assert_eq!(
                    result.input_levels[2].table_infos[0].sst_id,
                    (16 + pending_count) as u64
                );
                assert_eq!(result.select_input_size, 20);
                assert_eq!(result.target_input_size, 10);
                assert_eq!(result.total_file_count, 3);
            }
        }
    }

    // Run explicitly with cargo test --release -p risingwave_meta bench_level_to_base -- --ignored --nocapture --test-threads=1.
    #[test]
    #[ignore = "microbenchmark"]
    fn bench_level_to_base() {
        use std::hint::black_box;
        use std::time::Duration;

        let mut criterion = criterion::Criterion::default()
            .sample_size(30)
            .warm_up_time(Duration::from_secs(1))
            .measurement_time(Duration::from_secs(3));
        let picker = create_compaction_picker_for_test();
        for (name, count, target_size, pending_count) in [
            ("small", 8, 10, 0),
            ("first_valid", 256, 10, 0),
            ("late_valid", 256, 10, 255),
            ("all_pending", 256, 10, 256),
            ("all_rejected", 256, 100, 0),
        ] {
            let (l0, target) = candidate_fixture(count, target_size);
            let mut handlers = vec![LevelHandler::new(0), LevelHandler::new(1)];
            for sst in target.table_infos.iter().take(pending_count) {
                handlers[1].test_add_pending_sst(sst.sst_id, 1);
            }
            criterion.bench_function(&format!("level_to_base/{name}"), |b| {
                b.iter(|| {
                    let mut stats = LocalPickerStatistic::default();
                    let result = picker.pick_multi_level_to_base(
                        black_box(&l0),
                        black_box(&target),
                        0,
                        black_box(&handlers),
                        &mut stats,
                    );
                    black_box((result, stats))
                })
            });
        }
        criterion.final_summary();
    }

    fn create_compaction_picker_for_test() -> LevelCompactionPicker {
        let config = Arc::new(
            CompactionConfigBuilder::new()
                .level0_tier_compact_file_number(2)
                .level0_sub_level_compact_level_count(1)
                .build(),
        );
        LevelCompactionPicker::new(1, config, Arc::new(CompactionDeveloperConfig::default()))
    }

    #[test]
    fn test_compact_l0_to_l1() {
        let mut picker = create_compaction_picker_for_test();
        let l0 = generate_level(
            0,
            vec![
                generate_table(5, 1, 100, 200, 2),
                generate_table(4, 1, 201, 300, 2),
            ],
        );
        let mut levels = Levels {
            l0: OverlappingLevel {
                total_file_size: l0.total_file_size,
                uncompressed_file_size: l0.total_file_size,
                sub_levels: vec![l0],
            },
            levels: vec![generate_level(
                1,
                vec![
                    generate_table(3, 1, 1, 100, 1),
                    generate_table(2, 1, 101, 150, 1),
                    generate_table(1, 1, 201, 210, 1),
                ],
            )],
            ..Default::default()
        };
        let mut local_stats = LocalPickerStatistic::default();
        let mut levels_handler = vec![LevelHandler::new(0), LevelHandler::new(1)];

        let ret = picker
            .pick_compaction(&levels, &levels_handler, &mut local_stats)
            .unwrap();
        assert_eq!(ret.input_levels[0].table_infos.len(), 1);
        assert_eq!(ret.input_levels[0].table_infos[0].sst_id, 4);
        assert_eq!(ret.input_levels[1].table_infos[0].sst_id, 1);

        ret.add_pending_task(0, &mut levels_handler);
        {
            push_table_level0_nonoverlapping(&mut levels, generate_table(6, 1, 100, 200, 2));
            push_table_level0_nonoverlapping(&mut levels, generate_table(7, 1, 301, 333, 4));
            let ret2 = picker
                .pick_compaction(&levels, &levels_handler, &mut local_stats)
                .unwrap();

            assert_eq!(ret2.input_levels[0].table_infos.len(), 1);
            assert_eq!(ret2.input_levels[0].table_infos[0].sst_id, 6);
            assert_eq!(ret2.input_levels[1].table_infos[0].sst_id, 5);
        }

        levels.l0.sub_levels[0]
            .table_infos
            .retain(|table| table.sst_id != 4);
        levels.l0.total_file_size -= ret.input_levels[0].table_infos[0].file_size;

        levels_handler[0].remove_task(0);
        levels_handler[1].remove_task(0);

        let ret = picker
            .pick_compaction(&levels, &levels_handler, &mut local_stats)
            .unwrap();
        assert_eq!(ret.input_levels.len(), 3);
        assert_eq!(ret.input_levels[0].table_infos[0].sst_id, 6);
        assert_eq!(ret.input_levels[1].table_infos[0].sst_id, 5);
        assert_eq!(ret.input_levels[2].table_infos.len(), 2);
        assert_eq!(ret.input_levels[2].table_infos[0].sst_id, 3);
        assert_eq!(ret.input_levels[2].table_infos[1].sst_id, 2);
        ret.add_pending_task(1, &mut levels_handler);

        let mut local_stats = LocalPickerStatistic::default();
        // Cannot pick because no idle table in sub-level[0]. (And sub-level[0] is pending
        // actually).
        push_table_level0_overlapping(&mut levels, generate_table(8, 1, 199, 233, 3));
        let ret = picker.pick_compaction(&levels, &levels_handler, &mut local_stats);
        assert!(ret.is_none());

        // Don't pick overlapping sub-level 8
        levels_handler[0].remove_task(1);
        levels_handler[1].remove_task(1);
        let ret = picker
            .pick_compaction(&levels, &levels_handler, &mut local_stats)
            .unwrap();
        assert_eq!(ret.input_levels.len(), 3);
        assert_eq!(ret.input_levels[0].table_infos[0].sst_id, 6);
        assert_eq!(ret.input_levels[1].table_infos[0].sst_id, 5);
        assert_eq!(ret.input_levels[2].table_infos.len(), 2);
    }

    #[test]
    fn test_selecting_key_range_overlap() {
        // When picking L0->L1, all L1 files overlapped with selecting_key_range should be picked.
        let config = Arc::new(
            CompactionConfigBuilder::new()
                .level0_tier_compact_file_number(2)
                .compaction_mode(CompactionMode::Range as i32)
                .level0_sub_level_compact_level_count(1)
                .enable_optimize_l0_interval_selection(Some(false))
                .build(),
        );

        let config_enable_optimize_l0_interval_selection = Arc::new(
            CompactionConfigBuilder::new()
                .level0_tier_compact_file_number(2)
                .compaction_mode(CompactionMode::Range as i32)
                .level0_sub_level_compact_level_count(1)
                .enable_optimize_l0_interval_selection(Some(true))
                .build(),
        );

        let mut picker =
            LevelCompactionPicker::new(1, config, Arc::new(CompactionDeveloperConfig::default()));

        let mut picker_enable_optimize_l0_interval_selection = LevelCompactionPicker::new(
            1,
            config_enable_optimize_l0_interval_selection,
            Arc::new(CompactionDeveloperConfig::default()),
        );

        let levels = vec![Level {
            level_idx: 1,
            level_type: LevelType::Nonoverlapping,
            table_infos: vec![
                generate_table(3, 1, 0, 50, 1),
                generate_table(4, 1, 150, 180, 1),
                generate_table(5, 1, 250, 300, 1),
            ],
            ..Default::default()
        }];
        let mut levels = Levels {
            levels,
            l0: OverlappingLevel {
                sub_levels: vec![],
                total_file_size: 0,
                uncompressed_file_size: 0,
            },
            ..Default::default()
        };
        push_tables_level0_nonoverlapping(&mut levels, vec![generate_table(1, 1, 50, 140, 2)]);
        push_tables_level0_nonoverlapping(
            &mut levels,
            vec![
                generate_table(7, 1, 200, 250, 2),
                generate_table(8, 1, 400, 500, 2),
            ],
        );

        {
            let mut levels_handler = vec![LevelHandler::new(0), LevelHandler::new(1)];

            let mut local_stats = LocalPickerStatistic::default();
            let ret = picker_enable_optimize_l0_interval_selection
                .pick_compaction(&levels, &levels_handler, &mut local_stats)
                .unwrap();

            // pick
            // l0 [sst_8]
            assert_eq!(ret.input_levels.len(), 2);
            assert_eq!(
                ret.input_levels[0]
                    .table_infos
                    .iter()
                    .map(|t| t.sst_id)
                    .collect_vec(),
                vec![8]
            );
            // trivial_move
            assert!(ret.input_levels[1].table_infos.is_empty());

            ret.add_pending_task(0, &mut levels_handler);

            let ret = picker_enable_optimize_l0_interval_selection
                .pick_compaction(&levels, &levels_handler, &mut local_stats)
                .unwrap();

            assert_eq!(ret.input_levels.len(), 2);
            assert_eq!(
                ret.input_levels[0]
                    .table_infos
                    .iter()
                    .map(|t| t.sst_id)
                    .collect_vec(),
                vec![7]
            );

            assert_eq!(
                ret.input_levels[1]
                    .table_infos
                    .iter()
                    .map(|t| t.sst_id)
                    .collect_vec(),
                vec![5]
            );
        }

        {
            let levels_handler = vec![LevelHandler::new(0), LevelHandler::new(1)];

            let mut local_stats = LocalPickerStatistic::default();
            let ret = picker
                .pick_compaction(&levels, &levels_handler, &mut local_stats)
                .unwrap();

            // pick
            assert_eq!(ret.input_levels.len(), 2);
            assert_eq!(
                ret.input_levels[0]
                    .table_infos
                    .iter()
                    .map(|t| t.sst_id.as_raw_id())
                    .collect_vec(),
                vec![1]
            );

            assert_eq!(
                ret.input_levels[1]
                    .table_infos
                    .iter()
                    .map(|t| t.sst_id.as_raw_id())
                    .collect_vec(),
                vec![3]
            );
        }
    }

    #[test]
    fn test_l0_to_l1_compact_conflict() {
        // When picking L0->L1, L0's selecting_key_range should not be overlapped with L0's
        // compacting_key_range.
        let mut picker = create_compaction_picker_for_test();
        let levels = vec![Level {
            level_idx: 1,
            level_type: LevelType::Nonoverlapping,
            table_infos: vec![],
            total_file_size: 0,
            sub_level_id: 0,
            uncompressed_file_size: 0,
            ..Default::default()
        }];
        let mut levels = Levels {
            levels,
            l0: OverlappingLevel {
                sub_levels: vec![],
                total_file_size: 0,
                uncompressed_file_size: 0,
            },
            ..Default::default()
        };
        push_tables_level0_nonoverlapping(
            &mut levels,
            vec![
                generate_table(1, 1, 100, 300, 2),
                generate_table(2, 1, 350, 500, 2),
            ],
        );
        let mut levels_handler = vec![LevelHandler::new(0), LevelHandler::new(1)];

        let mut local_stats = LocalPickerStatistic::default();
        let ret = picker
            .pick_compaction(&levels, &levels_handler, &mut local_stats)
            .unwrap();
        // trivial_move
        ret.add_pending_task(0, &mut levels_handler); // pending only for test
        push_tables_level0_nonoverlapping(&mut levels, vec![generate_table(3, 1, 250, 300, 3)]);
        let config: CompactionConfig = CompactionConfigBuilder::new()
            .level0_tier_compact_file_number(2)
            .max_compaction_bytes(1000)
            .sub_level_max_compaction_bytes(150)
            .max_bytes_for_level_multiplier(1)
            .level0_sub_level_compact_level_count(3)
            .build();
        let mut picker = TierCompactionPicker::new(Arc::new(config));

        let ret: Option<CompactionInput> =
            picker.pick_compaction(&levels, &levels_handler, &mut local_stats);
        assert!(ret.is_none());
    }

    #[test]
    fn test_skip_compact_write_amplification_limit() {
        let config: CompactionConfig = CompactionConfigBuilder::new()
            .level0_tier_compact_file_number(2)
            .max_compaction_bytes(1000)
            .sub_level_max_compaction_bytes(150)
            .max_bytes_for_level_multiplier(1)
            .level0_sub_level_compact_level_count(2)
            .build();
        let mut picker = LevelCompactionPicker::new(
            1,
            Arc::new(config),
            Arc::new(CompactionDeveloperConfig::default()),
        );

        let mut levels = Levels {
            levels: vec![Level {
                level_idx: 1,
                level_type: LevelType::Nonoverlapping,
                table_infos: vec![
                    generate_table(1, 1, 100, 399, 2),
                    generate_table(2, 1, 400, 699, 2),
                    generate_table(3, 1, 700, 999, 2),
                ],
                total_file_size: 900,
                sub_level_id: 0,
                uncompressed_file_size: 900,
                ..Default::default()
            }],
            l0: generate_l0_nonoverlapping_sublevels(vec![]),
            ..Default::default()
        };
        push_tables_level0_nonoverlapping(
            &mut levels,
            vec![
                generate_table(4, 1, 100, 180, 2),
                generate_table(5, 1, 400, 450, 2),
                generate_table(6, 1, 600, 700, 2),
            ],
        );

        let mut levels_handler = vec![LevelHandler::new(0), LevelHandler::new(1)];
        let mut local_stats = LocalPickerStatistic::default();
        levels_handler[0].add_pending_task(1, 4, &levels.l0.sub_levels[0].table_infos);
        let ret = picker.pick_compaction(&levels, &levels_handler, &mut local_stats);
        // Skip this compaction because the write amplification is too large.
        assert!(ret.is_none());
    }

    #[test]
    fn test_l0_to_l1_break_on_exceed_compaction_size() {
        let mut local_stats = LocalPickerStatistic::default();
        let mut l0 = generate_l0_overlapping_sublevels(vec![
            vec![
                generate_table(4, 1, 10, 90, 1),
                generate_table(5, 1, 210, 220, 1),
            ],
            vec![generate_table(6, 1, 0, 100000, 1)],
            vec![generate_table(7, 1, 0, 100000, 1)],
        ]);
        // We can set level_type only because the input above is valid.
        for s in &mut l0.sub_levels {
            s.level_type = LevelType::Nonoverlapping;
        }
        let levels = Levels {
            l0,
            levels: vec![generate_level(1, vec![generate_table(3, 1, 0, 100000, 1)])],
            ..Default::default()
        };
        let levels_handler = vec![LevelHandler::new(0), LevelHandler::new(1)];

        // Pick with large max_compaction_bytes results all sub levels included in input.
        let config = Arc::new(
            CompactionConfigBuilder::new()
                .max_compaction_bytes(500000)
                .sub_level_max_compaction_bytes(50000)
                .max_bytes_for_level_base(500000)
                .level0_sub_level_compact_level_count(1)
                .build(),
        );
        // Only include sub-level 0 results will violate MAX_WRITE_AMPLIFICATION.
        // So all sub-levels are included to make write amplification < MAX_WRITE_AMPLIFICATION.
        let mut picker =
            LevelCompactionPicker::new(1, config, Arc::new(CompactionDeveloperConfig::default()));
        let ret = picker
            .pick_compaction(&levels, &levels_handler, &mut local_stats)
            .unwrap();
        assert_eq!(ret.input_levels[0].table_infos[0].sst_id, 7);
        assert_eq!(
            3,
            ret.input_levels.iter().filter(|l| l.level_idx == 0).count()
        );
        assert_eq!(
            4,
            ret.input_levels
                .iter()
                .filter(|l| l.level_idx == 0)
                .map(|l| l.table_infos.len())
                .sum::<usize>()
        );

        // Pick with small max_compaction_bytes results partial sub levels included in input.
        let config = Arc::new(
            CompactionConfigBuilder::new()
                .max_compaction_bytes(100010)
                .max_bytes_for_level_base(512)
                .level0_sub_level_compact_level_count(1)
                .build(),
        );
        let mut picker =
            LevelCompactionPicker::new(1, config, Arc::new(CompactionDeveloperConfig::default()));

        let ret = picker
            .pick_compaction(&levels, &levels_handler, &mut local_stats)
            .unwrap();
        assert_eq!(ret.input_levels[0].table_infos[0].sst_id, 6);
        assert_eq!(
            2,
            ret.input_levels.iter().filter(|l| l.level_idx == 0).count()
        );
        assert_eq!(
            3,
            ret.input_levels
                .iter()
                .filter(|l| l.level_idx == 0)
                .map(|l| l.table_infos.len())
                .sum::<usize>()
        );
    }

    #[test]
    fn test_l0_to_l1_break_on_pending_sub_level() {
        let l0 = generate_l0_nonoverlapping_multi_sublevels(vec![
            vec![
                generate_table(4, 1, 10, 90, 1),
                generate_table(5, 1, 210, 220, 1),
            ],
            vec![generate_table(6, 1, 0, 100000, 1)],
            vec![generate_table(7, 1, 0, 100000, 1)],
        ]);

        let levels = Levels {
            l0,
            levels: vec![generate_level(1, vec![generate_table(3, 1, 0, 100000, 1)])],
            ..Default::default()
        };
        let mut levels_handler = vec![LevelHandler::new(0), LevelHandler::new(1)];
        let mut local_stats = LocalPickerStatistic::default();

        // Create a pending sub-level.
        let pending_level = levels.l0.sub_levels[1].clone();
        assert_eq!(pending_level.sub_level_id, 1);
        let tier_task_input = CompactionInput {
            input_levels: vec![InputLevel {
                level_idx: 0,
                level_type: pending_level.level_type,
                table_infos: pending_level.table_infos.clone(),
            }],
            target_level: 1,
            target_sub_level_id: pending_level.sub_level_id,
            ..Default::default()
        };
        assert!(!levels_handler[0].is_level_pending_compact(&pending_level));
        tier_task_input.add_pending_task(1, &mut levels_handler);
        assert!(levels_handler[0].is_level_pending_compact(&pending_level));

        // Pick with large max_compaction_bytes results all sub levels included in input.
        let config = Arc::new(
            CompactionConfigBuilder::new()
                .max_compaction_bytes(500000)
                .level0_sub_level_compact_level_count(2)
                .build(),
        );

        // Only include sub-level 0 results will violate MAX_WRITE_AMPLIFICATION.
        // But stopped by pending sub-level when trying to include more sub-levels.
        let mut picker = LevelCompactionPicker::new(
            1,
            config.clone(),
            Arc::new(CompactionDeveloperConfig::default()),
        );
        let ret = picker.pick_compaction(&levels, &levels_handler, &mut local_stats);
        assert!(ret.is_none());

        // Free the pending sub-level.
        for pending_task_id in &levels_handler[0].pending_tasks_ids() {
            levels_handler[0].remove_task(*pending_task_id);
        }

        // No more pending sub-level so we can get a task now.
        let mut picker =
            LevelCompactionPicker::new(1, config, Arc::new(CompactionDeveloperConfig::default()));
        picker
            .pick_compaction(&levels, &levels_handler, &mut local_stats)
            .unwrap();
    }

    #[test]
    fn test_l0_to_base_when_all_base_pending() {
        let l0 = generate_l0_nonoverlapping_multi_sublevels(vec![
            vec![
                generate_table(4, 1, 10, 90, 1),
                generate_table(5, 1, 1000, 2000, 1),
            ],
            vec![generate_table(6, 1, 10, 90, 1)],
        ]);

        let levels = Levels {
            l0,
            levels: vec![generate_level(1, vec![generate_table(3, 1, 1, 100, 1)])],
            ..Default::default()
        };
        let mut levels_handler = vec![LevelHandler::new(0), LevelHandler::new(1)];
        let mut local_stats = LocalPickerStatistic::default();

        let config = Arc::new(
            CompactionConfigBuilder::new()
                .max_compaction_bytes(500000)
                .level0_sub_level_compact_level_count(2)
                .sub_level_max_compaction_bytes(1000)
                .build(),
        );

        let mut picker =
            LevelCompactionPicker::new(1, config, Arc::new(CompactionDeveloperConfig::default()));
        let ret = picker
            .pick_compaction(&levels, &levels_handler, &mut local_stats)
            .unwrap();
        // 1. trivial_move
        assert_eq!(2, ret.input_levels.len());
        assert!(ret.input_levels[1].table_infos.is_empty());
        assert_eq!(5, ret.input_levels[0].table_infos[0].sst_id);
        ret.add_pending_task(0, &mut levels_handler);

        let ret = picker
            .pick_compaction(&levels, &levels_handler, &mut local_stats)
            .unwrap();
        assert_eq!(3, ret.input_levels.len());
        assert_eq!(6, ret.input_levels[0].table_infos[0].sst_id);
    }
}