risingwave_meta/hummock/
level_handler.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::collections::HashMap;
16
17use itertools::Itertools;
18use risingwave_hummock_sdk::level::Level;
19use risingwave_hummock_sdk::sstable_info::SstableInfo;
20use risingwave_hummock_sdk::{HummockCompactionTaskId, HummockSstableId};
21use risingwave_pb::hummock::level_handler::RunningCompactTask;
22
23#[derive(Clone, Debug, PartialEq)]
24pub struct LevelHandler {
25    level: u32,
26    compacting_files: HashMap<HummockSstableId, HummockCompactionTaskId>,
27    pending_tasks: Vec<RunningCompactTask>,
28}
29
30impl LevelHandler {
31    pub fn new(level: u32) -> Self {
32        Self {
33            level,
34            compacting_files: HashMap::default(),
35            pending_tasks: vec![],
36        }
37    }
38
39    pub fn get_level(&self) -> u32 {
40        self.level
41    }
42
43    pub fn remove_task(&mut self, target_task_id: u64) {
44        for task in &self.pending_tasks {
45            if task.task_id == target_task_id {
46                for sst in &task.ssts {
47                    self.compacting_files.remove(sst);
48                }
49            }
50        }
51        self.pending_tasks
52            .retain(|task| task.task_id != target_task_id);
53    }
54
55    pub fn is_pending_compact(&self, sst_id: &HummockSstableId) -> bool {
56        self.compacting_files.contains_key(sst_id)
57    }
58
59    pub fn is_level_pending_compact(&self, level: &Level) -> bool {
60        level
61            .table_infos
62            .iter()
63            .any(|table| self.compacting_files.contains_key(&table.sst_id))
64    }
65
66    pub fn is_level_all_pending_compact(&self, level: &Level) -> bool {
67        if level.table_infos.is_empty() {
68            return false;
69        }
70
71        level
72            .table_infos
73            .iter()
74            .all(|table| self.compacting_files.contains_key(&table.sst_id))
75    }
76
77    pub fn add_pending_task<'a>(
78        &mut self,
79        task_id: u64,
80        target_level: usize,
81        ssts: impl IntoIterator<Item = &'a SstableInfo>,
82    ) {
83        let target_level = target_level as u32;
84        let mut table_ids = vec![];
85        let mut total_file_size = 0;
86        for sst in ssts {
87            self.compacting_files.insert(sst.sst_id, task_id);
88            total_file_size += sst.sst_size;
89            table_ids.push(sst.sst_id);
90        }
91
92        self.pending_tasks.push(RunningCompactTask {
93            task_id,
94            target_level,
95            total_file_size,
96            ssts: table_ids,
97        });
98    }
99
100    pub fn pending_file_count(&self) -> usize {
101        self.compacting_files.len()
102    }
103
104    pub fn pending_output_file_size(&self, target_level: u32) -> u64 {
105        self.pending_tasks
106            .iter()
107            .filter(|task| task.target_level == target_level)
108            .map(|task| task.total_file_size)
109            .sum::<u64>()
110    }
111
112    pub fn pending_tasks_ids(&self) -> Vec<u64> {
113        self.pending_tasks
114            .iter()
115            .map(|task| task.task_id)
116            .collect_vec()
117    }
118
119    pub fn pending_tasks(&self) -> &[RunningCompactTask] {
120        &self.pending_tasks
121    }
122
123    pub fn compacting_files(&self) -> &HashMap<HummockSstableId, HummockCompactionTaskId> {
124        &self.compacting_files
125    }
126
127    #[cfg(test)]
128    pub(crate) fn test_add_pending_sst(&mut self, sst_id: HummockSstableId, task_id: u64) {
129        self.compacting_files.insert(sst_id, task_id);
130    }
131}
132
133impl From<&LevelHandler> for risingwave_pb::hummock::LevelHandler {
134    fn from(lh: &LevelHandler) -> Self {
135        risingwave_pb::hummock::LevelHandler {
136            level: lh.level,
137            tasks: lh.pending_tasks.clone(),
138        }
139    }
140}
141
142impl From<&risingwave_pb::hummock::LevelHandler> for LevelHandler {
143    fn from(lh: &risingwave_pb::hummock::LevelHandler) -> Self {
144        let mut pending_tasks = vec![];
145        let mut compacting_files = HashMap::new();
146        for task in &lh.tasks {
147            pending_tasks.push(task.clone());
148            for s in &task.ssts {
149                compacting_files.insert(*s, task.task_id);
150            }
151        }
152        Self {
153            pending_tasks,
154            compacting_files,
155            level: lh.level,
156        }
157    }
158}
risingwave_meta/hummock/level_handler.rs

risingwave_meta/hummock/
level_handler.rs
