1pub mod opener;
18pub mod options;
19pub(crate) mod version;
20
21use std::collections::hash_map::Entry;
22use std::collections::HashMap;
23use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
24use std::sync::{Arc, RwLock};
25
26use common_telemetry::{error, info, warn};
27use crossbeam_utils::atomic::AtomicCell;
28use snafu::{ensure, OptionExt};
29use store_api::codec::PrimaryKeyEncoding;
30use store_api::logstore::provider::Provider;
31use store_api::metadata::RegionMetadataRef;
32use store_api::region_engine::{
33 RegionManifestInfo, RegionRole, RegionStatistic, SettableRegionRoleState,
34};
35use store_api::storage::{RegionId, SequenceNumber};
36use store_api::ManifestVersion;
37
38use crate::access_layer::AccessLayerRef;
39use crate::error::{
40 FlushableRegionStateSnafu, RegionNotFoundSnafu, RegionStateSnafu, RegionTruncatedSnafu, Result,
41 UpdateManifestSnafu,
42};
43use crate::manifest::action::{RegionManifest, RegionMetaAction, RegionMetaActionList};
44use crate::manifest::manager::RegionManifestManager;
45use crate::memtable::MemtableBuilderRef;
46use crate::region::version::{VersionControlRef, VersionRef};
47use crate::request::{OnFailure, OptionOutputTx};
48use crate::sst::file_purger::FilePurgerRef;
49use crate::time_provider::TimeProviderRef;
50
51const ESTIMATED_WAL_FACTOR: f32 = 0.42825;
53
54#[derive(Debug)]
56pub struct RegionUsage {
57 pub region_id: RegionId,
58 pub wal_usage: u64,
59 pub sst_usage: u64,
60 pub manifest_usage: u64,
61}
62
63impl RegionUsage {
64 pub fn disk_usage(&self) -> u64 {
65 self.wal_usage + self.sst_usage + self.manifest_usage
66 }
67}
68
69#[derive(Debug, Clone, Copy, PartialEq, Eq)]
70pub enum RegionLeaderState {
71 Writable,
73 Altering,
75 Dropping,
77 Truncating,
79 Editing,
81 Downgrading,
83}
84
85#[derive(Debug, Clone, Copy, PartialEq, Eq)]
86pub enum RegionRoleState {
87 Leader(RegionLeaderState),
88 Follower,
89}
90
91#[derive(Debug)]
97pub struct MitoRegion {
98 pub(crate) region_id: RegionId,
103
104 pub(crate) version_control: VersionControlRef,
108 pub(crate) access_layer: AccessLayerRef,
110 pub(crate) manifest_ctx: ManifestContextRef,
112 pub(crate) file_purger: FilePurgerRef,
114 pub(crate) provider: Provider,
116 last_flush_millis: AtomicI64,
118 last_compaction_millis: AtomicI64,
120 time_provider: TimeProviderRef,
122 pub(crate) topic_latest_entry_id: AtomicU64,
132 pub(crate) memtable_builder: MemtableBuilderRef,
134 stats: ManifestStats,
136}
137
138pub type MitoRegionRef = Arc<MitoRegion>;
139
140impl MitoRegion {
141 pub(crate) async fn stop(&self) {
143 self.manifest_ctx
144 .manifest_manager
145 .write()
146 .await
147 .stop()
148 .await;
149
150 info!(
151 "Stopped region manifest manager, region_id: {}",
152 self.region_id
153 );
154 }
155
156 pub(crate) fn metadata(&self) -> RegionMetadataRef {
158 let version_data = self.version_control.current();
159 version_data.version.metadata.clone()
160 }
161
162 pub(crate) fn primary_key_encoding(&self) -> PrimaryKeyEncoding {
164 let version_data = self.version_control.current();
165 version_data.version.metadata.primary_key_encoding
166 }
167
168 pub(crate) fn version(&self) -> VersionRef {
170 let version_data = self.version_control.current();
171 version_data.version
172 }
173
174 pub(crate) fn last_flush_millis(&self) -> i64 {
176 self.last_flush_millis.load(Ordering::Relaxed)
177 }
178
179 pub(crate) fn update_flush_millis(&self) {
181 let now = self.time_provider.current_time_millis();
182 self.last_flush_millis.store(now, Ordering::Relaxed);
183 }
184
185 pub(crate) fn last_compaction_millis(&self) -> i64 {
187 self.last_compaction_millis.load(Ordering::Relaxed)
188 }
189
190 pub(crate) fn update_compaction_millis(&self) {
192 let now = self.time_provider.current_time_millis();
193 self.last_compaction_millis.store(now, Ordering::Relaxed);
194 }
195
196 pub(crate) fn region_dir(&self) -> &str {
198 self.access_layer.region_dir()
199 }
200
201 pub(crate) fn is_writable(&self) -> bool {
203 self.manifest_ctx.state.load() == RegionRoleState::Leader(RegionLeaderState::Writable)
204 }
205
206 pub(crate) fn is_flushable(&self) -> bool {
208 matches!(
209 self.manifest_ctx.state.load(),
210 RegionRoleState::Leader(RegionLeaderState::Writable)
211 | RegionRoleState::Leader(RegionLeaderState::Downgrading)
212 )
213 }
214
215 pub(crate) fn is_downgrading(&self) -> bool {
217 matches!(
218 self.manifest_ctx.state.load(),
219 RegionRoleState::Leader(RegionLeaderState::Downgrading)
220 )
221 }
222
223 pub fn region_id(&self) -> RegionId {
224 self.region_id
225 }
226
227 pub fn find_committed_sequence(&self) -> SequenceNumber {
228 self.version_control.committed_sequence()
229 }
230
231 pub fn is_follower(&self) -> bool {
233 self.manifest_ctx.state.load() == RegionRoleState::Follower
234 }
235
236 pub(crate) fn state(&self) -> RegionRoleState {
238 self.manifest_ctx.state.load()
239 }
240
241 pub(crate) fn set_role(&self, next_role: RegionRole) {
243 self.manifest_ctx.set_role(next_role, self.region_id);
244 }
245
246 pub(crate) fn set_altering(&self) -> Result<()> {
249 self.compare_exchange_state(
250 RegionLeaderState::Writable,
251 RegionRoleState::Leader(RegionLeaderState::Altering),
252 )
253 }
254
255 pub(crate) fn set_dropping(&self) -> Result<()> {
258 self.compare_exchange_state(
259 RegionLeaderState::Writable,
260 RegionRoleState::Leader(RegionLeaderState::Dropping),
261 )
262 }
263
264 pub(crate) fn set_truncating(&self) -> Result<()> {
267 self.compare_exchange_state(
268 RegionLeaderState::Writable,
269 RegionRoleState::Leader(RegionLeaderState::Truncating),
270 )
271 }
272
273 pub(crate) fn set_editing(&self) -> Result<()> {
276 self.compare_exchange_state(
277 RegionLeaderState::Writable,
278 RegionRoleState::Leader(RegionLeaderState::Editing),
279 )
280 }
281
282 pub(crate) async fn set_role_state_gracefully(&self, state: SettableRegionRoleState) {
284 let _manager = self.manifest_ctx.manifest_manager.write().await;
285 self.set_role(state.into());
288 }
289
290 pub(crate) fn switch_state_to_writable(&self, expect: RegionLeaderState) {
293 if let Err(e) = self
294 .compare_exchange_state(expect, RegionRoleState::Leader(RegionLeaderState::Writable))
295 {
296 error!(e; "failed to switch region state to writable, expect state is {:?}", expect);
297 }
298 }
299
300 pub(crate) fn region_statistic(&self) -> RegionStatistic {
302 let version = self.version();
303 let memtables = &version.memtables;
304 let memtable_usage = (memtables.mutable_usage() + memtables.immutables_usage()) as u64;
305
306 let sst_usage = version.ssts.sst_usage();
307 let index_usage = version.ssts.index_usage();
308 let flushed_entry_id = version.flushed_entry_id;
309
310 let wal_usage = self.estimated_wal_usage(memtable_usage);
311 let manifest_usage = self.stats.total_manifest_size();
312 let num_rows = version.ssts.num_rows() + version.memtables.num_rows();
313 let manifest_version = self.stats.manifest_version();
314
315 let topic_latest_entry_id = self.topic_latest_entry_id.load(Ordering::Relaxed);
316
317 RegionStatistic {
318 num_rows,
319 memtable_size: memtable_usage,
320 wal_size: wal_usage,
321 manifest_size: manifest_usage,
322 sst_size: sst_usage,
323 index_size: index_usage,
324 manifest: RegionManifestInfo::Mito {
325 manifest_version,
326 flushed_entry_id,
327 },
328 data_topic_latest_entry_id: topic_latest_entry_id,
329 metadata_topic_latest_entry_id: topic_latest_entry_id,
330 }
331 }
332
333 fn estimated_wal_usage(&self, memtable_usage: u64) -> u64 {
336 ((memtable_usage as f32) * ESTIMATED_WAL_FACTOR) as u64
337 }
338
339 fn compare_exchange_state(
342 &self,
343 expect: RegionLeaderState,
344 state: RegionRoleState,
345 ) -> Result<()> {
346 self.manifest_ctx
347 .state
348 .compare_exchange(RegionRoleState::Leader(expect), state)
349 .map_err(|actual| {
350 RegionStateSnafu {
351 region_id: self.region_id,
352 state: actual,
353 expect: RegionRoleState::Leader(expect),
354 }
355 .build()
356 })?;
357 Ok(())
358 }
359}
360
361#[derive(Debug)]
363pub(crate) struct ManifestContext {
364 manifest_manager: tokio::sync::RwLock<RegionManifestManager>,
366 state: AtomicCell<RegionRoleState>,
369}
370
371impl ManifestContext {
372 pub(crate) fn new(manager: RegionManifestManager, state: RegionRoleState) -> Self {
373 ManifestContext {
374 manifest_manager: tokio::sync::RwLock::new(manager),
375 state: AtomicCell::new(state),
376 }
377 }
378
379 pub(crate) async fn manifest_version(&self) -> ManifestVersion {
380 self.manifest_manager
381 .read()
382 .await
383 .manifest()
384 .manifest_version
385 }
386
387 pub(crate) async fn has_update(&self) -> Result<bool> {
388 self.manifest_manager.read().await.has_update().await
389 }
390
391 pub(crate) async fn install_manifest_to(
397 &self,
398 version: ManifestVersion,
399 ) -> Result<Arc<RegionManifest>> {
400 let mut manager = self.manifest_manager.write().await;
401 manager.install_manifest_to(version).await?;
402
403 Ok(manager.manifest())
404 }
405
406 pub(crate) async fn update_manifest(
408 &self,
409 expect_state: RegionLeaderState,
410 action_list: RegionMetaActionList,
411 ) -> Result<ManifestVersion> {
412 let mut manager = self.manifest_manager.write().await;
414 let manifest = manager.manifest();
416 let current_state = self.state.load();
419
420 if expect_state != RegionLeaderState::Downgrading {
425 if current_state == RegionRoleState::Leader(RegionLeaderState::Downgrading) {
426 info!(
427 "Region {} is in downgrading leader state, updating manifest. state is {:?}",
428 manifest.metadata.region_id, expect_state
429 );
430 }
431 ensure!(
432 current_state == RegionRoleState::Leader(expect_state)
433 || current_state == RegionRoleState::Leader(RegionLeaderState::Downgrading),
434 UpdateManifestSnafu {
435 region_id: manifest.metadata.region_id,
436 state: current_state,
437 }
438 );
439 } else {
440 ensure!(
441 current_state == RegionRoleState::Leader(expect_state),
442 RegionStateSnafu {
443 region_id: manifest.metadata.region_id,
444 state: current_state,
445 expect: RegionRoleState::Leader(expect_state),
446 }
447 );
448 }
449
450 for action in &action_list.actions {
451 let RegionMetaAction::Edit(edit) = &action else {
453 continue;
454 };
455
456 let Some(truncated_entry_id) = manifest.truncated_entry_id else {
458 continue;
459 };
460
461 if let Some(flushed_entry_id) = edit.flushed_entry_id {
463 ensure!(
464 truncated_entry_id < flushed_entry_id,
465 RegionTruncatedSnafu {
466 region_id: manifest.metadata.region_id,
467 }
468 );
469 }
470
471 if !edit.files_to_remove.is_empty() {
473 for file in &edit.files_to_remove {
475 ensure!(
476 manifest.files.contains_key(&file.file_id),
477 RegionTruncatedSnafu {
478 region_id: manifest.metadata.region_id,
479 }
480 );
481 }
482 }
483 }
484
485 let version = manager.update(action_list).await.inspect_err(
487 |e| error!(e; "Failed to update manifest, region_id: {}", manifest.metadata.region_id),
488 )?;
489
490 if self.state.load() == RegionRoleState::Follower {
491 warn!(
492 "Region {} becomes follower while updating manifest which may cause inconsistency, manifest version: {version}",
493 manifest.metadata.region_id
494 );
495 }
496
497 Ok(version)
498 }
499
500 pub(crate) fn set_role(&self, next_role: RegionRole, region_id: RegionId) {
522 match next_role {
523 RegionRole::Follower => {
524 match self.state.fetch_update(|state| {
525 if !matches!(state, RegionRoleState::Follower) {
526 Some(RegionRoleState::Follower)
527 } else {
528 None
529 }
530 }) {
531 Ok(state) => info!(
532 "Convert region {} to follower, previous role state: {:?}",
533 region_id, state
534 ),
535 Err(state) => {
536 if state != RegionRoleState::Follower {
537 warn!(
538 "Failed to convert region {} to follower, current role state: {:?}",
539 region_id, state
540 )
541 }
542 }
543 }
544 }
545 RegionRole::Leader => {
546 match self.state.fetch_update(|state| {
547 if matches!(
548 state,
549 RegionRoleState::Follower
550 | RegionRoleState::Leader(RegionLeaderState::Downgrading)
551 ) {
552 Some(RegionRoleState::Leader(RegionLeaderState::Writable))
553 } else {
554 None
555 }
556 }) {
557 Ok(state) => info!(
558 "Convert region {} to leader, previous role state: {:?}",
559 region_id, state
560 ),
561 Err(state) => {
562 if state != RegionRoleState::Leader(RegionLeaderState::Writable) {
563 warn!(
564 "Failed to convert region {} to leader, current role state: {:?}",
565 region_id, state
566 )
567 }
568 }
569 }
570 }
571 RegionRole::DowngradingLeader => {
572 match self.state.compare_exchange(
573 RegionRoleState::Leader(RegionLeaderState::Writable),
574 RegionRoleState::Leader(RegionLeaderState::Downgrading),
575 ) {
576 Ok(state) => info!(
577 "Convert region {} to downgrading region, previous role state: {:?}",
578 region_id, state
579 ),
580 Err(state) => {
581 if state != RegionRoleState::Leader(RegionLeaderState::Downgrading) {
582 warn!(
583 "Failed to convert region {} to downgrading leader, current role state: {:?}",
584 region_id, state
585 )
586 }
587 }
588 }
589 }
590 }
591 }
592}
593
594#[cfg(test)]
595impl ManifestContext {
596 pub(crate) async fn manifest(&self) -> Arc<crate::manifest::action::RegionManifest> {
597 self.manifest_manager.read().await.manifest()
598 }
599}
600
601pub(crate) type ManifestContextRef = Arc<ManifestContext>;
602
603#[derive(Debug, Default)]
605pub(crate) struct RegionMap {
606 regions: RwLock<HashMap<RegionId, MitoRegionRef>>,
607}
608
609impl RegionMap {
610 pub(crate) fn is_region_exists(&self, region_id: RegionId) -> bool {
612 let regions = self.regions.read().unwrap();
613 regions.contains_key(®ion_id)
614 }
615
616 pub(crate) fn insert_region(&self, region: MitoRegionRef) {
618 let mut regions = self.regions.write().unwrap();
619 regions.insert(region.region_id, region);
620 }
621
622 pub(crate) fn get_region(&self, region_id: RegionId) -> Option<MitoRegionRef> {
624 let regions = self.regions.read().unwrap();
625 regions.get(®ion_id).cloned()
626 }
627
628 pub(crate) fn writable_region(&self, region_id: RegionId) -> Result<MitoRegionRef> {
632 let region = self
633 .get_region(region_id)
634 .context(RegionNotFoundSnafu { region_id })?;
635 ensure!(
636 region.is_writable(),
637 RegionStateSnafu {
638 region_id,
639 state: region.state(),
640 expect: RegionRoleState::Leader(RegionLeaderState::Writable),
641 }
642 );
643 Ok(region)
644 }
645
646 pub(crate) fn follower_region(&self, region_id: RegionId) -> Result<MitoRegionRef> {
650 let region = self
651 .get_region(region_id)
652 .context(RegionNotFoundSnafu { region_id })?;
653 ensure!(
654 region.is_follower(),
655 RegionStateSnafu {
656 region_id,
657 state: region.state(),
658 expect: RegionRoleState::Follower,
659 }
660 );
661
662 Ok(region)
663 }
664
665 pub(crate) fn get_region_or<F: OnFailure>(
669 &self,
670 region_id: RegionId,
671 cb: &mut F,
672 ) -> Option<MitoRegionRef> {
673 match self
674 .get_region(region_id)
675 .context(RegionNotFoundSnafu { region_id })
676 {
677 Ok(region) => Some(region),
678 Err(e) => {
679 cb.on_failure(e);
680 None
681 }
682 }
683 }
684
685 pub(crate) fn writable_region_or<F: OnFailure>(
689 &self,
690 region_id: RegionId,
691 cb: &mut F,
692 ) -> Option<MitoRegionRef> {
693 match self.writable_region(region_id) {
694 Ok(region) => Some(region),
695 Err(e) => {
696 cb.on_failure(e);
697 None
698 }
699 }
700 }
701
702 fn flushable_region(&self, region_id: RegionId) -> Result<MitoRegionRef> {
706 let region = self
707 .get_region(region_id)
708 .context(RegionNotFoundSnafu { region_id })?;
709 ensure!(
710 region.is_flushable(),
711 FlushableRegionStateSnafu {
712 region_id,
713 state: region.state(),
714 }
715 );
716 Ok(region)
717 }
718
719 pub(crate) fn flushable_region_or<F: OnFailure>(
723 &self,
724 region_id: RegionId,
725 cb: &mut F,
726 ) -> Option<MitoRegionRef> {
727 match self.flushable_region(region_id) {
728 Ok(region) => Some(region),
729 Err(e) => {
730 cb.on_failure(e);
731 None
732 }
733 }
734 }
735
736 pub(crate) fn remove_region(&self, region_id: RegionId) {
738 let mut regions = self.regions.write().unwrap();
739 regions.remove(®ion_id);
740 }
741
742 pub(crate) fn list_regions(&self) -> Vec<MitoRegionRef> {
744 let regions = self.regions.read().unwrap();
745 regions.values().cloned().collect()
746 }
747
748 pub(crate) fn clear(&self) {
750 self.regions.write().unwrap().clear();
751 }
752}
753
754pub(crate) type RegionMapRef = Arc<RegionMap>;
755
756#[derive(Debug, Default)]
758pub(crate) struct OpeningRegions {
759 regions: RwLock<HashMap<RegionId, Vec<OptionOutputTx>>>,
760}
761
762impl OpeningRegions {
763 pub(crate) fn wait_for_opening_region(
765 &self,
766 region_id: RegionId,
767 sender: OptionOutputTx,
768 ) -> Option<OptionOutputTx> {
769 let mut regions = self.regions.write().unwrap();
770 match regions.entry(region_id) {
771 Entry::Occupied(mut senders) => {
772 senders.get_mut().push(sender);
773 None
774 }
775 Entry::Vacant(_) => Some(sender),
776 }
777 }
778
779 pub(crate) fn is_region_exists(&self, region_id: RegionId) -> bool {
781 let regions = self.regions.read().unwrap();
782 regions.contains_key(®ion_id)
783 }
784
785 pub(crate) fn insert_sender(&self, region: RegionId, sender: OptionOutputTx) {
787 let mut regions = self.regions.write().unwrap();
788 regions.insert(region, vec![sender]);
789 }
790
791 pub(crate) fn remove_sender(&self, region_id: RegionId) -> Vec<OptionOutputTx> {
793 let mut regions = self.regions.write().unwrap();
794 regions.remove(®ion_id).unwrap_or_default()
795 }
796
797 #[cfg(test)]
798 pub(crate) fn sender_len(&self, region_id: RegionId) -> usize {
799 let regions = self.regions.read().unwrap();
800 if let Some(senders) = regions.get(®ion_id) {
801 senders.len()
802 } else {
803 0
804 }
805 }
806}
807
808pub(crate) type OpeningRegionsRef = Arc<OpeningRegions>;
809
810#[derive(Default, Debug, Clone)]
812pub(crate) struct ManifestStats {
813 total_manifest_size: Arc<AtomicU64>,
814 manifest_version: Arc<AtomicU64>,
815}
816
817impl ManifestStats {
818 fn total_manifest_size(&self) -> u64 {
819 self.total_manifest_size.load(Ordering::Relaxed)
820 }
821
822 fn manifest_version(&self) -> u64 {
823 self.manifest_version.load(Ordering::Relaxed)
824 }
825}
826
827#[cfg(test)]
828mod tests {
829 use std::sync::Arc;
830
831 use crossbeam_utils::atomic::AtomicCell;
832 use store_api::region_engine::RegionRole;
833 use store_api::storage::RegionId;
834
835 use crate::region::{RegionLeaderState, RegionRoleState};
836 use crate::test_util::scheduler_util::SchedulerEnv;
837 use crate::test_util::version_util::VersionControlBuilder;
838
839 #[test]
840 fn test_region_state_lock_free() {
841 assert!(AtomicCell::<RegionRoleState>::is_lock_free());
842 }
843
844 #[tokio::test]
845 async fn test_set_region_state() {
846 let env = SchedulerEnv::new().await;
847 let builder = VersionControlBuilder::new();
848 let version_control = Arc::new(builder.build());
849 let manifest_ctx = env
850 .mock_manifest_context(version_control.current().version.metadata.clone())
851 .await;
852
853 let region_id = RegionId::new(1024, 0);
854 manifest_ctx.set_role(RegionRole::Follower, region_id);
856 assert_eq!(manifest_ctx.state.load(), RegionRoleState::Follower);
857
858 manifest_ctx.set_role(RegionRole::Leader, region_id);
860 assert_eq!(
861 manifest_ctx.state.load(),
862 RegionRoleState::Leader(RegionLeaderState::Writable)
863 );
864
865 manifest_ctx.set_role(RegionRole::DowngradingLeader, region_id);
867 assert_eq!(
868 manifest_ctx.state.load(),
869 RegionRoleState::Leader(RegionLeaderState::Downgrading)
870 );
871
872 manifest_ctx.set_role(RegionRole::Follower, region_id);
874 assert_eq!(manifest_ctx.state.load(), RegionRoleState::Follower);
875
876 manifest_ctx.set_role(RegionRole::DowngradingLeader, region_id);
878 assert_eq!(manifest_ctx.state.load(), RegionRoleState::Follower);
879
880 manifest_ctx.set_role(RegionRole::Leader, region_id);
882 manifest_ctx.set_role(RegionRole::DowngradingLeader, region_id);
883 assert_eq!(
884 manifest_ctx.state.load(),
885 RegionRoleState::Leader(RegionLeaderState::Downgrading)
886 );
887
888 manifest_ctx.set_role(RegionRole::Leader, region_id);
890 assert_eq!(
891 manifest_ctx.state.load(),
892 RegionRoleState::Leader(RegionLeaderState::Writable)
893 );
894 }
895}