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 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::memory_manager::{CompactionMemoryManager, new_compaction_memory_manager};
65use crate::compaction::{CompactionCoordinator, CompactionScheduler};
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::lease::NodeLease;
71use crate::memtable::MemtableBuilderProvider;
72use crate::memtable::bulk::CompactDispatcher;
73use crate::metrics::{REGION_COUNT, REQUEST_WAIT_TIME, WRITE_STALLING};
74use crate::region::opener::PartitionExprFetcherRef;
75use crate::region::{
76 CatchupRegions, CatchupRegionsRef, MitoRegionRef, OpeningRegions, OpeningRegionsRef, RegionMap,
77 RegionMapRef,
78};
79use crate::request::{
80 BackgroundNotify, BulkInsertRequest, DdlRequest, OptionOutputTx, SenderBulkRequest,
81 SenderDdlRequest, SenderWriteRequest, WorkerRequest, WorkerRequestWithTime,
82};
83use crate::schedule::scheduler::{LocalScheduler, SchedulerRef};
84use crate::series_index::{
85 IndexFilePurger, SeriesIndexTaskState, series_index_channel, spawn_series_index_tasks,
86};
87use crate::sst::file::RegionFileId;
88use crate::sst::file_ref::FileReferenceManagerRef;
89use crate::sst::index::IndexBuildScheduler;
90use crate::sst::index::intermediate::IntermediateManager;
91use crate::sst::index::puffin_manager::PuffinManagerFactory;
92#[cfg(feature = "hdfs-object-store")]
93use crate::temp_file_sweeper;
94use crate::time_provider::{StdTimeProvider, TimeProviderRef};
95use crate::wal::Wal;
96use crate::worker::handle_flush::StagingRegionsRef;
97use crate::worker::handle_manifest::RegionEditQueues;
98
99pub(crate) type WorkerId = u32;
101
102pub(crate) const DROPPING_MARKER_FILE: &str = ".dropping";
103
104pub(crate) const CHECK_REGION_INTERVAL: Duration = Duration::from_secs(60);
106pub(crate) const MAX_INITIAL_CHECK_DELAY_SECS: u64 = 60 * 3;
108
109#[cfg_attr(doc, aquamarine::aquamarine)]
110pub(crate) struct WorkerGroup {
147 workers: Vec<RegionWorker>,
149 flush_job_pool: SchedulerRef,
151 compact_job_pool: Arc<CompactionCoordinator>,
153 index_build_job_pool: SchedulerRef,
155 purge_scheduler: SchedulerRef,
157 cache_manager: CacheManagerRef,
159 file_ref_manager: FileReferenceManagerRef,
161 gc_limiter: GcLimiterRef,
163 node_lease: Arc<NodeLease>,
164 object_store_manager: ObjectStoreManagerRef,
166 puffin_manager_factory: PuffinManagerFactory,
168 intermediate_manager: IntermediateManager,
170 schema_metadata_manager: SchemaMetadataManagerRef,
172}
173
174impl WorkerGroup {
175 #[allow(clippy::too_many_arguments)]
179 pub(crate) async fn start<S: LogStore>(
180 data_home: &str,
181 config: Arc<MitoConfig>,
182 log_store: Arc<S>,
183 object_store_manager: ObjectStoreManagerRef,
184 schema_metadata_manager: SchemaMetadataManagerRef,
185 file_ref_manager: FileReferenceManagerRef,
186 partition_expr_fetcher: PartitionExprFetcherRef,
187 plugins: Plugins,
188 ) -> Result<WorkerGroup> {
189 let (flush_sender, flush_receiver) = watch::channel(());
190 let write_buffer_manager = Arc::new(
191 WriteBufferManagerImpl::new(config.global_write_buffer_size.as_bytes() as usize)
192 .with_notifier(flush_sender.clone()),
193 );
194 let puffin_manager_factory = PuffinManagerFactory::new(
195 &config.index.aux_path,
196 config.index.staging_size.as_bytes(),
197 Some(config.index.write_buffer_size.as_bytes() as _),
198 config.index.staging_ttl,
199 )
200 .await?;
201 let intermediate_manager = IntermediateManager::init_fs(&config.index.aux_path)
202 .await?
203 .with_buffer_size(Some(config.index.write_buffer_size.as_bytes() as _));
204 let index_build_job_pool =
205 Arc::new(LocalScheduler::new(config.max_background_index_builds));
206 let series_index_store = series_index_store_from_config(&config, data_home).await?;
207 let series_index_disk_usage = Arc::new(AtomicU64::new(0));
208 let flush_job_pool = Arc::new(LocalScheduler::new(config.max_background_flushes));
209 let compact_job_pool = Arc::new(CompactionCoordinator::new_local(
210 config.max_background_compactions,
211 ));
212 let memtable_compact_dispatcher =
213 Arc::new(CompactDispatcher::new(config.max_background_compactions));
214 let flush_semaphore = Arc::new(Semaphore::new(config.max_background_flushes));
215 let purge_scheduler = Arc::new(LocalScheduler::new(config.max_background_purges));
217 let upload_store_wrapper = plugins.get::<WriteCacheUploadStoreWrapperRef>();
218 let write_cache = write_cache_from_config(
219 &config,
220 puffin_manager_factory.clone(),
221 intermediate_manager.clone(),
222 upload_store_wrapper,
223 )
224 .await?;
225 let cache_manager = Arc::new(
226 CacheManager::builder()
227 .sst_meta_cache_size(config.sst_meta_cache_size.as_bytes())
228 .vector_cache_size(config.vector_cache_size.as_bytes())
229 .page_cache_size(config.page_cache_size.as_bytes())
230 .selector_result_cache_size(config.selector_result_cache_size.as_bytes())
231 .range_result_cache_size(config.range_result_cache_size.as_bytes())
232 .prefilter_result_cache_size(config.prefilter_result_cache_size.as_bytes())
233 .index_metadata_size(config.index.metadata_cache_size.as_bytes())
234 .index_content_size(config.index.content_cache_size.as_bytes())
235 .index_content_page_size(config.index.content_cache_page_size.as_bytes())
236 .index_result_cache_size(config.index.result_cache_size.as_bytes())
237 .puffin_metadata_size(config.index.metadata_cache_size.as_bytes())
238 .write_cache(write_cache)
239 .build(),
240 );
241 let time_provider = Arc::new(StdTimeProvider);
242 let total_memory = get_total_memory_bytes();
243 let total_memory = if total_memory > 0 {
244 total_memory as u64
245 } else {
246 0
247 };
248 let compaction_limit_bytes = config
249 .experimental_compaction_memory_limit
250 .resolve(total_memory);
251 let compaction_memory_manager =
252 Arc::new(new_compaction_memory_manager(compaction_limit_bytes));
253 let gc_limiter = Arc::new(GcLimiter::new(config.gc.max_concurrent_gc_job));
254 #[cfg(feature = "hdfs-object-store")]
259 let sweep_limiter = (config.gc.max_concurrent_gc_job > 0)
260 .then(|| Arc::new(Semaphore::new(config.gc.max_concurrent_gc_job)));
261 let node_lease = Arc::new(NodeLease::default());
262 let staging_regions = StagingRegionsRef::default();
263
264 let workers = (0..config.num_workers)
265 .map(|id| {
266 WorkerStarter {
267 id: id as WorkerId,
268 config: config.clone(),
269 log_store: log_store.clone(),
270 object_store_manager: object_store_manager.clone(),
271 write_buffer_manager: write_buffer_manager.clone(),
272 index_build_job_pool: index_build_job_pool.clone(),
273 series_index_store: series_index_store.clone(),
274 series_index_disk_usage: series_index_disk_usage.clone(),
275 flush_job_pool: flush_job_pool.clone(),
276 compact_job_pool: compact_job_pool.clone(),
277 purge_scheduler: purge_scheduler.clone(),
278 memtable_compact_dispatcher: memtable_compact_dispatcher.clone(),
279 listener: WorkerListener::default(),
280 cache_manager: cache_manager.clone(),
281 compaction_memory_manager: compaction_memory_manager.clone(),
282 puffin_manager_factory: puffin_manager_factory.clone(),
283 intermediate_manager: intermediate_manager.clone(),
284 time_provider: time_provider.clone(),
285 flush_sender: flush_sender.clone(),
286 flush_receiver: flush_receiver.clone(),
287 plugins: plugins.clone(),
288 schema_metadata_manager: schema_metadata_manager.clone(),
289 file_ref_manager: file_ref_manager.clone(),
290 partition_expr_fetcher: partition_expr_fetcher.clone(),
291 flush_semaphore: flush_semaphore.clone(),
292 #[cfg(feature = "hdfs-object-store")]
293 sweep_limiter: sweep_limiter.clone(),
294 node_lease: node_lease.clone(),
295 staging_regions: staging_regions.clone(),
296 }
297 .start()
298 })
299 .collect::<Result<Vec<_>>>()?;
300
301 Ok(WorkerGroup {
302 workers,
303 flush_job_pool,
304 compact_job_pool,
305 index_build_job_pool,
306 purge_scheduler,
307 cache_manager,
308 file_ref_manager,
309 gc_limiter,
310 node_lease,
311 object_store_manager,
312 puffin_manager_factory,
313 intermediate_manager,
314 schema_metadata_manager,
315 })
316 }
317
318 pub(crate) async fn stop(&self) -> Result<()> {
320 info!("Stop region worker group");
321
322 self.compact_job_pool.stop().await?;
325 self.flush_job_pool.stop(true).await?;
327 self.purge_scheduler.stop(true).await?;
329 self.index_build_job_pool.stop(true).await?;
331
332 try_join_all(self.workers.iter().map(|worker| worker.stop())).await?;
333
334 Ok(())
335 }
336
337 pub(crate) async fn submit_to_worker(
339 &self,
340 region_id: RegionId,
341 request: WorkerRequest,
342 ) -> Result<()> {
343 self.worker(region_id).submit_request(request).await
344 }
345
346 pub(crate) async fn submit_writes(&self, writes: Vec<SenderWriteRequest>) {
351 let mut writes_by_worker: HashMap<usize, Vec<SenderWriteRequest>> = HashMap::new();
352 for write in writes {
353 let index = region_id_to_index(write.request.region_id, self.workers.len());
354 writes_by_worker.entry(index).or_default().push(write);
355 }
356 join_all(
357 writes_by_worker
358 .into_iter()
359 .map(|(index, writes)| async move {
360 let worker = &self.workers[index];
361 let Err(request) = worker
362 .try_submit_request(WorkerRequest::BatchWrite(writes))
363 .await
364 else {
365 return;
366 };
367 let WorkerRequest::BatchWrite(writes) = request else {
368 unreachable!("the worker returns the request it was given");
369 };
370 for write in writes {
371 write
372 .sender
373 .send(WorkerStoppedSnafu { id: worker.id }.fail());
374 }
375 }),
376 )
377 .await;
378 }
379
380 pub(crate) fn is_region_exists(&self, region_id: RegionId) -> bool {
382 self.worker(region_id).is_region_exists(region_id)
383 }
384
385 pub(crate) fn is_region_opening(&self, region_id: RegionId) -> bool {
387 self.worker(region_id).is_region_opening(region_id)
388 }
389
390 pub(crate) fn is_region_catching_up(&self, region_id: RegionId) -> bool {
392 self.worker(region_id).is_region_catching_up(region_id)
393 }
394
395 pub(crate) fn get_region(&self, region_id: RegionId) -> Option<MitoRegionRef> {
399 self.worker(region_id).get_region(region_id)
400 }
401
402 pub(crate) fn cache_manager(&self) -> CacheManagerRef {
404 self.cache_manager.clone()
405 }
406
407 pub(crate) fn file_ref_manager(&self) -> FileReferenceManagerRef {
408 self.file_ref_manager.clone()
409 }
410
411 pub(crate) fn gc_limiter(&self) -> GcLimiterRef {
412 self.gc_limiter.clone()
413 }
414
415 pub(crate) fn node_lease(&self) -> &NodeLease {
416 &self.node_lease
417 }
418
419 pub(crate) fn worker(&self, region_id: RegionId) -> &RegionWorker {
421 let index = region_id_to_index(region_id, self.workers.len());
422
423 &self.workers[index]
424 }
425
426 pub(crate) fn all_regions(&self) -> impl Iterator<Item = MitoRegionRef> + use<'_> {
427 self.workers
428 .iter()
429 .flat_map(|worker| worker.regions.list_regions())
430 }
431
432 pub(crate) fn object_store_manager(&self) -> &ObjectStoreManagerRef {
433 &self.object_store_manager
434 }
435
436 pub(crate) fn puffin_manager_factory(&self) -> &PuffinManagerFactory {
437 &self.puffin_manager_factory
438 }
439
440 pub(crate) fn intermediate_manager(&self) -> &IntermediateManager {
441 &self.intermediate_manager
442 }
443
444 pub(crate) fn schema_metadata_manager(&self) -> &SchemaMetadataManagerRef {
445 &self.schema_metadata_manager
446 }
447}
448
449#[cfg(any(test, feature = "test"))]
451impl WorkerGroup {
452 #[allow(clippy::too_many_arguments)]
456 pub(crate) async fn start_for_test<S: LogStore>(
457 data_home: &str,
458 config: Arc<MitoConfig>,
459 log_store: Arc<S>,
460 object_store_manager: ObjectStoreManagerRef,
461 write_buffer_manager: Option<WriteBufferManagerRef>,
462 listener: Option<crate::engine::listener::EventListenerRef>,
463 schema_metadata_manager: SchemaMetadataManagerRef,
464 file_ref_manager: FileReferenceManagerRef,
465 time_provider: TimeProviderRef,
466 partition_expr_fetcher: PartitionExprFetcherRef,
467 ) -> Result<WorkerGroup> {
468 let (flush_sender, flush_receiver) = watch::channel(());
469 let write_buffer_manager = write_buffer_manager.unwrap_or_else(|| {
470 Arc::new(
471 WriteBufferManagerImpl::new(config.global_write_buffer_size.as_bytes() as usize)
472 .with_notifier(flush_sender.clone()),
473 )
474 });
475 let index_build_job_pool =
476 Arc::new(LocalScheduler::new(config.max_background_index_builds));
477 let series_index_store = series_index_store_from_config(&config, data_home).await?;
478 let series_index_disk_usage = Arc::new(AtomicU64::new(0));
479 let flush_job_pool = Arc::new(LocalScheduler::new(config.max_background_flushes));
480 let memtable_compact_dispatcher =
481 Arc::new(CompactDispatcher::new(config.max_background_compactions));
482 let compact_job_pool = Arc::new(CompactionCoordinator::new_local(
483 config.max_background_compactions,
484 ));
485 let flush_semaphore = Arc::new(Semaphore::new(config.max_background_flushes));
486 let purge_scheduler = Arc::new(LocalScheduler::new(config.max_background_flushes));
487 let puffin_manager_factory = PuffinManagerFactory::new(
488 &config.index.aux_path,
489 config.index.staging_size.as_bytes(),
490 Some(config.index.write_buffer_size.as_bytes() as _),
491 config.index.staging_ttl,
492 )
493 .await?;
494 let intermediate_manager = IntermediateManager::init_fs(&config.index.aux_path)
495 .await?
496 .with_buffer_size(Some(config.index.write_buffer_size.as_bytes() as _));
497 let write_cache = write_cache_from_config(
498 &config,
499 puffin_manager_factory.clone(),
500 intermediate_manager.clone(),
501 None,
502 )
503 .await?;
504 let cache_manager = Arc::new(
505 CacheManager::builder()
506 .sst_meta_cache_size(config.sst_meta_cache_size.as_bytes())
507 .vector_cache_size(config.vector_cache_size.as_bytes())
508 .page_cache_size(config.page_cache_size.as_bytes())
509 .selector_result_cache_size(config.selector_result_cache_size.as_bytes())
510 .range_result_cache_size(config.range_result_cache_size.as_bytes())
511 .prefilter_result_cache_size(config.prefilter_result_cache_size.as_bytes())
512 .write_cache(write_cache)
513 .build(),
514 );
515 let total_memory = get_total_memory_bytes();
516 let total_memory = if total_memory > 0 {
517 total_memory as u64
518 } else {
519 0
520 };
521 let compaction_limit_bytes = config
522 .experimental_compaction_memory_limit
523 .resolve(total_memory);
524 let compaction_memory_manager =
525 Arc::new(new_compaction_memory_manager(compaction_limit_bytes));
526 let gc_limiter = Arc::new(GcLimiter::new(config.gc.max_concurrent_gc_job));
527 #[cfg(feature = "hdfs-object-store")]
532 let sweep_limiter = (config.gc.max_concurrent_gc_job > 0)
533 .then(|| Arc::new(Semaphore::new(config.gc.max_concurrent_gc_job)));
534 let node_lease = Arc::new(NodeLease::default());
535 let staging_regions = StagingRegionsRef::default();
536 let workers = (0..config.num_workers)
537 .map(|id| {
538 WorkerStarter {
539 id: id as WorkerId,
540 config: config.clone(),
541 log_store: log_store.clone(),
542 object_store_manager: object_store_manager.clone(),
543 write_buffer_manager: write_buffer_manager.clone(),
544 index_build_job_pool: index_build_job_pool.clone(),
545 series_index_store: series_index_store.clone(),
546 series_index_disk_usage: series_index_disk_usage.clone(),
547 flush_job_pool: flush_job_pool.clone(),
548 compact_job_pool: compact_job_pool.clone(),
549 purge_scheduler: purge_scheduler.clone(),
550 memtable_compact_dispatcher: memtable_compact_dispatcher.clone(),
551 listener: WorkerListener::new(listener.clone()),
552 cache_manager: cache_manager.clone(),
553 compaction_memory_manager: compaction_memory_manager.clone(),
554 puffin_manager_factory: puffin_manager_factory.clone(),
555 intermediate_manager: intermediate_manager.clone(),
556 time_provider: time_provider.clone(),
557 flush_sender: flush_sender.clone(),
558 flush_receiver: flush_receiver.clone(),
559 plugins: Plugins::new(),
560 schema_metadata_manager: schema_metadata_manager.clone(),
561 file_ref_manager: file_ref_manager.clone(),
562 partition_expr_fetcher: partition_expr_fetcher.clone(),
563 flush_semaphore: flush_semaphore.clone(),
564 #[cfg(feature = "hdfs-object-store")]
565 sweep_limiter: sweep_limiter.clone(),
566 node_lease: node_lease.clone(),
567 staging_regions: staging_regions.clone(),
568 }
569 .start()
570 })
571 .collect::<Result<Vec<_>>>()?;
572
573 Ok(WorkerGroup {
574 workers,
575 flush_job_pool,
576 compact_job_pool,
577 index_build_job_pool,
578 purge_scheduler,
579 cache_manager,
580 file_ref_manager,
581 gc_limiter,
582 node_lease,
583 object_store_manager,
584 puffin_manager_factory,
585 intermediate_manager,
586 schema_metadata_manager,
587 })
588 }
589
590 pub(crate) fn purge_scheduler(&self) -> &SchedulerRef {
592 &self.purge_scheduler
593 }
594}
595
596fn region_id_to_index(id: RegionId, num_workers: usize) -> usize {
597 ((id.table_id() as usize % num_workers) + (id.region_number() as usize % num_workers))
598 % num_workers
599}
600
601async fn series_index_store_from_config(
603 config: &MitoConfig,
604 data_home: &str,
605) -> Result<Option<ObjectStore>> {
606 if !config.experimental_enable_series_index {
607 return Ok(None);
608 }
609
610 let root = Path::new(data_home).join("series_index");
611 new_fs_cache_store(&root.to_string_lossy()).await.map(Some)
612}
613
614pub async fn write_cache_from_config(
615 config: &MitoConfig,
616 puffin_manager_factory: PuffinManagerFactory,
617 intermediate_manager: IntermediateManager,
618 upload_store_wrapper: Option<WriteCacheUploadStoreWrapperRef>,
619) -> Result<Option<WriteCacheRef>> {
620 if !config.enable_write_cache {
621 return Ok(None);
622 }
623
624 tokio::fs::create_dir_all(Path::new(&config.write_cache_path))
625 .await
626 .context(CreateDirSnafu {
627 dir: &config.write_cache_path,
628 })?;
629
630 let cache = WriteCache::new_fs(
631 &config.write_cache_path,
632 config.write_cache_size,
633 config.write_cache_ttl,
634 Some(config.index_cache_percent),
635 config.enable_refill_cache_on_read,
636 puffin_manager_factory,
637 intermediate_manager,
638 config.manifest_cache_size,
639 )
640 .await?
641 .with_upload_store_wrapper(upload_store_wrapper);
642 Ok(Some(Arc::new(cache)))
643}
644
645pub(crate) fn worker_init_check_delay() -> Duration {
647 let init_check_delay = rng().random_range(0..MAX_INITIAL_CHECK_DELAY_SECS);
648 Duration::from_secs(init_check_delay)
649}
650
651struct WorkerStarter<S> {
653 id: WorkerId,
654 config: Arc<MitoConfig>,
655 log_store: Arc<S>,
656 object_store_manager: ObjectStoreManagerRef,
657 write_buffer_manager: WriteBufferManagerRef,
658 compact_job_pool: Arc<CompactionCoordinator>,
659 index_build_job_pool: SchedulerRef,
660 series_index_store: Option<ObjectStore>,
661 series_index_disk_usage: Arc<AtomicU64>,
662 flush_job_pool: SchedulerRef,
663 purge_scheduler: SchedulerRef,
664 memtable_compact_dispatcher: Arc<CompactDispatcher>,
665 listener: WorkerListener,
666 cache_manager: CacheManagerRef,
667 compaction_memory_manager: Arc<CompactionMemoryManager>,
668 puffin_manager_factory: PuffinManagerFactory,
669 intermediate_manager: IntermediateManager,
670 time_provider: TimeProviderRef,
671 flush_sender: watch::Sender<()>,
673 flush_receiver: watch::Receiver<()>,
675 plugins: Plugins,
676 schema_metadata_manager: SchemaMetadataManagerRef,
677 file_ref_manager: FileReferenceManagerRef,
678 partition_expr_fetcher: PartitionExprFetcherRef,
679 flush_semaphore: Arc<Semaphore>,
680 #[cfg(feature = "hdfs-object-store")]
681 sweep_limiter: Option<Arc<Semaphore>>,
684 node_lease: Arc<NodeLease>,
685 staging_regions: StagingRegionsRef,
686}
687
688impl<S: LogStore> WorkerStarter<S> {
689 fn start(self) -> Result<RegionWorker> {
691 let regions = Arc::new(RegionMap::default());
692 let opening_regions = Arc::new(OpeningRegions::default());
693 let catchup_regions = Arc::new(CatchupRegions::default());
694 let (sender, receiver) = mpsc::channel(self.config.worker_channel_size);
695
696 let running = Arc::new(AtomicBool::new(true));
697 let (series_index_task_state, series_index_receiver) = if self.series_index_store.is_some()
698 {
699 let (state, receiver) =
700 SeriesIndexTaskState::new(self.id, self.config.worker_channel_size);
701 (Some(Arc::new(state)), Some(receiver))
702 } else {
703 (None, None)
704 };
705 let mut series_index_purger = None;
706 let series_index_handle = self
707 .series_index_store
708 .clone()
709 .zip(series_index_task_state.clone())
710 .zip(series_index_receiver)
711 .map(|((store, state), receiver)| {
712 let (purger, purge_receiver) = series_index_channel(store.clone());
713 series_index_purger = Some(purger.clone());
714 spawn_series_index_tasks(
715 self.id,
716 store,
717 regions.clone(),
718 state,
719 receiver,
720 self.config.experimental_series_index_bucket_width,
721 purger,
722 purge_receiver,
723 self.series_index_disk_usage.clone(),
724 self.config.experimental_series_index_max_size.as_bytes(),
725 self.config.experimental_series_index_maintenance_interval,
726 self.time_provider.clone(),
727 self.config.experimental_enable_range_index,
728 )
729 });
730 let now = self.time_provider.current_time_millis();
731 let id_string = self.id.to_string();
732 let mut worker_thread = RegionWorkerLoop {
733 id: self.id,
734 config: self.config.clone(),
735 regions: regions.clone(),
736 catchup_regions: catchup_regions.clone(),
737 dropping_regions: Arc::new(RegionMap::default()),
738 opening_regions: opening_regions.clone(),
739 sender: sender.clone(),
740 receiver,
741 wal: Wal::new(self.log_store),
742 object_store_manager: self.object_store_manager.clone(),
743 running: running.clone(),
744 memtable_builder_provider: MemtableBuilderProvider::new(
745 Some(self.write_buffer_manager.clone()),
746 self.config.clone(),
747 self.memtable_compact_dispatcher.clone(),
748 ),
749 purge_scheduler: self.purge_scheduler.clone(),
750 write_buffer_manager: self.write_buffer_manager,
751 index_build_scheduler: IndexBuildScheduler::new(
752 self.index_build_job_pool,
753 self.config.max_background_index_builds,
754 ),
755 series_index_task_state: series_index_task_state.clone(),
756 series_index_store: self.series_index_store,
757 series_index_purger,
758 flush_scheduler: FlushScheduler::new(self.flush_job_pool),
759 compaction_scheduler: CompactionScheduler::new(
760 self.compact_job_pool,
761 sender.clone(),
762 self.cache_manager.clone(),
763 self.config.clone(),
764 self.listener.clone(),
765 self.plugins.clone(),
766 self.compaction_memory_manager.clone(),
767 self.config.experimental_compaction_on_exhausted,
768 ),
769 stalled_requests: StalledRequests::default(),
770 listener: self.listener,
771 cache_manager: self.cache_manager,
772 puffin_manager_factory: self.puffin_manager_factory,
773 intermediate_manager: self.intermediate_manager,
774 time_provider: self.time_provider,
775 last_periodical_check_millis: now,
776 flush_sender: self.flush_sender,
777 flush_receiver: self.flush_receiver,
778 stalling_count: WRITE_STALLING.with_label_values(&[&id_string]),
779 region_count: REGION_COUNT.with_label_values(&[&id_string]),
780 request_wait_time: REQUEST_WAIT_TIME.with_label_values(&[&id_string]),
781 region_edit_queues: RegionEditQueues::default(),
782 pending_alters: Default::default(),
783 schema_metadata_manager: self.schema_metadata_manager,
784 file_ref_manager: self.file_ref_manager.clone(),
785 partition_expr_fetcher: self.partition_expr_fetcher,
786 flush_semaphore: self.flush_semaphore,
787 plugins: self.plugins,
788 #[cfg(feature = "hdfs-object-store")]
789 sweep_limiter: self.sweep_limiter,
790 node_lease: self.node_lease,
791 staging_regions: self.staging_regions,
792 };
793 let handle = common_runtime::spawn_global(async move {
794 worker_thread.run().await;
795 });
796
797 Ok(RegionWorker {
798 id: self.id,
799 regions,
800 opening_regions,
801 catchup_regions,
802 sender,
803 handle: Mutex::new(Some(handle)),
804 series_index_handle: Mutex::new(series_index_handle),
805 series_index_task_state,
806 running,
807 #[cfg(test)]
808 memtable_compact_dispatcher: self.memtable_compact_dispatcher.clone(),
809 })
810 }
811}
812
813pub(crate) struct RegionWorker {
815 id: WorkerId,
817 regions: RegionMapRef,
819 opening_regions: OpeningRegionsRef,
821 catchup_regions: CatchupRegionsRef,
823 sender: Sender<WorkerRequestWithTime>,
825 handle: Mutex<Option<JoinHandle<()>>>,
827 series_index_handle: Mutex<Option<JoinHandle<()>>>,
829 series_index_task_state: Option<Arc<SeriesIndexTaskState>>,
831 running: Arc<AtomicBool>,
833 #[cfg(test)]
834 memtable_compact_dispatcher: Arc<CompactDispatcher>,
837}
838
839impl RegionWorker {
840 async fn submit_request(&self, request: WorkerRequest) -> Result<()> {
842 self.try_submit_request(request)
843 .await
844 .map_err(|_| WorkerStoppedSnafu { id: self.id }.build())
845 }
846
847 async fn try_submit_request(
849 &self,
850 request: WorkerRequest,
851 ) -> std::result::Result<(), WorkerRequest> {
852 if !self.is_running() {
853 return Err(request);
854 }
855 let request_with_time = WorkerRequestWithTime::new(request);
856 if let Err(mpsc::error::SendError(request_with_time)) =
857 self.sender.send(request_with_time).await
858 {
859 warn!(
860 "Worker {} is already exited but the running flag is still true",
861 self.id
862 );
863 self.set_running(false);
865 if let Some(state) = &self.series_index_task_state {
866 state.stop();
867 }
868 return Err(request_with_time.request);
869 }
870
871 Ok(())
872 }
873
874 async fn stop(&self) -> Result<()> {
878 let handle = self.handle.lock().await.take();
879 self.set_running(false);
880 if let Some(state) = &self.series_index_task_state {
881 state.stop();
882 }
883
884 let mut worker_result = Ok(());
885 if let Some(handle) = handle {
886 info!("Stop region worker {}", self.id);
887
888 if self
889 .sender
890 .send(WorkerRequestWithTime::new(WorkerRequest::Stop))
891 .await
892 .is_err()
893 {
894 warn!("Worker {} is already exited before stop", self.id);
895 }
896
897 worker_result = handle.await.context(JoinSnafu);
898 }
899 let series_index_result = if let Some(handle) = self.series_index_handle.lock().await.take()
900 {
901 handle.await.context(JoinSnafu)
902 } else {
903 Ok(())
904 };
905 worker_result?;
906 series_index_result?;
907
908 Ok(())
909 }
910
911 fn is_running(&self) -> bool {
913 self.running.load(Ordering::Relaxed)
914 }
915
916 fn set_running(&self, value: bool) {
918 self.running.store(value, Ordering::Relaxed)
919 }
920
921 fn is_region_exists(&self, region_id: RegionId) -> bool {
923 self.regions.is_region_exists(region_id)
924 }
925
926 fn is_region_opening(&self, region_id: RegionId) -> bool {
928 self.opening_regions.is_region_exists(region_id)
929 }
930
931 fn is_region_catching_up(&self, region_id: RegionId) -> bool {
933 self.catchup_regions.is_region_exists(region_id)
934 }
935
936 fn get_region(&self, region_id: RegionId) -> Option<MitoRegionRef> {
938 self.regions.get_region(region_id)
939 }
940
941 #[cfg(test)]
942 pub(crate) fn opening_regions(&self) -> &OpeningRegionsRef {
944 &self.opening_regions
945 }
946
947 #[cfg(test)]
948 pub(crate) fn catchup_regions(&self) -> &CatchupRegionsRef {
950 &self.catchup_regions
951 }
952}
953
954impl Drop for RegionWorker {
955 fn drop(&mut self) {
956 if let Some(state) = &self.series_index_task_state {
957 state.stop();
958 }
959 if self.is_running() {
960 self.set_running(false);
961 }
963 }
964}
965
966type RequestBuffer = Vec<WorkerRequest>;
967
968#[derive(Default)]
972pub(crate) struct StalledRequests {
973 pub(crate) requests:
980 HashMap<RegionId, (usize, Vec<SenderWriteRequest>, Vec<SenderBulkRequest>)>,
981 pub(crate) estimated_size: usize,
983}
984
985impl StalledRequests {
986 pub(crate) fn estimated_size(&self, region_id: &RegionId) -> usize {
988 self.requests
989 .get(region_id)
990 .map(|(size, _, _)| *size)
991 .unwrap_or_default()
992 }
993
994 pub(crate) fn append(
996 &mut self,
997 requests: &mut Vec<SenderWriteRequest>,
998 bulk_requests: &mut Vec<SenderBulkRequest>,
999 ) {
1000 for req in requests.drain(..) {
1001 self.push(req);
1002 }
1003 for req in bulk_requests.drain(..) {
1004 self.push_bulk(req);
1005 }
1006 }
1007
1008 pub(crate) fn push(&mut self, req: SenderWriteRequest) {
1010 let (size, requests, _) = self.requests.entry(req.request.region_id).or_default();
1011 let req_size = req.request.estimated_size();
1012 *size += req_size;
1013 self.estimated_size += req_size;
1014 requests.push(req);
1015 }
1016
1017 pub(crate) fn push_bulk(&mut self, req: SenderBulkRequest) {
1018 let region_id = req.region_id;
1019 let (size, _, requests) = self.requests.entry(region_id).or_default();
1020 let req_size = req.request.estimated_size();
1021 *size += req_size;
1022 self.estimated_size += req_size;
1023 requests.push(req);
1024 }
1025
1026 pub(crate) fn remove(
1028 &mut self,
1029 region_id: &RegionId,
1030 ) -> (Vec<SenderWriteRequest>, Vec<SenderBulkRequest>) {
1031 if let Some((size, write_reqs, bulk_reqs)) = self.requests.remove(region_id) {
1032 self.estimated_size -= size;
1033 (write_reqs, bulk_reqs)
1034 } else {
1035 (vec![], vec![])
1036 }
1037 }
1038
1039 pub(crate) fn stalled_count(&self) -> usize {
1041 self.requests
1042 .values()
1043 .map(|(_, reqs, bulk_reqs)| reqs.len() + bulk_reqs.len())
1044 .sum()
1045 }
1046}
1047
1048struct RegionWorkerLoop<S> {
1050 id: WorkerId,
1052 config: Arc<MitoConfig>,
1054 regions: RegionMapRef,
1056 dropping_regions: RegionMapRef,
1058 opening_regions: OpeningRegionsRef,
1060 catchup_regions: CatchupRegionsRef,
1062 sender: Sender<WorkerRequestWithTime>,
1064 receiver: Receiver<WorkerRequestWithTime>,
1066 wal: Wal<S>,
1068 object_store_manager: ObjectStoreManagerRef,
1070 running: Arc<AtomicBool>,
1072 memtable_builder_provider: MemtableBuilderProvider,
1074 purge_scheduler: SchedulerRef,
1076 write_buffer_manager: WriteBufferManagerRef,
1078 index_build_scheduler: IndexBuildScheduler,
1080 series_index_task_state: Option<Arc<SeriesIndexTaskState>>,
1082 series_index_store: Option<ObjectStore>,
1084 series_index_purger: Option<IndexFilePurger>,
1085 flush_scheduler: FlushScheduler,
1087 compaction_scheduler: CompactionScheduler,
1089 stalled_requests: StalledRequests,
1091 listener: WorkerListener,
1093 cache_manager: CacheManagerRef,
1095 puffin_manager_factory: PuffinManagerFactory,
1097 intermediate_manager: IntermediateManager,
1099 time_provider: TimeProviderRef,
1101 last_periodical_check_millis: i64,
1103 flush_sender: watch::Sender<()>,
1105 flush_receiver: watch::Receiver<()>,
1107 stalling_count: IntGauge,
1109 region_count: IntGauge,
1111 request_wait_time: Histogram,
1113 region_edit_queues: RegionEditQueues,
1115 pending_alters: handle_alter::PendingAlters,
1117 schema_metadata_manager: SchemaMetadataManagerRef,
1119 file_ref_manager: FileReferenceManagerRef,
1121 partition_expr_fetcher: PartitionExprFetcherRef,
1123 flush_semaphore: Arc<Semaphore>,
1125 plugins: Plugins,
1127 #[cfg(feature = "hdfs-object-store")]
1128 sweep_limiter: Option<Arc<Semaphore>>,
1131 node_lease: Arc<NodeLease>,
1132 staging_regions: StagingRegionsRef,
1134}
1135
1136impl<S: LogStore> RegionWorkerLoop<S> {
1137 fn register_region(&mut self, region: MitoRegionRef) {
1139 if let Some(old) = self.regions.get_region(region.region_id) {
1140 old.deactivate_region_manifest_size_metric();
1141 }
1142
1143 self.compaction_scheduler.register_region(
1144 ®ion.version_control,
1145 ®ion.access_layer,
1146 ®ion.compaction_state,
1147 );
1148 self.regions.insert_region(region.clone());
1149 region.activate_region_manifest_size_metric();
1150 }
1151
1152 async fn run(&mut self) {
1154 let init_check_delay = worker_init_check_delay();
1155 info!(
1156 "Start region worker thread {}, init_check_delay: {:?}",
1157 self.id, init_check_delay
1158 );
1159 self.last_periodical_check_millis += init_check_delay.as_millis() as i64;
1160
1161 let mut write_req_buffer: Vec<SenderWriteRequest> =
1163 Vec::with_capacity(self.config.worker_request_batch_size);
1164 let mut bulk_req_buffer: Vec<SenderBulkRequest> =
1165 Vec::with_capacity(self.config.worker_request_batch_size);
1166 let mut ddl_req_buffer: Vec<SenderDdlRequest> =
1167 Vec::with_capacity(self.config.worker_request_batch_size);
1168 let mut general_req_buffer: Vec<WorkerRequest> =
1169 RequestBuffer::with_capacity(self.config.worker_request_batch_size);
1170
1171 while self.running.load(Ordering::Relaxed) {
1172 write_req_buffer.clear();
1174 ddl_req_buffer.clear();
1175 general_req_buffer.clear();
1176 let mut bulk_insert_req_num = 0;
1177
1178 let max_wait_time = self.time_provider.wait_duration(CHECK_REGION_INTERVAL);
1179 let sleep = tokio::time::sleep(max_wait_time);
1180 tokio::pin!(sleep);
1181
1182 tokio::select! {
1183 request_opt = self.receiver.recv() => {
1184 match request_opt {
1185 Some(request_with_time) => {
1186 let wait_time = request_with_time.created_at.elapsed();
1188 self.request_wait_time.observe(wait_time.as_secs_f64());
1189
1190 match request_with_time.request {
1191 WorkerRequest::Write(sender_req) => write_req_buffer.push(sender_req),
1192 WorkerRequest::BatchWrite(sender_reqs) => write_req_buffer.extend(sender_reqs),
1193 WorkerRequest::Ddl(sender_req) => ddl_req_buffer.push(sender_req),
1194 WorkerRequest::BulkInserts(bulk_insert) => {
1195 bulk_insert_req_num += 1;
1196 self.buffer_bulk_insert_request(bulk_insert, &mut bulk_req_buffer)
1197 .await;
1198 }
1199 req => general_req_buffer.push(req),
1200 }
1201 },
1202 None => break,
1204 }
1205 }
1206 recv_res = self.flush_receiver.changed() => {
1207 if recv_res.is_err() {
1208 break;
1210 } else {
1211 self.maybe_flush_worker();
1216 self.handle_stalled_requests().await;
1218 continue;
1219 }
1220 }
1221 _ = &mut sleep => {
1222 self.handle_periodical_tasks();
1224 continue;
1225 }
1226 }
1227
1228 if self.flush_receiver.has_changed().unwrap_or(false) {
1229 self.handle_stalled_requests().await;
1233 }
1234
1235 for _ in 1..self.config.worker_request_batch_size {
1237 match self.receiver.try_recv() {
1239 Ok(request_with_time) => {
1240 let wait_time = request_with_time.created_at.elapsed();
1242 self.request_wait_time.observe(wait_time.as_secs_f64());
1243
1244 match request_with_time.request {
1245 WorkerRequest::Write(sender_req) => write_req_buffer.push(sender_req),
1246 WorkerRequest::BatchWrite(sender_reqs) => {
1247 write_req_buffer.extend(sender_reqs)
1248 }
1249 WorkerRequest::Ddl(sender_req) => ddl_req_buffer.push(sender_req),
1250 WorkerRequest::BulkInserts(bulk_insert) => {
1251 bulk_insert_req_num += 1;
1252 self.buffer_bulk_insert_request(bulk_insert, &mut bulk_req_buffer)
1253 .await
1254 }
1255 req => general_req_buffer.push(req),
1256 }
1257 }
1258 Err(_) => break,
1260 }
1261 }
1262
1263 self.listener.on_recv_requests(
1264 write_req_buffer.len()
1265 + ddl_req_buffer.len()
1266 + general_req_buffer.len()
1267 + bulk_insert_req_num,
1268 );
1269
1270 self.handle_requests(
1271 &mut write_req_buffer,
1272 &mut ddl_req_buffer,
1273 &mut general_req_buffer,
1274 &mut bulk_req_buffer,
1275 )
1276 .await;
1277
1278 self.handle_periodical_tasks();
1279 }
1280
1281 self.clean().await;
1282
1283 info!("Exit region worker thread {}", self.id);
1284 }
1285
1286 async fn buffer_bulk_insert_request(
1287 &mut self,
1288 bulk_insert: BulkInsertRequest,
1289 bulk_requests: &mut Vec<SenderBulkRequest>,
1290 ) {
1291 let BulkInsertRequest {
1292 metadata,
1293 request,
1294 sender,
1295 } = bulk_insert;
1296
1297 if let Some(region_metadata) = metadata {
1298 self.handle_bulk_insert_batch(region_metadata, request, bulk_requests, sender)
1299 .await;
1300 } else {
1301 error!("Cannot find region metadata for {}", request.region_id);
1302 sender.send(
1303 error::RegionNotFoundSnafu {
1304 region_id: request.region_id,
1305 }
1306 .fail(),
1307 );
1308 }
1309 }
1310
1311 async fn handle_requests(
1315 &mut self,
1316 write_requests: &mut Vec<SenderWriteRequest>,
1317 ddl_requests: &mut Vec<SenderDdlRequest>,
1318 general_requests: &mut Vec<WorkerRequest>,
1319 bulk_requests: &mut Vec<SenderBulkRequest>,
1320 ) {
1321 for worker_req in general_requests.drain(..) {
1322 match worker_req {
1323 WorkerRequest::Write(_) | WorkerRequest::BatchWrite(_) | WorkerRequest::Ddl(_) => {
1324 continue;
1326 }
1327 WorkerRequest::BulkInserts(_) => unreachable!("bulk inserts are buffered"),
1328 WorkerRequest::Background { region_id, notify } => {
1329 if matches!(
1330 ¬ify,
1331 BackgroundNotify::RegionEdit(edit_result)
1332 if edit_result.update_region_state
1333 ) {
1334 self.handle_buffered_region_write_requests(
1343 ®ion_id,
1344 write_requests,
1345 bulk_requests,
1346 )
1347 .await;
1348 }
1349 self.handle_background_notify(region_id, notify).await;
1351 }
1352 WorkerRequest::SetRegionRoleStateGracefully {
1353 region_id,
1354 region_role_state,
1355 sender,
1356 } => {
1357 self.set_role_state_gracefully(region_id, region_role_state, sender)
1358 .await;
1359 }
1360 WorkerRequest::EditRegion(request) => {
1361 self.handle_region_edit(request);
1362 }
1363 WorkerRequest::Stop => {
1364 debug_assert!(!self.running.load(Ordering::Relaxed));
1365 }
1366 WorkerRequest::SyncRegion(req) => {
1367 self.handle_region_sync(req).await;
1368 }
1369 WorkerRequest::RemapManifests(req) => {
1370 self.handle_remap_manifests_request(req);
1371 }
1372 WorkerRequest::CopyRegionFrom(req) => {
1373 self.handle_copy_region_from_request(req);
1374 }
1375 }
1376 }
1377
1378 self.handle_write_requests(write_requests, bulk_requests, true)
1381 .await;
1382
1383 self.handle_ddl_requests(ddl_requests).await;
1384 }
1385
1386 async fn handle_ddl_requests(&mut self, ddl_requests: &mut Vec<SenderDdlRequest>) {
1388 if ddl_requests.is_empty() {
1389 return;
1390 }
1391
1392 for ddl in ddl_requests.drain(..) {
1393 let res = match ddl.request {
1394 DdlRequest::Create(req) => self.handle_create_request(ddl.region_id, req).await,
1395 DdlRequest::Drop(req) => {
1396 self.handle_drop_request(ddl.region_id, req, ddl.sender)
1397 .await;
1398 continue;
1399 }
1400 DdlRequest::Open((req, wal_entry_receiver)) => {
1401 self.handle_open_request(ddl.region_id, req, wal_entry_receiver, ddl.sender)
1402 .await;
1403 continue;
1404 }
1405 DdlRequest::OfflineCleanup(req) => {
1406 self.handle_offline_cleanup_request(ddl.region_id, req)
1407 .await
1408 }
1409 DdlRequest::Close(req) => {
1410 self.handle_close_request(ddl.region_id, req, ddl.sender)
1411 .await;
1412 continue;
1413 }
1414 DdlRequest::Alter(req) => {
1415 self.handle_alter_request(ddl.region_id, req, ddl.sender)
1416 .await;
1417 continue;
1418 }
1419 DdlRequest::Flush(req) => {
1420 self.handle_flush_request(ddl.region_id, req, ddl.sender);
1421 continue;
1422 }
1423 DdlRequest::Compact(req) => {
1424 self.handle_compaction_request(ddl.region_id, req, ddl.sender)
1425 .await;
1426 continue;
1427 }
1428 DdlRequest::BuildIndex(req) => {
1429 self.handle_build_index_request(ddl.region_id, req, ddl.sender)
1430 .await;
1431 continue;
1432 }
1433 DdlRequest::Truncate(req) => {
1434 self.handle_truncate_request(ddl.region_id, req, ddl.sender)
1435 .await;
1436 continue;
1437 }
1438 DdlRequest::Catchup((req, wal_entry_receiver)) => {
1439 self.handle_catchup_request(ddl.region_id, req, wal_entry_receiver, ddl.sender)
1440 .await;
1441 continue;
1442 }
1443 DdlRequest::EnterStaging(req) => {
1444 self.handle_enter_staging_request(
1445 ddl.region_id,
1446 req.partition_directive,
1447 ddl.sender,
1448 )
1449 .await;
1450 continue;
1451 }
1452 DdlRequest::ApplyStagingManifest(req) => {
1453 self.handle_apply_staging_manifest_request(ddl.region_id, req, ddl.sender)
1454 .await;
1455 continue;
1456 }
1457 };
1458
1459 ddl.sender.send(res);
1460 }
1461 }
1462
1463 fn handle_periodical_tasks(&mut self) {
1465 let interval = CHECK_REGION_INTERVAL.as_millis() as i64;
1466 if self
1467 .time_provider
1468 .elapsed_since(self.last_periodical_check_millis)
1469 < interval
1470 {
1471 return;
1472 }
1473
1474 self.last_periodical_check_millis = self.time_provider.current_time_millis();
1475
1476 if let Err(e) = self.flush_periodically() {
1477 error!(e; "Failed to flush regions periodically");
1478 }
1479 }
1480
1481 async fn handle_background_notify(&mut self, region_id: RegionId, notify: BackgroundNotify) {
1483 match notify {
1484 BackgroundNotify::RegionOpened { region, registered } => {
1485 self.register_region(region);
1486 self.region_count.inc();
1487
1488 let senders = self.opening_regions.remove_sender(region_id);
1491 if let Some(state) = &self.series_index_task_state {
1492 state.wake();
1493 }
1494 let _ = registered.send(());
1495 for sender in senders {
1496 sender.send(Ok(0));
1497 }
1498 }
1499 BackgroundNotify::CompactionUnit(req) => {
1500 self.handle_compaction_unit(region_id, req).await;
1501 }
1502 BackgroundNotify::CompactionDdlComplete {
1503 generation,
1504 request_index,
1505 } => {
1506 self.compaction_scheduler
1507 .on_ddl_complete(region_id, generation, request_index);
1508 }
1509 BackgroundNotify::CompactionPickFinished(req) => {
1510 self.handle_compaction_pick_finished(region_id, req).await
1511 }
1512 BackgroundNotify::FlushFinished(req) => {
1513 self.handle_flush_finished(region_id, req).await
1514 }
1515 BackgroundNotify::FlushFailed(req) => self.handle_flush_failed(region_id, req).await,
1516 BackgroundNotify::IndexBuildFinished(req) => {
1517 self.handle_index_build_finished(region_id, req).await
1518 }
1519 BackgroundNotify::IndexBuildStopped(req) => {
1520 self.handle_index_build_stopped(region_id, req).await
1521 }
1522 BackgroundNotify::IndexBuildRetry(req) => {
1523 self.handle_rebuild_index(req, OptionOutputTx::new(None))
1524 .await
1525 }
1526 BackgroundNotify::CompactionFinished(req) => {
1527 self.handle_compaction_finished(region_id, req).await
1528 }
1529 BackgroundNotify::CompactionFailed(req) => self.handle_compaction_failure(req).await,
1530 BackgroundNotify::Truncate(req) => self.handle_truncate_result(req).await,
1531 BackgroundNotify::DiscardUnflushed(req) => {
1532 self.handle_discard_unflushed_result(req).await
1533 }
1534 BackgroundNotify::RegionChange(req) => {
1535 self.handle_manifest_region_change_result(req).await
1536 }
1537 BackgroundNotify::EnterStaging(req) => self.handle_enter_staging_result(req).await,
1538 BackgroundNotify::RegionEdit(req) => self.handle_region_edit_result(req).await,
1539 BackgroundNotify::CopyRegionFromFinished(req) => {
1540 self.handle_copy_region_from_finished(req)
1541 }
1542 }
1543 }
1544
1545 async fn set_role_state_gracefully(
1547 &mut self,
1548 region_id: RegionId,
1549 region_role_state: SettableRegionRoleState,
1550 sender: oneshot::Sender<SetRegionRoleStateResponse>,
1551 ) {
1552 if let Some(region) = self.regions.get_region(region_id) {
1553 if matches!(
1554 region_role_state,
1555 SettableRegionRoleState::Follower | SettableRegionRoleState::DowngradingLeader
1556 ) {
1557 region.manifest_ctx.begin_downgrade(region_id);
1562 }
1563 #[cfg(feature = "hdfs-object-store")]
1564 let sweep_older_than = self.config.gc.unknown_file_lingering_time;
1565 #[cfg(feature = "hdfs-object-store")]
1566 let sweep_limiter = self.sweep_limiter.clone();
1567 common_runtime::spawn_global(async move {
1569 match region.set_role_state_gracefully(region_role_state).await {
1570 Ok(()) => {
1571 let last_entry_id = region.version_control.current().last_entry_id;
1572 let _ = sender.send(SetRegionRoleStateResponse::success(
1573 SetRegionRoleStateSuccess::mito(last_entry_id),
1574 ));
1575 #[cfg(feature = "hdfs-object-store")]
1578 if region_role_state == SettableRegionRoleState::Leader {
1579 temp_file_sweeper::spawn_region_sweep(
1580 region.access_layer(),
1581 region_id,
1582 sweep_older_than,
1583 sweep_limiter,
1584 );
1585 }
1586 }
1587 Err(e) => {
1588 error!(e; "Failed to set region {} role state to {:?}", region_id, region_role_state);
1589 let _ = sender.send(SetRegionRoleStateResponse::invalid_transition(
1590 BoxedError::new(e),
1591 ));
1592 }
1593 }
1594 });
1595 } else {
1596 let _ = sender.send(SetRegionRoleStateResponse::NotFound);
1597 }
1598 }
1599}
1600
1601impl<S> RegionWorkerLoop<S> {
1602 async fn clean(&self) {
1604 let regions = self.regions.list_regions();
1606 for region in regions {
1607 region.deactivate_region_manifest_size_metric();
1608 region.stop().await;
1609 }
1610
1611 self.regions.clear();
1612 }
1613
1614 fn notify_group(&mut self) {
1617 let _ = self.flush_sender.send(());
1619 self.flush_receiver.borrow_and_update();
1621 }
1622}
1623
1624#[derive(Default, Clone)]
1626pub(crate) struct WorkerListener {
1627 #[cfg(any(test, feature = "test"))]
1628 listener: Option<crate::engine::listener::EventListenerRef>,
1629}
1630
1631impl WorkerListener {
1632 #[cfg(any(test, feature = "test"))]
1633 pub(crate) fn new(
1634 listener: Option<crate::engine::listener::EventListenerRef>,
1635 ) -> WorkerListener {
1636 WorkerListener { listener }
1637 }
1638
1639 pub(crate) async fn on_region_open_registered(&self, _region_id: RegionId) {
1641 #[cfg(any(test, feature = "test"))]
1642 if let Some(listener) = &self.listener {
1643 listener.on_region_open_registered(_region_id).await;
1644 }
1645 }
1646
1647 pub(crate) fn on_flush_success(&self, region_id: RegionId) {
1649 #[cfg(any(test, feature = "test"))]
1650 if let Some(listener) = &self.listener {
1651 listener.on_flush_success(region_id);
1652 }
1653 let _ = region_id;
1655 }
1656
1657 pub(crate) fn on_write_stall(&self) {
1659 #[cfg(any(test, feature = "test"))]
1660 if let Some(listener) = &self.listener {
1661 listener.on_write_stall();
1662 }
1663 }
1664
1665 pub(crate) async fn on_write_wal_end(&self) {
1667 #[cfg(any(test, feature = "test"))]
1668 if let Some(listener) = &self.listener {
1669 listener.on_write_wal_end().await;
1670 }
1671 }
1672
1673 pub(crate) async fn on_flush_begin(&self, region_id: RegionId) {
1674 #[cfg(any(test, feature = "test"))]
1675 if let Some(listener) = &self.listener {
1676 listener.on_flush_begin(region_id).await;
1677 }
1678 let _ = region_id;
1680 }
1681
1682 pub(crate) async fn on_flush_commit_begin(&self, _region_id: RegionId) {
1683 #[cfg(any(test, feature = "test"))]
1684 if let Some(listener) = &self.listener {
1685 listener.on_flush_commit_begin(_region_id).await;
1686 }
1687 }
1688
1689 pub(crate) fn on_flush_cancel_requested(&self, _region_id: RegionId) {
1690 #[cfg(any(test, feature = "test"))]
1691 if let Some(listener) = &self.listener {
1692 listener.on_flush_cancel_requested(_region_id);
1693 }
1694 }
1695
1696 pub(crate) fn on_later_drop_begin(&self, region_id: RegionId) -> Option<Duration> {
1697 #[cfg(any(test, feature = "test"))]
1698 if let Some(listener) = &self.listener {
1699 return listener.on_later_drop_begin(region_id);
1700 }
1701 let _ = region_id;
1703 None
1704 }
1705
1706 pub(crate) fn on_later_drop_end(&self, region_id: RegionId, removed: bool) {
1708 #[cfg(any(test, feature = "test"))]
1709 if let Some(listener) = &self.listener {
1710 listener.on_later_drop_end(region_id, removed);
1711 }
1712 let _ = region_id;
1714 let _ = removed;
1715 }
1716
1717 pub(crate) async fn on_merge_ssts_finished(&self, region_id: RegionId) {
1718 #[cfg(any(test, feature = "test"))]
1719 if let Some(listener) = &self.listener {
1720 listener.on_merge_ssts_finished(region_id).await;
1721 }
1722 let _ = region_id;
1724 }
1725
1726 pub(crate) fn on_recv_requests(&self, request_num: usize) {
1727 #[cfg(any(test, feature = "test"))]
1728 if let Some(listener) = &self.listener {
1729 listener.on_recv_requests(request_num);
1730 }
1731 let _ = request_num;
1733 }
1734
1735 pub(crate) fn on_file_cache_filled(&self, _file_id: FileId) {
1736 #[cfg(any(test, feature = "test"))]
1737 if let Some(listener) = &self.listener {
1738 listener.on_file_cache_filled(_file_id);
1739 }
1740 }
1741
1742 pub(crate) fn on_compaction_scheduled(&self, _region_id: RegionId) {
1743 #[cfg(any(test, feature = "test"))]
1744 if let Some(listener) = &self.listener {
1745 listener.on_compaction_scheduled(_region_id);
1746 }
1747 }
1748
1749 pub(crate) async fn on_compaction_pick_begin(&self, _region_id: RegionId) {
1750 #[cfg(any(test, feature = "test"))]
1751 if let Some(listener) = &self.listener {
1752 listener.on_compaction_pick_begin(_region_id).await;
1753 }
1754 }
1755
1756 pub(crate) async fn on_compaction_commit_begin(&self, _region_id: RegionId) {
1757 #[cfg(any(test, feature = "test"))]
1758 if let Some(listener) = &self.listener {
1759 listener.on_compaction_commit_begin(_region_id).await;
1760 }
1761 }
1762
1763 pub(crate) async fn on_compaction_unit_merge_begin(&self, _region_id: RegionId, _plan_id: u64) {
1765 #[cfg(any(test, feature = "test"))]
1766 if let Some(listener) = &self.listener {
1767 listener
1768 .on_compaction_unit_merge_begin(_region_id, _plan_id)
1769 .await;
1770 }
1771 }
1772
1773 pub(crate) async fn on_compaction_unit_committed(&self, _region_id: RegionId, _plan_id: u64) {
1775 #[cfg(any(test, feature = "test"))]
1776 if let Some(listener) = &self.listener {
1777 listener
1778 .on_compaction_unit_committed(_region_id, _plan_id)
1779 .await;
1780 }
1781 }
1782
1783 pub(crate) fn on_compaction_unit_applied(&self, _region_id: RegionId, _plan_id: u64) {
1785 #[cfg(any(test, feature = "test"))]
1786 if let Some(listener) = &self.listener {
1787 listener.on_compaction_unit_applied(_region_id, _plan_id);
1788 }
1789 }
1790
1791 pub(crate) async fn on_compaction_result_notified(&self, _region_id: RegionId) {
1792 #[cfg(any(test, feature = "test"))]
1793 if let Some(listener) = &self.listener {
1794 listener.on_compaction_result_notified(_region_id).await;
1795 }
1796 }
1797
1798 pub(crate) fn on_compaction_cancel_requested(&self, _region_id: RegionId) {
1799 #[cfg(any(test, feature = "test"))]
1800 if let Some(listener) = &self.listener {
1801 listener.on_compaction_cancel_requested(_region_id);
1802 }
1803 }
1804
1805 pub(crate) async fn on_notify_region_change_result_begin(&self, _region_id: RegionId) {
1806 #[cfg(any(test, feature = "test"))]
1807 if let Some(listener) = &self.listener {
1808 listener
1809 .on_notify_region_change_result_begin(_region_id)
1810 .await;
1811 }
1812 }
1813
1814 pub(crate) async fn on_enter_staging_result_begin(&self, _region_id: RegionId) {
1815 #[cfg(any(test, feature = "test"))]
1816 if let Some(listener) = &self.listener {
1817 listener.on_enter_staging_result_begin(_region_id).await;
1818 }
1819 }
1820
1821 pub(crate) async fn on_index_build_finish(&self, _region_file_id: RegionFileId) {
1822 #[cfg(any(test, feature = "test"))]
1823 if let Some(listener) = &self.listener {
1824 listener.on_index_build_finish(_region_file_id).await;
1825 }
1826 }
1827
1828 pub(crate) async fn on_index_build_begin(&self, _region_file_id: RegionFileId) {
1829 #[cfg(any(test, feature = "test"))]
1830 if let Some(listener) = &self.listener {
1831 listener.on_index_build_begin(_region_file_id).await;
1832 }
1833 }
1834
1835 pub(crate) async fn on_index_build_abort(&self, _region_file_id: RegionFileId) {
1836 #[cfg(any(test, feature = "test"))]
1837 if let Some(listener) = &self.listener {
1838 listener.on_index_build_abort(_region_file_id).await;
1839 }
1840 }
1841
1842 pub(crate) async fn on_index_build_before_manifest_commit(
1843 &self,
1844 _region_file_id: RegionFileId,
1845 ) {
1846 #[cfg(any(test, feature = "test"))]
1847 if let Some(listener) = &self.listener {
1848 listener
1849 .on_index_build_before_manifest_commit(_region_file_id)
1850 .await;
1851 }
1852 }
1853
1854 pub(crate) async fn on_index_build_manifest_committed(&self, _region_file_id: RegionFileId) {
1855 #[cfg(any(test, feature = "test"))]
1856 if let Some(listener) = &self.listener {
1857 listener
1858 .on_index_build_manifest_committed(_region_file_id)
1859 .await;
1860 }
1861 }
1862}
1863
1864#[cfg(test)]
1865mod tests {
1866 use super::*;
1867 use crate::test_util::TestEnv;
1868
1869 #[test]
1870 fn test_region_id_to_index() {
1871 let num_workers = 4;
1872
1873 let region_id = RegionId::new(1, 2);
1874 let index = region_id_to_index(region_id, num_workers);
1875 assert_eq!(index, 3);
1876
1877 let region_id = RegionId::new(2, 3);
1878 let index = region_id_to_index(region_id, num_workers);
1879 assert_eq!(index, 1);
1880 }
1881
1882 #[tokio::test]
1883 async fn test_worker_group_start_stop() {
1884 for (enabled, existing_index) in [(false, false), (true, false), (false, true)] {
1885 let env = TestEnv::with_prefix("group-stop").await;
1887 let root = env.data_home().join("series_index");
1888 if existing_index {
1889 tokio::fs::create_dir_all(&root).await.unwrap();
1890 tokio::fs::write(root.join("retained"), b"index")
1891 .await
1892 .unwrap();
1893 }
1894 let group = env
1895 .create_worker_group(MitoConfig {
1896 num_workers: 4,
1897 experimental_enable_series_index: enabled,
1898 experimental_series_index_maintenance_interval: Duration::from_secs(3600),
1899 ..Default::default()
1900 })
1901 .await;
1902
1903 for worker in &group.workers {
1904 assert_eq!(worker.series_index_task_state.is_some(), enabled);
1905 assert_eq!(worker.series_index_handle.lock().await.is_some(), enabled);
1906 }
1907 assert_eq!(root.is_dir(), enabled || existing_index);
1908 if existing_index {
1909 assert_eq!(
1910 tokio::fs::read(root.join("retained")).await.unwrap(),
1911 b"index"
1912 );
1913 }
1914
1915 tokio::time::timeout(Duration::from_secs(5), group.stop())
1916 .await
1917 .expect("series-index tasks should stop without waiting for their interval")
1918 .unwrap();
1919 }
1920 }
1921
1922 #[tokio::test]
1924 async fn test_workers_share_memtable_compact_dispatcher() {
1925 let env = TestEnv::with_prefix("share-dispatcher").await;
1926 let group = env
1927 .create_worker_group(MitoConfig {
1928 num_workers: 4,
1929 ..Default::default()
1930 })
1931 .await;
1932 let first = &group.workers[0].memtable_compact_dispatcher;
1933 for (index, worker) in group.workers.iter().enumerate() {
1934 let current = &worker.memtable_compact_dispatcher;
1935 assert!(
1936 Arc::ptr_eq(first, current),
1937 "worker {index} has its own memtable compaction dispatcher"
1938 );
1939 }
1940
1941 group.stop().await.unwrap();
1942 }
1943}