Skip to main content

mito2/worker/
handle_rebuild_index.rs

1// Copyright 2023 Greptime Team
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//! Handles index build requests.
16
17use std::collections::HashMap;
18use std::sync::Arc;
19
20use common_telemetry::{debug, error, warn};
21use store_api::region_request::RegionBuildIndexRequest;
22use store_api::storage::{FileId, RegionId};
23use tokio::sync::mpsc;
24
25use crate::cache::CacheStrategy;
26use crate::error::Result;
27use crate::manifest::action::RegionEdit;
28use crate::metrics::INDEX_PUBLICATION_STALE_TOTAL;
29use crate::region::version::VersionRef;
30use crate::region::{IndexBuildSource, MitoRegionRef};
31use crate::request::{
32    BuildIndexRequest, IndexBuildFailed, IndexBuildFinished, IndexBuildStopped, OptionOutputTx,
33};
34use crate::sst::file::{FileHandle, FileMeta, RegionFileId, RegionIndexId};
35use crate::sst::index::{
36    IndexBuildOutcome, IndexBuildTask, IndexBuildType, IndexerBuilderImpl, ResultMpscSender,
37    cleanup_stale_index_caches,
38};
39use crate::worker::RegionWorkerLoop;
40
41impl<S> RegionWorkerLoop<S> {
42    pub(crate) fn new_index_build_task(
43        &self,
44        region: &MitoRegionRef,
45        version: &VersionRef,
46        file: FileHandle,
47        file_meta: FileMeta,
48        build_type: IndexBuildType,
49        result_sender: ResultMpscSender,
50    ) -> IndexBuildTask {
51        let access_layer = region.access_layer.clone();
52
53        let puffin_manager = if let Some(write_cache) = self.cache_manager.write_cache() {
54            write_cache.build_puffin_manager()
55        } else {
56            access_layer.build_puffin_manager()
57        };
58
59        let intermediate_manager = if let Some(write_cache) = self.cache_manager.write_cache() {
60            write_cache.intermediate_manager().clone()
61        } else {
62            access_layer.intermediate_manager().clone()
63        };
64
65        let indexer_builder_ref = Arc::new(IndexerBuilderImpl {
66            build_type: build_type.clone(),
67            metadata: version.metadata.clone(),
68            inverted_index_config: self.config.inverted_index.clone(),
69            fulltext_index_config: self.config.fulltext_index.clone(),
70            bloom_filter_index_config: self.config.bloom_filter_index.clone(),
71            #[cfg(feature = "vector_index")]
72            vector_index_config: self.config.vector_index.clone(),
73            index_options: version.options.index_options.clone(),
74            intermediate_manager,
75            puffin_manager,
76            write_cache_enabled: self.cache_manager.write_cache().is_some(),
77        });
78
79        IndexBuildTask {
80            region_id: region.region_id,
81            file: file.clone(),
82            source: IndexBuildSource::new(file_meta, version.metadata.schema_version),
83            reason: build_type,
84            access_layer: access_layer.clone(),
85            listener: self.listener.clone(),
86            manifest_ctx: region.manifest_ctx.clone(),
87            write_cache: self.cache_manager.write_cache().cloned(),
88            cache_manager: Some(self.cache_manager.clone()),
89            file_purger: file.file_purger(),
90            request_sender: self.sender.clone(),
91            indexer_builder: indexer_builder_ref.clone(),
92            result_sender,
93        }
94    }
95
96    /// Handles manual build index requests.
97    pub(crate) async fn handle_build_index_request(
98        &mut self,
99        region_id: RegionId,
100        _req: RegionBuildIndexRequest,
101        sender: OptionOutputTx,
102    ) {
103        self.handle_rebuild_index(
104            BuildIndexRequest {
105                region_id,
106                build_type: IndexBuildType::Manual,
107                file_metas: Vec::new(),
108            },
109            sender,
110        )
111        .await;
112    }
113
114    pub(crate) async fn handle_rebuild_index(
115        &mut self,
116        request: BuildIndexRequest,
117        mut sender: OptionOutputTx,
118    ) {
119        let region_id = request.region_id;
120        let Some(region) = self.regions.writable_region_or(region_id, &mut sender) else {
121            return;
122        };
123
124        let version_control = region.version_control.clone();
125        let version = version_control.current().version;
126        // A committed index publication may not have reached version control
127        // yet. Use the manifest's FileMeta as the conditional publication
128        // source while keeping the builder's schema generation from `version`.
129        let manifest = region.manifest_ctx.manifest().await;
130        let current_file_meta = |file_id| {
131            let file_meta = manifest.files.get(&file_id);
132            if file_meta.is_none() {
133                debug!(
134                    "Skipping index build because file is absent from manifest, region: {}, file_id: {}",
135                    region_id, file_id
136                );
137            }
138            file_meta.cloned()
139        };
140
141        let all_files: HashMap<FileId, FileHandle> = version
142            .ssts
143            .levels()
144            .iter()
145            .flat_map(|level| level.files.iter())
146            .filter(|(_, handle)| !handle.is_deleted() && !handle.compacting())
147            .map(|(id, handle)| (*id, handle.clone()))
148            .collect();
149
150        let build_tasks = if request.file_metas.is_empty() {
151            // If no specific files are provided, find files whose index is inconsistent with the region metadata.
152            all_files
153                .values()
154                .filter_map(|file| {
155                    let file_meta = current_file_meta(file.meta_ref().file_id)?;
156                    (!file_meta.is_index_consistent_with_region(&version.metadata.column_metadatas))
157                        .then(|| (file.clone(), file_meta))
158                })
159                .collect::<Vec<_>>()
160        } else {
161            request
162                .file_metas
163                .iter()
164                .filter_map(|meta| {
165                    let file = all_files.get(&meta.file_id)?;
166                    let file_meta = current_file_meta(meta.file_id)?;
167                    Some((file.clone(), file_meta))
168                })
169                .collect::<Vec<_>>()
170        };
171
172        if build_tasks.is_empty() {
173            debug!(
174                "No files need to build index for region {}, request: {:?}",
175                region_id, request
176            );
177            sender.send(Ok(0));
178            return;
179        }
180
181        let num_tasks = build_tasks.len();
182        let (tx, mut rx) = mpsc::channel::<Result<IndexBuildOutcome>>(num_tasks);
183
184        for (file_handle, file_meta) in build_tasks {
185            debug!(
186                "Scheduling index build for region {}, file_id {}",
187                region_id,
188                file_handle.meta_ref().file_id
189            );
190
191            if region.should_abort_index() {
192                warn!(
193                    "Region {} is in state {:?}, abort index rebuild process for file_id {}",
194                    region_id,
195                    region.state(),
196                    file_handle.meta_ref().file_id
197                );
198                break;
199            }
200
201            let task = self.new_index_build_task(
202                &region,
203                &version,
204                file_handle.clone(),
205                file_meta,
206                request.build_type.clone(),
207                tx.clone(),
208            );
209            let _ = self
210                .index_build_scheduler
211                .schedule_build(&region.version_control, task)
212                .await;
213        }
214        // Wait for all index build tasks to finish and notify the caller.
215        common_runtime::spawn_global(async move {
216            for _ in 0..num_tasks {
217                if let Some(Err(e)) = rx.recv().await {
218                    warn!(e; "Index build task failed for region: {}", region_id);
219                    sender.send(Err(e));
220                    return;
221                }
222            }
223            sender.send(Ok(0));
224        });
225    }
226
227    pub(crate) async fn handle_index_build_finished(
228        &mut self,
229        region_id: RegionId,
230        request: IndexBuildFinished,
231    ) {
232        let region = match self.regions.get_region(region_id) {
233            Some(region) => region,
234            None => {
235                warn!(
236                    "Region not found for index build finished, region_id: {}",
237                    region_id
238                );
239                return;
240            }
241        };
242
243        let file_meta = request.file_meta;
244        let region_file_id = RegionFileId::new(file_meta.region_id, file_meta.file_id);
245        let index_id = RegionIndexId::new(region_file_id, file_meta.index_version);
246
247        // Clean old puffin-related cache before making the new metadata visible.
248        let cache_strategy = CacheStrategy::EnableAll(self.cache_manager.clone());
249        cache_strategy.evict_puffin_cache(index_id).await;
250
251        let manager = region.manifest_ctx.manifest_manager.read().await;
252        let manifest = manager.manifest();
253        let is_current = manifest.files.get(&file_meta.file_id) == Some(&file_meta);
254        if is_current {
255            region.version_control.apply_edit(
256                Some(RegionEdit {
257                    files_to_add: vec![file_meta.clone()],
258                    files_to_remove: Vec::new(),
259                    timestamp_ms: None,
260                    flushed_sequence: None,
261                    flushed_entry_id: None,
262                    committed_sequence: None,
263                    compaction_time_window: None,
264                }),
265                &[],
266                region.file_purger.clone(),
267            );
268        }
269        drop(manager);
270
271        if !is_current {
272            INDEX_PUBLICATION_STALE_TOTAL
273                .with_label_values(&["worker_apply"])
274                .inc();
275            warn!(
276                "Ignores stale index build result, region: {}, file_id: {}, index_version: {}, committed manifest version: {}",
277                region_id, file_meta.file_id, file_meta.index_version, request.manifest_version
278            );
279            cleanup_stale_index_caches(
280                index_id,
281                &region.access_layer,
282                Some(&self.cache_manager),
283                self.cache_manager.write_cache(),
284            )
285            .await;
286            self.listener.on_index_build_abort(region_file_id).await;
287            return;
288        }
289
290        self.listener.on_index_build_finish(region_file_id).await;
291    }
292
293    pub(crate) async fn handle_index_build_failed(
294        &mut self,
295        region_id: RegionId,
296        request: IndexBuildFailed,
297    ) {
298        error!(request.err; "Index build failed for region: {}", region_id);
299        self.index_build_scheduler
300            .on_failure(region_id, request.err.clone())
301            .await;
302    }
303
304    pub(crate) async fn handle_index_build_stopped(
305        &mut self,
306        region_id: RegionId,
307        request: IndexBuildStopped,
308    ) {
309        self.index_build_scheduler
310            .on_task_stopped(region_id, request.file_id);
311    }
312}