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    SetRegionRoleStateResponse, SetRegionRoleStateSuccess, SettableRegionRoleState,
56};
57use store_api::storage::{FileId, RegionId};
58use tokio::sync::mpsc::{Receiver, Sender};
59use tokio::sync::{Mutex, Semaphore, mpsc, oneshot, watch};
60
61use crate::access_layer::new_fs_cache_store;
62use crate::cache::write_cache::{WriteCache, WriteCacheRef};
63use crate::cache::{CacheManager, CacheManagerRef, WriteCacheUploadStoreWrapperRef};
64use crate::compaction::memory_manager::{CompactionMemoryManager, new_compaction_memory_manager};
65use crate::compaction::{CompactionCoordinator, CompactionScheduler};
66use crate::config::MitoConfig;
67use crate::error::{self, CreateDirSnafu, JoinSnafu, Result, WorkerStoppedSnafu};
68use crate::flush::{FlushScheduler, WriteBufferManagerImpl, WriteBufferManagerRef};
69use crate::gc::{GcLimiter, GcLimiterRef};
70use crate::lease::NodeLease;
71use crate::memtable::MemtableBuilderProvider;
72use crate::memtable::bulk::CompactDispatcher;
73use crate::metrics::{REGION_COUNT, REQUEST_WAIT_TIME, WRITE_STALLING};
74use crate::region::opener::PartitionExprFetcherRef;
75use crate::region::{
76    CatchupRegions, CatchupRegionsRef, MitoRegionRef, OpeningRegions, OpeningRegionsRef, RegionMap,
77    RegionMapRef,
78};
79use crate::request::{
80    BackgroundNotify, BulkInsertRequest, DdlRequest, OptionOutputTx, SenderBulkRequest,
81    SenderDdlRequest, SenderWriteRequest, WorkerRequest, WorkerRequestWithTime,
82};
83use crate::schedule::scheduler::{LocalScheduler, SchedulerRef};
84use crate::series_index::{
85    IndexFilePurger, SeriesIndexTaskState, series_index_channel, spawn_series_index_tasks,
86};
87use crate::sst::file::RegionFileId;
88use crate::sst::file_ref::FileReferenceManagerRef;
89use crate::sst::index::IndexBuildScheduler;
90use crate::sst::index::intermediate::IntermediateManager;
91use crate::sst::index::puffin_manager::PuffinManagerFactory;
92#[cfg(feature = "hdfs-object-store")]
93use crate::temp_file_sweeper;
94use crate::time_provider::{StdTimeProvider, TimeProviderRef};
95use crate::wal::Wal;
96use crate::worker::handle_flush::StagingRegionsRef;
97use crate::worker::handle_manifest::RegionEditQueues;
98
99/// Identifier for a worker.
100pub(crate) type WorkerId = u32;
101
102pub(crate) const DROPPING_MARKER_FILE: &str = ".dropping";
103
104/// Interval to check whether regions should flush.
105pub(crate) const CHECK_REGION_INTERVAL: Duration = Duration::from_secs(60);
106/// Max delay to check region periodical tasks.
107pub(crate) const MAX_INITIAL_CHECK_DELAY_SECS: u64 = 60 * 3;
108
109#[cfg_attr(doc, aquamarine::aquamarine)]
110/// A fixed size group of [RegionWorkers](RegionWorker).
111///
112/// A worker group binds each region to a specific [RegionWorker] and sends
113/// requests to region's dedicated worker.
114///
115/// ```mermaid
116/// graph LR
117///
118/// RegionRequest -- Route by region id --> Worker0 & Worker1
119///
120/// subgraph MitoEngine
121///     subgraph WorkerGroup
122///         Worker0["RegionWorker 0"]
123///         Worker1["RegionWorker 1"]
124///     end
125/// end
126///
127/// Chan0[" Request channel 0"]
128/// Chan1[" Request channel 1"]
129/// WorkerThread1["RegionWorkerLoop 1"]
130///
131/// subgraph WorkerThread0["RegionWorkerLoop 0"]
132///     subgraph RegionMap["RegionMap (regions bound to worker 0)"]
133///         Region0["Region 0"]
134///         Region2["Region 2"]
135///     end
136///     Buffer0["RequestBuffer"]
137///
138///     Buffer0 -- modify regions --> RegionMap
139/// end
140///
141/// Worker0 --> Chan0
142/// Worker1 --> Chan1
143/// Chan0 --> Buffer0
144/// Chan1 --> WorkerThread1
145/// ```
146pub(crate) struct WorkerGroup {
147    /// Workers of the group.
148    workers: Vec<RegionWorker>,
149    /// Flush background job pool.
150    flush_job_pool: SchedulerRef,
151    /// Compaction background job pool.
152    compact_job_pool: Arc<CompactionCoordinator>,
153    /// Scheduler for index build jobs.
154    index_build_job_pool: SchedulerRef,
155    /// Scheduler for file purgers.
156    purge_scheduler: SchedulerRef,
157    /// Cache.
158    cache_manager: CacheManagerRef,
159    /// File reference manager.
160    file_ref_manager: FileReferenceManagerRef,
161    /// Gc limiter to limit concurrent gc jobs.
162    gc_limiter: GcLimiterRef,
163    node_lease: Arc<NodeLease>,
164    /// Object store manager.
165    object_store_manager: ObjectStoreManagerRef,
166    /// Puffin manager factory.
167    puffin_manager_factory: PuffinManagerFactory,
168    /// Intermediate manager.
169    intermediate_manager: IntermediateManager,
170    /// Schema metadata manager.
171    schema_metadata_manager: SchemaMetadataManagerRef,
172}
173
174impl WorkerGroup {
175    /// Starts a worker group.
176    ///
177    /// The number of workers should be power of two.
178    #[allow(clippy::too_many_arguments)]
179    pub(crate) async fn start<S: LogStore>(
180        data_home: &str,
181        config: Arc<MitoConfig>,
182        log_store: Arc<S>,
183        object_store_manager: ObjectStoreManagerRef,
184        schema_metadata_manager: SchemaMetadataManagerRef,
185        file_ref_manager: FileReferenceManagerRef,
186        partition_expr_fetcher: PartitionExprFetcherRef,
187        plugins: Plugins,
188    ) -> Result<WorkerGroup> {
189        let (flush_sender, flush_receiver) = watch::channel(());
190        let write_buffer_manager = Arc::new(
191            WriteBufferManagerImpl::new(config.global_write_buffer_size.as_bytes() as usize)
192                .with_notifier(flush_sender.clone()),
193        );
194        let puffin_manager_factory = PuffinManagerFactory::new(
195            &config.index.aux_path,
196            config.index.staging_size.as_bytes(),
197            Some(config.index.write_buffer_size.as_bytes() as _),
198            config.index.staging_ttl,
199        )
200        .await?;
201        let intermediate_manager = IntermediateManager::init_fs(&config.index.aux_path)
202            .await?
203            .with_buffer_size(Some(config.index.write_buffer_size.as_bytes() as _));
204        let index_build_job_pool =
205            Arc::new(LocalScheduler::new(config.max_background_index_builds));
206        let series_index_store = series_index_store_from_config(&config, data_home).await?;
207        let series_index_disk_usage = Arc::new(AtomicU64::new(0));
208        let flush_job_pool = Arc::new(LocalScheduler::new(config.max_background_flushes));
209        let compact_job_pool = Arc::new(CompactionCoordinator::new_local(
210            config.max_background_compactions,
211        ));
212        let memtable_compact_dispatcher =
213            Arc::new(CompactDispatcher::new(config.max_background_compactions));
214        let flush_semaphore = Arc::new(Semaphore::new(config.max_background_flushes));
215        // We use another scheduler to avoid purge jobs blocking other jobs.
216        let purge_scheduler = Arc::new(LocalScheduler::new(config.max_background_purges));
217        let upload_store_wrapper = plugins.get::<WriteCacheUploadStoreWrapperRef>();
218        let write_cache = write_cache_from_config(
219            &config,
220            puffin_manager_factory.clone(),
221            intermediate_manager.clone(),
222            upload_store_wrapper,
223        )
224        .await?;
225        let cache_manager = Arc::new(
226            CacheManager::builder()
227                .sst_meta_cache_size(config.sst_meta_cache_size.as_bytes())
228                .vector_cache_size(config.vector_cache_size.as_bytes())
229                .page_cache_size(config.page_cache_size.as_bytes())
230                .selector_result_cache_size(config.selector_result_cache_size.as_bytes())
231                .range_result_cache_size(config.range_result_cache_size.as_bytes())
232                .prefilter_result_cache_size(config.prefilter_result_cache_size.as_bytes())
233                .index_metadata_size(config.index.metadata_cache_size.as_bytes())
234                .index_content_size(config.index.content_cache_size.as_bytes())
235                .index_content_page_size(config.index.content_cache_page_size.as_bytes())
236                .index_result_cache_size(config.index.result_cache_size.as_bytes())
237                .puffin_metadata_size(config.index.metadata_cache_size.as_bytes())
238                .write_cache(write_cache)
239                .build(),
240        );
241        let time_provider = Arc::new(StdTimeProvider);
242        let total_memory = get_total_memory_bytes();
243        let total_memory = if total_memory > 0 {
244            total_memory as u64
245        } else {
246            0
247        };
248        let compaction_limit_bytes = config
249            .experimental_compaction_memory_limit
250            .resolve(total_memory);
251        let compaction_memory_manager =
252            Arc::new(new_compaction_memory_manager(compaction_limit_bytes));
253        let gc_limiter = Arc::new(GcLimiter::new(config.gc.max_concurrent_gc_job));
254        // A sweep waits for a permit while a GC job takes one without waiting, so
255        // sharing the GC limiter would fail GC jobs and report sweeps as GC jobs.
256        // A datanode that allows no concurrent GC job runs no sweep either: a
257        // limiter of zero would never hand out the permit the task waits for.
258        #[cfg(feature = "hdfs-object-store")]
259        let sweep_limiter = (config.gc.max_concurrent_gc_job > 0)
260            .then(|| Arc::new(Semaphore::new(config.gc.max_concurrent_gc_job)));
261        let node_lease = Arc::new(NodeLease::default());
262        let staging_regions = StagingRegionsRef::default();
263
264        let workers = (0..config.num_workers)
265            .map(|id| {
266                WorkerStarter {
267                    id: id as WorkerId,
268                    config: config.clone(),
269                    log_store: log_store.clone(),
270                    object_store_manager: object_store_manager.clone(),
271                    write_buffer_manager: write_buffer_manager.clone(),
272                    index_build_job_pool: index_build_job_pool.clone(),
273                    series_index_store: series_index_store.clone(),
274                    series_index_disk_usage: series_index_disk_usage.clone(),
275                    flush_job_pool: flush_job_pool.clone(),
276                    compact_job_pool: compact_job_pool.clone(),
277                    purge_scheduler: purge_scheduler.clone(),
278                    memtable_compact_dispatcher: memtable_compact_dispatcher.clone(),
279                    listener: WorkerListener::default(),
280                    cache_manager: cache_manager.clone(),
281                    compaction_memory_manager: compaction_memory_manager.clone(),
282                    puffin_manager_factory: puffin_manager_factory.clone(),
283                    intermediate_manager: intermediate_manager.clone(),
284                    time_provider: time_provider.clone(),
285                    flush_sender: flush_sender.clone(),
286                    flush_receiver: flush_receiver.clone(),
287                    plugins: plugins.clone(),
288                    schema_metadata_manager: schema_metadata_manager.clone(),
289                    file_ref_manager: file_ref_manager.clone(),
290                    partition_expr_fetcher: partition_expr_fetcher.clone(),
291                    flush_semaphore: flush_semaphore.clone(),
292                    #[cfg(feature = "hdfs-object-store")]
293                    sweep_limiter: sweep_limiter.clone(),
294                    node_lease: node_lease.clone(),
295                    staging_regions: staging_regions.clone(),
296                }
297                .start()
298            })
299            .collect::<Result<Vec<_>>>()?;
300
301        Ok(WorkerGroup {
302            workers,
303            flush_job_pool,
304            compact_job_pool,
305            index_build_job_pool,
306            purge_scheduler,
307            cache_manager,
308            file_ref_manager,
309            gc_limiter,
310            node_lease,
311            object_store_manager,
312            puffin_manager_factory,
313            intermediate_manager,
314            schema_metadata_manager,
315        })
316    }
317
318    /// Stops the worker group.
319    pub(crate) async fn stop(&self) -> Result<()> {
320        info!("Stop region worker group");
321
322        // TODO(yingwen): Do we need to stop gracefully?
323        // Stops the scheduler gracefully.
324        self.compact_job_pool.stop().await?;
325        // Stops the scheduler gracefully.
326        self.flush_job_pool.stop(true).await?;
327        // Stops the purge scheduler gracefully.
328        self.purge_scheduler.stop(true).await?;
329        // Stops the index build job pool gracefully.
330        self.index_build_job_pool.stop(true).await?;
331
332        try_join_all(self.workers.iter().map(|worker| worker.stop())).await?;
333
334        Ok(())
335    }
336
337    /// Submits a request to a worker in the group.
338    pub(crate) async fn submit_to_worker(
339        &self,
340        region_id: RegionId,
341        request: WorkerRequest,
342    ) -> Result<()> {
343        self.worker(region_id).submit_request(request).await
344    }
345
346    /// Submits the writes of one request, those of each worker in one message, so that the
347    /// worker handles them in one batch. The messages are sent concurrently, so a full worker
348    /// queue holds back no other worker. A write whose worker has stopped receives the error a
349    /// write submitted on its own gets.
350    pub(crate) async fn submit_writes(&self, writes: Vec<SenderWriteRequest>) {
351        let mut writes_by_worker: HashMap<usize, Vec<SenderWriteRequest>> = HashMap::new();
352        for write in writes {
353            let index = region_id_to_index(write.request.region_id, self.workers.len());
354            writes_by_worker.entry(index).or_default().push(write);
355        }
356        join_all(
357            writes_by_worker
358                .into_iter()
359                .map(|(index, writes)| async move {
360                    let worker = &self.workers[index];
361                    let Err(request) = worker
362                        .try_submit_request(WorkerRequest::BatchWrite(writes))
363                        .await
364                    else {
365                        return;
366                    };
367                    let WorkerRequest::BatchWrite(writes) = request else {
368                        unreachable!("the worker returns the request it was given");
369                    };
370                    for write in writes {
371                        write
372                            .sender
373                            .send(WorkerStoppedSnafu { id: worker.id }.fail());
374                    }
375                }),
376        )
377        .await;
378    }
379
380    /// Returns true if the specific region exists.
381    pub(crate) fn is_region_exists(&self, region_id: RegionId) -> bool {
382        self.worker(region_id).is_region_exists(region_id)
383    }
384
385    /// Returns true if the specific region is opening.
386    pub(crate) fn is_region_opening(&self, region_id: RegionId) -> bool {
387        self.worker(region_id).is_region_opening(region_id)
388    }
389
390    /// Returns true if the specific region is catching up.
391    pub(crate) fn is_region_catching_up(&self, region_id: RegionId) -> bool {
392        self.worker(region_id).is_region_catching_up(region_id)
393    }
394
395    /// Returns region of specific `region_id`.
396    ///
397    /// This method should not be public.
398    pub(crate) fn get_region(&self, region_id: RegionId) -> Option<MitoRegionRef> {
399        self.worker(region_id).get_region(region_id)
400    }
401
402    /// Returns cache of the group.
403    pub(crate) fn cache_manager(&self) -> CacheManagerRef {
404        self.cache_manager.clone()
405    }
406
407    pub(crate) fn file_ref_manager(&self) -> FileReferenceManagerRef {
408        self.file_ref_manager.clone()
409    }
410
411    pub(crate) fn gc_limiter(&self) -> GcLimiterRef {
412        self.gc_limiter.clone()
413    }
414
415    pub(crate) fn node_lease(&self) -> &NodeLease {
416        &self.node_lease
417    }
418
419    /// Get worker for specific `region_id`.
420    pub(crate) fn worker(&self, region_id: RegionId) -> &RegionWorker {
421        let index = region_id_to_index(region_id, self.workers.len());
422
423        &self.workers[index]
424    }
425
426    pub(crate) fn all_regions(&self) -> impl Iterator<Item = MitoRegionRef> + use<'_> {
427        self.workers
428            .iter()
429            .flat_map(|worker| worker.regions.list_regions())
430    }
431
432    pub(crate) fn object_store_manager(&self) -> &ObjectStoreManagerRef {
433        &self.object_store_manager
434    }
435
436    pub(crate) fn puffin_manager_factory(&self) -> &PuffinManagerFactory {
437        &self.puffin_manager_factory
438    }
439
440    pub(crate) fn intermediate_manager(&self) -> &IntermediateManager {
441        &self.intermediate_manager
442    }
443
444    pub(crate) fn schema_metadata_manager(&self) -> &SchemaMetadataManagerRef {
445        &self.schema_metadata_manager
446    }
447}
448
449// Tests methods.
450#[cfg(any(test, feature = "test"))]
451impl WorkerGroup {
452    /// Starts a worker group with `write_buffer_manager` and `listener` for tests.
453    ///
454    /// The number of workers should be power of two.
455    #[allow(clippy::too_many_arguments)]
456    pub(crate) async fn start_for_test<S: LogStore>(
457        data_home: &str,
458        config: Arc<MitoConfig>,
459        log_store: Arc<S>,
460        object_store_manager: ObjectStoreManagerRef,
461        write_buffer_manager: Option<WriteBufferManagerRef>,
462        listener: Option<crate::engine::listener::EventListenerRef>,
463        schema_metadata_manager: SchemaMetadataManagerRef,
464        file_ref_manager: FileReferenceManagerRef,
465        time_provider: TimeProviderRef,
466        partition_expr_fetcher: PartitionExprFetcherRef,
467    ) -> Result<WorkerGroup> {
468        let (flush_sender, flush_receiver) = watch::channel(());
469        let write_buffer_manager = write_buffer_manager.unwrap_or_else(|| {
470            Arc::new(
471                WriteBufferManagerImpl::new(config.global_write_buffer_size.as_bytes() as usize)
472                    .with_notifier(flush_sender.clone()),
473            )
474        });
475        let index_build_job_pool =
476            Arc::new(LocalScheduler::new(config.max_background_index_builds));
477        let series_index_store = series_index_store_from_config(&config, data_home).await?;
478        let series_index_disk_usage = Arc::new(AtomicU64::new(0));
479        let flush_job_pool = Arc::new(LocalScheduler::new(config.max_background_flushes));
480        let memtable_compact_dispatcher =
481            Arc::new(CompactDispatcher::new(config.max_background_compactions));
482        let compact_job_pool = Arc::new(CompactionCoordinator::new_local(
483            config.max_background_compactions,
484        ));
485        let flush_semaphore = Arc::new(Semaphore::new(config.max_background_flushes));
486        let purge_scheduler = Arc::new(LocalScheduler::new(config.max_background_flushes));
487        let puffin_manager_factory = PuffinManagerFactory::new(
488            &config.index.aux_path,
489            config.index.staging_size.as_bytes(),
490            Some(config.index.write_buffer_size.as_bytes() as _),
491            config.index.staging_ttl,
492        )
493        .await?;
494        let intermediate_manager = IntermediateManager::init_fs(&config.index.aux_path)
495            .await?
496            .with_buffer_size(Some(config.index.write_buffer_size.as_bytes() as _));
497        let write_cache = write_cache_from_config(
498            &config,
499            puffin_manager_factory.clone(),
500            intermediate_manager.clone(),
501            None,
502        )
503        .await?;
504        let cache_manager = Arc::new(
505            CacheManager::builder()
506                .sst_meta_cache_size(config.sst_meta_cache_size.as_bytes())
507                .vector_cache_size(config.vector_cache_size.as_bytes())
508                .page_cache_size(config.page_cache_size.as_bytes())
509                .selector_result_cache_size(config.selector_result_cache_size.as_bytes())
510                .range_result_cache_size(config.range_result_cache_size.as_bytes())
511                .prefilter_result_cache_size(config.prefilter_result_cache_size.as_bytes())
512                .write_cache(write_cache)
513                .build(),
514        );
515        let total_memory = get_total_memory_bytes();
516        let total_memory = if total_memory > 0 {
517            total_memory as u64
518        } else {
519            0
520        };
521        let compaction_limit_bytes = config
522            .experimental_compaction_memory_limit
523            .resolve(total_memory);
524        let compaction_memory_manager =
525            Arc::new(new_compaction_memory_manager(compaction_limit_bytes));
526        let gc_limiter = Arc::new(GcLimiter::new(config.gc.max_concurrent_gc_job));
527        // A sweep waits for a permit while a GC job takes one without waiting, so
528        // sharing the GC limiter would fail GC jobs and report sweeps as GC jobs.
529        // A datanode that allows no concurrent GC job runs no sweep either: a
530        // limiter of zero would never hand out the permit the task waits for.
531        #[cfg(feature = "hdfs-object-store")]
532        let sweep_limiter = (config.gc.max_concurrent_gc_job > 0)
533            .then(|| Arc::new(Semaphore::new(config.gc.max_concurrent_gc_job)));
534        let node_lease = Arc::new(NodeLease::default());
535        let staging_regions = StagingRegionsRef::default();
536        let workers = (0..config.num_workers)
537            .map(|id| {
538                WorkerStarter {
539                    id: id as WorkerId,
540                    config: config.clone(),
541                    log_store: log_store.clone(),
542                    object_store_manager: object_store_manager.clone(),
543                    write_buffer_manager: write_buffer_manager.clone(),
544                    index_build_job_pool: index_build_job_pool.clone(),
545                    series_index_store: series_index_store.clone(),
546                    series_index_disk_usage: series_index_disk_usage.clone(),
547                    flush_job_pool: flush_job_pool.clone(),
548                    compact_job_pool: compact_job_pool.clone(),
549                    purge_scheduler: purge_scheduler.clone(),
550                    memtable_compact_dispatcher: memtable_compact_dispatcher.clone(),
551                    listener: WorkerListener::new(listener.clone()),
552                    cache_manager: cache_manager.clone(),
553                    compaction_memory_manager: compaction_memory_manager.clone(),
554                    puffin_manager_factory: puffin_manager_factory.clone(),
555                    intermediate_manager: intermediate_manager.clone(),
556                    time_provider: time_provider.clone(),
557                    flush_sender: flush_sender.clone(),
558                    flush_receiver: flush_receiver.clone(),
559                    plugins: Plugins::new(),
560                    schema_metadata_manager: schema_metadata_manager.clone(),
561                    file_ref_manager: file_ref_manager.clone(),
562                    partition_expr_fetcher: partition_expr_fetcher.clone(),
563                    flush_semaphore: flush_semaphore.clone(),
564                    #[cfg(feature = "hdfs-object-store")]
565                    sweep_limiter: sweep_limiter.clone(),
566                    node_lease: node_lease.clone(),
567                    staging_regions: staging_regions.clone(),
568                }
569                .start()
570            })
571            .collect::<Result<Vec<_>>>()?;
572
573        Ok(WorkerGroup {
574            workers,
575            flush_job_pool,
576            compact_job_pool,
577            index_build_job_pool,
578            purge_scheduler,
579            cache_manager,
580            file_ref_manager,
581            gc_limiter,
582            node_lease,
583            object_store_manager,
584            puffin_manager_factory,
585            intermediate_manager,
586            schema_metadata_manager,
587        })
588    }
589
590    /// Returns the purge scheduler.
591    pub(crate) fn purge_scheduler(&self) -> &SchedulerRef {
592        &self.purge_scheduler
593    }
594}
595
596fn region_id_to_index(id: RegionId, num_workers: usize) -> usize {
597    ((id.table_id() as usize % num_workers) + (id.region_number() as usize % num_workers))
598        % num_workers
599}
600
601/// Opens the fixed local series-index store only when the feature is enabled.
602async fn series_index_store_from_config(
603    config: &MitoConfig,
604    data_home: &str,
605) -> Result<Option<ObjectStore>> {
606    if !config.experimental_enable_series_index {
607        return Ok(None);
608    }
609
610    let root = Path::new(data_home).join("series_index");
611    new_fs_cache_store(&root.to_string_lossy()).await.map(Some)
612}
613
614pub async fn write_cache_from_config(
615    config: &MitoConfig,
616    puffin_manager_factory: PuffinManagerFactory,
617    intermediate_manager: IntermediateManager,
618    upload_store_wrapper: Option<WriteCacheUploadStoreWrapperRef>,
619) -> Result<Option<WriteCacheRef>> {
620    if !config.enable_write_cache {
621        return Ok(None);
622    }
623
624    tokio::fs::create_dir_all(Path::new(&config.write_cache_path))
625        .await
626        .context(CreateDirSnafu {
627            dir: &config.write_cache_path,
628        })?;
629
630    let cache = WriteCache::new_fs(
631        &config.write_cache_path,
632        config.write_cache_size,
633        config.write_cache_ttl,
634        Some(config.index_cache_percent),
635        config.enable_refill_cache_on_read,
636        puffin_manager_factory,
637        intermediate_manager,
638        config.manifest_cache_size,
639    )
640    .await?
641    .with_upload_store_wrapper(upload_store_wrapper);
642    Ok(Some(Arc::new(cache)))
643}
644
645/// Computes a initial check delay for a worker.
646pub(crate) fn worker_init_check_delay() -> Duration {
647    let init_check_delay = rng().random_range(0..MAX_INITIAL_CHECK_DELAY_SECS);
648    Duration::from_secs(init_check_delay)
649}
650
651/// Worker start config.
652struct WorkerStarter<S> {
653    id: WorkerId,
654    config: Arc<MitoConfig>,
655    log_store: Arc<S>,
656    object_store_manager: ObjectStoreManagerRef,
657    write_buffer_manager: WriteBufferManagerRef,
658    compact_job_pool: Arc<CompactionCoordinator>,
659    index_build_job_pool: SchedulerRef,
660    series_index_store: Option<ObjectStore>,
661    series_index_disk_usage: Arc<AtomicU64>,
662    flush_job_pool: SchedulerRef,
663    purge_scheduler: SchedulerRef,
664    memtable_compact_dispatcher: Arc<CompactDispatcher>,
665    listener: WorkerListener,
666    cache_manager: CacheManagerRef,
667    compaction_memory_manager: Arc<CompactionMemoryManager>,
668    puffin_manager_factory: PuffinManagerFactory,
669    intermediate_manager: IntermediateManager,
670    time_provider: TimeProviderRef,
671    /// Watch channel sender to notify workers to handle stalled requests.
672    flush_sender: watch::Sender<()>,
673    /// Watch channel receiver to wait for background flush job.
674    flush_receiver: watch::Receiver<()>,
675    plugins: Plugins,
676    schema_metadata_manager: SchemaMetadataManagerRef,
677    file_ref_manager: FileReferenceManagerRef,
678    partition_expr_fetcher: PartitionExprFetcherRef,
679    flush_semaphore: Arc<Semaphore>,
680    #[cfg(feature = "hdfs-object-store")]
681    /// Limiter for the temporary file sweeps, or `None` if the datanode allows
682    /// no concurrent GC job.
683    sweep_limiter: Option<Arc<Semaphore>>,
684    node_lease: Arc<NodeLease>,
685    staging_regions: StagingRegionsRef,
686}
687
688impl<S: LogStore> WorkerStarter<S> {
689    /// Starts a region worker and its background thread.
690    fn start(self) -> Result<RegionWorker> {
691        let regions = Arc::new(RegionMap::default());
692        let opening_regions = Arc::new(OpeningRegions::default());
693        let catchup_regions = Arc::new(CatchupRegions::default());
694        let (sender, receiver) = mpsc::channel(self.config.worker_channel_size);
695
696        let running = Arc::new(AtomicBool::new(true));
697        let (series_index_task_state, series_index_receiver) = if self.series_index_store.is_some()
698        {
699            let (state, receiver) =
700                SeriesIndexTaskState::new(self.id, self.config.worker_channel_size);
701            (Some(Arc::new(state)), Some(receiver))
702        } else {
703            (None, None)
704        };
705        let mut series_index_purger = None;
706        let series_index_handle = self
707            .series_index_store
708            .clone()
709            .zip(series_index_task_state.clone())
710            .zip(series_index_receiver)
711            .map(|((store, state), receiver)| {
712                let (purger, purge_receiver) = series_index_channel(store.clone());
713                series_index_purger = Some(purger.clone());
714                spawn_series_index_tasks(
715                    self.id,
716                    store,
717                    regions.clone(),
718                    state,
719                    receiver,
720                    self.config.experimental_series_index_bucket_width,
721                    purger,
722                    purge_receiver,
723                    self.series_index_disk_usage.clone(),
724                    self.config.experimental_series_index_max_size.as_bytes(),
725                    self.config.experimental_series_index_maintenance_interval,
726                    self.time_provider.clone(),
727                    self.config.experimental_enable_range_index,
728                )
729            });
730        let now = self.time_provider.current_time_millis();
731        let id_string = self.id.to_string();
732        let mut worker_thread = RegionWorkerLoop {
733            id: self.id,
734            config: self.config.clone(),
735            regions: regions.clone(),
736            catchup_regions: catchup_regions.clone(),
737            dropping_regions: Arc::new(RegionMap::default()),
738            opening_regions: opening_regions.clone(),
739            sender: sender.clone(),
740            receiver,
741            wal: Wal::new(self.log_store),
742            object_store_manager: self.object_store_manager.clone(),
743            running: running.clone(),
744            memtable_builder_provider: MemtableBuilderProvider::new(
745                Some(self.write_buffer_manager.clone()),
746                self.config.clone(),
747                self.memtable_compact_dispatcher.clone(),
748            ),
749            purge_scheduler: self.purge_scheduler.clone(),
750            write_buffer_manager: self.write_buffer_manager,
751            index_build_scheduler: IndexBuildScheduler::new(
752                self.index_build_job_pool,
753                self.config.max_background_index_builds,
754            ),
755            series_index_task_state: series_index_task_state.clone(),
756            series_index_store: self.series_index_store,
757            series_index_purger,
758            flush_scheduler: FlushScheduler::new(self.flush_job_pool),
759            compaction_scheduler: CompactionScheduler::new(
760                self.compact_job_pool,
761                sender.clone(),
762                self.cache_manager.clone(),
763                self.config.clone(),
764                self.listener.clone(),
765                self.plugins.clone(),
766                self.compaction_memory_manager.clone(),
767                self.config.experimental_compaction_on_exhausted,
768            ),
769            stalled_requests: StalledRequests::default(),
770            listener: self.listener,
771            cache_manager: self.cache_manager,
772            puffin_manager_factory: self.puffin_manager_factory,
773            intermediate_manager: self.intermediate_manager,
774            time_provider: self.time_provider,
775            last_periodical_check_millis: now,
776            flush_sender: self.flush_sender,
777            flush_receiver: self.flush_receiver,
778            stalling_count: WRITE_STALLING.with_label_values(&[&id_string]),
779            region_count: REGION_COUNT.with_label_values(&[&id_string]),
780            request_wait_time: REQUEST_WAIT_TIME.with_label_values(&[&id_string]),
781            region_edit_queues: RegionEditQueues::default(),
782            pending_alters: Default::default(),
783            schema_metadata_manager: self.schema_metadata_manager,
784            file_ref_manager: self.file_ref_manager.clone(),
785            partition_expr_fetcher: self.partition_expr_fetcher,
786            flush_semaphore: self.flush_semaphore,
787            plugins: self.plugins,
788            #[cfg(feature = "hdfs-object-store")]
789            sweep_limiter: self.sweep_limiter,
790            node_lease: self.node_lease,
791            staging_regions: self.staging_regions,
792        };
793        let handle = common_runtime::spawn_global(async move {
794            worker_thread.run().await;
795        });
796
797        Ok(RegionWorker {
798            id: self.id,
799            regions,
800            opening_regions,
801            catchup_regions,
802            sender,
803            handle: Mutex::new(Some(handle)),
804            series_index_handle: Mutex::new(series_index_handle),
805            series_index_task_state,
806            running,
807            #[cfg(test)]
808            memtable_compact_dispatcher: self.memtable_compact_dispatcher.clone(),
809        })
810    }
811}
812
813/// Worker to write and alter regions bound to it.
814pub(crate) struct RegionWorker {
815    /// Id of the worker.
816    id: WorkerId,
817    /// Regions bound to the worker.
818    regions: RegionMapRef,
819    /// The opening regions.
820    opening_regions: OpeningRegionsRef,
821    /// The catching up regions.
822    catchup_regions: CatchupRegionsRef,
823    /// Request sender.
824    sender: Sender<WorkerRequestWithTime>,
825    /// Handle to the worker thread.
826    handle: Mutex<Option<JoinHandle<()>>>,
827    /// Handle to the sequential series-index task.
828    series_index_handle: Mutex<Option<JoinHandle<()>>>,
829    /// Controls the worker-owned series-index task.
830    series_index_task_state: Option<Arc<SeriesIndexTaskState>>,
831    /// Whether to run the worker thread.
832    running: Arc<AtomicBool>,
833    #[cfg(test)]
834    /// Dispatcher for bulk memtable compactions, shared by all workers in the group.
835    /// Kept only so tests can check that every worker holds the same one.
836    memtable_compact_dispatcher: Arc<CompactDispatcher>,
837}
838
839impl RegionWorker {
840    /// Submits request to background worker thread.
841    async fn submit_request(&self, request: WorkerRequest) -> Result<()> {
842        self.try_submit_request(request)
843            .await
844            .map_err(|_| WorkerStoppedSnafu { id: self.id }.build())
845    }
846
847    /// Submits request to background worker thread, or returns it if the worker is stopped.
848    async fn try_submit_request(
849        &self,
850        request: WorkerRequest,
851    ) -> std::result::Result<(), WorkerRequest> {
852        if !self.is_running() {
853            return Err(request);
854        }
855        let request_with_time = WorkerRequestWithTime::new(request);
856        if let Err(mpsc::error::SendError(request_with_time)) =
857            self.sender.send(request_with_time).await
858        {
859            warn!(
860                "Worker {} is already exited but the running flag is still true",
861                self.id
862            );
863            // Manually set the running flag to false to avoid printing more warning logs.
864            self.set_running(false);
865            if let Some(state) = &self.series_index_task_state {
866                state.stop();
867            }
868            return Err(request_with_time.request);
869        }
870
871        Ok(())
872    }
873
874    /// Stop the worker.
875    ///
876    /// This method waits until the worker thread exists.
877    async fn stop(&self) -> Result<()> {
878        let handle = self.handle.lock().await.take();
879        self.set_running(false);
880        if let Some(state) = &self.series_index_task_state {
881            state.stop();
882        }
883
884        let mut worker_result = Ok(());
885        if let Some(handle) = handle {
886            info!("Stop region worker {}", self.id);
887
888            if self
889                .sender
890                .send(WorkerRequestWithTime::new(WorkerRequest::Stop))
891                .await
892                .is_err()
893            {
894                warn!("Worker {} is already exited before stop", self.id);
895            }
896
897            worker_result = handle.await.context(JoinSnafu);
898        }
899        let series_index_result = if let Some(handle) = self.series_index_handle.lock().await.take()
900        {
901            handle.await.context(JoinSnafu)
902        } else {
903            Ok(())
904        };
905        worker_result?;
906        series_index_result?;
907
908        Ok(())
909    }
910
911    /// Returns true if the worker is still running.
912    fn is_running(&self) -> bool {
913        self.running.load(Ordering::Relaxed)
914    }
915
916    /// Sets whether the worker is still running.
917    fn set_running(&self, value: bool) {
918        self.running.store(value, Ordering::Relaxed)
919    }
920
921    /// Returns true if the worker contains specific region.
922    fn is_region_exists(&self, region_id: RegionId) -> bool {
923        self.regions.is_region_exists(region_id)
924    }
925
926    /// Returns true if the region is opening.
927    fn is_region_opening(&self, region_id: RegionId) -> bool {
928        self.opening_regions.is_region_exists(region_id)
929    }
930
931    /// Returns true if the region is catching up.
932    fn is_region_catching_up(&self, region_id: RegionId) -> bool {
933        self.catchup_regions.is_region_exists(region_id)
934    }
935
936    /// Returns region of specific `region_id`.
937    fn get_region(&self, region_id: RegionId) -> Option<MitoRegionRef> {
938        self.regions.get_region(region_id)
939    }
940
941    #[cfg(test)]
942    /// Returns the [OpeningRegionsRef].
943    pub(crate) fn opening_regions(&self) -> &OpeningRegionsRef {
944        &self.opening_regions
945    }
946
947    #[cfg(test)]
948    /// Returns the [CatchupRegionsRef].
949    pub(crate) fn catchup_regions(&self) -> &CatchupRegionsRef {
950        &self.catchup_regions
951    }
952}
953
954impl Drop for RegionWorker {
955    fn drop(&mut self) {
956        if let Some(state) = &self.series_index_task_state {
957            state.stop();
958        }
959        if self.is_running() {
960            self.set_running(false);
961            // Once we drop the sender, the worker thread will receive a disconnected error.
962        }
963    }
964}
965
966type RequestBuffer = Vec<WorkerRequest>;
967
968/// Buffer for stalled write requests.
969///
970/// Maintains stalled write requests and their estimated size.
971#[derive(Default)]
972pub(crate) struct StalledRequests {
973    /// Stalled requests.
974    /// Remember to use `StalledRequests::stalled_count()` to get the total number of stalled requests
975    /// instead of `StalledRequests::requests.len()`.
976    ///
977    /// Key: RegionId
978    /// Value: (estimated size, stalled requests)
979    pub(crate) requests:
980        HashMap<RegionId, (usize, Vec<SenderWriteRequest>, Vec<SenderBulkRequest>)>,
981    /// Estimated size of all stalled requests.
982    pub(crate) estimated_size: usize,
983}
984
985impl StalledRequests {
986    /// Returns the estimated size of stalled requests for a region.
987    pub(crate) fn estimated_size(&self, region_id: &RegionId) -> usize {
988        self.requests
989            .get(region_id)
990            .map(|(size, _, _)| *size)
991            .unwrap_or_default()
992    }
993
994    /// Appends stalled requests.
995    pub(crate) fn append(
996        &mut self,
997        requests: &mut Vec<SenderWriteRequest>,
998        bulk_requests: &mut Vec<SenderBulkRequest>,
999    ) {
1000        for req in requests.drain(..) {
1001            self.push(req);
1002        }
1003        for req in bulk_requests.drain(..) {
1004            self.push_bulk(req);
1005        }
1006    }
1007
1008    /// Pushes a stalled request to the buffer.
1009    pub(crate) fn push(&mut self, req: SenderWriteRequest) {
1010        let (size, requests, _) = self.requests.entry(req.request.region_id).or_default();
1011        let req_size = req.request.estimated_size();
1012        *size += req_size;
1013        self.estimated_size += req_size;
1014        requests.push(req);
1015    }
1016
1017    pub(crate) fn push_bulk(&mut self, req: SenderBulkRequest) {
1018        let region_id = req.region_id;
1019        let (size, _, requests) = self.requests.entry(region_id).or_default();
1020        let req_size = req.request.estimated_size();
1021        *size += req_size;
1022        self.estimated_size += req_size;
1023        requests.push(req);
1024    }
1025
1026    /// Removes stalled requests of specific region.
1027    pub(crate) fn remove(
1028        &mut self,
1029        region_id: &RegionId,
1030    ) -> (Vec<SenderWriteRequest>, Vec<SenderBulkRequest>) {
1031        if let Some((size, write_reqs, bulk_reqs)) = self.requests.remove(region_id) {
1032            self.estimated_size -= size;
1033            (write_reqs, bulk_reqs)
1034        } else {
1035            (vec![], vec![])
1036        }
1037    }
1038
1039    /// Returns the total number of all stalled requests.
1040    pub(crate) fn stalled_count(&self) -> usize {
1041        self.requests
1042            .values()
1043            .map(|(_, reqs, bulk_reqs)| reqs.len() + bulk_reqs.len())
1044            .sum()
1045    }
1046}
1047
1048/// Background worker loop to handle requests.
1049struct RegionWorkerLoop<S> {
1050    /// Id of the worker.
1051    id: WorkerId,
1052    /// Engine config.
1053    config: Arc<MitoConfig>,
1054    /// Regions bound to the worker.
1055    regions: RegionMapRef,
1056    /// Regions that are not yet fully dropped.
1057    dropping_regions: RegionMapRef,
1058    /// Regions that are opening.
1059    opening_regions: OpeningRegionsRef,
1060    /// Regions that are catching up.
1061    catchup_regions: CatchupRegionsRef,
1062    /// Request sender.
1063    sender: Sender<WorkerRequestWithTime>,
1064    /// Request receiver.
1065    receiver: Receiver<WorkerRequestWithTime>,
1066    /// WAL of the engine.
1067    wal: Wal<S>,
1068    /// Manages object stores for manifest and SSTs.
1069    object_store_manager: ObjectStoreManagerRef,
1070    /// Whether the worker thread is still running.
1071    running: Arc<AtomicBool>,
1072    /// Memtable builder provider for each region.
1073    memtable_builder_provider: MemtableBuilderProvider,
1074    /// Background purge job scheduler.
1075    purge_scheduler: SchedulerRef,
1076    /// Engine write buffer manager.
1077    write_buffer_manager: WriteBufferManagerRef,
1078    /// Scheduler for index build task.
1079    index_build_scheduler: IndexBuildScheduler,
1080    /// Controls the worker-owned series-index task.
1081    series_index_task_state: Option<Arc<SeriesIndexTaskState>>,
1082    /// Local store for series and range indexes managed by reconciliation.
1083    series_index_store: Option<ObjectStore>,
1084    series_index_purger: Option<IndexFilePurger>,
1085    /// Schedules background flush requests.
1086    flush_scheduler: FlushScheduler,
1087    /// Scheduler for compaction tasks.
1088    compaction_scheduler: CompactionScheduler,
1089    /// Stalled write requests.
1090    stalled_requests: StalledRequests,
1091    /// Event listener for tests.
1092    listener: WorkerListener,
1093    /// Cache.
1094    cache_manager: CacheManagerRef,
1095    /// Puffin manager factory for index.
1096    puffin_manager_factory: PuffinManagerFactory,
1097    /// Intermediate manager for inverted index.
1098    intermediate_manager: IntermediateManager,
1099    /// Provider to get current time.
1100    time_provider: TimeProviderRef,
1101    /// Last time to check regions periodically.
1102    last_periodical_check_millis: i64,
1103    /// Watch channel sender to notify workers to handle stalled requests.
1104    flush_sender: watch::Sender<()>,
1105    /// Watch channel receiver to wait for background flush job.
1106    flush_receiver: watch::Receiver<()>,
1107    /// Gauge of stalling request count.
1108    stalling_count: IntGauge,
1109    /// Gauge of regions in the worker.
1110    region_count: IntGauge,
1111    /// Histogram of request wait time for this worker.
1112    request_wait_time: Histogram,
1113    /// Queues for region edit requests.
1114    region_edit_queues: RegionEditQueues,
1115    /// ALTER requests waiting for the current edit or ALTER to finish.
1116    pending_alters: handle_alter::PendingAlters,
1117    /// Database level metadata manager.
1118    schema_metadata_manager: SchemaMetadataManagerRef,
1119    /// Datanode level file references manager.
1120    file_ref_manager: FileReferenceManagerRef,
1121    /// Partition expr fetcher used to backfill partition expr on open for compatibility.
1122    partition_expr_fetcher: PartitionExprFetcherRef,
1123    /// Semaphore to control flush concurrency.
1124    flush_semaphore: Arc<Semaphore>,
1125    /// Plugins for flush hooks.
1126    plugins: Plugins,
1127    #[cfg(feature = "hdfs-object-store")]
1128    /// Limiter for the temporary file sweeps, or `None` if the datanode allows
1129    /// no concurrent GC job.
1130    sweep_limiter: Option<Arc<Semaphore>>,
1131    node_lease: Arc<NodeLease>,
1132    /// Staging regions of this node that have received writes.
1133    staging_regions: StagingRegionsRef,
1134}
1135
1136impl<S: LogStore> RegionWorkerLoop<S> {
1137    /// Publishes a loaded region together with its resident compaction state.
1138    fn register_region(&mut self, region: MitoRegionRef) {
1139        if let Some(old) = self.regions.get_region(region.region_id) {
1140            old.deactivate_region_manifest_size_metric();
1141        }
1142
1143        self.compaction_scheduler.register_region(
1144            &region.version_control,
1145            &region.access_layer,
1146            &region.compaction_state,
1147        );
1148        self.regions.insert_region(region.clone());
1149        region.activate_region_manifest_size_metric();
1150    }
1151
1152    /// Starts the worker loop.
1153    async fn run(&mut self) {
1154        let init_check_delay = worker_init_check_delay();
1155        info!(
1156            "Start region worker thread {}, init_check_delay: {:?}",
1157            self.id, init_check_delay
1158        );
1159        self.last_periodical_check_millis += init_check_delay.as_millis() as i64;
1160
1161        // Buffer to retrieve requests from receiver.
1162        let mut write_req_buffer: Vec<SenderWriteRequest> =
1163            Vec::with_capacity(self.config.worker_request_batch_size);
1164        let mut bulk_req_buffer: Vec<SenderBulkRequest> =
1165            Vec::with_capacity(self.config.worker_request_batch_size);
1166        let mut ddl_req_buffer: Vec<SenderDdlRequest> =
1167            Vec::with_capacity(self.config.worker_request_batch_size);
1168        let mut general_req_buffer: Vec<WorkerRequest> =
1169            RequestBuffer::with_capacity(self.config.worker_request_batch_size);
1170
1171        while self.running.load(Ordering::Relaxed) {
1172            // Clear the buffer before handling next batch of requests.
1173            write_req_buffer.clear();
1174            ddl_req_buffer.clear();
1175            general_req_buffer.clear();
1176            let mut bulk_insert_req_num = 0;
1177
1178            let max_wait_time = self.time_provider.wait_duration(CHECK_REGION_INTERVAL);
1179            let sleep = tokio::time::sleep(max_wait_time);
1180            tokio::pin!(sleep);
1181
1182            tokio::select! {
1183                request_opt = self.receiver.recv() => {
1184                    match request_opt {
1185                        Some(request_with_time) => {
1186                            // Observe the wait time
1187                            let wait_time = request_with_time.created_at.elapsed();
1188                            self.request_wait_time.observe(wait_time.as_secs_f64());
1189
1190                            match request_with_time.request {
1191                                WorkerRequest::Write(sender_req) => write_req_buffer.push(sender_req),
1192                                WorkerRequest::BatchWrite(sender_reqs) => write_req_buffer.extend(sender_reqs),
1193                                WorkerRequest::Ddl(sender_req) => ddl_req_buffer.push(sender_req),
1194                                WorkerRequest::BulkInserts(bulk_insert) => {
1195                                    bulk_insert_req_num += 1;
1196                                    self.buffer_bulk_insert_request(bulk_insert, &mut bulk_req_buffer)
1197                                        .await;
1198                                }
1199                                req => general_req_buffer.push(req),
1200                            }
1201                        },
1202                        // The channel is disconnected.
1203                        None => break,
1204                    }
1205                }
1206                recv_res = self.flush_receiver.changed() => {
1207                    if recv_res.is_err() {
1208                        // The channel is disconnected.
1209                        break;
1210                    } else {
1211                        // Also flush this worker if other workers trigger flush as this worker may have
1212                        // a large memtable to flush. We may not have chance to flush that memtable if we
1213                        // never write to this worker. So only flushing other workers may not release enough
1214                        // memory.
1215                        self.maybe_flush_worker();
1216                        // A flush job is finished, handles stalled requests.
1217                        self.handle_stalled_requests().await;
1218                        continue;
1219                    }
1220                }
1221                _ = &mut sleep => {
1222                    // Timeout. Checks periodical tasks.
1223                    self.handle_periodical_tasks();
1224                    continue;
1225                }
1226            }
1227
1228            if self.flush_receiver.has_changed().unwrap_or(false) {
1229                // Always checks whether we could process stalled requests to avoid a request
1230                // hangs too long.
1231                // If the channel is closed, do nothing.
1232                self.handle_stalled_requests().await;
1233            }
1234
1235            // Try to recv more requests from the channel.
1236            for _ in 1..self.config.worker_request_batch_size {
1237                // We have received one request so we start from 1.
1238                match self.receiver.try_recv() {
1239                    Ok(request_with_time) => {
1240                        // Observe the wait time
1241                        let wait_time = request_with_time.created_at.elapsed();
1242                        self.request_wait_time.observe(wait_time.as_secs_f64());
1243
1244                        match request_with_time.request {
1245                            WorkerRequest::Write(sender_req) => write_req_buffer.push(sender_req),
1246                            WorkerRequest::BatchWrite(sender_reqs) => {
1247                                write_req_buffer.extend(sender_reqs)
1248                            }
1249                            WorkerRequest::Ddl(sender_req) => ddl_req_buffer.push(sender_req),
1250                            WorkerRequest::BulkInserts(bulk_insert) => {
1251                                bulk_insert_req_num += 1;
1252                                self.buffer_bulk_insert_request(bulk_insert, &mut bulk_req_buffer)
1253                                    .await
1254                            }
1255                            req => general_req_buffer.push(req),
1256                        }
1257                    }
1258                    // We still need to handle remaining requests.
1259                    Err(_) => break,
1260                }
1261            }
1262
1263            self.listener.on_recv_requests(
1264                write_req_buffer.len()
1265                    + ddl_req_buffer.len()
1266                    + general_req_buffer.len()
1267                    + bulk_insert_req_num,
1268            );
1269
1270            self.handle_requests(
1271                &mut write_req_buffer,
1272                &mut ddl_req_buffer,
1273                &mut general_req_buffer,
1274                &mut bulk_req_buffer,
1275            )
1276            .await;
1277
1278            self.handle_periodical_tasks();
1279        }
1280
1281        self.clean().await;
1282
1283        info!("Exit region worker thread {}", self.id);
1284    }
1285
1286    async fn buffer_bulk_insert_request(
1287        &mut self,
1288        bulk_insert: BulkInsertRequest,
1289        bulk_requests: &mut Vec<SenderBulkRequest>,
1290    ) {
1291        let BulkInsertRequest {
1292            metadata,
1293            request,
1294            sender,
1295        } = bulk_insert;
1296
1297        if let Some(region_metadata) = metadata {
1298            self.handle_bulk_insert_batch(region_metadata, request, bulk_requests, sender)
1299                .await;
1300        } else {
1301            error!("Cannot find region metadata for {}", request.region_id);
1302            sender.send(
1303                error::RegionNotFoundSnafu {
1304                    region_id: request.region_id,
1305                }
1306                .fail(),
1307            );
1308        }
1309    }
1310
1311    /// Dispatches and processes requests.
1312    ///
1313    /// `general_requests` should not contain categorized write, ddl, or bulk insert requests.
1314    async fn handle_requests(
1315        &mut self,
1316        write_requests: &mut Vec<SenderWriteRequest>,
1317        ddl_requests: &mut Vec<SenderDdlRequest>,
1318        general_requests: &mut Vec<WorkerRequest>,
1319        bulk_requests: &mut Vec<SenderBulkRequest>,
1320    ) {
1321        for worker_req in general_requests.drain(..) {
1322            match worker_req {
1323                WorkerRequest::Write(_) | WorkerRequest::BatchWrite(_) | WorkerRequest::Ddl(_) => {
1324                    // These requests are categorized before dispatching general requests.
1325                    continue;
1326                }
1327                WorkerRequest::BulkInserts(_) => unreachable!("bulk inserts are buffered"),
1328                WorkerRequest::Background { region_id, notify } => {
1329                    if matches!(
1330                        &notify,
1331                        BackgroundNotify::RegionEdit(edit_result)
1332                            if edit_result.update_region_state
1333                    ) {
1334                        // Region state must be Editing when reach here.
1335                        // This call only moves write/bulk write request into stall queue. When region edit result
1336                        // is processed inside handle_background_notify and region state is switched back to Writable,
1337                        // stalled request will be processed before the next region edit is dequeued from
1338                        // RegionEditQueue immediately in handle_region_edit_result. It not only ensured pending writes
1339                        // are processed in time, but also prevents them from starvation.
1340                        // TODO(hl): maybe we need to merge those queues for pending requests like pending_ddl,
1341                        // region edits and stalled request, so we can simplify the coordination between these queues.
1342                        self.handle_buffered_region_write_requests(
1343                            &region_id,
1344                            write_requests,
1345                            bulk_requests,
1346                        )
1347                        .await;
1348                    }
1349                    // For background notify, we handle it directly.
1350                    self.handle_background_notify(region_id, notify).await;
1351                }
1352                WorkerRequest::SetRegionRoleStateGracefully {
1353                    region_id,
1354                    region_role_state,
1355                    sender,
1356                } => {
1357                    self.set_role_state_gracefully(region_id, region_role_state, sender)
1358                        .await;
1359                }
1360                WorkerRequest::EditRegion(request) => {
1361                    self.handle_region_edit(request);
1362                }
1363                WorkerRequest::Stop => {
1364                    debug_assert!(!self.running.load(Ordering::Relaxed));
1365                }
1366                WorkerRequest::SyncRegion(req) => {
1367                    self.handle_region_sync(req).await;
1368                }
1369                WorkerRequest::RemapManifests(req) => {
1370                    self.handle_remap_manifests_request(req);
1371                }
1372                WorkerRequest::CopyRegionFrom(req) => {
1373                    self.handle_copy_region_from_request(req);
1374                }
1375            }
1376        }
1377
1378        // Handles all write requests first. So we can alter regions without
1379        // considering existing write requests.
1380        self.handle_write_requests(write_requests, bulk_requests, true)
1381            .await;
1382
1383        self.handle_ddl_requests(ddl_requests).await;
1384    }
1385
1386    /// Takes and handles all ddl requests.
1387    async fn handle_ddl_requests(&mut self, ddl_requests: &mut Vec<SenderDdlRequest>) {
1388        if ddl_requests.is_empty() {
1389            return;
1390        }
1391
1392        for ddl in ddl_requests.drain(..) {
1393            let res = match ddl.request {
1394                DdlRequest::Create(req) => self.handle_create_request(ddl.region_id, req).await,
1395                DdlRequest::Drop(req) => {
1396                    self.handle_drop_request(ddl.region_id, req, ddl.sender)
1397                        .await;
1398                    continue;
1399                }
1400                DdlRequest::Open((req, wal_entry_receiver)) => {
1401                    self.handle_open_request(ddl.region_id, req, wal_entry_receiver, ddl.sender)
1402                        .await;
1403                    continue;
1404                }
1405                DdlRequest::OfflineCleanup(req) => {
1406                    self.handle_offline_cleanup_request(ddl.region_id, req)
1407                        .await
1408                }
1409                DdlRequest::Close(req) => {
1410                    self.handle_close_request(ddl.region_id, req, ddl.sender)
1411                        .await;
1412                    continue;
1413                }
1414                DdlRequest::Alter(req) => {
1415                    self.handle_alter_request(ddl.region_id, req, ddl.sender)
1416                        .await;
1417                    continue;
1418                }
1419                DdlRequest::Flush(req) => {
1420                    self.handle_flush_request(ddl.region_id, req, ddl.sender);
1421                    continue;
1422                }
1423                DdlRequest::Compact(req) => {
1424                    self.handle_compaction_request(ddl.region_id, req, ddl.sender)
1425                        .await;
1426                    continue;
1427                }
1428                DdlRequest::BuildIndex(req) => {
1429                    self.handle_build_index_request(ddl.region_id, req, ddl.sender)
1430                        .await;
1431                    continue;
1432                }
1433                DdlRequest::Truncate(req) => {
1434                    self.handle_truncate_request(ddl.region_id, req, ddl.sender)
1435                        .await;
1436                    continue;
1437                }
1438                DdlRequest::Catchup((req, wal_entry_receiver)) => {
1439                    self.handle_catchup_request(ddl.region_id, req, wal_entry_receiver, ddl.sender)
1440                        .await;
1441                    continue;
1442                }
1443                DdlRequest::EnterStaging(req) => {
1444                    self.handle_enter_staging_request(
1445                        ddl.region_id,
1446                        req.partition_directive,
1447                        ddl.sender,
1448                    )
1449                    .await;
1450                    continue;
1451                }
1452                DdlRequest::ApplyStagingManifest(req) => {
1453                    self.handle_apply_staging_manifest_request(ddl.region_id, req, ddl.sender)
1454                        .await;
1455                    continue;
1456                }
1457            };
1458
1459            ddl.sender.send(res);
1460        }
1461    }
1462
1463    /// Handle periodical tasks such as region auto flush.
1464    fn handle_periodical_tasks(&mut self) {
1465        let interval = CHECK_REGION_INTERVAL.as_millis() as i64;
1466        if self
1467            .time_provider
1468            .elapsed_since(self.last_periodical_check_millis)
1469            < interval
1470        {
1471            return;
1472        }
1473
1474        self.last_periodical_check_millis = self.time_provider.current_time_millis();
1475
1476        if let Err(e) = self.flush_periodically() {
1477            error!(e; "Failed to flush regions periodically");
1478        }
1479    }
1480
1481    /// Handles region background request
1482    async fn handle_background_notify(&mut self, region_id: RegionId, notify: BackgroundNotify) {
1483        match notify {
1484            BackgroundNotify::RegionOpened { region, registered } => {
1485                self.register_region(region);
1486                self.region_count.inc();
1487
1488                // Finish this opening batch before the worker can process a close/reopen,
1489                // without waiting for the background opener to resume after the ACK.
1490                let senders = self.opening_regions.remove_sender(region_id);
1491                if let Some(state) = &self.series_index_task_state {
1492                    state.wake();
1493                }
1494                let _ = registered.send(());
1495                for sender in senders {
1496                    sender.send(Ok(0));
1497                }
1498            }
1499            BackgroundNotify::CompactionUnit(req) => {
1500                self.handle_compaction_unit(region_id, req).await;
1501            }
1502            BackgroundNotify::CompactionDdlComplete {
1503                generation,
1504                request_index,
1505            } => {
1506                self.compaction_scheduler
1507                    .on_ddl_complete(region_id, generation, request_index);
1508            }
1509            BackgroundNotify::CompactionPickFinished(req) => {
1510                self.handle_compaction_pick_finished(region_id, req).await
1511            }
1512            BackgroundNotify::FlushFinished(req) => {
1513                self.handle_flush_finished(region_id, req).await
1514            }
1515            BackgroundNotify::FlushFailed(req) => self.handle_flush_failed(region_id, req).await,
1516            BackgroundNotify::IndexBuildFinished(req) => {
1517                self.handle_index_build_finished(region_id, req).await
1518            }
1519            BackgroundNotify::IndexBuildStopped(req) => {
1520                self.handle_index_build_stopped(region_id, req).await
1521            }
1522            BackgroundNotify::IndexBuildRetry(req) => {
1523                self.handle_rebuild_index(req, OptionOutputTx::new(None))
1524                    .await
1525            }
1526            BackgroundNotify::CompactionFinished(req) => {
1527                self.handle_compaction_finished(region_id, req).await
1528            }
1529            BackgroundNotify::CompactionFailed(req) => self.handle_compaction_failure(req).await,
1530            BackgroundNotify::Truncate(req) => self.handle_truncate_result(req).await,
1531            BackgroundNotify::DiscardUnflushed(req) => {
1532                self.handle_discard_unflushed_result(req).await
1533            }
1534            BackgroundNotify::RegionChange(req) => {
1535                self.handle_manifest_region_change_result(req).await
1536            }
1537            BackgroundNotify::EnterStaging(req) => self.handle_enter_staging_result(req).await,
1538            BackgroundNotify::RegionEdit(req) => self.handle_region_edit_result(req).await,
1539            BackgroundNotify::CopyRegionFromFinished(req) => {
1540                self.handle_copy_region_from_finished(req)
1541            }
1542        }
1543    }
1544
1545    /// Handles `set_region_role_gracefully`.
1546    async fn set_role_state_gracefully(
1547        &mut self,
1548        region_id: RegionId,
1549        region_role_state: SettableRegionRoleState,
1550        sender: oneshot::Sender<SetRegionRoleStateResponse>,
1551    ) {
1552        if let Some(region) = self.regions.get_region(region_id) {
1553            if matches!(
1554                region_role_state,
1555                SettableRegionRoleState::Follower | SettableRegionRoleState::DowngradingLeader
1556            ) {
1557                // Fence writes here: writes the worker handles later see
1558                // Downgrading, and earlier ones have finished. This only
1559                // changes a writable or staging leader, but always fences
1560                // promotions issued before.
1561                region.manifest_ctx.begin_downgrade(region_id);
1562            }
1563            #[cfg(feature = "hdfs-object-store")]
1564            let sweep_older_than = self.config.gc.unknown_file_lingering_time;
1565            #[cfg(feature = "hdfs-object-store")]
1566            let sweep_limiter = self.sweep_limiter.clone();
1567            // We need to do this in background as we need the manifest lock.
1568            common_runtime::spawn_global(async move {
1569                match region.set_role_state_gracefully(region_role_state).await {
1570                    Ok(()) => {
1571                        let last_entry_id = region.version_control.current().last_entry_id;
1572                        let _ = sender.send(SetRegionRoleStateResponse::success(
1573                            SetRegionRoleStateSuccess::mito(last_entry_id),
1574                        ));
1575                        // Only a writable leader reclaims what a writer that died
1576                        // before it renamed its temporary file left behind.
1577                        #[cfg(feature = "hdfs-object-store")]
1578                        if region_role_state == SettableRegionRoleState::Leader {
1579                            temp_file_sweeper::spawn_region_sweep(
1580                                region.access_layer(),
1581                                region_id,
1582                                sweep_older_than,
1583                                sweep_limiter,
1584                            );
1585                        }
1586                    }
1587                    Err(e) => {
1588                        error!(e; "Failed to set region {} role state to {:?}", region_id, region_role_state);
1589                        let _ = sender.send(SetRegionRoleStateResponse::invalid_transition(
1590                            BoxedError::new(e),
1591                        ));
1592                    }
1593                }
1594            });
1595        } else {
1596            let _ = sender.send(SetRegionRoleStateResponse::NotFound);
1597        }
1598    }
1599}
1600
1601impl<S> RegionWorkerLoop<S> {
1602    /// Cleans up the worker.
1603    async fn clean(&self) {
1604        // Closes remaining regions.
1605        let regions = self.regions.list_regions();
1606        for region in regions {
1607            region.deactivate_region_manifest_size_metric();
1608            region.stop().await;
1609        }
1610
1611        self.regions.clear();
1612    }
1613
1614    /// Notifies the whole group that a flush job is finished so other
1615    /// workers can handle stalled requests.
1616    fn notify_group(&mut self) {
1617        // Notifies all receivers.
1618        let _ = self.flush_sender.send(());
1619        // Marks the receiver in current worker as seen so the loop won't be waked up immediately.
1620        self.flush_receiver.borrow_and_update();
1621    }
1622}
1623
1624/// Wrapper that only calls event listener in tests.
1625#[derive(Default, Clone)]
1626pub(crate) struct WorkerListener {
1627    #[cfg(any(test, feature = "test"))]
1628    listener: Option<crate::engine::listener::EventListenerRef>,
1629}
1630
1631impl WorkerListener {
1632    #[cfg(any(test, feature = "test"))]
1633    pub(crate) fn new(
1634        listener: Option<crate::engine::listener::EventListenerRef>,
1635    ) -> WorkerListener {
1636        WorkerListener { listener }
1637    }
1638
1639    /// Notifies the background opener after the worker acknowledges region registration.
1640    pub(crate) async fn on_region_open_registered(&self, _region_id: RegionId) {
1641        #[cfg(any(test, feature = "test"))]
1642        if let Some(listener) = &self.listener {
1643            listener.on_region_open_registered(_region_id).await;
1644        }
1645    }
1646
1647    /// Flush is finished successfully.
1648    pub(crate) fn on_flush_success(&self, region_id: RegionId) {
1649        #[cfg(any(test, feature = "test"))]
1650        if let Some(listener) = &self.listener {
1651            listener.on_flush_success(region_id);
1652        }
1653        // Avoid compiler warning.
1654        let _ = region_id;
1655    }
1656
1657    /// Engine is stalled.
1658    pub(crate) fn on_write_stall(&self) {
1659        #[cfg(any(test, feature = "test"))]
1660        if let Some(listener) = &self.listener {
1661            listener.on_write_stall();
1662        }
1663    }
1664
1665    /// A batch of writes is appended to the WAL.
1666    pub(crate) async fn on_write_wal_end(&self) {
1667        #[cfg(any(test, feature = "test"))]
1668        if let Some(listener) = &self.listener {
1669            listener.on_write_wal_end().await;
1670        }
1671    }
1672
1673    pub(crate) async fn on_flush_begin(&self, region_id: RegionId) {
1674        #[cfg(any(test, feature = "test"))]
1675        if let Some(listener) = &self.listener {
1676            listener.on_flush_begin(region_id).await;
1677        }
1678        // Avoid compiler warning.
1679        let _ = region_id;
1680    }
1681
1682    pub(crate) async fn on_flush_commit_begin(&self, _region_id: RegionId) {
1683        #[cfg(any(test, feature = "test"))]
1684        if let Some(listener) = &self.listener {
1685            listener.on_flush_commit_begin(_region_id).await;
1686        }
1687    }
1688
1689    pub(crate) fn on_flush_cancel_requested(&self, _region_id: RegionId) {
1690        #[cfg(any(test, feature = "test"))]
1691        if let Some(listener) = &self.listener {
1692            listener.on_flush_cancel_requested(_region_id);
1693        }
1694    }
1695
1696    pub(crate) fn on_later_drop_begin(&self, region_id: RegionId) -> Option<Duration> {
1697        #[cfg(any(test, feature = "test"))]
1698        if let Some(listener) = &self.listener {
1699            return listener.on_later_drop_begin(region_id);
1700        }
1701        // Avoid compiler warning.
1702        let _ = region_id;
1703        None
1704    }
1705
1706    /// On later drop task is finished.
1707    pub(crate) fn on_later_drop_end(&self, region_id: RegionId, removed: bool) {
1708        #[cfg(any(test, feature = "test"))]
1709        if let Some(listener) = &self.listener {
1710            listener.on_later_drop_end(region_id, removed);
1711        }
1712        // Avoid compiler warning.
1713        let _ = region_id;
1714        let _ = removed;
1715    }
1716
1717    pub(crate) async fn on_merge_ssts_finished(&self, region_id: RegionId) {
1718        #[cfg(any(test, feature = "test"))]
1719        if let Some(listener) = &self.listener {
1720            listener.on_merge_ssts_finished(region_id).await;
1721        }
1722        // Avoid compiler warning.
1723        let _ = region_id;
1724    }
1725
1726    pub(crate) fn on_recv_requests(&self, request_num: usize) {
1727        #[cfg(any(test, feature = "test"))]
1728        if let Some(listener) = &self.listener {
1729            listener.on_recv_requests(request_num);
1730        }
1731        // Avoid compiler warning.
1732        let _ = request_num;
1733    }
1734
1735    pub(crate) fn on_file_cache_filled(&self, _file_id: FileId) {
1736        #[cfg(any(test, feature = "test"))]
1737        if let Some(listener) = &self.listener {
1738            listener.on_file_cache_filled(_file_id);
1739        }
1740    }
1741
1742    pub(crate) fn on_compaction_scheduled(&self, _region_id: RegionId) {
1743        #[cfg(any(test, feature = "test"))]
1744        if let Some(listener) = &self.listener {
1745            listener.on_compaction_scheduled(_region_id);
1746        }
1747    }
1748
1749    pub(crate) async fn on_compaction_pick_begin(&self, _region_id: RegionId) {
1750        #[cfg(any(test, feature = "test"))]
1751        if let Some(listener) = &self.listener {
1752            listener.on_compaction_pick_begin(_region_id).await;
1753        }
1754    }
1755
1756    pub(crate) async fn on_compaction_commit_begin(&self, _region_id: RegionId) {
1757        #[cfg(any(test, feature = "test"))]
1758        if let Some(listener) = &self.listener {
1759            listener.on_compaction_commit_begin(_region_id).await;
1760        }
1761    }
1762
1763    /// Exposes a per-unit test gate before merge starts.
1764    pub(crate) async fn on_compaction_unit_merge_begin(&self, _region_id: RegionId, _plan_id: u64) {
1765        #[cfg(any(test, feature = "test"))]
1766        if let Some(listener) = &self.listener {
1767            listener
1768                .on_compaction_unit_merge_begin(_region_id, _plan_id)
1769                .await;
1770        }
1771    }
1772
1773    /// Exposes a per-unit test gate between manifest commit and worker notification.
1774    pub(crate) async fn on_compaction_unit_committed(&self, _region_id: RegionId, _plan_id: u64) {
1775        #[cfg(any(test, feature = "test"))]
1776        if let Some(listener) = &self.listener {
1777            listener
1778                .on_compaction_unit_committed(_region_id, _plan_id)
1779                .await;
1780        }
1781    }
1782
1783    /// Reports unit visibility to the test listener independently of request completion.
1784    pub(crate) fn on_compaction_unit_applied(&self, _region_id: RegionId, _plan_id: u64) {
1785        #[cfg(any(test, feature = "test"))]
1786        if let Some(listener) = &self.listener {
1787            listener.on_compaction_unit_applied(_region_id, _plan_id);
1788        }
1789    }
1790
1791    pub(crate) async fn on_compaction_result_notified(&self, _region_id: RegionId) {
1792        #[cfg(any(test, feature = "test"))]
1793        if let Some(listener) = &self.listener {
1794            listener.on_compaction_result_notified(_region_id).await;
1795        }
1796    }
1797
1798    pub(crate) fn on_compaction_cancel_requested(&self, _region_id: RegionId) {
1799        #[cfg(any(test, feature = "test"))]
1800        if let Some(listener) = &self.listener {
1801            listener.on_compaction_cancel_requested(_region_id);
1802        }
1803    }
1804
1805    pub(crate) async fn on_notify_region_change_result_begin(&self, _region_id: RegionId) {
1806        #[cfg(any(test, feature = "test"))]
1807        if let Some(listener) = &self.listener {
1808            listener
1809                .on_notify_region_change_result_begin(_region_id)
1810                .await;
1811        }
1812    }
1813
1814    pub(crate) async fn on_enter_staging_result_begin(&self, _region_id: RegionId) {
1815        #[cfg(any(test, feature = "test"))]
1816        if let Some(listener) = &self.listener {
1817            listener.on_enter_staging_result_begin(_region_id).await;
1818        }
1819    }
1820
1821    pub(crate) async fn on_index_build_finish(&self, _region_file_id: RegionFileId) {
1822        #[cfg(any(test, feature = "test"))]
1823        if let Some(listener) = &self.listener {
1824            listener.on_index_build_finish(_region_file_id).await;
1825        }
1826    }
1827
1828    pub(crate) async fn on_index_build_begin(&self, _region_file_id: RegionFileId) {
1829        #[cfg(any(test, feature = "test"))]
1830        if let Some(listener) = &self.listener {
1831            listener.on_index_build_begin(_region_file_id).await;
1832        }
1833    }
1834
1835    pub(crate) async fn on_index_build_abort(&self, _region_file_id: RegionFileId) {
1836        #[cfg(any(test, feature = "test"))]
1837        if let Some(listener) = &self.listener {
1838            listener.on_index_build_abort(_region_file_id).await;
1839        }
1840    }
1841
1842    pub(crate) async fn on_index_build_before_manifest_commit(
1843        &self,
1844        _region_file_id: RegionFileId,
1845    ) {
1846        #[cfg(any(test, feature = "test"))]
1847        if let Some(listener) = &self.listener {
1848            listener
1849                .on_index_build_before_manifest_commit(_region_file_id)
1850                .await;
1851        }
1852    }
1853
1854    pub(crate) async fn on_index_build_manifest_committed(&self, _region_file_id: RegionFileId) {
1855        #[cfg(any(test, feature = "test"))]
1856        if let Some(listener) = &self.listener {
1857            listener
1858                .on_index_build_manifest_committed(_region_file_id)
1859                .await;
1860        }
1861    }
1862}
1863
1864#[cfg(test)]
1865mod tests {
1866    use super::*;
1867    use crate::test_util::TestEnv;
1868
1869    #[test]
1870    fn test_region_id_to_index() {
1871        let num_workers = 4;
1872
1873        let region_id = RegionId::new(1, 2);
1874        let index = region_id_to_index(region_id, num_workers);
1875        assert_eq!(index, 3);
1876
1877        let region_id = RegionId::new(2, 3);
1878        let index = region_id_to_index(region_id, num_workers);
1879        assert_eq!(index, 1);
1880    }
1881
1882    #[tokio::test]
1883    async fn test_worker_group_start_stop() {
1884        for (enabled, existing_index) in [(false, false), (true, false), (false, true)] {
1885            // Use a fresh WAL directory for each case: stopping workers does not stop the log store.
1886            let env = TestEnv::with_prefix("group-stop").await;
1887            let root = env.data_home().join("series_index");
1888            if existing_index {
1889                tokio::fs::create_dir_all(&root).await.unwrap();
1890                tokio::fs::write(root.join("retained"), b"index")
1891                    .await
1892                    .unwrap();
1893            }
1894            let group = env
1895                .create_worker_group(MitoConfig {
1896                    num_workers: 4,
1897                    experimental_enable_series_index: enabled,
1898                    experimental_series_index_maintenance_interval: Duration::from_secs(3600),
1899                    ..Default::default()
1900                })
1901                .await;
1902
1903            for worker in &group.workers {
1904                assert_eq!(worker.series_index_task_state.is_some(), enabled);
1905                assert_eq!(worker.series_index_handle.lock().await.is_some(), enabled);
1906            }
1907            assert_eq!(root.is_dir(), enabled || existing_index);
1908            if existing_index {
1909                assert_eq!(
1910                    tokio::fs::read(root.join("retained")).await.unwrap(),
1911                    b"index"
1912                );
1913            }
1914
1915            tokio::time::timeout(Duration::from_secs(5), group.stop())
1916                .await
1917                .expect("series-index tasks should stop without waiting for their interval")
1918                .unwrap();
1919        }
1920    }
1921
1922    /// Checks that all workers in a group share one bulk memtable compaction dispatcher.
1923    #[tokio::test]
1924    async fn test_workers_share_memtable_compact_dispatcher() {
1925        let env = TestEnv::with_prefix("share-dispatcher").await;
1926        let group = env
1927            .create_worker_group(MitoConfig {
1928                num_workers: 4,
1929                ..Default::default()
1930            })
1931            .await;
1932        let first = &group.workers[0].memtable_compact_dispatcher;
1933        for (index, worker) in group.workers.iter().enumerate() {
1934            let current = &worker.memtable_compact_dispatcher;
1935            assert!(
1936                Arc::ptr_eq(first, current),
1937                "worker {index} has its own memtable compaction dispatcher"
1938            );
1939        }
1940
1941        group.stop().await.unwrap();
1942    }
1943}