1mod handle_alter;
18mod handle_apply_staging;
19mod handle_bulk_insert;
20mod handle_catchup;
21mod handle_close;
22mod handle_compaction;
23mod handle_copy_region;
24mod handle_create;
25mod handle_drop;
26mod handle_enter_staging;
27mod handle_flush;
28mod handle_manifest;
29mod handle_open;
30mod handle_rebuild_index;
31mod handle_remap;
32mod handle_truncate;
33mod handle_write;
34
35use std::collections::HashMap;
36use std::path::Path;
37use std::sync::Arc;
38use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
39use std::time::Duration;
40
41use common_base::Plugins;
42use common_error::ext::BoxedError;
43use common_meta::key::SchemaMetadataManagerRef;
44use common_runtime::JoinHandle;
45use common_stat::get_total_memory_bytes;
46use common_telemetry::{error, info, warn};
47use futures::future::{join_all, try_join_all};
48use object_store::ObjectStore;
49use object_store::manager::ObjectStoreManagerRef;
50use prometheus::{Histogram, IntGauge};
51use rand::{Rng, rng};
52use snafu::ResultExt;
53use store_api::logstore::LogStore;
54use store_api::region_engine::{
55 RegionRole, SetRegionRoleStateResponse, SetRegionRoleStateSuccess, SettableRegionRoleState,
56};
57use store_api::storage::{FileId, RegionId};
58use tokio::sync::mpsc::{Receiver, Sender};
59use tokio::sync::{Mutex, Semaphore, mpsc, oneshot, watch};
60
61use crate::access_layer::new_fs_cache_store;
62use crate::cache::write_cache::{WriteCache, WriteCacheRef};
63use crate::cache::{CacheManager, CacheManagerRef, WriteCacheUploadStoreWrapperRef};
64use crate::compaction::CompactionScheduler;
65use crate::compaction::memory_manager::{CompactionMemoryManager, new_compaction_memory_manager};
66use crate::config::MitoConfig;
67use crate::error::{self, CreateDirSnafu, JoinSnafu, Result, WorkerStoppedSnafu};
68use crate::flush::{FlushScheduler, WriteBufferManagerImpl, WriteBufferManagerRef};
69use crate::gc::{GcLimiter, GcLimiterRef};
70use crate::memtable::MemtableBuilderProvider;
71use crate::metrics::{REGION_COUNT, REQUEST_WAIT_TIME, WRITE_STALLING};
72use crate::region::opener::PartitionExprFetcherRef;
73use crate::region::{
74 CatchupRegions, CatchupRegionsRef, MitoRegionRef, NodeLease, OpeningRegions, OpeningRegionsRef,
75 RegionMap, RegionMapRef,
76};
77use crate::request::{
78 BackgroundNotify, BulkInsertRequest, DdlRequest, OptionOutputTx, SenderBulkRequest,
79 SenderDdlRequest, SenderWriteRequest, WorkerRequest, WorkerRequestWithTime,
80};
81use crate::schedule::scheduler::{LocalScheduler, SchedulerRef};
82use crate::series_index::{
83 IndexFilePurger, SeriesIndexTaskState, series_index_channel, spawn_series_index_tasks,
84};
85use crate::sst::file::RegionFileId;
86use crate::sst::file_ref::FileReferenceManagerRef;
87use crate::sst::index::IndexBuildScheduler;
88use crate::sst::index::intermediate::IntermediateManager;
89use crate::sst::index::puffin_manager::PuffinManagerFactory;
90#[cfg(feature = "hdfs-object-store")]
91use crate::temp_file_sweeper;
92use crate::time_provider::{StdTimeProvider, TimeProviderRef};
93use crate::wal::Wal;
94use crate::worker::handle_manifest::RegionEditQueues;
95
96pub(crate) type WorkerId = u32;
98
99pub(crate) const DROPPING_MARKER_FILE: &str = ".dropping";
100
101pub(crate) const CHECK_REGION_INTERVAL: Duration = Duration::from_secs(60);
103pub(crate) const MAX_INITIAL_CHECK_DELAY_SECS: u64 = 60 * 3;
105
106#[cfg_attr(doc, aquamarine::aquamarine)]
107pub(crate) struct WorkerGroup {
144 workers: Vec<RegionWorker>,
146 flush_job_pool: SchedulerRef,
148 compact_job_pool: SchedulerRef,
150 index_build_job_pool: SchedulerRef,
152 purge_scheduler: SchedulerRef,
154 cache_manager: CacheManagerRef,
156 file_ref_manager: FileReferenceManagerRef,
158 gc_limiter: GcLimiterRef,
160 node_lease: Arc<NodeLease>,
161 object_store_manager: ObjectStoreManagerRef,
163 puffin_manager_factory: PuffinManagerFactory,
165 intermediate_manager: IntermediateManager,
167 schema_metadata_manager: SchemaMetadataManagerRef,
169}
170
171impl WorkerGroup {
172 #[allow(clippy::too_many_arguments)]
176 pub(crate) async fn start<S: LogStore>(
177 data_home: &str,
178 config: Arc<MitoConfig>,
179 log_store: Arc<S>,
180 object_store_manager: ObjectStoreManagerRef,
181 schema_metadata_manager: SchemaMetadataManagerRef,
182 file_ref_manager: FileReferenceManagerRef,
183 partition_expr_fetcher: PartitionExprFetcherRef,
184 plugins: Plugins,
185 ) -> Result<WorkerGroup> {
186 let (flush_sender, flush_receiver) = watch::channel(());
187 let write_buffer_manager = Arc::new(
188 WriteBufferManagerImpl::new(config.global_write_buffer_size.as_bytes() as usize)
189 .with_notifier(flush_sender.clone()),
190 );
191 let puffin_manager_factory = PuffinManagerFactory::new(
192 &config.index.aux_path,
193 config.index.staging_size.as_bytes(),
194 Some(config.index.write_buffer_size.as_bytes() as _),
195 config.index.staging_ttl,
196 )
197 .await?;
198 let intermediate_manager = IntermediateManager::init_fs(&config.index.aux_path)
199 .await?
200 .with_buffer_size(Some(config.index.write_buffer_size.as_bytes() as _));
201 let index_build_job_pool =
202 Arc::new(LocalScheduler::new(config.max_background_index_builds));
203 let series_index_store = series_index_store_from_config(&config, data_home).await?;
204 let series_index_disk_usage = Arc::new(AtomicU64::new(0));
205 let flush_job_pool = Arc::new(LocalScheduler::new(config.max_background_flushes));
206 let compact_job_pool = Arc::new(LocalScheduler::new(config.max_background_compactions));
207 let flush_semaphore = Arc::new(Semaphore::new(config.max_background_flushes));
208 let purge_scheduler = Arc::new(LocalScheduler::new(config.max_background_purges));
210 let upload_store_wrapper = plugins.get::<WriteCacheUploadStoreWrapperRef>();
211 let write_cache = write_cache_from_config(
212 &config,
213 puffin_manager_factory.clone(),
214 intermediate_manager.clone(),
215 upload_store_wrapper,
216 )
217 .await?;
218 let cache_manager = Arc::new(
219 CacheManager::builder()
220 .sst_meta_cache_size(config.sst_meta_cache_size.as_bytes())
221 .vector_cache_size(config.vector_cache_size.as_bytes())
222 .page_cache_size(config.page_cache_size.as_bytes())
223 .selector_result_cache_size(config.selector_result_cache_size.as_bytes())
224 .range_result_cache_size(config.range_result_cache_size.as_bytes())
225 .prefilter_result_cache_size(config.prefilter_result_cache_size.as_bytes())
226 .index_metadata_size(config.index.metadata_cache_size.as_bytes())
227 .index_content_size(config.index.content_cache_size.as_bytes())
228 .index_content_page_size(config.index.content_cache_page_size.as_bytes())
229 .index_result_cache_size(config.index.result_cache_size.as_bytes())
230 .puffin_metadata_size(config.index.metadata_cache_size.as_bytes())
231 .write_cache(write_cache)
232 .build(),
233 );
234 let time_provider = Arc::new(StdTimeProvider);
235 let total_memory = get_total_memory_bytes();
236 let total_memory = if total_memory > 0 {
237 total_memory as u64
238 } else {
239 0
240 };
241 let compaction_limit_bytes = config
242 .experimental_compaction_memory_limit
243 .resolve(total_memory);
244 let compaction_memory_manager =
245 Arc::new(new_compaction_memory_manager(compaction_limit_bytes));
246 let gc_limiter = Arc::new(GcLimiter::new(config.gc.max_concurrent_gc_job));
247 #[cfg(feature = "hdfs-object-store")]
252 let sweep_limiter = (config.gc.max_concurrent_gc_job > 0)
253 .then(|| Arc::new(Semaphore::new(config.gc.max_concurrent_gc_job)));
254 let node_lease = Arc::new(NodeLease::default());
255
256 let workers = (0..config.num_workers)
257 .map(|id| {
258 WorkerStarter {
259 id: id as WorkerId,
260 config: config.clone(),
261 log_store: log_store.clone(),
262 object_store_manager: object_store_manager.clone(),
263 write_buffer_manager: write_buffer_manager.clone(),
264 index_build_job_pool: index_build_job_pool.clone(),
265 series_index_store: series_index_store.clone(),
266 series_index_disk_usage: series_index_disk_usage.clone(),
267 flush_job_pool: flush_job_pool.clone(),
268 compact_job_pool: compact_job_pool.clone(),
269 purge_scheduler: purge_scheduler.clone(),
270 listener: WorkerListener::default(),
271 cache_manager: cache_manager.clone(),
272 compaction_memory_manager: compaction_memory_manager.clone(),
273 puffin_manager_factory: puffin_manager_factory.clone(),
274 intermediate_manager: intermediate_manager.clone(),
275 time_provider: time_provider.clone(),
276 flush_sender: flush_sender.clone(),
277 flush_receiver: flush_receiver.clone(),
278 plugins: plugins.clone(),
279 schema_metadata_manager: schema_metadata_manager.clone(),
280 file_ref_manager: file_ref_manager.clone(),
281 partition_expr_fetcher: partition_expr_fetcher.clone(),
282 flush_semaphore: flush_semaphore.clone(),
283 #[cfg(feature = "hdfs-object-store")]
284 sweep_limiter: sweep_limiter.clone(),
285 node_lease: node_lease.clone(),
286 }
287 .start()
288 })
289 .collect::<Result<Vec<_>>>()?;
290
291 Ok(WorkerGroup {
292 workers,
293 flush_job_pool,
294 compact_job_pool,
295 index_build_job_pool,
296 purge_scheduler,
297 cache_manager,
298 file_ref_manager,
299 gc_limiter,
300 node_lease,
301 object_store_manager,
302 puffin_manager_factory,
303 intermediate_manager,
304 schema_metadata_manager,
305 })
306 }
307
308 pub(crate) async fn stop(&self) -> Result<()> {
310 info!("Stop region worker group");
311
312 self.compact_job_pool.stop(true).await?;
315 self.flush_job_pool.stop(true).await?;
317 self.purge_scheduler.stop(true).await?;
319 self.index_build_job_pool.stop(true).await?;
321
322 try_join_all(self.workers.iter().map(|worker| worker.stop())).await?;
323
324 Ok(())
325 }
326
327 pub(crate) async fn submit_to_worker(
329 &self,
330 region_id: RegionId,
331 request: WorkerRequest,
332 ) -> Result<()> {
333 self.worker(region_id).submit_request(request).await
334 }
335
336 pub(crate) async fn submit_writes(&self, writes: Vec<SenderWriteRequest>) {
341 let mut writes_by_worker: HashMap<usize, Vec<SenderWriteRequest>> = HashMap::new();
342 for write in writes {
343 let index = region_id_to_index(write.request.region_id, self.workers.len());
344 writes_by_worker.entry(index).or_default().push(write);
345 }
346 join_all(
347 writes_by_worker
348 .into_iter()
349 .map(|(index, writes)| async move {
350 let worker = &self.workers[index];
351 let Err(request) = worker
352 .try_submit_request(WorkerRequest::BatchWrite(writes))
353 .await
354 else {
355 return;
356 };
357 let WorkerRequest::BatchWrite(writes) = request else {
358 unreachable!("the worker returns the request it was given");
359 };
360 for write in writes {
361 write
362 .sender
363 .send(WorkerStoppedSnafu { id: worker.id }.fail());
364 }
365 }),
366 )
367 .await;
368 }
369
370 pub(crate) fn is_region_exists(&self, region_id: RegionId) -> bool {
372 self.worker(region_id).is_region_exists(region_id)
373 }
374
375 pub(crate) fn is_region_opening(&self, region_id: RegionId) -> bool {
377 self.worker(region_id).is_region_opening(region_id)
378 }
379
380 pub(crate) fn is_region_catching_up(&self, region_id: RegionId) -> bool {
382 self.worker(region_id).is_region_catching_up(region_id)
383 }
384
385 pub(crate) fn get_region(&self, region_id: RegionId) -> Option<MitoRegionRef> {
389 self.worker(region_id).get_region(region_id)
390 }
391
392 pub(crate) fn cache_manager(&self) -> CacheManagerRef {
394 self.cache_manager.clone()
395 }
396
397 pub(crate) fn file_ref_manager(&self) -> FileReferenceManagerRef {
398 self.file_ref_manager.clone()
399 }
400
401 pub(crate) fn gc_limiter(&self) -> GcLimiterRef {
402 self.gc_limiter.clone()
403 }
404
405 pub(crate) fn node_lease(&self) -> &NodeLease {
406 &self.node_lease
407 }
408
409 pub(crate) fn worker(&self, region_id: RegionId) -> &RegionWorker {
411 let index = region_id_to_index(region_id, self.workers.len());
412
413 &self.workers[index]
414 }
415
416 pub(crate) fn all_regions(&self) -> impl Iterator<Item = MitoRegionRef> + use<'_> {
417 self.workers
418 .iter()
419 .flat_map(|worker| worker.regions.list_regions())
420 }
421
422 pub(crate) fn object_store_manager(&self) -> &ObjectStoreManagerRef {
423 &self.object_store_manager
424 }
425
426 pub(crate) fn puffin_manager_factory(&self) -> &PuffinManagerFactory {
427 &self.puffin_manager_factory
428 }
429
430 pub(crate) fn intermediate_manager(&self) -> &IntermediateManager {
431 &self.intermediate_manager
432 }
433
434 pub(crate) fn schema_metadata_manager(&self) -> &SchemaMetadataManagerRef {
435 &self.schema_metadata_manager
436 }
437}
438
439#[cfg(any(test, feature = "test"))]
441impl WorkerGroup {
442 #[allow(clippy::too_many_arguments)]
446 pub(crate) async fn start_for_test<S: LogStore>(
447 data_home: &str,
448 config: Arc<MitoConfig>,
449 log_store: Arc<S>,
450 object_store_manager: ObjectStoreManagerRef,
451 write_buffer_manager: Option<WriteBufferManagerRef>,
452 listener: Option<crate::engine::listener::EventListenerRef>,
453 schema_metadata_manager: SchemaMetadataManagerRef,
454 file_ref_manager: FileReferenceManagerRef,
455 time_provider: TimeProviderRef,
456 partition_expr_fetcher: PartitionExprFetcherRef,
457 ) -> Result<WorkerGroup> {
458 let (flush_sender, flush_receiver) = watch::channel(());
459 let write_buffer_manager = write_buffer_manager.unwrap_or_else(|| {
460 Arc::new(
461 WriteBufferManagerImpl::new(config.global_write_buffer_size.as_bytes() as usize)
462 .with_notifier(flush_sender.clone()),
463 )
464 });
465 let index_build_job_pool =
466 Arc::new(LocalScheduler::new(config.max_background_index_builds));
467 let series_index_store = series_index_store_from_config(&config, data_home).await?;
468 let series_index_disk_usage = Arc::new(AtomicU64::new(0));
469 let flush_job_pool = Arc::new(LocalScheduler::new(config.max_background_flushes));
470 let compact_job_pool = Arc::new(LocalScheduler::new(config.max_background_compactions));
471 let flush_semaphore = Arc::new(Semaphore::new(config.max_background_flushes));
472 let purge_scheduler = Arc::new(LocalScheduler::new(config.max_background_flushes));
473 let puffin_manager_factory = PuffinManagerFactory::new(
474 &config.index.aux_path,
475 config.index.staging_size.as_bytes(),
476 Some(config.index.write_buffer_size.as_bytes() as _),
477 config.index.staging_ttl,
478 )
479 .await?;
480 let intermediate_manager = IntermediateManager::init_fs(&config.index.aux_path)
481 .await?
482 .with_buffer_size(Some(config.index.write_buffer_size.as_bytes() as _));
483 let write_cache = write_cache_from_config(
484 &config,
485 puffin_manager_factory.clone(),
486 intermediate_manager.clone(),
487 None,
488 )
489 .await?;
490 let cache_manager = Arc::new(
491 CacheManager::builder()
492 .sst_meta_cache_size(config.sst_meta_cache_size.as_bytes())
493 .vector_cache_size(config.vector_cache_size.as_bytes())
494 .page_cache_size(config.page_cache_size.as_bytes())
495 .selector_result_cache_size(config.selector_result_cache_size.as_bytes())
496 .range_result_cache_size(config.range_result_cache_size.as_bytes())
497 .prefilter_result_cache_size(config.prefilter_result_cache_size.as_bytes())
498 .write_cache(write_cache)
499 .build(),
500 );
501 let total_memory = get_total_memory_bytes();
502 let total_memory = if total_memory > 0 {
503 total_memory as u64
504 } else {
505 0
506 };
507 let compaction_limit_bytes = config
508 .experimental_compaction_memory_limit
509 .resolve(total_memory);
510 let compaction_memory_manager =
511 Arc::new(new_compaction_memory_manager(compaction_limit_bytes));
512 let gc_limiter = Arc::new(GcLimiter::new(config.gc.max_concurrent_gc_job));
513 #[cfg(feature = "hdfs-object-store")]
518 let sweep_limiter = (config.gc.max_concurrent_gc_job > 0)
519 .then(|| Arc::new(Semaphore::new(config.gc.max_concurrent_gc_job)));
520 let node_lease = Arc::new(NodeLease::default());
521 let workers = (0..config.num_workers)
522 .map(|id| {
523 WorkerStarter {
524 id: id as WorkerId,
525 config: config.clone(),
526 log_store: log_store.clone(),
527 object_store_manager: object_store_manager.clone(),
528 write_buffer_manager: write_buffer_manager.clone(),
529 index_build_job_pool: index_build_job_pool.clone(),
530 series_index_store: series_index_store.clone(),
531 series_index_disk_usage: series_index_disk_usage.clone(),
532 flush_job_pool: flush_job_pool.clone(),
533 compact_job_pool: compact_job_pool.clone(),
534 purge_scheduler: purge_scheduler.clone(),
535 listener: WorkerListener::new(listener.clone()),
536 cache_manager: cache_manager.clone(),
537 compaction_memory_manager: compaction_memory_manager.clone(),
538 puffin_manager_factory: puffin_manager_factory.clone(),
539 intermediate_manager: intermediate_manager.clone(),
540 time_provider: time_provider.clone(),
541 flush_sender: flush_sender.clone(),
542 flush_receiver: flush_receiver.clone(),
543 plugins: Plugins::new(),
544 schema_metadata_manager: schema_metadata_manager.clone(),
545 file_ref_manager: file_ref_manager.clone(),
546 partition_expr_fetcher: partition_expr_fetcher.clone(),
547 flush_semaphore: flush_semaphore.clone(),
548 #[cfg(feature = "hdfs-object-store")]
549 sweep_limiter: sweep_limiter.clone(),
550 node_lease: node_lease.clone(),
551 }
552 .start()
553 })
554 .collect::<Result<Vec<_>>>()?;
555
556 Ok(WorkerGroup {
557 workers,
558 flush_job_pool,
559 compact_job_pool,
560 index_build_job_pool,
561 purge_scheduler,
562 cache_manager,
563 file_ref_manager,
564 gc_limiter,
565 node_lease,
566 object_store_manager,
567 puffin_manager_factory,
568 intermediate_manager,
569 schema_metadata_manager,
570 })
571 }
572
573 pub(crate) fn purge_scheduler(&self) -> &SchedulerRef {
575 &self.purge_scheduler
576 }
577}
578
579fn region_id_to_index(id: RegionId, num_workers: usize) -> usize {
580 ((id.table_id() as usize % num_workers) + (id.region_number() as usize % num_workers))
581 % num_workers
582}
583
584async fn series_index_store_from_config(
586 config: &MitoConfig,
587 data_home: &str,
588) -> Result<Option<ObjectStore>> {
589 if !config.experimental_enable_series_index {
590 return Ok(None);
591 }
592
593 let root = Path::new(data_home).join("series_index");
594 new_fs_cache_store(&root.to_string_lossy()).await.map(Some)
595}
596
597pub async fn write_cache_from_config(
598 config: &MitoConfig,
599 puffin_manager_factory: PuffinManagerFactory,
600 intermediate_manager: IntermediateManager,
601 upload_store_wrapper: Option<WriteCacheUploadStoreWrapperRef>,
602) -> Result<Option<WriteCacheRef>> {
603 if !config.enable_write_cache {
604 return Ok(None);
605 }
606
607 tokio::fs::create_dir_all(Path::new(&config.write_cache_path))
608 .await
609 .context(CreateDirSnafu {
610 dir: &config.write_cache_path,
611 })?;
612
613 let cache = WriteCache::new_fs(
614 &config.write_cache_path,
615 config.write_cache_size,
616 config.write_cache_ttl,
617 Some(config.index_cache_percent),
618 config.enable_refill_cache_on_read,
619 puffin_manager_factory,
620 intermediate_manager,
621 config.manifest_cache_size,
622 )
623 .await?
624 .with_upload_store_wrapper(upload_store_wrapper);
625 Ok(Some(Arc::new(cache)))
626}
627
628pub(crate) fn worker_init_check_delay() -> Duration {
630 let init_check_delay = rng().random_range(0..MAX_INITIAL_CHECK_DELAY_SECS);
631 Duration::from_secs(init_check_delay)
632}
633
634struct WorkerStarter<S> {
636 id: WorkerId,
637 config: Arc<MitoConfig>,
638 log_store: Arc<S>,
639 object_store_manager: ObjectStoreManagerRef,
640 write_buffer_manager: WriteBufferManagerRef,
641 compact_job_pool: SchedulerRef,
642 index_build_job_pool: SchedulerRef,
643 series_index_store: Option<ObjectStore>,
644 series_index_disk_usage: Arc<AtomicU64>,
645 flush_job_pool: SchedulerRef,
646 purge_scheduler: SchedulerRef,
647 listener: WorkerListener,
648 cache_manager: CacheManagerRef,
649 compaction_memory_manager: Arc<CompactionMemoryManager>,
650 puffin_manager_factory: PuffinManagerFactory,
651 intermediate_manager: IntermediateManager,
652 time_provider: TimeProviderRef,
653 flush_sender: watch::Sender<()>,
655 flush_receiver: watch::Receiver<()>,
657 plugins: Plugins,
658 schema_metadata_manager: SchemaMetadataManagerRef,
659 file_ref_manager: FileReferenceManagerRef,
660 partition_expr_fetcher: PartitionExprFetcherRef,
661 flush_semaphore: Arc<Semaphore>,
662 #[cfg(feature = "hdfs-object-store")]
663 sweep_limiter: Option<Arc<Semaphore>>,
666 node_lease: Arc<NodeLease>,
667}
668
669impl<S: LogStore> WorkerStarter<S> {
670 fn start(self) -> Result<RegionWorker> {
672 let regions = Arc::new(RegionMap::default());
673 let opening_regions = Arc::new(OpeningRegions::default());
674 let catchup_regions = Arc::new(CatchupRegions::default());
675 let (sender, receiver) = mpsc::channel(self.config.worker_channel_size);
676
677 let running = Arc::new(AtomicBool::new(true));
678 let (series_index_task_state, series_index_receiver) = if self.series_index_store.is_some()
679 {
680 let (state, receiver) =
681 SeriesIndexTaskState::new(self.id, self.config.worker_channel_size);
682 (Some(Arc::new(state)), Some(receiver))
683 } else {
684 (None, None)
685 };
686 let mut series_index_purger = None;
687 let series_index_handle = self
688 .series_index_store
689 .clone()
690 .zip(series_index_task_state.clone())
691 .zip(series_index_receiver)
692 .map(|((store, state), receiver)| {
693 let (purger, purge_receiver) = series_index_channel(store.clone());
694 series_index_purger = Some(purger.clone());
695 spawn_series_index_tasks(
696 self.id,
697 store,
698 regions.clone(),
699 state,
700 receiver,
701 self.config.experimental_series_index_bucket_width,
702 purger,
703 purge_receiver,
704 self.series_index_disk_usage.clone(),
705 self.config.experimental_series_index_max_size.as_bytes(),
706 self.config.experimental_series_index_maintenance_interval,
707 self.time_provider.clone(),
708 self.config.experimental_enable_range_index,
709 )
710 });
711 let now = self.time_provider.current_time_millis();
712 let id_string = self.id.to_string();
713 let mut worker_thread = RegionWorkerLoop {
714 id: self.id,
715 config: self.config.clone(),
716 regions: regions.clone(),
717 catchup_regions: catchup_regions.clone(),
718 dropping_regions: Arc::new(RegionMap::default()),
719 opening_regions: opening_regions.clone(),
720 sender: sender.clone(),
721 receiver,
722 wal: Wal::new(self.log_store),
723 object_store_manager: self.object_store_manager.clone(),
724 running: running.clone(),
725 memtable_builder_provider: MemtableBuilderProvider::new(
726 Some(self.write_buffer_manager.clone()),
727 self.config.clone(),
728 ),
729 purge_scheduler: self.purge_scheduler.clone(),
730 write_buffer_manager: self.write_buffer_manager,
731 index_build_scheduler: IndexBuildScheduler::new(
732 self.index_build_job_pool,
733 self.config.max_background_index_builds,
734 ),
735 series_index_task_state: series_index_task_state.clone(),
736 series_index_store: self.series_index_store,
737 series_index_purger,
738 flush_scheduler: FlushScheduler::new(self.flush_job_pool),
739 compaction_scheduler: CompactionScheduler::new(
740 self.compact_job_pool,
741 sender.clone(),
742 self.cache_manager.clone(),
743 self.config.clone(),
744 self.listener.clone(),
745 self.plugins.clone(),
746 self.compaction_memory_manager.clone(),
747 self.config.experimental_compaction_on_exhausted,
748 ),
749 stalled_requests: StalledRequests::default(),
750 listener: self.listener,
751 cache_manager: self.cache_manager,
752 puffin_manager_factory: self.puffin_manager_factory,
753 intermediate_manager: self.intermediate_manager,
754 time_provider: self.time_provider,
755 last_periodical_check_millis: now,
756 flush_sender: self.flush_sender,
757 flush_receiver: self.flush_receiver,
758 stalling_count: WRITE_STALLING.with_label_values(&[&id_string]),
759 region_count: REGION_COUNT.with_label_values(&[&id_string]),
760 request_wait_time: REQUEST_WAIT_TIME.with_label_values(&[&id_string]),
761 region_edit_queues: RegionEditQueues::default(),
762 schema_metadata_manager: self.schema_metadata_manager,
763 file_ref_manager: self.file_ref_manager.clone(),
764 partition_expr_fetcher: self.partition_expr_fetcher,
765 flush_semaphore: self.flush_semaphore,
766 plugins: self.plugins,
767 #[cfg(feature = "hdfs-object-store")]
768 sweep_limiter: self.sweep_limiter,
769 node_lease: self.node_lease,
770 };
771 let handle = common_runtime::spawn_global(async move {
772 worker_thread.run().await;
773 });
774
775 Ok(RegionWorker {
776 id: self.id,
777 regions,
778 opening_regions,
779 catchup_regions,
780 sender,
781 handle: Mutex::new(Some(handle)),
782 series_index_handle: Mutex::new(series_index_handle),
783 series_index_task_state,
784 running,
785 })
786 }
787}
788
789pub(crate) struct RegionWorker {
791 id: WorkerId,
793 regions: RegionMapRef,
795 opening_regions: OpeningRegionsRef,
797 catchup_regions: CatchupRegionsRef,
799 sender: Sender<WorkerRequestWithTime>,
801 handle: Mutex<Option<JoinHandle<()>>>,
803 series_index_handle: Mutex<Option<JoinHandle<()>>>,
805 series_index_task_state: Option<Arc<SeriesIndexTaskState>>,
807 running: Arc<AtomicBool>,
809}
810
811impl RegionWorker {
812 async fn submit_request(&self, request: WorkerRequest) -> Result<()> {
814 self.try_submit_request(request)
815 .await
816 .map_err(|_| WorkerStoppedSnafu { id: self.id }.build())
817 }
818
819 async fn try_submit_request(
821 &self,
822 request: WorkerRequest,
823 ) -> std::result::Result<(), WorkerRequest> {
824 if !self.is_running() {
825 return Err(request);
826 }
827 let request_with_time = WorkerRequestWithTime::new(request);
828 if let Err(mpsc::error::SendError(request_with_time)) =
829 self.sender.send(request_with_time).await
830 {
831 warn!(
832 "Worker {} is already exited but the running flag is still true",
833 self.id
834 );
835 self.set_running(false);
837 if let Some(state) = &self.series_index_task_state {
838 state.stop();
839 }
840 return Err(request_with_time.request);
841 }
842
843 Ok(())
844 }
845
846 async fn stop(&self) -> Result<()> {
850 let handle = self.handle.lock().await.take();
851 self.set_running(false);
852 if let Some(state) = &self.series_index_task_state {
853 state.stop();
854 }
855
856 let mut worker_result = Ok(());
857 if let Some(handle) = handle {
858 info!("Stop region worker {}", self.id);
859
860 if self
861 .sender
862 .send(WorkerRequestWithTime::new(WorkerRequest::Stop))
863 .await
864 .is_err()
865 {
866 warn!("Worker {} is already exited before stop", self.id);
867 }
868
869 worker_result = handle.await.context(JoinSnafu);
870 }
871 let series_index_result = if let Some(handle) = self.series_index_handle.lock().await.take()
872 {
873 handle.await.context(JoinSnafu)
874 } else {
875 Ok(())
876 };
877 worker_result?;
878 series_index_result?;
879
880 Ok(())
881 }
882
883 fn is_running(&self) -> bool {
885 self.running.load(Ordering::Relaxed)
886 }
887
888 fn set_running(&self, value: bool) {
890 self.running.store(value, Ordering::Relaxed)
891 }
892
893 fn is_region_exists(&self, region_id: RegionId) -> bool {
895 self.regions.is_region_exists(region_id)
896 }
897
898 fn is_region_opening(&self, region_id: RegionId) -> bool {
900 self.opening_regions.is_region_exists(region_id)
901 }
902
903 fn is_region_catching_up(&self, region_id: RegionId) -> bool {
905 self.catchup_regions.is_region_exists(region_id)
906 }
907
908 fn get_region(&self, region_id: RegionId) -> Option<MitoRegionRef> {
910 self.regions.get_region(region_id)
911 }
912
913 #[cfg(test)]
914 pub(crate) fn opening_regions(&self) -> &OpeningRegionsRef {
916 &self.opening_regions
917 }
918
919 #[cfg(test)]
920 pub(crate) fn catchup_regions(&self) -> &CatchupRegionsRef {
922 &self.catchup_regions
923 }
924}
925
926impl Drop for RegionWorker {
927 fn drop(&mut self) {
928 if let Some(state) = &self.series_index_task_state {
929 state.stop();
930 }
931 if self.is_running() {
932 self.set_running(false);
933 }
935 }
936}
937
938type RequestBuffer = Vec<WorkerRequest>;
939
940#[derive(Default)]
944pub(crate) struct StalledRequests {
945 pub(crate) requests:
952 HashMap<RegionId, (usize, Vec<SenderWriteRequest>, Vec<SenderBulkRequest>)>,
953 pub(crate) estimated_size: usize,
955}
956
957impl StalledRequests {
958 pub(crate) fn estimated_size(&self, region_id: &RegionId) -> usize {
960 self.requests
961 .get(region_id)
962 .map(|(size, _, _)| *size)
963 .unwrap_or_default()
964 }
965
966 pub(crate) fn append(
968 &mut self,
969 requests: &mut Vec<SenderWriteRequest>,
970 bulk_requests: &mut Vec<SenderBulkRequest>,
971 ) {
972 for req in requests.drain(..) {
973 self.push(req);
974 }
975 for req in bulk_requests.drain(..) {
976 self.push_bulk(req);
977 }
978 }
979
980 pub(crate) fn push(&mut self, req: SenderWriteRequest) {
982 let (size, requests, _) = self.requests.entry(req.request.region_id).or_default();
983 let req_size = req.request.estimated_size();
984 *size += req_size;
985 self.estimated_size += req_size;
986 requests.push(req);
987 }
988
989 pub(crate) fn push_bulk(&mut self, req: SenderBulkRequest) {
990 let region_id = req.region_id;
991 let (size, _, requests) = self.requests.entry(region_id).or_default();
992 let req_size = req.request.estimated_size();
993 *size += req_size;
994 self.estimated_size += req_size;
995 requests.push(req);
996 }
997
998 pub(crate) fn remove(
1000 &mut self,
1001 region_id: &RegionId,
1002 ) -> (Vec<SenderWriteRequest>, Vec<SenderBulkRequest>) {
1003 if let Some((size, write_reqs, bulk_reqs)) = self.requests.remove(region_id) {
1004 self.estimated_size -= size;
1005 (write_reqs, bulk_reqs)
1006 } else {
1007 (vec![], vec![])
1008 }
1009 }
1010
1011 pub(crate) fn stalled_count(&self) -> usize {
1013 self.requests
1014 .values()
1015 .map(|(_, reqs, bulk_reqs)| reqs.len() + bulk_reqs.len())
1016 .sum()
1017 }
1018}
1019
1020struct RegionWorkerLoop<S> {
1022 id: WorkerId,
1024 config: Arc<MitoConfig>,
1026 regions: RegionMapRef,
1028 dropping_regions: RegionMapRef,
1030 opening_regions: OpeningRegionsRef,
1032 catchup_regions: CatchupRegionsRef,
1034 sender: Sender<WorkerRequestWithTime>,
1036 receiver: Receiver<WorkerRequestWithTime>,
1038 wal: Wal<S>,
1040 object_store_manager: ObjectStoreManagerRef,
1042 running: Arc<AtomicBool>,
1044 memtable_builder_provider: MemtableBuilderProvider,
1046 purge_scheduler: SchedulerRef,
1048 write_buffer_manager: WriteBufferManagerRef,
1050 index_build_scheduler: IndexBuildScheduler,
1052 series_index_task_state: Option<Arc<SeriesIndexTaskState>>,
1054 series_index_store: Option<ObjectStore>,
1056 series_index_purger: Option<IndexFilePurger>,
1057 flush_scheduler: FlushScheduler,
1059 compaction_scheduler: CompactionScheduler,
1061 stalled_requests: StalledRequests,
1063 listener: WorkerListener,
1065 cache_manager: CacheManagerRef,
1067 puffin_manager_factory: PuffinManagerFactory,
1069 intermediate_manager: IntermediateManager,
1071 time_provider: TimeProviderRef,
1073 last_periodical_check_millis: i64,
1075 flush_sender: watch::Sender<()>,
1077 flush_receiver: watch::Receiver<()>,
1079 stalling_count: IntGauge,
1081 region_count: IntGauge,
1083 request_wait_time: Histogram,
1085 region_edit_queues: RegionEditQueues,
1087 schema_metadata_manager: SchemaMetadataManagerRef,
1089 file_ref_manager: FileReferenceManagerRef,
1091 partition_expr_fetcher: PartitionExprFetcherRef,
1093 flush_semaphore: Arc<Semaphore>,
1095 plugins: Plugins,
1097 #[cfg(feature = "hdfs-object-store")]
1098 sweep_limiter: Option<Arc<Semaphore>>,
1101 node_lease: Arc<NodeLease>,
1102}
1103
1104impl<S: LogStore> RegionWorkerLoop<S> {
1105 fn register_region(&mut self, region: MitoRegionRef) {
1107 self.compaction_scheduler
1108 .register_region(®ion.version_control, ®ion.access_layer);
1109 self.regions.insert_region(region);
1110 }
1111
1112 async fn run(&mut self) {
1114 let init_check_delay = worker_init_check_delay();
1115 info!(
1116 "Start region worker thread {}, init_check_delay: {:?}",
1117 self.id, init_check_delay
1118 );
1119 self.last_periodical_check_millis += init_check_delay.as_millis() as i64;
1120
1121 let mut write_req_buffer: Vec<SenderWriteRequest> =
1123 Vec::with_capacity(self.config.worker_request_batch_size);
1124 let mut bulk_req_buffer: Vec<SenderBulkRequest> =
1125 Vec::with_capacity(self.config.worker_request_batch_size);
1126 let mut ddl_req_buffer: Vec<SenderDdlRequest> =
1127 Vec::with_capacity(self.config.worker_request_batch_size);
1128 let mut general_req_buffer: Vec<WorkerRequest> =
1129 RequestBuffer::with_capacity(self.config.worker_request_batch_size);
1130
1131 while self.running.load(Ordering::Relaxed) {
1132 write_req_buffer.clear();
1134 ddl_req_buffer.clear();
1135 general_req_buffer.clear();
1136 let mut bulk_insert_req_num = 0;
1137
1138 let max_wait_time = self.time_provider.wait_duration(CHECK_REGION_INTERVAL);
1139 let sleep = tokio::time::sleep(max_wait_time);
1140 tokio::pin!(sleep);
1141
1142 tokio::select! {
1143 request_opt = self.receiver.recv() => {
1144 match request_opt {
1145 Some(request_with_time) => {
1146 let wait_time = request_with_time.created_at.elapsed();
1148 self.request_wait_time.observe(wait_time.as_secs_f64());
1149
1150 match request_with_time.request {
1151 WorkerRequest::Write(sender_req) => write_req_buffer.push(sender_req),
1152 WorkerRequest::BatchWrite(sender_reqs) => write_req_buffer.extend(sender_reqs),
1153 WorkerRequest::Ddl(sender_req) => ddl_req_buffer.push(sender_req),
1154 WorkerRequest::BulkInserts(bulk_insert) => {
1155 bulk_insert_req_num += 1;
1156 self.buffer_bulk_insert_request(bulk_insert, &mut bulk_req_buffer)
1157 .await;
1158 }
1159 req => general_req_buffer.push(req),
1160 }
1161 },
1162 None => break,
1164 }
1165 }
1166 recv_res = self.flush_receiver.changed() => {
1167 if recv_res.is_err() {
1168 break;
1170 } else {
1171 self.maybe_flush_worker();
1176 self.handle_stalled_requests().await;
1178 continue;
1179 }
1180 }
1181 _ = &mut sleep => {
1182 self.handle_periodical_tasks();
1184 continue;
1185 }
1186 }
1187
1188 if self.flush_receiver.has_changed().unwrap_or(false) {
1189 self.handle_stalled_requests().await;
1193 }
1194
1195 for _ in 1..self.config.worker_request_batch_size {
1197 match self.receiver.try_recv() {
1199 Ok(request_with_time) => {
1200 let wait_time = request_with_time.created_at.elapsed();
1202 self.request_wait_time.observe(wait_time.as_secs_f64());
1203
1204 match request_with_time.request {
1205 WorkerRequest::Write(sender_req) => write_req_buffer.push(sender_req),
1206 WorkerRequest::BatchWrite(sender_reqs) => {
1207 write_req_buffer.extend(sender_reqs)
1208 }
1209 WorkerRequest::Ddl(sender_req) => ddl_req_buffer.push(sender_req),
1210 WorkerRequest::BulkInserts(bulk_insert) => {
1211 bulk_insert_req_num += 1;
1212 self.buffer_bulk_insert_request(bulk_insert, &mut bulk_req_buffer)
1213 .await
1214 }
1215 req => general_req_buffer.push(req),
1216 }
1217 }
1218 Err(_) => break,
1220 }
1221 }
1222
1223 self.listener.on_recv_requests(
1224 write_req_buffer.len()
1225 + ddl_req_buffer.len()
1226 + general_req_buffer.len()
1227 + bulk_insert_req_num,
1228 );
1229
1230 self.handle_requests(
1231 &mut write_req_buffer,
1232 &mut ddl_req_buffer,
1233 &mut general_req_buffer,
1234 &mut bulk_req_buffer,
1235 )
1236 .await;
1237
1238 self.handle_periodical_tasks();
1239 }
1240
1241 self.clean().await;
1242
1243 info!("Exit region worker thread {}", self.id);
1244 }
1245
1246 async fn buffer_bulk_insert_request(
1247 &mut self,
1248 bulk_insert: BulkInsertRequest,
1249 bulk_requests: &mut Vec<SenderBulkRequest>,
1250 ) {
1251 let BulkInsertRequest {
1252 metadata,
1253 request,
1254 sender,
1255 } = bulk_insert;
1256
1257 if let Some(region_metadata) = metadata {
1258 self.handle_bulk_insert_batch(region_metadata, request, bulk_requests, sender)
1259 .await;
1260 } else {
1261 error!("Cannot find region metadata for {}", request.region_id);
1262 sender.send(
1263 error::RegionNotFoundSnafu {
1264 region_id: request.region_id,
1265 }
1266 .fail(),
1267 );
1268 }
1269 }
1270
1271 async fn handle_requests(
1275 &mut self,
1276 write_requests: &mut Vec<SenderWriteRequest>,
1277 ddl_requests: &mut Vec<SenderDdlRequest>,
1278 general_requests: &mut Vec<WorkerRequest>,
1279 bulk_requests: &mut Vec<SenderBulkRequest>,
1280 ) {
1281 for worker_req in general_requests.drain(..) {
1282 match worker_req {
1283 WorkerRequest::Write(_) | WorkerRequest::BatchWrite(_) | WorkerRequest::Ddl(_) => {
1284 continue;
1286 }
1287 WorkerRequest::BulkInserts(_) => unreachable!("bulk inserts are buffered"),
1288 WorkerRequest::Background { region_id, notify } => {
1289 if matches!(
1290 ¬ify,
1291 BackgroundNotify::RegionEdit(edit_result)
1292 if edit_result.update_region_state
1293 ) {
1294 self.handle_buffered_region_write_requests(
1303 ®ion_id,
1304 write_requests,
1305 bulk_requests,
1306 )
1307 .await;
1308 }
1309 self.handle_background_notify(region_id, notify).await;
1311 }
1312 WorkerRequest::SetRegionRoleStateGracefully {
1313 region_id,
1314 region_role_state,
1315 sender,
1316 } => {
1317 self.set_role_state_gracefully(region_id, region_role_state, sender)
1318 .await;
1319 }
1320 WorkerRequest::EditRegion(request) => {
1321 self.handle_region_edit(request);
1322 }
1323 WorkerRequest::Stop => {
1324 debug_assert!(!self.running.load(Ordering::Relaxed));
1325 }
1326 WorkerRequest::SyncRegion(req) => {
1327 self.handle_region_sync(req).await;
1328 }
1329 WorkerRequest::RemapManifests(req) => {
1330 self.handle_remap_manifests_request(req);
1331 }
1332 WorkerRequest::CopyRegionFrom(req) => {
1333 self.handle_copy_region_from_request(req);
1334 }
1335 }
1336 }
1337
1338 self.handle_write_requests(write_requests, bulk_requests, true)
1341 .await;
1342
1343 self.handle_ddl_requests(ddl_requests).await;
1344 }
1345
1346 async fn handle_ddl_requests(&mut self, ddl_requests: &mut Vec<SenderDdlRequest>) {
1348 if ddl_requests.is_empty() {
1349 return;
1350 }
1351
1352 for ddl in ddl_requests.drain(..) {
1353 let res = match ddl.request {
1354 DdlRequest::Create(req) => self.handle_create_request(ddl.region_id, req).await,
1355 DdlRequest::Drop(req) => {
1356 self.handle_drop_request(ddl.region_id, req, ddl.sender)
1357 .await;
1358 continue;
1359 }
1360 DdlRequest::Open((req, wal_entry_receiver)) => {
1361 self.handle_open_request(ddl.region_id, req, wal_entry_receiver, ddl.sender)
1362 .await;
1363 continue;
1364 }
1365 DdlRequest::OfflineCleanup(req) => {
1366 self.handle_offline_cleanup_request(ddl.region_id, req)
1367 .await
1368 }
1369 DdlRequest::Close(req) => {
1370 self.handle_close_request(ddl.region_id, req, ddl.sender)
1371 .await;
1372 continue;
1373 }
1374 DdlRequest::Alter(req) => {
1375 self.handle_alter_request(ddl.region_id, req, ddl.sender)
1376 .await;
1377 continue;
1378 }
1379 DdlRequest::Flush(req) => {
1380 self.handle_flush_request(ddl.region_id, req, ddl.sender);
1381 continue;
1382 }
1383 DdlRequest::Compact(req) => {
1384 self.handle_compaction_request(ddl.region_id, req, ddl.sender)
1385 .await;
1386 continue;
1387 }
1388 DdlRequest::BuildIndex(req) => {
1389 self.handle_build_index_request(ddl.region_id, req, ddl.sender)
1390 .await;
1391 continue;
1392 }
1393 DdlRequest::Truncate(req) => {
1394 self.handle_truncate_request(ddl.region_id, req, ddl.sender)
1395 .await;
1396 continue;
1397 }
1398 DdlRequest::Catchup((req, wal_entry_receiver)) => {
1399 self.handle_catchup_request(ddl.region_id, req, wal_entry_receiver, ddl.sender)
1400 .await;
1401 continue;
1402 }
1403 DdlRequest::EnterStaging(req) => {
1404 self.handle_enter_staging_request(
1405 ddl.region_id,
1406 req.partition_directive,
1407 ddl.sender,
1408 )
1409 .await;
1410 continue;
1411 }
1412 DdlRequest::ApplyStagingManifest(req) => {
1413 self.handle_apply_staging_manifest_request(ddl.region_id, req, ddl.sender)
1414 .await;
1415 continue;
1416 }
1417 };
1418
1419 ddl.sender.send(res);
1420 }
1421 }
1422
1423 fn handle_periodical_tasks(&mut self) {
1425 let interval = CHECK_REGION_INTERVAL.as_millis() as i64;
1426 if self
1427 .time_provider
1428 .elapsed_since(self.last_periodical_check_millis)
1429 < interval
1430 {
1431 return;
1432 }
1433
1434 self.last_periodical_check_millis = self.time_provider.current_time_millis();
1435
1436 if let Err(e) = self.flush_periodically() {
1437 error!(e; "Failed to flush regions periodically");
1438 }
1439 }
1440
1441 async fn handle_background_notify(&mut self, region_id: RegionId, notify: BackgroundNotify) {
1443 match notify {
1444 BackgroundNotify::RegionOpened { region, registered } => {
1445 self.register_region(region);
1446 self.region_count.inc();
1447
1448 let senders = self.opening_regions.remove_sender(region_id);
1451 if let Some(state) = &self.series_index_task_state {
1452 state.wake();
1453 }
1454 let _ = registered.send(());
1455 for sender in senders {
1456 sender.send(Ok(0));
1457 }
1458 }
1459 BackgroundNotify::CompactionUnit(req) => {
1460 self.handle_compaction_unit(region_id, req).await;
1461 }
1462 BackgroundNotify::CompactionDdlComplete {
1463 generation,
1464 request_index,
1465 } => {
1466 self.compaction_scheduler
1467 .on_ddl_complete(region_id, generation, request_index);
1468 }
1469 BackgroundNotify::CompactionPickFinished(req) => {
1470 self.handle_compaction_pick_finished(region_id, req).await
1471 }
1472 BackgroundNotify::FlushFinished(req) => {
1473 self.handle_flush_finished(region_id, req).await
1474 }
1475 BackgroundNotify::FlushFailed(req) => self.handle_flush_failed(region_id, req).await,
1476 BackgroundNotify::IndexBuildFinished(req) => {
1477 self.handle_index_build_finished(region_id, req).await
1478 }
1479 BackgroundNotify::IndexBuildStopped(req) => {
1480 self.handle_index_build_stopped(region_id, req).await
1481 }
1482 BackgroundNotify::IndexBuildRetry(req) => {
1483 self.handle_rebuild_index(req, OptionOutputTx::new(None))
1484 .await
1485 }
1486 BackgroundNotify::CompactionFinished(req) => {
1487 self.handle_compaction_finished(region_id, req).await
1488 }
1489 BackgroundNotify::CompactionFailed(req) => self.handle_compaction_failure(req).await,
1490 BackgroundNotify::Truncate(req) => self.handle_truncate_result(req).await,
1491 BackgroundNotify::DiscardUnflushed(req) => {
1492 self.handle_discard_unflushed_result(req).await
1493 }
1494 BackgroundNotify::RegionChange(req) => {
1495 self.handle_manifest_region_change_result(req).await
1496 }
1497 BackgroundNotify::EnterStaging(req) => self.handle_enter_staging_result(req).await,
1498 BackgroundNotify::RegionEdit(req) => self.handle_region_edit_result(req).await,
1499 BackgroundNotify::CopyRegionFromFinished(req) => {
1500 self.handle_copy_region_from_finished(req)
1501 }
1502 }
1503 }
1504
1505 async fn set_role_state_gracefully(
1507 &mut self,
1508 region_id: RegionId,
1509 region_role_state: SettableRegionRoleState,
1510 sender: oneshot::Sender<SetRegionRoleStateResponse>,
1511 ) {
1512 if let Some(region) = self.regions.get_region(region_id) {
1513 if matches!(
1514 region_role_state,
1515 SettableRegionRoleState::Follower | SettableRegionRoleState::DowngradingLeader
1516 ) {
1517 region.set_role(RegionRole::DowngradingLeader);
1521 }
1522 #[cfg(feature = "hdfs-object-store")]
1523 let sweep_older_than = self.config.gc.unknown_file_lingering_time;
1524 #[cfg(feature = "hdfs-object-store")]
1525 let sweep_limiter = self.sweep_limiter.clone();
1526 common_runtime::spawn_global(async move {
1528 match region.set_role_state_gracefully(region_role_state).await {
1529 Ok(()) => {
1530 let last_entry_id = region.version_control.current().last_entry_id;
1531 let _ = sender.send(SetRegionRoleStateResponse::success(
1532 SetRegionRoleStateSuccess::mito(last_entry_id),
1533 ));
1534 #[cfg(feature = "hdfs-object-store")]
1537 if region_role_state == SettableRegionRoleState::Leader {
1538 temp_file_sweeper::spawn_region_sweep(
1539 region.access_layer(),
1540 region_id,
1541 sweep_older_than,
1542 sweep_limiter,
1543 );
1544 }
1545 }
1546 Err(e) => {
1547 error!(e; "Failed to set region {} role state to {:?}", region_id, region_role_state);
1548 let _ = sender.send(SetRegionRoleStateResponse::invalid_transition(
1549 BoxedError::new(e),
1550 ));
1551 }
1552 }
1553 });
1554 } else {
1555 let _ = sender.send(SetRegionRoleStateResponse::NotFound);
1556 }
1557 }
1558}
1559
1560impl<S> RegionWorkerLoop<S> {
1561 async fn clean(&self) {
1563 let regions = self.regions.list_regions();
1565 for region in regions {
1566 region.stop().await;
1567 }
1568
1569 self.regions.clear();
1570 }
1571
1572 fn notify_group(&mut self) {
1575 let _ = self.flush_sender.send(());
1577 self.flush_receiver.borrow_and_update();
1579 }
1580}
1581
1582#[derive(Default, Clone)]
1584pub(crate) struct WorkerListener {
1585 #[cfg(any(test, feature = "test"))]
1586 listener: Option<crate::engine::listener::EventListenerRef>,
1587}
1588
1589impl WorkerListener {
1590 #[cfg(any(test, feature = "test"))]
1591 pub(crate) fn new(
1592 listener: Option<crate::engine::listener::EventListenerRef>,
1593 ) -> WorkerListener {
1594 WorkerListener { listener }
1595 }
1596
1597 pub(crate) async fn on_region_open_registered(&self, _region_id: RegionId) {
1599 #[cfg(any(test, feature = "test"))]
1600 if let Some(listener) = &self.listener {
1601 listener.on_region_open_registered(_region_id).await;
1602 }
1603 }
1604
1605 pub(crate) fn on_flush_success(&self, region_id: RegionId) {
1607 #[cfg(any(test, feature = "test"))]
1608 if let Some(listener) = &self.listener {
1609 listener.on_flush_success(region_id);
1610 }
1611 let _ = region_id;
1613 }
1614
1615 pub(crate) fn on_write_stall(&self) {
1617 #[cfg(any(test, feature = "test"))]
1618 if let Some(listener) = &self.listener {
1619 listener.on_write_stall();
1620 }
1621 }
1622
1623 pub(crate) async fn on_write_wal_end(&self) {
1625 #[cfg(any(test, feature = "test"))]
1626 if let Some(listener) = &self.listener {
1627 listener.on_write_wal_end().await;
1628 }
1629 }
1630
1631 pub(crate) async fn on_flush_begin(&self, region_id: RegionId) {
1632 #[cfg(any(test, feature = "test"))]
1633 if let Some(listener) = &self.listener {
1634 listener.on_flush_begin(region_id).await;
1635 }
1636 let _ = region_id;
1638 }
1639
1640 pub(crate) async fn on_flush_commit_begin(&self, _region_id: RegionId) {
1641 #[cfg(any(test, feature = "test"))]
1642 if let Some(listener) = &self.listener {
1643 listener.on_flush_commit_begin(_region_id).await;
1644 }
1645 }
1646
1647 pub(crate) fn on_flush_cancel_requested(&self, _region_id: RegionId) {
1648 #[cfg(any(test, feature = "test"))]
1649 if let Some(listener) = &self.listener {
1650 listener.on_flush_cancel_requested(_region_id);
1651 }
1652 }
1653
1654 pub(crate) fn on_later_drop_begin(&self, region_id: RegionId) -> Option<Duration> {
1655 #[cfg(any(test, feature = "test"))]
1656 if let Some(listener) = &self.listener {
1657 return listener.on_later_drop_begin(region_id);
1658 }
1659 let _ = region_id;
1661 None
1662 }
1663
1664 pub(crate) fn on_later_drop_end(&self, region_id: RegionId, removed: bool) {
1666 #[cfg(any(test, feature = "test"))]
1667 if let Some(listener) = &self.listener {
1668 listener.on_later_drop_end(region_id, removed);
1669 }
1670 let _ = region_id;
1672 let _ = removed;
1673 }
1674
1675 pub(crate) async fn on_merge_ssts_finished(&self, region_id: RegionId) {
1676 #[cfg(any(test, feature = "test"))]
1677 if let Some(listener) = &self.listener {
1678 listener.on_merge_ssts_finished(region_id).await;
1679 }
1680 let _ = region_id;
1682 }
1683
1684 pub(crate) fn on_recv_requests(&self, request_num: usize) {
1685 #[cfg(any(test, feature = "test"))]
1686 if let Some(listener) = &self.listener {
1687 listener.on_recv_requests(request_num);
1688 }
1689 let _ = request_num;
1691 }
1692
1693 pub(crate) fn on_file_cache_filled(&self, _file_id: FileId) {
1694 #[cfg(any(test, feature = "test"))]
1695 if let Some(listener) = &self.listener {
1696 listener.on_file_cache_filled(_file_id);
1697 }
1698 }
1699
1700 pub(crate) fn on_compaction_scheduled(&self, _region_id: RegionId) {
1701 #[cfg(any(test, feature = "test"))]
1702 if let Some(listener) = &self.listener {
1703 listener.on_compaction_scheduled(_region_id);
1704 }
1705 }
1706
1707 pub(crate) async fn on_compaction_pick_begin(&self, _region_id: RegionId) {
1708 #[cfg(any(test, feature = "test"))]
1709 if let Some(listener) = &self.listener {
1710 listener.on_compaction_pick_begin(_region_id).await;
1711 }
1712 }
1713
1714 pub(crate) async fn on_compaction_commit_begin(&self, _region_id: RegionId) {
1715 #[cfg(any(test, feature = "test"))]
1716 if let Some(listener) = &self.listener {
1717 listener.on_compaction_commit_begin(_region_id).await;
1718 }
1719 }
1720
1721 pub(crate) async fn on_compaction_unit_merge_begin(&self, _region_id: RegionId, _plan_id: u64) {
1723 #[cfg(any(test, feature = "test"))]
1724 if let Some(listener) = &self.listener {
1725 listener
1726 .on_compaction_unit_merge_begin(_region_id, _plan_id)
1727 .await;
1728 }
1729 }
1730
1731 pub(crate) async fn on_compaction_unit_committed(&self, _region_id: RegionId, _plan_id: u64) {
1733 #[cfg(any(test, feature = "test"))]
1734 if let Some(listener) = &self.listener {
1735 listener
1736 .on_compaction_unit_committed(_region_id, _plan_id)
1737 .await;
1738 }
1739 }
1740
1741 pub(crate) fn on_compaction_unit_applied(&self, _region_id: RegionId, _plan_id: u64) {
1743 #[cfg(any(test, feature = "test"))]
1744 if let Some(listener) = &self.listener {
1745 listener.on_compaction_unit_applied(_region_id, _plan_id);
1746 }
1747 }
1748
1749 pub(crate) async fn on_compaction_result_notified(&self, _region_id: RegionId) {
1750 #[cfg(any(test, feature = "test"))]
1751 if let Some(listener) = &self.listener {
1752 listener.on_compaction_result_notified(_region_id).await;
1753 }
1754 }
1755
1756 pub(crate) fn on_compaction_cancel_requested(&self, _region_id: RegionId) {
1757 #[cfg(any(test, feature = "test"))]
1758 if let Some(listener) = &self.listener {
1759 listener.on_compaction_cancel_requested(_region_id);
1760 }
1761 }
1762
1763 pub(crate) async fn on_notify_region_change_result_begin(&self, _region_id: RegionId) {
1764 #[cfg(any(test, feature = "test"))]
1765 if let Some(listener) = &self.listener {
1766 listener
1767 .on_notify_region_change_result_begin(_region_id)
1768 .await;
1769 }
1770 }
1771
1772 pub(crate) async fn on_enter_staging_result_begin(&self, _region_id: RegionId) {
1773 #[cfg(any(test, feature = "test"))]
1774 if let Some(listener) = &self.listener {
1775 listener.on_enter_staging_result_begin(_region_id).await;
1776 }
1777 }
1778
1779 pub(crate) async fn on_index_build_finish(&self, _region_file_id: RegionFileId) {
1780 #[cfg(any(test, feature = "test"))]
1781 if let Some(listener) = &self.listener {
1782 listener.on_index_build_finish(_region_file_id).await;
1783 }
1784 }
1785
1786 pub(crate) async fn on_index_build_begin(&self, _region_file_id: RegionFileId) {
1787 #[cfg(any(test, feature = "test"))]
1788 if let Some(listener) = &self.listener {
1789 listener.on_index_build_begin(_region_file_id).await;
1790 }
1791 }
1792
1793 pub(crate) async fn on_index_build_abort(&self, _region_file_id: RegionFileId) {
1794 #[cfg(any(test, feature = "test"))]
1795 if let Some(listener) = &self.listener {
1796 listener.on_index_build_abort(_region_file_id).await;
1797 }
1798 }
1799
1800 pub(crate) async fn on_index_build_before_manifest_commit(
1801 &self,
1802 _region_file_id: RegionFileId,
1803 ) {
1804 #[cfg(any(test, feature = "test"))]
1805 if let Some(listener) = &self.listener {
1806 listener
1807 .on_index_build_before_manifest_commit(_region_file_id)
1808 .await;
1809 }
1810 }
1811
1812 pub(crate) async fn on_index_build_manifest_committed(&self, _region_file_id: RegionFileId) {
1813 #[cfg(any(test, feature = "test"))]
1814 if let Some(listener) = &self.listener {
1815 listener
1816 .on_index_build_manifest_committed(_region_file_id)
1817 .await;
1818 }
1819 }
1820}
1821
1822#[cfg(test)]
1823mod tests {
1824 use super::*;
1825 use crate::test_util::TestEnv;
1826
1827 #[test]
1828 fn test_region_id_to_index() {
1829 let num_workers = 4;
1830
1831 let region_id = RegionId::new(1, 2);
1832 let index = region_id_to_index(region_id, num_workers);
1833 assert_eq!(index, 3);
1834
1835 let region_id = RegionId::new(2, 3);
1836 let index = region_id_to_index(region_id, num_workers);
1837 assert_eq!(index, 1);
1838 }
1839
1840 #[tokio::test]
1841 async fn test_worker_group_start_stop() {
1842 for (enabled, existing_index) in [(false, false), (true, false), (false, true)] {
1843 let env = TestEnv::with_prefix("group-stop").await;
1845 let root = env.data_home().join("series_index");
1846 if existing_index {
1847 tokio::fs::create_dir_all(&root).await.unwrap();
1848 tokio::fs::write(root.join("retained"), b"index")
1849 .await
1850 .unwrap();
1851 }
1852 let group = env
1853 .create_worker_group(MitoConfig {
1854 num_workers: 4,
1855 experimental_enable_series_index: enabled,
1856 experimental_series_index_maintenance_interval: Duration::from_secs(3600),
1857 ..Default::default()
1858 })
1859 .await;
1860
1861 for worker in &group.workers {
1862 assert_eq!(worker.series_index_task_state.is_some(), enabled);
1863 assert_eq!(worker.series_index_handle.lock().await.is_some(), enabled);
1864 }
1865 assert_eq!(root.is_dir(), enabled || existing_index);
1866 if existing_index {
1867 assert_eq!(
1868 tokio::fs::read(root.join("retained")).await.unwrap(),
1869 b"index"
1870 );
1871 }
1872
1873 tokio::time::timeout(Duration::from_secs(5), group.stop())
1874 .await
1875 .expect("series-index tasks should stop without waiting for their interval")
1876 .unwrap();
1877 }
1878 }
1879}