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