1use 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 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 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 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 ®ion,
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(®ion.version_control, task)
212 .await;
213 }
214 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 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 ®ion.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}