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, OpeningRegions, OpeningRegionsRef, RegionMap,
75 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;
90use crate::time_provider::{StdTimeProvider, TimeProviderRef};
91use crate::wal::Wal;
92use crate::worker::handle_manifest::RegionEditQueues;
93
94pub(crate) type WorkerId = u32;
96
97pub(crate) const DROPPING_MARKER_FILE: &str = ".dropping";
98
99pub(crate) const CHECK_REGION_INTERVAL: Duration = Duration::from_secs(60);
101pub(crate) const MAX_INITIAL_CHECK_DELAY_SECS: u64 = 60 * 3;
103
104#[cfg_attr(doc, aquamarine::aquamarine)]
105pub(crate) struct WorkerGroup {
142 workers: Vec<RegionWorker>,
144 flush_job_pool: SchedulerRef,
146 compact_job_pool: SchedulerRef,
148 index_build_job_pool: SchedulerRef,
150 purge_scheduler: SchedulerRef,
152 cache_manager: CacheManagerRef,
154 file_ref_manager: FileReferenceManagerRef,
156 gc_limiter: GcLimiterRef,
158 object_store_manager: ObjectStoreManagerRef,
160 puffin_manager_factory: PuffinManagerFactory,
162 intermediate_manager: IntermediateManager,
164 schema_metadata_manager: SchemaMetadataManagerRef,
166}
167
168impl WorkerGroup {
169 #[allow(clippy::too_many_arguments)]
173 pub(crate) async fn start<S: LogStore>(
174 data_home: &str,
175 config: Arc<MitoConfig>,
176 log_store: Arc<S>,
177 object_store_manager: ObjectStoreManagerRef,
178 schema_metadata_manager: SchemaMetadataManagerRef,
179 file_ref_manager: FileReferenceManagerRef,
180 partition_expr_fetcher: PartitionExprFetcherRef,
181 plugins: Plugins,
182 ) -> Result<WorkerGroup> {
183 let (flush_sender, flush_receiver) = watch::channel(());
184 let write_buffer_manager = Arc::new(
185 WriteBufferManagerImpl::new(config.global_write_buffer_size.as_bytes() as usize)
186 .with_notifier(flush_sender.clone()),
187 );
188 let puffin_manager_factory = PuffinManagerFactory::new(
189 &config.index.aux_path,
190 config.index.staging_size.as_bytes(),
191 Some(config.index.write_buffer_size.as_bytes() as _),
192 config.index.staging_ttl,
193 )
194 .await?;
195 let intermediate_manager = IntermediateManager::init_fs(&config.index.aux_path)
196 .await?
197 .with_buffer_size(Some(config.index.write_buffer_size.as_bytes() as _));
198 let index_build_job_pool =
199 Arc::new(LocalScheduler::new(config.max_background_index_builds));
200 let series_index_store = series_index_store_from_config(&config, data_home).await?;
201 let series_index_disk_usage = Arc::new(AtomicU64::new(0));
202 let flush_job_pool = Arc::new(LocalScheduler::new(config.max_background_flushes));
203 let compact_job_pool = Arc::new(LocalScheduler::new(config.max_background_compactions));
204 let flush_semaphore = Arc::new(Semaphore::new(config.max_background_flushes));
205 let purge_scheduler = Arc::new(LocalScheduler::new(config.max_background_purges));
207 let upload_store_wrapper = plugins.get::<WriteCacheUploadStoreWrapperRef>();
208 let write_cache = write_cache_from_config(
209 &config,
210 puffin_manager_factory.clone(),
211 intermediate_manager.clone(),
212 upload_store_wrapper,
213 )
214 .await?;
215 let cache_manager = Arc::new(
216 CacheManager::builder()
217 .sst_meta_cache_size(config.sst_meta_cache_size.as_bytes())
218 .vector_cache_size(config.vector_cache_size.as_bytes())
219 .page_cache_size(config.page_cache_size.as_bytes())
220 .selector_result_cache_size(config.selector_result_cache_size.as_bytes())
221 .range_result_cache_size(config.range_result_cache_size.as_bytes())
222 .prefilter_result_cache_size(config.prefilter_result_cache_size.as_bytes())
223 .index_metadata_size(config.index.metadata_cache_size.as_bytes())
224 .index_content_size(config.index.content_cache_size.as_bytes())
225 .index_content_page_size(config.index.content_cache_page_size.as_bytes())
226 .index_result_cache_size(config.index.result_cache_size.as_bytes())
227 .puffin_metadata_size(config.index.metadata_cache_size.as_bytes())
228 .write_cache(write_cache)
229 .build(),
230 );
231 let time_provider = Arc::new(StdTimeProvider);
232 let total_memory = get_total_memory_bytes();
233 let total_memory = if total_memory > 0 {
234 total_memory as u64
235 } else {
236 0
237 };
238 let compaction_limit_bytes = config
239 .experimental_compaction_memory_limit
240 .resolve(total_memory);
241 let compaction_memory_manager =
242 Arc::new(new_compaction_memory_manager(compaction_limit_bytes));
243 let gc_limiter = Arc::new(GcLimiter::new(config.gc.max_concurrent_gc_job));
244
245 let workers = (0..config.num_workers)
246 .map(|id| {
247 WorkerStarter {
248 id: id as WorkerId,
249 config: config.clone(),
250 log_store: log_store.clone(),
251 object_store_manager: object_store_manager.clone(),
252 write_buffer_manager: write_buffer_manager.clone(),
253 index_build_job_pool: index_build_job_pool.clone(),
254 series_index_store: series_index_store.clone(),
255 series_index_disk_usage: series_index_disk_usage.clone(),
256 flush_job_pool: flush_job_pool.clone(),
257 compact_job_pool: compact_job_pool.clone(),
258 purge_scheduler: purge_scheduler.clone(),
259 listener: WorkerListener::default(),
260 cache_manager: cache_manager.clone(),
261 compaction_memory_manager: compaction_memory_manager.clone(),
262 puffin_manager_factory: puffin_manager_factory.clone(),
263 intermediate_manager: intermediate_manager.clone(),
264 time_provider: time_provider.clone(),
265 flush_sender: flush_sender.clone(),
266 flush_receiver: flush_receiver.clone(),
267 plugins: plugins.clone(),
268 schema_metadata_manager: schema_metadata_manager.clone(),
269 file_ref_manager: file_ref_manager.clone(),
270 partition_expr_fetcher: partition_expr_fetcher.clone(),
271 flush_semaphore: flush_semaphore.clone(),
272 }
273 .start()
274 })
275 .collect::<Result<Vec<_>>>()?;
276
277 Ok(WorkerGroup {
278 workers,
279 flush_job_pool,
280 compact_job_pool,
281 index_build_job_pool,
282 purge_scheduler,
283 cache_manager,
284 file_ref_manager,
285 gc_limiter,
286 object_store_manager,
287 puffin_manager_factory,
288 intermediate_manager,
289 schema_metadata_manager,
290 })
291 }
292
293 pub(crate) async fn stop(&self) -> Result<()> {
295 info!("Stop region worker group");
296
297 self.compact_job_pool.stop(true).await?;
300 self.flush_job_pool.stop(true).await?;
302 self.purge_scheduler.stop(true).await?;
304 self.index_build_job_pool.stop(true).await?;
306
307 try_join_all(self.workers.iter().map(|worker| worker.stop())).await?;
308
309 Ok(())
310 }
311
312 pub(crate) async fn submit_to_worker(
314 &self,
315 region_id: RegionId,
316 request: WorkerRequest,
317 ) -> Result<()> {
318 self.worker(region_id).submit_request(request).await
319 }
320
321 pub(crate) async fn submit_writes(&self, writes: Vec<SenderWriteRequest>) {
326 let mut writes_by_worker: HashMap<usize, Vec<SenderWriteRequest>> = HashMap::new();
327 for write in writes {
328 let index = region_id_to_index(write.request.region_id, self.workers.len());
329 writes_by_worker.entry(index).or_default().push(write);
330 }
331 join_all(
332 writes_by_worker
333 .into_iter()
334 .map(|(index, writes)| async move {
335 let worker = &self.workers[index];
336 let Err(request) = worker
337 .try_submit_request(WorkerRequest::BatchWrite(writes))
338 .await
339 else {
340 return;
341 };
342 let WorkerRequest::BatchWrite(writes) = request else {
343 unreachable!("the worker returns the request it was given");
344 };
345 for write in writes {
346 write
347 .sender
348 .send(WorkerStoppedSnafu { id: worker.id }.fail());
349 }
350 }),
351 )
352 .await;
353 }
354
355 pub(crate) fn is_region_exists(&self, region_id: RegionId) -> bool {
357 self.worker(region_id).is_region_exists(region_id)
358 }
359
360 pub(crate) fn is_region_opening(&self, region_id: RegionId) -> bool {
362 self.worker(region_id).is_region_opening(region_id)
363 }
364
365 pub(crate) fn is_region_catching_up(&self, region_id: RegionId) -> bool {
367 self.worker(region_id).is_region_catching_up(region_id)
368 }
369
370 pub(crate) fn get_region(&self, region_id: RegionId) -> Option<MitoRegionRef> {
374 self.worker(region_id).get_region(region_id)
375 }
376
377 pub(crate) fn cache_manager(&self) -> CacheManagerRef {
379 self.cache_manager.clone()
380 }
381
382 pub(crate) fn file_ref_manager(&self) -> FileReferenceManagerRef {
383 self.file_ref_manager.clone()
384 }
385
386 pub(crate) fn gc_limiter(&self) -> GcLimiterRef {
387 self.gc_limiter.clone()
388 }
389
390 pub(crate) fn worker(&self, region_id: RegionId) -> &RegionWorker {
392 let index = region_id_to_index(region_id, self.workers.len());
393
394 &self.workers[index]
395 }
396
397 pub(crate) fn all_regions(&self) -> impl Iterator<Item = MitoRegionRef> + use<'_> {
398 self.workers
399 .iter()
400 .flat_map(|worker| worker.regions.list_regions())
401 }
402
403 pub(crate) fn object_store_manager(&self) -> &ObjectStoreManagerRef {
404 &self.object_store_manager
405 }
406
407 pub(crate) fn puffin_manager_factory(&self) -> &PuffinManagerFactory {
408 &self.puffin_manager_factory
409 }
410
411 pub(crate) fn intermediate_manager(&self) -> &IntermediateManager {
412 &self.intermediate_manager
413 }
414
415 pub(crate) fn schema_metadata_manager(&self) -> &SchemaMetadataManagerRef {
416 &self.schema_metadata_manager
417 }
418}
419
420#[cfg(any(test, feature = "test"))]
422impl WorkerGroup {
423 #[allow(clippy::too_many_arguments)]
427 pub(crate) async fn start_for_test<S: LogStore>(
428 data_home: &str,
429 config: Arc<MitoConfig>,
430 log_store: Arc<S>,
431 object_store_manager: ObjectStoreManagerRef,
432 write_buffer_manager: Option<WriteBufferManagerRef>,
433 listener: Option<crate::engine::listener::EventListenerRef>,
434 schema_metadata_manager: SchemaMetadataManagerRef,
435 file_ref_manager: FileReferenceManagerRef,
436 time_provider: TimeProviderRef,
437 partition_expr_fetcher: PartitionExprFetcherRef,
438 ) -> Result<WorkerGroup> {
439 let (flush_sender, flush_receiver) = watch::channel(());
440 let write_buffer_manager = write_buffer_manager.unwrap_or_else(|| {
441 Arc::new(
442 WriteBufferManagerImpl::new(config.global_write_buffer_size.as_bytes() as usize)
443 .with_notifier(flush_sender.clone()),
444 )
445 });
446 let index_build_job_pool =
447 Arc::new(LocalScheduler::new(config.max_background_index_builds));
448 let series_index_store = series_index_store_from_config(&config, data_home).await?;
449 let series_index_disk_usage = Arc::new(AtomicU64::new(0));
450 let flush_job_pool = Arc::new(LocalScheduler::new(config.max_background_flushes));
451 let compact_job_pool = Arc::new(LocalScheduler::new(config.max_background_compactions));
452 let flush_semaphore = Arc::new(Semaphore::new(config.max_background_flushes));
453 let purge_scheduler = Arc::new(LocalScheduler::new(config.max_background_flushes));
454 let puffin_manager_factory = PuffinManagerFactory::new(
455 &config.index.aux_path,
456 config.index.staging_size.as_bytes(),
457 Some(config.index.write_buffer_size.as_bytes() as _),
458 config.index.staging_ttl,
459 )
460 .await?;
461 let intermediate_manager = IntermediateManager::init_fs(&config.index.aux_path)
462 .await?
463 .with_buffer_size(Some(config.index.write_buffer_size.as_bytes() as _));
464 let write_cache = write_cache_from_config(
465 &config,
466 puffin_manager_factory.clone(),
467 intermediate_manager.clone(),
468 None,
469 )
470 .await?;
471 let cache_manager = Arc::new(
472 CacheManager::builder()
473 .sst_meta_cache_size(config.sst_meta_cache_size.as_bytes())
474 .vector_cache_size(config.vector_cache_size.as_bytes())
475 .page_cache_size(config.page_cache_size.as_bytes())
476 .selector_result_cache_size(config.selector_result_cache_size.as_bytes())
477 .range_result_cache_size(config.range_result_cache_size.as_bytes())
478 .prefilter_result_cache_size(config.prefilter_result_cache_size.as_bytes())
479 .write_cache(write_cache)
480 .build(),
481 );
482 let total_memory = get_total_memory_bytes();
483 let total_memory = if total_memory > 0 {
484 total_memory as u64
485 } else {
486 0
487 };
488 let compaction_limit_bytes = config
489 .experimental_compaction_memory_limit
490 .resolve(total_memory);
491 let compaction_memory_manager =
492 Arc::new(new_compaction_memory_manager(compaction_limit_bytes));
493 let gc_limiter = Arc::new(GcLimiter::new(config.gc.max_concurrent_gc_job));
494 let workers = (0..config.num_workers)
495 .map(|id| {
496 WorkerStarter {
497 id: id as WorkerId,
498 config: config.clone(),
499 log_store: log_store.clone(),
500 object_store_manager: object_store_manager.clone(),
501 write_buffer_manager: write_buffer_manager.clone(),
502 index_build_job_pool: index_build_job_pool.clone(),
503 series_index_store: series_index_store.clone(),
504 series_index_disk_usage: series_index_disk_usage.clone(),
505 flush_job_pool: flush_job_pool.clone(),
506 compact_job_pool: compact_job_pool.clone(),
507 purge_scheduler: purge_scheduler.clone(),
508 listener: WorkerListener::new(listener.clone()),
509 cache_manager: cache_manager.clone(),
510 compaction_memory_manager: compaction_memory_manager.clone(),
511 puffin_manager_factory: puffin_manager_factory.clone(),
512 intermediate_manager: intermediate_manager.clone(),
513 time_provider: time_provider.clone(),
514 flush_sender: flush_sender.clone(),
515 flush_receiver: flush_receiver.clone(),
516 plugins: Plugins::new(),
517 schema_metadata_manager: schema_metadata_manager.clone(),
518 file_ref_manager: file_ref_manager.clone(),
519 partition_expr_fetcher: partition_expr_fetcher.clone(),
520 flush_semaphore: flush_semaphore.clone(),
521 }
522 .start()
523 })
524 .collect::<Result<Vec<_>>>()?;
525
526 Ok(WorkerGroup {
527 workers,
528 flush_job_pool,
529 compact_job_pool,
530 index_build_job_pool,
531 purge_scheduler,
532 cache_manager,
533 file_ref_manager,
534 gc_limiter,
535 object_store_manager,
536 puffin_manager_factory,
537 intermediate_manager,
538 schema_metadata_manager,
539 })
540 }
541
542 pub(crate) fn purge_scheduler(&self) -> &SchedulerRef {
544 &self.purge_scheduler
545 }
546}
547
548fn region_id_to_index(id: RegionId, num_workers: usize) -> usize {
549 ((id.table_id() as usize % num_workers) + (id.region_number() as usize % num_workers))
550 % num_workers
551}
552
553async fn series_index_store_from_config(
555 config: &MitoConfig,
556 data_home: &str,
557) -> Result<Option<ObjectStore>> {
558 if !config.experimental_enable_series_index {
559 return Ok(None);
560 }
561
562 let root = Path::new(data_home).join("series_index");
563 new_fs_cache_store(&root.to_string_lossy()).await.map(Some)
564}
565
566pub async fn write_cache_from_config(
567 config: &MitoConfig,
568 puffin_manager_factory: PuffinManagerFactory,
569 intermediate_manager: IntermediateManager,
570 upload_store_wrapper: Option<WriteCacheUploadStoreWrapperRef>,
571) -> Result<Option<WriteCacheRef>> {
572 if !config.enable_write_cache {
573 return Ok(None);
574 }
575
576 tokio::fs::create_dir_all(Path::new(&config.write_cache_path))
577 .await
578 .context(CreateDirSnafu {
579 dir: &config.write_cache_path,
580 })?;
581
582 let cache = WriteCache::new_fs(
583 &config.write_cache_path,
584 config.write_cache_size,
585 config.write_cache_ttl,
586 Some(config.index_cache_percent),
587 config.enable_refill_cache_on_read,
588 puffin_manager_factory,
589 intermediate_manager,
590 config.manifest_cache_size,
591 )
592 .await?
593 .with_upload_store_wrapper(upload_store_wrapper);
594 Ok(Some(Arc::new(cache)))
595}
596
597pub(crate) fn worker_init_check_delay() -> Duration {
599 let init_check_delay = rng().random_range(0..MAX_INITIAL_CHECK_DELAY_SECS);
600 Duration::from_secs(init_check_delay)
601}
602
603struct WorkerStarter<S> {
605 id: WorkerId,
606 config: Arc<MitoConfig>,
607 log_store: Arc<S>,
608 object_store_manager: ObjectStoreManagerRef,
609 write_buffer_manager: WriteBufferManagerRef,
610 compact_job_pool: SchedulerRef,
611 index_build_job_pool: SchedulerRef,
612 series_index_store: Option<ObjectStore>,
613 series_index_disk_usage: Arc<AtomicU64>,
614 flush_job_pool: SchedulerRef,
615 purge_scheduler: SchedulerRef,
616 listener: WorkerListener,
617 cache_manager: CacheManagerRef,
618 compaction_memory_manager: Arc<CompactionMemoryManager>,
619 puffin_manager_factory: PuffinManagerFactory,
620 intermediate_manager: IntermediateManager,
621 time_provider: TimeProviderRef,
622 flush_sender: watch::Sender<()>,
624 flush_receiver: watch::Receiver<()>,
626 plugins: Plugins,
627 schema_metadata_manager: SchemaMetadataManagerRef,
628 file_ref_manager: FileReferenceManagerRef,
629 partition_expr_fetcher: PartitionExprFetcherRef,
630 flush_semaphore: Arc<Semaphore>,
631}
632
633impl<S: LogStore> WorkerStarter<S> {
634 fn start(self) -> Result<RegionWorker> {
636 let regions = Arc::new(RegionMap::default());
637 let opening_regions = Arc::new(OpeningRegions::default());
638 let catchup_regions = Arc::new(CatchupRegions::default());
639 let (sender, receiver) = mpsc::channel(self.config.worker_channel_size);
640
641 let running = Arc::new(AtomicBool::new(true));
642 let (series_index_task_state, series_index_receiver) = if self.series_index_store.is_some()
643 {
644 let (state, receiver) =
645 SeriesIndexTaskState::new(self.id, self.config.worker_channel_size);
646 (Some(Arc::new(state)), Some(receiver))
647 } else {
648 (None, None)
649 };
650 let mut series_index_purger = None;
651 let series_index_handle = self
652 .series_index_store
653 .clone()
654 .zip(series_index_task_state.clone())
655 .zip(series_index_receiver)
656 .map(|((store, state), receiver)| {
657 let (purger, purge_receiver) = series_index_channel(store.clone());
658 series_index_purger = Some(purger.clone());
659 spawn_series_index_tasks(
660 self.id,
661 store,
662 regions.clone(),
663 state,
664 receiver,
665 self.config.experimental_series_index_bucket_width,
666 purger,
667 purge_receiver,
668 self.series_index_disk_usage.clone(),
669 self.config.experimental_series_index_max_size.as_bytes(),
670 self.config.experimental_series_index_maintenance_interval,
671 self.time_provider.clone(),
672 self.config.experimental_enable_range_index,
673 )
674 });
675 let now = self.time_provider.current_time_millis();
676 let id_string = self.id.to_string();
677 let mut worker_thread = RegionWorkerLoop {
678 id: self.id,
679 config: self.config.clone(),
680 regions: regions.clone(),
681 catchup_regions: catchup_regions.clone(),
682 dropping_regions: Arc::new(RegionMap::default()),
683 opening_regions: opening_regions.clone(),
684 sender: sender.clone(),
685 receiver,
686 wal: Wal::new(self.log_store),
687 object_store_manager: self.object_store_manager.clone(),
688 running: running.clone(),
689 memtable_builder_provider: MemtableBuilderProvider::new(
690 Some(self.write_buffer_manager.clone()),
691 self.config.clone(),
692 ),
693 purge_scheduler: self.purge_scheduler.clone(),
694 write_buffer_manager: self.write_buffer_manager,
695 index_build_scheduler: IndexBuildScheduler::new(
696 self.index_build_job_pool,
697 self.config.max_background_index_builds,
698 ),
699 series_index_task_state: series_index_task_state.clone(),
700 series_index_store: self.series_index_store,
701 series_index_purger,
702 flush_scheduler: FlushScheduler::new(self.flush_job_pool),
703 compaction_scheduler: CompactionScheduler::new(
704 self.compact_job_pool,
705 sender.clone(),
706 self.cache_manager.clone(),
707 self.config.clone(),
708 self.listener.clone(),
709 self.plugins.clone(),
710 self.compaction_memory_manager.clone(),
711 self.config.experimental_compaction_on_exhausted,
712 ),
713 stalled_requests: StalledRequests::default(),
714 listener: self.listener,
715 cache_manager: self.cache_manager,
716 puffin_manager_factory: self.puffin_manager_factory,
717 intermediate_manager: self.intermediate_manager,
718 time_provider: self.time_provider,
719 last_periodical_check_millis: now,
720 flush_sender: self.flush_sender,
721 flush_receiver: self.flush_receiver,
722 stalling_count: WRITE_STALLING.with_label_values(&[&id_string]),
723 region_count: REGION_COUNT.with_label_values(&[&id_string]),
724 request_wait_time: REQUEST_WAIT_TIME.with_label_values(&[&id_string]),
725 region_edit_queues: RegionEditQueues::default(),
726 schema_metadata_manager: self.schema_metadata_manager,
727 file_ref_manager: self.file_ref_manager.clone(),
728 partition_expr_fetcher: self.partition_expr_fetcher,
729 flush_semaphore: self.flush_semaphore,
730 plugins: self.plugins,
731 };
732 let handle = common_runtime::spawn_global(async move {
733 worker_thread.run().await;
734 });
735
736 Ok(RegionWorker {
737 id: self.id,
738 regions,
739 opening_regions,
740 catchup_regions,
741 sender,
742 handle: Mutex::new(Some(handle)),
743 series_index_handle: Mutex::new(series_index_handle),
744 series_index_task_state,
745 running,
746 })
747 }
748}
749
750pub(crate) struct RegionWorker {
752 id: WorkerId,
754 regions: RegionMapRef,
756 opening_regions: OpeningRegionsRef,
758 catchup_regions: CatchupRegionsRef,
760 sender: Sender<WorkerRequestWithTime>,
762 handle: Mutex<Option<JoinHandle<()>>>,
764 series_index_handle: Mutex<Option<JoinHandle<()>>>,
766 series_index_task_state: Option<Arc<SeriesIndexTaskState>>,
768 running: Arc<AtomicBool>,
770}
771
772impl RegionWorker {
773 async fn submit_request(&self, request: WorkerRequest) -> Result<()> {
775 self.try_submit_request(request)
776 .await
777 .map_err(|_| WorkerStoppedSnafu { id: self.id }.build())
778 }
779
780 async fn try_submit_request(
782 &self,
783 request: WorkerRequest,
784 ) -> std::result::Result<(), WorkerRequest> {
785 if !self.is_running() {
786 return Err(request);
787 }
788 let request_with_time = WorkerRequestWithTime::new(request);
789 if let Err(mpsc::error::SendError(request_with_time)) =
790 self.sender.send(request_with_time).await
791 {
792 warn!(
793 "Worker {} is already exited but the running flag is still true",
794 self.id
795 );
796 self.set_running(false);
798 if let Some(state) = &self.series_index_task_state {
799 state.stop();
800 }
801 return Err(request_with_time.request);
802 }
803
804 Ok(())
805 }
806
807 async fn stop(&self) -> Result<()> {
811 let handle = self.handle.lock().await.take();
812 self.set_running(false);
813 if let Some(state) = &self.series_index_task_state {
814 state.stop();
815 }
816
817 let mut worker_result = Ok(());
818 if let Some(handle) = handle {
819 info!("Stop region worker {}", self.id);
820
821 if self
822 .sender
823 .send(WorkerRequestWithTime::new(WorkerRequest::Stop))
824 .await
825 .is_err()
826 {
827 warn!("Worker {} is already exited before stop", self.id);
828 }
829
830 worker_result = handle.await.context(JoinSnafu);
831 }
832 let series_index_result = if let Some(handle) = self.series_index_handle.lock().await.take()
833 {
834 handle.await.context(JoinSnafu)
835 } else {
836 Ok(())
837 };
838 worker_result?;
839 series_index_result?;
840
841 Ok(())
842 }
843
844 fn is_running(&self) -> bool {
846 self.running.load(Ordering::Relaxed)
847 }
848
849 fn set_running(&self, value: bool) {
851 self.running.store(value, Ordering::Relaxed)
852 }
853
854 fn is_region_exists(&self, region_id: RegionId) -> bool {
856 self.regions.is_region_exists(region_id)
857 }
858
859 fn is_region_opening(&self, region_id: RegionId) -> bool {
861 self.opening_regions.is_region_exists(region_id)
862 }
863
864 fn is_region_catching_up(&self, region_id: RegionId) -> bool {
866 self.catchup_regions.is_region_exists(region_id)
867 }
868
869 fn get_region(&self, region_id: RegionId) -> Option<MitoRegionRef> {
871 self.regions.get_region(region_id)
872 }
873
874 #[cfg(test)]
875 pub(crate) fn opening_regions(&self) -> &OpeningRegionsRef {
877 &self.opening_regions
878 }
879
880 #[cfg(test)]
881 pub(crate) fn catchup_regions(&self) -> &CatchupRegionsRef {
883 &self.catchup_regions
884 }
885}
886
887impl Drop for RegionWorker {
888 fn drop(&mut self) {
889 if let Some(state) = &self.series_index_task_state {
890 state.stop();
891 }
892 if self.is_running() {
893 self.set_running(false);
894 }
896 }
897}
898
899type RequestBuffer = Vec<WorkerRequest>;
900
901#[derive(Default)]
905pub(crate) struct StalledRequests {
906 pub(crate) requests:
913 HashMap<RegionId, (usize, Vec<SenderWriteRequest>, Vec<SenderBulkRequest>)>,
914 pub(crate) estimated_size: usize,
916}
917
918impl StalledRequests {
919 pub(crate) fn estimated_size(&self, region_id: &RegionId) -> usize {
921 self.requests
922 .get(region_id)
923 .map(|(size, _, _)| *size)
924 .unwrap_or_default()
925 }
926
927 pub(crate) fn append(
929 &mut self,
930 requests: &mut Vec<SenderWriteRequest>,
931 bulk_requests: &mut Vec<SenderBulkRequest>,
932 ) {
933 for req in requests.drain(..) {
934 self.push(req);
935 }
936 for req in bulk_requests.drain(..) {
937 self.push_bulk(req);
938 }
939 }
940
941 pub(crate) fn push(&mut self, req: SenderWriteRequest) {
943 let (size, requests, _) = self.requests.entry(req.request.region_id).or_default();
944 let req_size = req.request.estimated_size();
945 *size += req_size;
946 self.estimated_size += req_size;
947 requests.push(req);
948 }
949
950 pub(crate) fn push_bulk(&mut self, req: SenderBulkRequest) {
951 let region_id = req.region_id;
952 let (size, _, requests) = self.requests.entry(region_id).or_default();
953 let req_size = req.request.estimated_size();
954 *size += req_size;
955 self.estimated_size += req_size;
956 requests.push(req);
957 }
958
959 pub(crate) fn remove(
961 &mut self,
962 region_id: &RegionId,
963 ) -> (Vec<SenderWriteRequest>, Vec<SenderBulkRequest>) {
964 if let Some((size, write_reqs, bulk_reqs)) = self.requests.remove(region_id) {
965 self.estimated_size -= size;
966 (write_reqs, bulk_reqs)
967 } else {
968 (vec![], vec![])
969 }
970 }
971
972 pub(crate) fn stalled_count(&self) -> usize {
974 self.requests
975 .values()
976 .map(|(_, reqs, bulk_reqs)| reqs.len() + bulk_reqs.len())
977 .sum()
978 }
979}
980
981struct RegionWorkerLoop<S> {
983 id: WorkerId,
985 config: Arc<MitoConfig>,
987 regions: RegionMapRef,
989 dropping_regions: RegionMapRef,
991 opening_regions: OpeningRegionsRef,
993 catchup_regions: CatchupRegionsRef,
995 sender: Sender<WorkerRequestWithTime>,
997 receiver: Receiver<WorkerRequestWithTime>,
999 wal: Wal<S>,
1001 object_store_manager: ObjectStoreManagerRef,
1003 running: Arc<AtomicBool>,
1005 memtable_builder_provider: MemtableBuilderProvider,
1007 purge_scheduler: SchedulerRef,
1009 write_buffer_manager: WriteBufferManagerRef,
1011 index_build_scheduler: IndexBuildScheduler,
1013 series_index_task_state: Option<Arc<SeriesIndexTaskState>>,
1015 series_index_store: Option<ObjectStore>,
1017 series_index_purger: Option<IndexFilePurger>,
1018 flush_scheduler: FlushScheduler,
1020 compaction_scheduler: CompactionScheduler,
1022 stalled_requests: StalledRequests,
1024 listener: WorkerListener,
1026 cache_manager: CacheManagerRef,
1028 puffin_manager_factory: PuffinManagerFactory,
1030 intermediate_manager: IntermediateManager,
1032 time_provider: TimeProviderRef,
1034 last_periodical_check_millis: i64,
1036 flush_sender: watch::Sender<()>,
1038 flush_receiver: watch::Receiver<()>,
1040 stalling_count: IntGauge,
1042 region_count: IntGauge,
1044 request_wait_time: Histogram,
1046 region_edit_queues: RegionEditQueues,
1048 schema_metadata_manager: SchemaMetadataManagerRef,
1050 file_ref_manager: FileReferenceManagerRef,
1052 partition_expr_fetcher: PartitionExprFetcherRef,
1054 flush_semaphore: Arc<Semaphore>,
1056 plugins: Plugins,
1058}
1059
1060impl<S: LogStore> RegionWorkerLoop<S> {
1061 fn register_region(&mut self, region: MitoRegionRef) {
1063 self.compaction_scheduler
1064 .register_region(®ion.version_control, ®ion.access_layer);
1065 self.regions.insert_region(region);
1066 }
1067
1068 async fn run(&mut self) {
1070 let init_check_delay = worker_init_check_delay();
1071 info!(
1072 "Start region worker thread {}, init_check_delay: {:?}",
1073 self.id, init_check_delay
1074 );
1075 self.last_periodical_check_millis += init_check_delay.as_millis() as i64;
1076
1077 let mut write_req_buffer: Vec<SenderWriteRequest> =
1079 Vec::with_capacity(self.config.worker_request_batch_size);
1080 let mut bulk_req_buffer: Vec<SenderBulkRequest> =
1081 Vec::with_capacity(self.config.worker_request_batch_size);
1082 let mut ddl_req_buffer: Vec<SenderDdlRequest> =
1083 Vec::with_capacity(self.config.worker_request_batch_size);
1084 let mut general_req_buffer: Vec<WorkerRequest> =
1085 RequestBuffer::with_capacity(self.config.worker_request_batch_size);
1086
1087 while self.running.load(Ordering::Relaxed) {
1088 write_req_buffer.clear();
1090 ddl_req_buffer.clear();
1091 general_req_buffer.clear();
1092 let mut bulk_insert_req_num = 0;
1093
1094 let max_wait_time = self.time_provider.wait_duration(CHECK_REGION_INTERVAL);
1095 let sleep = tokio::time::sleep(max_wait_time);
1096 tokio::pin!(sleep);
1097
1098 tokio::select! {
1099 request_opt = self.receiver.recv() => {
1100 match request_opt {
1101 Some(request_with_time) => {
1102 let wait_time = request_with_time.created_at.elapsed();
1104 self.request_wait_time.observe(wait_time.as_secs_f64());
1105
1106 match request_with_time.request {
1107 WorkerRequest::Write(sender_req) => write_req_buffer.push(sender_req),
1108 WorkerRequest::BatchWrite(sender_reqs) => write_req_buffer.extend(sender_reqs),
1109 WorkerRequest::Ddl(sender_req) => ddl_req_buffer.push(sender_req),
1110 WorkerRequest::BulkInserts(bulk_insert) => {
1111 bulk_insert_req_num += 1;
1112 self.buffer_bulk_insert_request(bulk_insert, &mut bulk_req_buffer)
1113 .await;
1114 }
1115 req => general_req_buffer.push(req),
1116 }
1117 },
1118 None => break,
1120 }
1121 }
1122 recv_res = self.flush_receiver.changed() => {
1123 if recv_res.is_err() {
1124 break;
1126 } else {
1127 self.maybe_flush_worker();
1132 self.handle_stalled_requests().await;
1134 continue;
1135 }
1136 }
1137 _ = &mut sleep => {
1138 self.handle_periodical_tasks();
1140 continue;
1141 }
1142 }
1143
1144 if self.flush_receiver.has_changed().unwrap_or(false) {
1145 self.handle_stalled_requests().await;
1149 }
1150
1151 for _ in 1..self.config.worker_request_batch_size {
1153 match self.receiver.try_recv() {
1155 Ok(request_with_time) => {
1156 let wait_time = request_with_time.created_at.elapsed();
1158 self.request_wait_time.observe(wait_time.as_secs_f64());
1159
1160 match request_with_time.request {
1161 WorkerRequest::Write(sender_req) => write_req_buffer.push(sender_req),
1162 WorkerRequest::BatchWrite(sender_reqs) => {
1163 write_req_buffer.extend(sender_reqs)
1164 }
1165 WorkerRequest::Ddl(sender_req) => ddl_req_buffer.push(sender_req),
1166 WorkerRequest::BulkInserts(bulk_insert) => {
1167 bulk_insert_req_num += 1;
1168 self.buffer_bulk_insert_request(bulk_insert, &mut bulk_req_buffer)
1169 .await
1170 }
1171 req => general_req_buffer.push(req),
1172 }
1173 }
1174 Err(_) => break,
1176 }
1177 }
1178
1179 self.listener.on_recv_requests(
1180 write_req_buffer.len()
1181 + ddl_req_buffer.len()
1182 + general_req_buffer.len()
1183 + bulk_insert_req_num,
1184 );
1185
1186 self.handle_requests(
1187 &mut write_req_buffer,
1188 &mut ddl_req_buffer,
1189 &mut general_req_buffer,
1190 &mut bulk_req_buffer,
1191 )
1192 .await;
1193
1194 self.handle_periodical_tasks();
1195 }
1196
1197 self.clean().await;
1198
1199 info!("Exit region worker thread {}", self.id);
1200 }
1201
1202 async fn buffer_bulk_insert_request(
1203 &mut self,
1204 bulk_insert: BulkInsertRequest,
1205 bulk_requests: &mut Vec<SenderBulkRequest>,
1206 ) {
1207 let BulkInsertRequest {
1208 metadata,
1209 request,
1210 sender,
1211 } = bulk_insert;
1212
1213 if let Some(region_metadata) = metadata {
1214 self.handle_bulk_insert_batch(region_metadata, request, bulk_requests, sender)
1215 .await;
1216 } else {
1217 error!("Cannot find region metadata for {}", request.region_id);
1218 sender.send(
1219 error::RegionNotFoundSnafu {
1220 region_id: request.region_id,
1221 }
1222 .fail(),
1223 );
1224 }
1225 }
1226
1227 async fn handle_requests(
1231 &mut self,
1232 write_requests: &mut Vec<SenderWriteRequest>,
1233 ddl_requests: &mut Vec<SenderDdlRequest>,
1234 general_requests: &mut Vec<WorkerRequest>,
1235 bulk_requests: &mut Vec<SenderBulkRequest>,
1236 ) {
1237 for worker_req in general_requests.drain(..) {
1238 match worker_req {
1239 WorkerRequest::Write(_) | WorkerRequest::BatchWrite(_) | WorkerRequest::Ddl(_) => {
1240 continue;
1242 }
1243 WorkerRequest::BulkInserts(_) => unreachable!("bulk inserts are buffered"),
1244 WorkerRequest::Background { region_id, notify } => {
1245 if matches!(
1246 ¬ify,
1247 BackgroundNotify::RegionEdit(edit_result)
1248 if edit_result.update_region_state
1249 ) {
1250 self.handle_buffered_region_write_requests(
1259 ®ion_id,
1260 write_requests,
1261 bulk_requests,
1262 )
1263 .await;
1264 }
1265 self.handle_background_notify(region_id, notify).await;
1267 }
1268 WorkerRequest::SetRegionRoleStateGracefully {
1269 region_id,
1270 region_role_state,
1271 sender,
1272 } => {
1273 self.set_role_state_gracefully(region_id, region_role_state, sender)
1274 .await;
1275 }
1276 WorkerRequest::EditRegion(request) => {
1277 self.handle_region_edit(request);
1278 }
1279 WorkerRequest::Stop => {
1280 debug_assert!(!self.running.load(Ordering::Relaxed));
1281 }
1282 WorkerRequest::SyncRegion(req) => {
1283 self.handle_region_sync(req).await;
1284 }
1285 WorkerRequest::RemapManifests(req) => {
1286 self.handle_remap_manifests_request(req);
1287 }
1288 WorkerRequest::CopyRegionFrom(req) => {
1289 self.handle_copy_region_from_request(req);
1290 }
1291 }
1292 }
1293
1294 self.handle_write_requests(write_requests, bulk_requests, true)
1297 .await;
1298
1299 self.handle_ddl_requests(ddl_requests).await;
1300 }
1301
1302 async fn handle_ddl_requests(&mut self, ddl_requests: &mut Vec<SenderDdlRequest>) {
1304 if ddl_requests.is_empty() {
1305 return;
1306 }
1307
1308 for ddl in ddl_requests.drain(..) {
1309 let res = match ddl.request {
1310 DdlRequest::Create(req) => self.handle_create_request(ddl.region_id, req).await,
1311 DdlRequest::Drop(req) => {
1312 self.handle_drop_request(ddl.region_id, req, ddl.sender)
1313 .await;
1314 continue;
1315 }
1316 DdlRequest::Open((req, wal_entry_receiver)) => {
1317 self.handle_open_request(ddl.region_id, req, wal_entry_receiver, ddl.sender)
1318 .await;
1319 continue;
1320 }
1321 DdlRequest::OfflineCleanup(req) => {
1322 self.handle_offline_cleanup_request(ddl.region_id, req)
1323 .await
1324 }
1325 DdlRequest::Close(req) => {
1326 self.handle_close_request(ddl.region_id, req, ddl.sender)
1327 .await;
1328 continue;
1329 }
1330 DdlRequest::Alter(req) => {
1331 self.handle_alter_request(ddl.region_id, req, ddl.sender)
1332 .await;
1333 continue;
1334 }
1335 DdlRequest::Flush(req) => {
1336 self.handle_flush_request(ddl.region_id, req, ddl.sender);
1337 continue;
1338 }
1339 DdlRequest::Compact(req) => {
1340 self.handle_compaction_request(ddl.region_id, req, ddl.sender)
1341 .await;
1342 continue;
1343 }
1344 DdlRequest::BuildIndex(req) => {
1345 self.handle_build_index_request(ddl.region_id, req, ddl.sender)
1346 .await;
1347 continue;
1348 }
1349 DdlRequest::Truncate(req) => {
1350 self.handle_truncate_request(ddl.region_id, req, ddl.sender)
1351 .await;
1352 continue;
1353 }
1354 DdlRequest::Catchup((req, wal_entry_receiver)) => {
1355 self.handle_catchup_request(ddl.region_id, req, wal_entry_receiver, ddl.sender)
1356 .await;
1357 continue;
1358 }
1359 DdlRequest::EnterStaging(req) => {
1360 self.handle_enter_staging_request(
1361 ddl.region_id,
1362 req.partition_directive,
1363 ddl.sender,
1364 )
1365 .await;
1366 continue;
1367 }
1368 DdlRequest::ApplyStagingManifest(req) => {
1369 self.handle_apply_staging_manifest_request(ddl.region_id, req, ddl.sender)
1370 .await;
1371 continue;
1372 }
1373 };
1374
1375 ddl.sender.send(res);
1376 }
1377 }
1378
1379 fn handle_periodical_tasks(&mut self) {
1381 let interval = CHECK_REGION_INTERVAL.as_millis() as i64;
1382 if self
1383 .time_provider
1384 .elapsed_since(self.last_periodical_check_millis)
1385 < interval
1386 {
1387 return;
1388 }
1389
1390 self.last_periodical_check_millis = self.time_provider.current_time_millis();
1391
1392 if let Err(e) = self.flush_periodically() {
1393 error!(e; "Failed to flush regions periodically");
1394 }
1395 }
1396
1397 async fn handle_background_notify(&mut self, region_id: RegionId, notify: BackgroundNotify) {
1399 match notify {
1400 BackgroundNotify::RegionOpened { region, registered } => {
1401 self.register_region(region);
1402 self.region_count.inc();
1403
1404 let senders = self.opening_regions.remove_sender(region_id);
1407 if let Some(state) = &self.series_index_task_state {
1408 state.wake();
1409 }
1410 let _ = registered.send(());
1411 for sender in senders {
1412 sender.send(Ok(0));
1413 }
1414 }
1415 BackgroundNotify::CompactionUnit(req) => {
1416 self.handle_compaction_unit(region_id, req).await;
1417 }
1418 BackgroundNotify::CompactionDdlComplete {
1419 generation,
1420 request_index,
1421 } => {
1422 self.compaction_scheduler
1423 .on_ddl_complete(region_id, generation, request_index);
1424 }
1425 BackgroundNotify::CompactionPickFinished(req) => {
1426 self.handle_compaction_pick_finished(region_id, req).await
1427 }
1428 BackgroundNotify::FlushFinished(req) => {
1429 self.handle_flush_finished(region_id, req).await
1430 }
1431 BackgroundNotify::FlushFailed(req) => self.handle_flush_failed(region_id, req).await,
1432 BackgroundNotify::IndexBuildFinished(req) => {
1433 self.handle_index_build_finished(region_id, req).await
1434 }
1435 BackgroundNotify::IndexBuildStopped(req) => {
1436 self.handle_index_build_stopped(region_id, req).await
1437 }
1438 BackgroundNotify::IndexBuildFailed(req) => {
1439 self.handle_index_build_failed(region_id, req).await
1440 }
1441 BackgroundNotify::IndexBuildRetry(req) => {
1442 self.handle_rebuild_index(req, OptionOutputTx::new(None))
1443 .await
1444 }
1445 BackgroundNotify::CompactionFinished(req) => {
1446 self.handle_compaction_finished(region_id, req).await
1447 }
1448 BackgroundNotify::CompactionFailed(req) => self.handle_compaction_failure(req).await,
1449 BackgroundNotify::Truncate(req) => self.handle_truncate_result(req).await,
1450 BackgroundNotify::DiscardUnflushed(req) => {
1451 self.handle_discard_unflushed_result(req).await
1452 }
1453 BackgroundNotify::RegionChange(req) => {
1454 self.handle_manifest_region_change_result(req).await
1455 }
1456 BackgroundNotify::EnterStaging(req) => self.handle_enter_staging_result(req).await,
1457 BackgroundNotify::RegionEdit(req) => self.handle_region_edit_result(req).await,
1458 BackgroundNotify::CopyRegionFromFinished(req) => {
1459 self.handle_copy_region_from_finished(req)
1460 }
1461 }
1462 }
1463
1464 async fn set_role_state_gracefully(
1466 &mut self,
1467 region_id: RegionId,
1468 region_role_state: SettableRegionRoleState,
1469 sender: oneshot::Sender<SetRegionRoleStateResponse>,
1470 ) {
1471 if let Some(region) = self.regions.get_region(region_id) {
1472 if matches!(
1473 region_role_state,
1474 SettableRegionRoleState::Follower | SettableRegionRoleState::DowngradingLeader
1475 ) {
1476 region.set_role(RegionRole::DowngradingLeader);
1480 }
1481 common_runtime::spawn_global(async move {
1483 match region.set_role_state_gracefully(region_role_state).await {
1484 Ok(()) => {
1485 let last_entry_id = region.version_control.current().last_entry_id;
1486 let _ = sender.send(SetRegionRoleStateResponse::success(
1487 SetRegionRoleStateSuccess::mito(last_entry_id),
1488 ));
1489 }
1490 Err(e) => {
1491 error!(e; "Failed to set region {} role state to {:?}", region_id, region_role_state);
1492 let _ = sender.send(SetRegionRoleStateResponse::invalid_transition(
1493 BoxedError::new(e),
1494 ));
1495 }
1496 }
1497 });
1498 } else {
1499 let _ = sender.send(SetRegionRoleStateResponse::NotFound);
1500 }
1501 }
1502}
1503
1504impl<S> RegionWorkerLoop<S> {
1505 async fn clean(&self) {
1507 let regions = self.regions.list_regions();
1509 for region in regions {
1510 region.stop().await;
1511 }
1512
1513 self.regions.clear();
1514 }
1515
1516 fn notify_group(&mut self) {
1519 let _ = self.flush_sender.send(());
1521 self.flush_receiver.borrow_and_update();
1523 }
1524}
1525
1526#[derive(Default, Clone)]
1528pub(crate) struct WorkerListener {
1529 #[cfg(any(test, feature = "test"))]
1530 listener: Option<crate::engine::listener::EventListenerRef>,
1531}
1532
1533impl WorkerListener {
1534 #[cfg(any(test, feature = "test"))]
1535 pub(crate) fn new(
1536 listener: Option<crate::engine::listener::EventListenerRef>,
1537 ) -> WorkerListener {
1538 WorkerListener { listener }
1539 }
1540
1541 pub(crate) async fn on_region_open_registered(&self, _region_id: RegionId) {
1543 #[cfg(any(test, feature = "test"))]
1544 if let Some(listener) = &self.listener {
1545 listener.on_region_open_registered(_region_id).await;
1546 }
1547 }
1548
1549 pub(crate) fn on_flush_success(&self, region_id: RegionId) {
1551 #[cfg(any(test, feature = "test"))]
1552 if let Some(listener) = &self.listener {
1553 listener.on_flush_success(region_id);
1554 }
1555 let _ = region_id;
1557 }
1558
1559 pub(crate) fn on_write_stall(&self) {
1561 #[cfg(any(test, feature = "test"))]
1562 if let Some(listener) = &self.listener {
1563 listener.on_write_stall();
1564 }
1565 }
1566
1567 pub(crate) async fn on_flush_begin(&self, region_id: RegionId) {
1568 #[cfg(any(test, feature = "test"))]
1569 if let Some(listener) = &self.listener {
1570 listener.on_flush_begin(region_id).await;
1571 }
1572 let _ = region_id;
1574 }
1575
1576 pub(crate) async fn on_flush_commit_begin(&self, _region_id: RegionId) {
1577 #[cfg(any(test, feature = "test"))]
1578 if let Some(listener) = &self.listener {
1579 listener.on_flush_commit_begin(_region_id).await;
1580 }
1581 }
1582
1583 pub(crate) fn on_flush_cancel_requested(&self, _region_id: RegionId) {
1584 #[cfg(any(test, feature = "test"))]
1585 if let Some(listener) = &self.listener {
1586 listener.on_flush_cancel_requested(_region_id);
1587 }
1588 }
1589
1590 pub(crate) fn on_later_drop_begin(&self, region_id: RegionId) -> Option<Duration> {
1591 #[cfg(any(test, feature = "test"))]
1592 if let Some(listener) = &self.listener {
1593 return listener.on_later_drop_begin(region_id);
1594 }
1595 let _ = region_id;
1597 None
1598 }
1599
1600 pub(crate) fn on_later_drop_end(&self, region_id: RegionId, removed: bool) {
1602 #[cfg(any(test, feature = "test"))]
1603 if let Some(listener) = &self.listener {
1604 listener.on_later_drop_end(region_id, removed);
1605 }
1606 let _ = region_id;
1608 let _ = removed;
1609 }
1610
1611 pub(crate) async fn on_merge_ssts_finished(&self, region_id: RegionId) {
1612 #[cfg(any(test, feature = "test"))]
1613 if let Some(listener) = &self.listener {
1614 listener.on_merge_ssts_finished(region_id).await;
1615 }
1616 let _ = region_id;
1618 }
1619
1620 pub(crate) fn on_recv_requests(&self, request_num: usize) {
1621 #[cfg(any(test, feature = "test"))]
1622 if let Some(listener) = &self.listener {
1623 listener.on_recv_requests(request_num);
1624 }
1625 let _ = request_num;
1627 }
1628
1629 pub(crate) fn on_file_cache_filled(&self, _file_id: FileId) {
1630 #[cfg(any(test, feature = "test"))]
1631 if let Some(listener) = &self.listener {
1632 listener.on_file_cache_filled(_file_id);
1633 }
1634 }
1635
1636 pub(crate) fn on_compaction_scheduled(&self, _region_id: RegionId) {
1637 #[cfg(any(test, feature = "test"))]
1638 if let Some(listener) = &self.listener {
1639 listener.on_compaction_scheduled(_region_id);
1640 }
1641 }
1642
1643 pub(crate) async fn on_compaction_pick_begin(&self, _region_id: RegionId) {
1644 #[cfg(any(test, feature = "test"))]
1645 if let Some(listener) = &self.listener {
1646 listener.on_compaction_pick_begin(_region_id).await;
1647 }
1648 }
1649
1650 pub(crate) async fn on_compaction_commit_begin(&self, _region_id: RegionId) {
1651 #[cfg(any(test, feature = "test"))]
1652 if let Some(listener) = &self.listener {
1653 listener.on_compaction_commit_begin(_region_id).await;
1654 }
1655 }
1656
1657 pub(crate) async fn on_compaction_unit_merge_begin(&self, _region_id: RegionId, _plan_id: u64) {
1659 #[cfg(any(test, feature = "test"))]
1660 if let Some(listener) = &self.listener {
1661 listener
1662 .on_compaction_unit_merge_begin(_region_id, _plan_id)
1663 .await;
1664 }
1665 }
1666
1667 pub(crate) async fn on_compaction_unit_committed(&self, _region_id: RegionId, _plan_id: u64) {
1669 #[cfg(any(test, feature = "test"))]
1670 if let Some(listener) = &self.listener {
1671 listener
1672 .on_compaction_unit_committed(_region_id, _plan_id)
1673 .await;
1674 }
1675 }
1676
1677 pub(crate) fn on_compaction_unit_applied(&self, _region_id: RegionId, _plan_id: u64) {
1679 #[cfg(any(test, feature = "test"))]
1680 if let Some(listener) = &self.listener {
1681 listener.on_compaction_unit_applied(_region_id, _plan_id);
1682 }
1683 }
1684
1685 pub(crate) async fn on_compaction_result_notified(&self, _region_id: RegionId) {
1686 #[cfg(any(test, feature = "test"))]
1687 if let Some(listener) = &self.listener {
1688 listener.on_compaction_result_notified(_region_id).await;
1689 }
1690 }
1691
1692 pub(crate) fn on_compaction_cancel_requested(&self, _region_id: RegionId) {
1693 #[cfg(any(test, feature = "test"))]
1694 if let Some(listener) = &self.listener {
1695 listener.on_compaction_cancel_requested(_region_id);
1696 }
1697 }
1698
1699 pub(crate) async fn on_notify_region_change_result_begin(&self, _region_id: RegionId) {
1700 #[cfg(any(test, feature = "test"))]
1701 if let Some(listener) = &self.listener {
1702 listener
1703 .on_notify_region_change_result_begin(_region_id)
1704 .await;
1705 }
1706 }
1707
1708 pub(crate) async fn on_enter_staging_result_begin(&self, _region_id: RegionId) {
1709 #[cfg(any(test, feature = "test"))]
1710 if let Some(listener) = &self.listener {
1711 listener.on_enter_staging_result_begin(_region_id).await;
1712 }
1713 }
1714
1715 pub(crate) async fn on_index_build_finish(&self, _region_file_id: RegionFileId) {
1716 #[cfg(any(test, feature = "test"))]
1717 if let Some(listener) = &self.listener {
1718 listener.on_index_build_finish(_region_file_id).await;
1719 }
1720 }
1721
1722 pub(crate) async fn on_index_build_begin(&self, _region_file_id: RegionFileId) {
1723 #[cfg(any(test, feature = "test"))]
1724 if let Some(listener) = &self.listener {
1725 listener.on_index_build_begin(_region_file_id).await;
1726 }
1727 }
1728
1729 pub(crate) async fn on_index_build_abort(&self, _region_file_id: RegionFileId) {
1730 #[cfg(any(test, feature = "test"))]
1731 if let Some(listener) = &self.listener {
1732 listener.on_index_build_abort(_region_file_id).await;
1733 }
1734 }
1735
1736 pub(crate) async fn on_index_build_before_manifest_commit(
1737 &self,
1738 _region_file_id: RegionFileId,
1739 ) {
1740 #[cfg(any(test, feature = "test"))]
1741 if let Some(listener) = &self.listener {
1742 listener
1743 .on_index_build_before_manifest_commit(_region_file_id)
1744 .await;
1745 }
1746 }
1747
1748 pub(crate) async fn on_index_build_manifest_committed(&self, _region_file_id: RegionFileId) {
1749 #[cfg(any(test, feature = "test"))]
1750 if let Some(listener) = &self.listener {
1751 listener
1752 .on_index_build_manifest_committed(_region_file_id)
1753 .await;
1754 }
1755 }
1756}
1757
1758#[cfg(test)]
1759mod tests {
1760 use super::*;
1761 use crate::test_util::TestEnv;
1762
1763 #[test]
1764 fn test_region_id_to_index() {
1765 let num_workers = 4;
1766
1767 let region_id = RegionId::new(1, 2);
1768 let index = region_id_to_index(region_id, num_workers);
1769 assert_eq!(index, 3);
1770
1771 let region_id = RegionId::new(2, 3);
1772 let index = region_id_to_index(region_id, num_workers);
1773 assert_eq!(index, 1);
1774 }
1775
1776 #[tokio::test]
1777 async fn test_worker_group_start_stop() {
1778 for (enabled, existing_index) in [(false, false), (true, false), (false, true)] {
1779 let env = TestEnv::with_prefix("group-stop").await;
1781 let root = env.data_home().join("series_index");
1782 if existing_index {
1783 tokio::fs::create_dir_all(&root).await.unwrap();
1784 tokio::fs::write(root.join("retained"), b"index")
1785 .await
1786 .unwrap();
1787 }
1788 let group = env
1789 .create_worker_group(MitoConfig {
1790 num_workers: 4,
1791 experimental_enable_series_index: enabled,
1792 experimental_series_index_maintenance_interval: Duration::from_secs(3600),
1793 ..Default::default()
1794 })
1795 .await;
1796
1797 for worker in &group.workers {
1798 assert_eq!(worker.series_index_task_state.is_some(), enabled);
1799 assert_eq!(worker.series_index_handle.lock().await.is_some(), enabled);
1800 }
1801 assert_eq!(root.is_dir(), enabled || existing_index);
1802 if existing_index {
1803 assert_eq!(
1804 tokio::fs::read(root.join("retained")).await.unwrap(),
1805 b"index"
1806 );
1807 }
1808
1809 tokio::time::timeout(Duration::from_secs(5), group.stop())
1810 .await
1811 .expect("series-index tasks should stop without waiting for their interval")
1812 .unwrap();
1813 }
1814 }
1815}