1use std::collections::{HashMap, VecDeque};
20use std::num::NonZeroU64;
21use std::sync::Arc;
22
23use common_telemetry::{debug, info, warn};
24use parquet::file::metadata::PageIndexPolicy;
25use snafu::ResultExt;
26use store_api::logstore::LogStore;
27use store_api::metadata::RegionMetadataRef;
28use store_api::storage::{RegionId, SequenceNumber};
29
30use crate::cache::CacheManagerRef;
31use crate::cache::file_cache::{FileType, IndexKey};
32use crate::config::IndexBuildMode;
33use crate::error::{
34 EditRegionSnafu, RegionBusySnafu, RegionNotFoundSnafu, RegionStateSnafu, Result,
35};
36use crate::manifest::action::{
37 RegionChange, RegionEdit, RegionMetaAction, RegionMetaActionList, RegionTruncate, TruncateKind,
38};
39use crate::memtable::MemtableBuilderProvider;
40use crate::metrics::WRITE_CACHE_INFLIGHT_DOWNLOAD;
41use crate::region::opener::{sanitize_region_options, version_builder_from_manifest};
42use crate::region::options::RegionOptions;
43use crate::region::version::VersionControlRef;
44use crate::region::{MitoRegionRef, RegionLeaderState, RegionRoleState};
45use crate::request::{
46 BackgroundNotify, BuildIndexRequest, DiscardUnflushedResult, OptionOutputTx,
47 RegionChangeResult, RegionEditRequest, RegionEditResult, RegionSyncRequest, TruncateResult,
48 WorkerRequest, WorkerRequestWithTime,
49};
50use crate::sst::index::IndexBuildType;
51use crate::sst::location;
52use crate::wal::EntryId;
53use crate::worker::{RegionWorkerLoop, WorkerListener};
54
55pub(crate) type RegionEditQueues = HashMap<RegionId, RegionEditQueue>;
56
57pub(crate) enum RegionEditDispatch {
59 Started,
60 Queued,
61 Rejected,
62}
63
64pub(crate) struct RegionEditQueue {
72 region_id: RegionId,
73 requests: VecDeque<RegionEditRequest>,
74}
75
76impl RegionEditQueue {
77 const QUEUE_MAX_LEN: usize = 128;
78
79 fn new(region_id: RegionId) -> Self {
80 Self {
81 region_id,
82 requests: VecDeque::new(),
83 }
84 }
85
86 fn enqueue(&mut self, request: RegionEditRequest) -> RegionEditDispatch {
87 if self.requests.len() > Self::QUEUE_MAX_LEN {
88 request.waiters.reply_with(|| {
89 RegionBusySnafu {
90 region_id: self.region_id,
91 }
92 .fail()
93 });
94 return RegionEditDispatch::Rejected;
95 };
96 self.requests.push_back(request);
97 RegionEditDispatch::Queued
98 }
99
100 fn dequeue(&mut self) -> Option<RegionEditRequest> {
101 fn can_merge(edit: &RegionEdit) -> bool {
102 edit.files_to_add.iter().all(|f| f.sequence.is_none())
112 && edit.files_to_remove.is_empty()
113 && edit.timestamp_ms.is_none()
114 && edit.compaction_time_window.is_none()
115 && edit.flushed_entry_id.is_none()
116 && edit.flushed_sequence.is_none()
117 && edit.committed_sequence.is_none()
118 }
119
120 let mut merged = self.requests.pop_front()?;
121 if !can_merge(&merged.edit) {
122 return Some(merged);
123 }
124
125 while let Some(request) = self
126 .requests
127 .pop_front_if(|request| can_merge(&request.edit))
128 {
129 merged.edit.files_to_add.extend(request.edit.files_to_add);
130 merged.waiters.merge(request.waiters);
131 }
132 debug!(
133 "the files to add: [{}] are merged in one edit",
134 merged
135 .edit
136 .files_to_add
137 .iter()
138 .map(|x| x.file_id.to_string())
139 .collect::<Vec<_>>()
140 .join(", ")
141 );
142 Some(merged)
143 }
144
145 fn is_empty(&self) -> bool {
146 self.requests.is_empty()
147 }
148
149 fn reject_all_as_not_found(mut self) {
150 while let Some(request) = self.requests.pop_front() {
151 request.waiters.reply_with(|| {
152 RegionNotFoundSnafu {
153 region_id: self.region_id,
154 }
155 .fail()
156 });
157 }
158 }
159}
160
161impl<S: LogStore> RegionWorkerLoop<S> {
162 pub(crate) fn reject_region_edit_queue_as_not_found(&mut self, region_id: RegionId) {
164 if let Some(queue) = self.pending_alters.remove(®ion_id) {
165 for (_, sender) in queue {
166 sender.send(RegionNotFoundSnafu { region_id }.fail());
167 }
168 }
169 if let Some(edit_queue) = self.region_edit_queues.remove(®ion_id) {
170 edit_queue.reject_all_as_not_found();
171 }
172 }
173
174 pub(crate) async fn handle_manifest_region_change_result(
176 &mut self,
177 change_result: RegionChangeResult,
178 ) {
179 let region = match self.regions.get_region(change_result.region_id) {
180 Some(region) => region,
181 None => {
182 self.reject_region_stalled_requests(&change_result.region_id);
183 change_result.sender.send(
184 RegionNotFoundSnafu {
185 region_id: change_result.region_id,
186 }
187 .fail(),
188 );
189 return;
190 }
191 };
192
193 if change_result.result.is_ok() {
194 Self::update_region_version(
196 ®ion.version_control,
197 change_result.new_meta,
198 change_result.new_options,
199 &self.memtable_builder_provider,
200 );
201 }
202
203 region.switch_state_to_writable(RegionLeaderState::Altering);
205 change_result.sender.send(change_result.result.map(|_| 0));
207
208 if self.config.index.build_mode == IndexBuildMode::Async && change_result.need_index {
210 self.handle_rebuild_index(
211 BuildIndexRequest {
212 region_id: region.region_id,
213 build_type: IndexBuildType::SchemaChange,
214 file_metas: Vec::new(),
215 result_senders: Vec::new(),
216 },
217 OptionOutputTx::new(None),
218 )
219 .await;
220 }
221 self.handle_region_stalled_requests(&change_result.region_id, true)
223 .await;
224 self.dispatch_pending_region_requests(change_result.region_id)
225 .await;
226 }
227
228 pub(crate) async fn handle_region_sync(&mut self, request: RegionSyncRequest) {
233 let region_id = request.region_id;
234 let sender = request.sender;
235 let region = match self.regions.follower_region(region_id) {
236 Ok(region) => region,
237 Err(e) => {
238 let _ = sender.send(Err(e));
239 return;
240 }
241 };
242
243 let original_manifest_version = region.manifest_ctx.manifest_version().await;
244 let manifest = match region
245 .manifest_ctx
246 .install_manifest_to(request.manifest_version)
247 .await
248 {
249 Ok(manifest) => manifest,
250 Err(e) => {
251 let _ = sender.send(Err(e));
252 return;
253 }
254 };
255 let version = region.version();
256 let mut region_options = version.options.clone();
257 let old_format = region_options.sst_format.unwrap_or_default();
258 sanitize_region_options(&manifest, &mut region_options);
260 if !version.memtables.is_empty() {
261 let current = region.version_control.current();
262 warn!(
263 "Region {} memtables is not empty, which should not happen, manifest version: {}, last entry id: {}",
264 region.region_id, manifest.manifest_version, current.last_entry_id
265 );
266 }
267
268 let memtable_builder = if old_format != region_options.sst_format.unwrap_or_default() {
270 Some(
272 self.memtable_builder_provider
273 .builder_for_options(®ion_options),
274 )
275 } else {
276 None
277 };
278 let new_mutable = version
279 .memtables
280 .mutable
281 .new_with_part_duration(version.compaction_time_window, memtable_builder);
282 let metadata = manifest.metadata.clone();
284
285 let version_builder = version_builder_from_manifest(
286 &manifest,
287 metadata,
288 region.file_purger.clone(),
289 new_mutable,
290 region_options,
291 );
292 let version = version_builder.build();
293 region.version_control.overwrite_current(Arc::new(version));
294 if let Some(sequence) = manifest.committed_sequence {
296 region.version_control.update_last_sequence(sequence);
297 }
298
299 let updated = manifest.manifest_version > original_manifest_version;
300 let _ = sender.send(Ok((manifest.manifest_version, updated)));
301 }
302}
303
304impl<S: LogStore> RegionWorkerLoop<S> {
305 pub(crate) fn handle_region_edit(&mut self, request: RegionEditRequest) -> RegionEditDispatch {
307 let region_id = request.region_id;
308 let Some(region) = self.regions.get_region(region_id) else {
309 request
310 .waiters
311 .reply_with(|| RegionNotFoundSnafu { region_id }.fail());
312 return RegionEditDispatch::Rejected;
313 };
314
315 if self.flush_scheduler.has_close_time_flush(region_id) {
318 request
319 .waiters
320 .reply_with(|| RegionBusySnafu { region_id }.fail());
321 return RegionEditDispatch::Rejected;
322 }
323
324 if let Some(expected) = &request.staging_partition_expr
326 && region
327 .manifest_ctx
328 .staging_partition_info()
329 .as_ref()
330 .and_then(|info| info.partition_expr())
331 != Some(expected.as_str())
332 {
333 let state = region.state();
334 request.waiters.reply_with(|| {
335 RegionStateSnafu {
336 region_id,
337 state,
338 expect: RegionRoleState::Leader(RegionLeaderState::Staging),
339 }
340 .fail()
341 });
342 return RegionEditDispatch::Rejected;
343 }
344
345 if !region.is_writable() || self.flush_scheduler.has_pending_alter(region_id) {
346 if matches!(
347 region.state(),
348 RegionRoleState::Leader(RegionLeaderState::Editing | RegionLeaderState::Altering)
349 ) || (region.is_writable() && self.flush_scheduler.has_pending_alter(region_id))
350 {
351 return self
352 .region_edit_queues
353 .entry(region_id)
354 .or_insert_with(|| RegionEditQueue::new(region_id))
355 .enqueue(request);
356 } else {
357 request
358 .waiters
359 .reply_with(|| RegionBusySnafu { region_id }.fail());
360 }
361 return RegionEditDispatch::Rejected;
362 }
363
364 let RegionEditRequest {
365 region_id: _,
366 mut edit,
367 waiters,
368 preload_sst_cache,
369 staging_partition_expr: _,
370 } = request;
371 let file_sequence = region.version_control.current().last_sequence + 1;
372 edit.committed_sequence = Some(file_sequence);
373
374 for file in &mut edit.files_to_add {
380 file.sequence = NonZeroU64::new(file_sequence);
381 file.preserve_row_sequence = false;
382 }
383
384 let is_staging = region.is_staging();
386 let expect_state = if is_staging {
387 RegionLeaderState::Staging
388 } else {
389 RegionLeaderState::Writable
390 };
391 if let Err(e) = region.set_editing(expect_state) {
393 let e = Arc::new(e);
394 waiters.reply_with(|| Err(e.clone()).context(EditRegionSnafu { region_id }));
395 return RegionEditDispatch::Rejected;
396 }
397 region.version_control.update_last_sequence(file_sequence);
398
399 let request_sender = self.sender.clone();
400 let cache_manager = self.cache_manager.clone();
401 let listener = self.listener.clone();
402 common_runtime::spawn_global(async move {
405 let result = edit_region(
406 ®ion,
407 edit.clone(),
408 cache_manager,
409 listener,
410 is_staging,
411 preload_sst_cache,
412 )
413 .await
414 .map_err(Arc::new);
415 let notify = WorkerRequest::Background {
416 region_id,
417 notify: BackgroundNotify::RegionEdit(RegionEditResult {
418 region_id,
419 waiters,
420 edit,
421 result,
422 update_region_state: true,
424 is_staging,
425 }),
426 };
427
428 if let Err(res) = request_sender
430 .send(WorkerRequestWithTime::new(notify))
431 .await
432 {
433 warn!(
434 "Failed to send region edit result back to the worker, region_id: {}, res: {:?}",
435 region_id, res
436 );
437 }
438 });
439 RegionEditDispatch::Started
440 }
441
442 pub(crate) async fn handle_region_edit_result(&mut self, edit_result: RegionEditResult) {
444 let region = match self.regions.get_region(edit_result.region_id) {
445 Some(region) => region,
446 None => {
447 self.fail_region_stalled_requests_as_not_found(&edit_result.region_id);
450 self.reject_region_edit_queue_as_not_found(edit_result.region_id);
451
452 edit_result.waiters.reply_with(|| {
453 RegionNotFoundSnafu {
454 region_id: edit_result.region_id,
455 }
456 .fail()
457 });
458 return;
459 }
460 };
461
462 let need_compaction = if edit_result.is_staging {
463 if edit_result.update_region_state
465 && region.state() == RegionRoleState::Leader(RegionLeaderState::Editing)
466 {
467 region.switch_state_to_staging(RegionLeaderState::Editing);
470 }
471
472 false
473 } else {
474 let need_compaction = self.config.schedule_compaction_after_edit
475 && edit_result.result.is_ok()
476 && !edit_result.edit.files_to_add.is_empty();
477
478 if edit_result.result.is_ok() {
480 region.version_control.apply_edit(
482 Some(edit_result.edit),
483 &[],
484 region.file_purger.clone(),
485 );
486 }
487 if edit_result.update_region_state {
488 region.switch_state_to_writable(RegionLeaderState::Editing);
489 }
490
491 need_compaction
492 };
493
494 edit_result
495 .waiters
496 .reply_with(|| match &edit_result.result {
497 Ok(()) => Ok(()),
498 Err(e) => Err(e.clone()).context(EditRegionSnafu {
499 region_id: edit_result.region_id,
500 }),
501 });
502
503 if edit_result.update_region_state {
504 self.handle_region_stalled_requests(&edit_result.region_id, false)
507 .await;
508 }
509
510 self.dispatch_pending_region_requests(edit_result.region_id)
511 .await;
512
513 if need_compaction {
514 self.schedule_compaction(®ion).await;
515 }
516 }
517
518 pub(crate) async fn dispatch_pending_region_requests(&mut self, region_id: RegionId) {
521 loop {
522 if self.flush_scheduler.has_pending_alter(region_id) {
523 return;
524 }
525 if self.regions.get_region(region_id).is_some_and(|region| {
526 matches!(
527 region.state(),
528 RegionRoleState::Leader(
529 RegionLeaderState::Editing | RegionLeaderState::Altering
530 )
531 )
532 }) {
533 return;
534 }
535 let next = self
536 .pending_alters
537 .get_mut(®ion_id)
538 .and_then(|q| q.pop_front());
539 let Some((request, sender)) = next else {
540 break;
541 };
542 self.execute_alter_request(region_id, request, sender).await;
543 }
544 self.pending_alters.remove(®ion_id);
545 loop {
546 let next_request = self
547 .region_edit_queues
548 .get_mut(®ion_id)
549 .and_then(|queue| queue.dequeue());
550 if self
551 .region_edit_queues
552 .get(®ion_id)
553 .is_some_and(|queue| queue.is_empty())
554 {
555 self.region_edit_queues.remove(®ion_id);
556 }
557 let Some(request) = next_request else {
558 return;
559 };
560 match self.handle_region_edit(request) {
566 RegionEditDispatch::Started | RegionEditDispatch::Queued => return,
567 RegionEditDispatch::Rejected => continue,
568 }
569 }
570 }
571
572 pub(crate) fn handle_manifest_truncate_action(
574 &self,
575 region: MitoRegionRef,
576 truncate: RegionTruncate,
577 sender: OptionOutputTx,
578 ) {
579 if let Err(e) = region.set_truncating() {
582 sender.send(Err(e));
583 return;
584 }
585 let request_sender = self.sender.clone();
588 let manifest_ctx = region.manifest_ctx.clone();
589 let is_staging = region.is_staging();
590 let durability_barrier = self
591 .wal
592 .durability_barrier(region.region_id, ®ion.provider);
593
594 common_runtime::spawn_global(async move {
596 let durable = match &truncate.kind {
599 TruncateKind::All {
600 truncated_entry_id, ..
601 } => durability_barrier.wait(*truncated_entry_id).await,
602 _ => Ok(()),
603 };
604 let action_list =
606 RegionMetaActionList::with_action(RegionMetaAction::Truncate(truncate.clone()));
607
608 let result = match durable {
609 Ok(()) => manifest_ctx
610 .update_manifest(RegionLeaderState::Truncating, action_list, is_staging)
611 .await
612 .map(|_| ()),
613 Err(e) => Err(e),
614 };
615
616 let truncate_result = TruncateResult {
618 region_id: truncate.region_id,
619 sender,
620 result,
621 kind: truncate.kind,
622 };
623 let _ = request_sender
624 .send(WorkerRequestWithTime::new(WorkerRequest::Background {
625 region_id: truncate.region_id,
626 notify: BackgroundNotify::Truncate(truncate_result),
627 }))
628 .await
629 .inspect_err(|_| warn!("failed to send truncate result"));
630 });
631 }
632
633 pub(crate) fn handle_manifest_discard_unflushed_action(
635 &self,
636 region: MitoRegionRef,
637 discarded_entry_id: EntryId,
638 discarded_sequence: SequenceNumber,
639 discarded_rows: u64,
640 discarded_bytes: u64,
641 sender: OptionOutputTx,
642 ) {
643 if let Err(e) = region.set_truncating() {
644 sender.send(Err(e));
645 return;
646 }
647
648 let region_id = region.region_id;
649 let request_sender = self.sender.clone();
650 let manifest_ctx = region.manifest_ctx.clone();
651 let durability_barrier = self.wal.durability_barrier(region_id, ®ion.provider);
652
653 common_runtime::spawn_global(async move {
654 let edit = RegionEdit {
658 files_to_add: Vec::new(),
659 files_to_remove: Vec::new(),
660 timestamp_ms: None,
661 compaction_time_window: None,
662 flushed_entry_id: Some(discarded_entry_id),
663 flushed_sequence: Some(discarded_sequence),
664 committed_sequence: None,
665 };
666 let action_list = RegionMetaActionList::with_action(RegionMetaAction::Edit(edit));
667 let result = match durability_barrier.wait(discarded_entry_id).await {
668 Ok(()) => manifest_ctx
669 .update_manifest(RegionLeaderState::Truncating, action_list, false)
670 .await
671 .map(|_| ()),
672 Err(e) => Err(e),
673 };
674
675 let result = DiscardUnflushedResult {
676 region_id,
677 sender,
678 result,
679 discarded_entry_id,
680 discarded_sequence,
681 discarded_rows,
682 discarded_bytes,
683 };
684 let _ = request_sender
685 .send(WorkerRequestWithTime::new(WorkerRequest::Background {
686 region_id,
687 notify: BackgroundNotify::DiscardUnflushed(result),
688 }))
689 .await
690 .inspect_err(|_| warn!("failed to send discard unflushed result"));
691 });
692 }
693
694 pub(crate) fn handle_manifest_region_change(
696 &self,
697 region: MitoRegionRef,
698 change: RegionChange,
699 need_index: bool,
700 new_options: Option<RegionOptions>,
701 sender: OptionOutputTx,
702 ) {
703 if let Err(e) = region.set_altering() {
705 sender.send(Err(e));
706 return;
707 }
708 let listener = self.listener.clone();
709 let request_sender = self.sender.clone();
710 let is_staging = region.is_staging();
711 common_runtime::spawn_global(async move {
713 let new_meta = change.metadata.clone();
714 let action_list = RegionMetaActionList::with_action(RegionMetaAction::Change(change));
715
716 let result = region
717 .manifest_ctx
718 .update_manifest(RegionLeaderState::Altering, action_list, is_staging)
719 .await
720 .map(|_| ());
721 let notify = WorkerRequest::Background {
722 region_id: region.region_id,
723 notify: BackgroundNotify::RegionChange(RegionChangeResult {
724 region_id: region.region_id,
725 sender,
726 result,
727 new_meta,
728 need_index,
729 new_options,
730 }),
731 };
732 listener
733 .on_notify_region_change_result_begin(region.region_id)
734 .await;
735
736 if let Err(res) = request_sender
737 .send(WorkerRequestWithTime::new(notify))
738 .await
739 {
740 warn!(
741 "Failed to send region change result back to the worker, region_id: {}, res: {:?}",
742 region.region_id, res
743 );
744 }
745 });
746 }
747
748 fn update_region_version(
749 version_control: &VersionControlRef,
750 new_meta: RegionMetadataRef,
751 new_options: Option<RegionOptions>,
752 memtable_builder_provider: &MemtableBuilderProvider,
753 ) {
754 let options_changed = new_options.is_some();
755 let region_id = new_meta.region_id;
756 if let Some(new_options) = new_options {
757 let new_memtable_builder = memtable_builder_provider.builder_for_options(&new_options);
760 version_control.alter_schema_and_format(new_meta, new_options, new_memtable_builder);
761 } else {
762 version_control.alter_schema(new_meta);
764 }
765
766 let version_data = version_control.current();
767 let version = version_data.version;
768 info!(
769 "Region {} is altered, metadata is {:?}, options: {:?}, options_changed: {}",
770 region_id, version.metadata, version.options, options_changed,
771 );
772 }
773}
774
775async fn edit_region(
777 region: &MitoRegionRef,
778 edit: RegionEdit,
779 cache_manager: CacheManagerRef,
780 listener: WorkerListener,
781 is_staging: bool,
782 preload_sst_cache: bool,
783) -> Result<()> {
784 let region_id = region.region_id;
785 if let Some(write_cache) = cache_manager.write_cache()
786 && preload_sst_cache
787 {
788 for file_meta in &edit.files_to_add {
789 let write_cache = write_cache.clone();
790 let layer = region.access_layer.clone();
791 let listener = listener.clone();
792
793 let index_key = IndexKey::new(region_id, file_meta.file_id, FileType::Parquet);
794 let remote_path =
795 location::sst_file_path(layer.table_dir(), file_meta.file_id(), layer.path_type());
796
797 let is_index_exist = file_meta.exists_index();
798 let index_file_size = file_meta.index_file_size();
799
800 let index_file_index_key = IndexKey::new(
801 region_id,
802 file_meta.index_id().file_id.file_id(),
803 FileType::Puffin(file_meta.index_version),
804 );
805 let index_remote_path = location::index_file_path(
806 layer.table_dir(),
807 file_meta.index_id(),
808 layer.path_type(),
809 );
810
811 let file_size = file_meta.file_size;
812 common_runtime::spawn_global(async move {
813 WRITE_CACHE_INFLIGHT_DOWNLOAD.add(1);
814
815 let parquet_cached = write_cache
816 .download_if_absent(index_key, &remote_path, layer.object_store(), file_size)
817 .await;
818
819 if parquet_cached.is_ok() {
820 let mut cache_metrics = Default::default();
823 let _ = write_cache
824 .file_cache()
825 .get_parquet_meta_data(
826 index_key,
827 &mut cache_metrics,
828 PageIndexPolicy::Optional,
829 )
830 .await;
831
832 if matches!(parquet_cached, Ok(true)) {
833 listener.on_file_cache_filled(index_key.file_id);
834 }
835 }
836 if is_index_exist {
837 if let Err(err) = write_cache
839 .download(
840 index_file_index_key,
841 &index_remote_path,
842 layer.object_store(),
843 index_file_size,
844 )
845 .await
846 {
847 common_telemetry::error!(
848 err; "Failed to download puffin file, region_id: {}, index_file_index_key: {:?}, index_remote_path: {}", region_id, index_file_index_key, index_remote_path
849 );
850 }
851 }
852
853 WRITE_CACHE_INFLIGHT_DOWNLOAD.sub(1);
854 });
855 }
856 }
857
858 info!(
859 "Applying {edit:?} to region {}, is_staging: {}",
860 region_id, is_staging
861 );
862
863 let action_list = RegionMetaActionList::with_action(RegionMetaAction::Edit(edit));
864 region
865 .manifest_ctx
866 .update_manifest(RegionLeaderState::Editing, action_list, is_staging)
867 .await
868 .map(|_| ())
869}