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