Skip to main content

mito2/
worker.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Structs and utilities for writing regions.
16
17mod 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
94/// Identifier for a worker.
95pub(crate) type WorkerId = u32;
96
97pub(crate) const DROPPING_MARKER_FILE: &str = ".dropping";
98
99/// Interval to check whether regions should flush.
100pub(crate) const CHECK_REGION_INTERVAL: Duration = Duration::from_secs(60);
101/// Max delay to check region periodical tasks.
102pub(crate) const MAX_INITIAL_CHECK_DELAY_SECS: u64 = 60 * 3;
103
104#[cfg_attr(doc, aquamarine::aquamarine)]
105/// A fixed size group of [RegionWorkers](RegionWorker).
106///
107/// A worker group binds each region to a specific [RegionWorker] and sends
108/// requests to region's dedicated worker.
109///
110/// ```mermaid
111/// graph LR
112///
113/// RegionRequest -- Route by region id --> Worker0 & Worker1
114///
115/// subgraph MitoEngine
116///     subgraph WorkerGroup
117///         Worker0["RegionWorker 0"]
118///         Worker1["RegionWorker 1"]
119///     end
120/// end
121///
122/// Chan0[" Request channel 0"]
123/// Chan1[" Request channel 1"]
124/// WorkerThread1["RegionWorkerLoop 1"]
125///
126/// subgraph WorkerThread0["RegionWorkerLoop 0"]
127///     subgraph RegionMap["RegionMap (regions bound to worker 0)"]
128///         Region0["Region 0"]
129///         Region2["Region 2"]
130///     end
131///     Buffer0["RequestBuffer"]
132///
133///     Buffer0 -- modify regions --> RegionMap
134/// end
135///
136/// Worker0 --> Chan0
137/// Worker1 --> Chan1
138/// Chan0 --> Buffer0
139/// Chan1 --> WorkerThread1
140/// ```
141pub(crate) struct WorkerGroup {
142    /// Workers of the group.
143    workers: Vec<RegionWorker>,
144    /// Flush background job pool.
145    flush_job_pool: SchedulerRef,
146    /// Compaction background job pool.
147    compact_job_pool: SchedulerRef,
148    /// Scheduler for index build jobs.
149    index_build_job_pool: SchedulerRef,
150    /// Scheduler for file purgers.
151    purge_scheduler: SchedulerRef,
152    /// Cache.
153    cache_manager: CacheManagerRef,
154    /// File reference manager.
155    file_ref_manager: FileReferenceManagerRef,
156    /// Gc limiter to limit concurrent gc jobs.
157    gc_limiter: GcLimiterRef,
158    /// Object store manager.
159    object_store_manager: ObjectStoreManagerRef,
160    /// Puffin manager factory.
161    puffin_manager_factory: PuffinManagerFactory,
162    /// Intermediate manager.
163    intermediate_manager: IntermediateManager,
164    /// Schema metadata manager.
165    schema_metadata_manager: SchemaMetadataManagerRef,
166}
167
168impl WorkerGroup {
169    /// Starts a worker group.
170    ///
171    /// The number of workers should be power of two.
172    #[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        // We use another scheduler to avoid purge jobs blocking other jobs.
206        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    /// Stops the worker group.
294    pub(crate) async fn stop(&self) -> Result<()> {
295        info!("Stop region worker group");
296
297        // TODO(yingwen): Do we need to stop gracefully?
298        // Stops the scheduler gracefully.
299        self.compact_job_pool.stop(true).await?;
300        // Stops the scheduler gracefully.
301        self.flush_job_pool.stop(true).await?;
302        // Stops the purge scheduler gracefully.
303        self.purge_scheduler.stop(true).await?;
304        // Stops the index build job pool gracefully.
305        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    /// Submits a request to a worker in the group.
313    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    /// Submits the writes of one request, those of each worker in one message, so that the
322    /// worker handles them in one batch. The messages are sent concurrently, so a full worker
323    /// queue holds back no other worker. A write whose worker has stopped receives the error a
324    /// write submitted on its own gets.
325    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    /// Returns true if the specific region exists.
356    pub(crate) fn is_region_exists(&self, region_id: RegionId) -> bool {
357        self.worker(region_id).is_region_exists(region_id)
358    }
359
360    /// Returns true if the specific region is opening.
361    pub(crate) fn is_region_opening(&self, region_id: RegionId) -> bool {
362        self.worker(region_id).is_region_opening(region_id)
363    }
364
365    /// Returns true if the specific region is catching up.
366    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    /// Returns region of specific `region_id`.
371    ///
372    /// This method should not be public.
373    pub(crate) fn get_region(&self, region_id: RegionId) -> Option<MitoRegionRef> {
374        self.worker(region_id).get_region(region_id)
375    }
376
377    /// Returns cache of the group.
378    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    /// Get worker for specific `region_id`.
391    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// Tests methods.
421#[cfg(any(test, feature = "test"))]
422impl WorkerGroup {
423    /// Starts a worker group with `write_buffer_manager` and `listener` for tests.
424    ///
425    /// The number of workers should be power of two.
426    #[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    /// Returns the purge scheduler.
543    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
553/// Opens the fixed local series-index store only when the feature is enabled.
554async 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
597/// Computes a initial check delay for a worker.
598pub(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
603/// Worker start config.
604struct 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    /// Watch channel sender to notify workers to handle stalled requests.
623    flush_sender: watch::Sender<()>,
624    /// Watch channel receiver to wait for background flush job.
625    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    /// Starts a region worker and its background thread.
635    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
750/// Worker to write and alter regions bound to it.
751pub(crate) struct RegionWorker {
752    /// Id of the worker.
753    id: WorkerId,
754    /// Regions bound to the worker.
755    regions: RegionMapRef,
756    /// The opening regions.
757    opening_regions: OpeningRegionsRef,
758    /// The catching up regions.
759    catchup_regions: CatchupRegionsRef,
760    /// Request sender.
761    sender: Sender<WorkerRequestWithTime>,
762    /// Handle to the worker thread.
763    handle: Mutex<Option<JoinHandle<()>>>,
764    /// Handle to the sequential series-index task.
765    series_index_handle: Mutex<Option<JoinHandle<()>>>,
766    /// Controls the worker-owned series-index task.
767    series_index_task_state: Option<Arc<SeriesIndexTaskState>>,
768    /// Whether to run the worker thread.
769    running: Arc<AtomicBool>,
770}
771
772impl RegionWorker {
773    /// Submits request to background worker thread.
774    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    /// Submits request to background worker thread, or returns it if the worker is stopped.
781    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            // Manually set the running flag to false to avoid printing more warning logs.
797            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    /// Stop the worker.
808    ///
809    /// This method waits until the worker thread exists.
810    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    /// Returns true if the worker is still running.
845    fn is_running(&self) -> bool {
846        self.running.load(Ordering::Relaxed)
847    }
848
849    /// Sets whether the worker is still running.
850    fn set_running(&self, value: bool) {
851        self.running.store(value, Ordering::Relaxed)
852    }
853
854    /// Returns true if the worker contains specific region.
855    fn is_region_exists(&self, region_id: RegionId) -> bool {
856        self.regions.is_region_exists(region_id)
857    }
858
859    /// Returns true if the region is opening.
860    fn is_region_opening(&self, region_id: RegionId) -> bool {
861        self.opening_regions.is_region_exists(region_id)
862    }
863
864    /// Returns true if the region is catching up.
865    fn is_region_catching_up(&self, region_id: RegionId) -> bool {
866        self.catchup_regions.is_region_exists(region_id)
867    }
868
869    /// Returns region of specific `region_id`.
870    fn get_region(&self, region_id: RegionId) -> Option<MitoRegionRef> {
871        self.regions.get_region(region_id)
872    }
873
874    #[cfg(test)]
875    /// Returns the [OpeningRegionsRef].
876    pub(crate) fn opening_regions(&self) -> &OpeningRegionsRef {
877        &self.opening_regions
878    }
879
880    #[cfg(test)]
881    /// Returns the [CatchupRegionsRef].
882    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            // Once we drop the sender, the worker thread will receive a disconnected error.
895        }
896    }
897}
898
899type RequestBuffer = Vec<WorkerRequest>;
900
901/// Buffer for stalled write requests.
902///
903/// Maintains stalled write requests and their estimated size.
904#[derive(Default)]
905pub(crate) struct StalledRequests {
906    /// Stalled requests.
907    /// Remember to use `StalledRequests::stalled_count()` to get the total number of stalled requests
908    /// instead of `StalledRequests::requests.len()`.
909    ///
910    /// Key: RegionId
911    /// Value: (estimated size, stalled requests)
912    pub(crate) requests:
913        HashMap<RegionId, (usize, Vec<SenderWriteRequest>, Vec<SenderBulkRequest>)>,
914    /// Estimated size of all stalled requests.
915    pub(crate) estimated_size: usize,
916}
917
918impl StalledRequests {
919    /// Returns the estimated size of stalled requests for a region.
920    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    /// Appends stalled requests.
928    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    /// Pushes a stalled request to the buffer.
942    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    /// Removes stalled requests of specific region.
960    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    /// Returns the total number of all stalled requests.
973    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
981/// Background worker loop to handle requests.
982struct RegionWorkerLoop<S> {
983    /// Id of the worker.
984    id: WorkerId,
985    /// Engine config.
986    config: Arc<MitoConfig>,
987    /// Regions bound to the worker.
988    regions: RegionMapRef,
989    /// Regions that are not yet fully dropped.
990    dropping_regions: RegionMapRef,
991    /// Regions that are opening.
992    opening_regions: OpeningRegionsRef,
993    /// Regions that are catching up.
994    catchup_regions: CatchupRegionsRef,
995    /// Request sender.
996    sender: Sender<WorkerRequestWithTime>,
997    /// Request receiver.
998    receiver: Receiver<WorkerRequestWithTime>,
999    /// WAL of the engine.
1000    wal: Wal<S>,
1001    /// Manages object stores for manifest and SSTs.
1002    object_store_manager: ObjectStoreManagerRef,
1003    /// Whether the worker thread is still running.
1004    running: Arc<AtomicBool>,
1005    /// Memtable builder provider for each region.
1006    memtable_builder_provider: MemtableBuilderProvider,
1007    /// Background purge job scheduler.
1008    purge_scheduler: SchedulerRef,
1009    /// Engine write buffer manager.
1010    write_buffer_manager: WriteBufferManagerRef,
1011    /// Scheduler for index build task.
1012    index_build_scheduler: IndexBuildScheduler,
1013    /// Controls the worker-owned series-index task.
1014    series_index_task_state: Option<Arc<SeriesIndexTaskState>>,
1015    /// Local store for series and range indexes managed by reconciliation.
1016    series_index_store: Option<ObjectStore>,
1017    series_index_purger: Option<IndexFilePurger>,
1018    /// Schedules background flush requests.
1019    flush_scheduler: FlushScheduler,
1020    /// Scheduler for compaction tasks.
1021    compaction_scheduler: CompactionScheduler,
1022    /// Stalled write requests.
1023    stalled_requests: StalledRequests,
1024    /// Event listener for tests.
1025    listener: WorkerListener,
1026    /// Cache.
1027    cache_manager: CacheManagerRef,
1028    /// Puffin manager factory for index.
1029    puffin_manager_factory: PuffinManagerFactory,
1030    /// Intermediate manager for inverted index.
1031    intermediate_manager: IntermediateManager,
1032    /// Provider to get current time.
1033    time_provider: TimeProviderRef,
1034    /// Last time to check regions periodically.
1035    last_periodical_check_millis: i64,
1036    /// Watch channel sender to notify workers to handle stalled requests.
1037    flush_sender: watch::Sender<()>,
1038    /// Watch channel receiver to wait for background flush job.
1039    flush_receiver: watch::Receiver<()>,
1040    /// Gauge of stalling request count.
1041    stalling_count: IntGauge,
1042    /// Gauge of regions in the worker.
1043    region_count: IntGauge,
1044    /// Histogram of request wait time for this worker.
1045    request_wait_time: Histogram,
1046    /// Queues for region edit requests.
1047    region_edit_queues: RegionEditQueues,
1048    /// Database level metadata manager.
1049    schema_metadata_manager: SchemaMetadataManagerRef,
1050    /// Datanode level file references manager.
1051    file_ref_manager: FileReferenceManagerRef,
1052    /// Partition expr fetcher used to backfill partition expr on open for compatibility.
1053    partition_expr_fetcher: PartitionExprFetcherRef,
1054    /// Semaphore to control flush concurrency.
1055    flush_semaphore: Arc<Semaphore>,
1056    /// Plugins for flush hooks.
1057    plugins: Plugins,
1058}
1059
1060impl<S: LogStore> RegionWorkerLoop<S> {
1061    /// Publishes a loaded region together with its resident compaction state.
1062    fn register_region(&mut self, region: MitoRegionRef) {
1063        self.compaction_scheduler
1064            .register_region(&region.version_control, &region.access_layer);
1065        self.regions.insert_region(region);
1066    }
1067
1068    /// Starts the worker loop.
1069    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        // Buffer to retrieve requests from receiver.
1078        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            // Clear the buffer before handling next batch of requests.
1089            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                            // Observe the wait time
1103                            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                        // The channel is disconnected.
1119                        None => break,
1120                    }
1121                }
1122                recv_res = self.flush_receiver.changed() => {
1123                    if recv_res.is_err() {
1124                        // The channel is disconnected.
1125                        break;
1126                    } else {
1127                        // Also flush this worker if other workers trigger flush as this worker may have
1128                        // a large memtable to flush. We may not have chance to flush that memtable if we
1129                        // never write to this worker. So only flushing other workers may not release enough
1130                        // memory.
1131                        self.maybe_flush_worker();
1132                        // A flush job is finished, handles stalled requests.
1133                        self.handle_stalled_requests().await;
1134                        continue;
1135                    }
1136                }
1137                _ = &mut sleep => {
1138                    // Timeout. Checks periodical tasks.
1139                    self.handle_periodical_tasks();
1140                    continue;
1141                }
1142            }
1143
1144            if self.flush_receiver.has_changed().unwrap_or(false) {
1145                // Always checks whether we could process stalled requests to avoid a request
1146                // hangs too long.
1147                // If the channel is closed, do nothing.
1148                self.handle_stalled_requests().await;
1149            }
1150
1151            // Try to recv more requests from the channel.
1152            for _ in 1..self.config.worker_request_batch_size {
1153                // We have received one request so we start from 1.
1154                match self.receiver.try_recv() {
1155                    Ok(request_with_time) => {
1156                        // Observe the wait time
1157                        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                    // We still need to handle remaining requests.
1175                    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    /// Dispatches and processes requests.
1228    ///
1229    /// `general_requests` should not contain categorized write, ddl, or bulk insert requests.
1230    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                    // These requests are categorized before dispatching general requests.
1241                    continue;
1242                }
1243                WorkerRequest::BulkInserts(_) => unreachable!("bulk inserts are buffered"),
1244                WorkerRequest::Background { region_id, notify } => {
1245                    if matches!(
1246                        &notify,
1247                        BackgroundNotify::RegionEdit(edit_result)
1248                            if edit_result.update_region_state
1249                    ) {
1250                        // Region state must be Editing when reach here.
1251                        // This call only moves write/bulk write request into stall queue. When region edit result
1252                        // is processed inside handle_background_notify and region state is switched back to Writable,
1253                        // stalled request will be processed before the next region edit is dequeued from
1254                        // RegionEditQueue immediately in handle_region_edit_result. It not only ensured pending writes
1255                        // are processed in time, but also prevents them from starvation.
1256                        // TODO(hl): maybe we need to merge those queues for pending requests like pending_ddl,
1257                        // region edits and stalled request, so we can simplify the coordination between these queues.
1258                        self.handle_buffered_region_write_requests(
1259                            &region_id,
1260                            write_requests,
1261                            bulk_requests,
1262                        )
1263                        .await;
1264                    }
1265                    // For background notify, we handle it directly.
1266                    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        // Handles all write requests first. So we can alter regions without
1295        // considering existing write requests.
1296        self.handle_write_requests(write_requests, bulk_requests, true)
1297            .await;
1298
1299        self.handle_ddl_requests(ddl_requests).await;
1300    }
1301
1302    /// Takes and handles all ddl requests.
1303    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    /// Handle periodical tasks such as region auto flush.
1380    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    /// Handles region background request
1398    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                // Finish this opening batch before the worker can process a close/reopen,
1405                // without waiting for the background opener to resume after the ACK.
1406                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    /// Handles `set_region_role_gracefully`.
1465    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                // Fence writes here: writes the worker handles later see
1477                // Downgrading, and earlier ones have finished. This only
1478                // changes a writable or staging leader.
1479                region.set_role(RegionRole::DowngradingLeader);
1480            }
1481            // We need to do this in background as we need the manifest lock.
1482            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    /// Cleans up the worker.
1506    async fn clean(&self) {
1507        // Closes remaining regions.
1508        let regions = self.regions.list_regions();
1509        for region in regions {
1510            region.stop().await;
1511        }
1512
1513        self.regions.clear();
1514    }
1515
1516    /// Notifies the whole group that a flush job is finished so other
1517    /// workers can handle stalled requests.
1518    fn notify_group(&mut self) {
1519        // Notifies all receivers.
1520        let _ = self.flush_sender.send(());
1521        // Marks the receiver in current worker as seen so the loop won't be waked up immediately.
1522        self.flush_receiver.borrow_and_update();
1523    }
1524}
1525
1526/// Wrapper that only calls event listener in tests.
1527#[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    /// Notifies the background opener after the worker acknowledges region registration.
1542    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    /// Flush is finished successfully.
1550    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        // Avoid compiler warning.
1556        let _ = region_id;
1557    }
1558
1559    /// Engine is stalled.
1560    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        // Avoid compiler warning.
1573        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        // Avoid compiler warning.
1596        let _ = region_id;
1597        None
1598    }
1599
1600    /// On later drop task is finished.
1601    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        // Avoid compiler warning.
1607        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        // Avoid compiler warning.
1617        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        // Avoid compiler warning.
1626        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    /// Exposes a per-unit test gate before merge starts.
1658    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    /// Exposes a per-unit test gate between manifest commit and worker notification.
1668    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    /// Reports unit visibility to the test listener independently of request completion.
1678    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            // Use a fresh WAL directory for each case: stopping workers does not stop the log store.
1780            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}