mito2/worker/

handle_flush.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Handling flush related requests.

use std::collections::{HashMap, HashSet};
use std::sync::atomic::Ordering;
use std::sync::{Arc, Mutex, Weak};

use common_base::readable_size::ReadableSize;
use common_telemetry::{debug, error, info};
use store_api::logstore::LogStore;
use store_api::region_request::{RegionFlushReason, RegionFlushRequest};
use store_api::storage::RegionId;

use crate::config::{IndexBuildMode, MitoConfig};
use crate::error::{RegionBusySnafu, RegionNotFoundSnafu, Result};
use crate::flush::{FlushReason, RegionFlushTask};
use crate::region::version::VersionRef;
use crate::region::{MitoRegion, MitoRegionRef};
use crate::request::{BuildIndexRequest, FlushFailed, FlushFinished, OnFailure, OptionOutputTx};
use crate::sst::index::IndexBuildType;
use crate::worker::RegionWorkerLoop;

#[derive(Debug, Default)]
pub(crate) struct RegionWriteBufferPressure {
    pub(crate) stalled_region_ids: HashSet<RegionId>,
    pub(crate) rejected_region_ids: HashSet<RegionId>,
}

/// Staging regions of a node that have received writes.
///
/// Workers share it to bound the memtables that stay unflushed until staging ends.
#[derive(Debug, Default)]
pub(crate) struct StagingRegions {
    regions: Mutex<HashMap<RegionId, Weak<MitoRegion>>>,
}

pub(crate) type StagingRegionsRef = Arc<StagingRegions>;

impl StagingRegions {
    /// Records `region` and returns the memtable usage of the regions still in staging.
    fn record_and_memtable_usage(&self, region: &MitoRegionRef) -> usize {
        let mut regions = self.regions.lock().unwrap();
        regions.insert(region.region_id, Arc::downgrade(region));
        let mut usage = 0;
        regions.retain(|_, region| match region.upgrade() {
            Some(region) if region.is_staging_or_staged_edit() => {
                usage += region_memtable_usage(&region.version());
                true
            }
            _ => false,
        });
        usage
    }
}

#[derive(Debug, Default, PartialEq, Eq)]
struct RegionWriteBufferStatus {
    should_flush: bool,
    should_stall: bool,
    should_reject: bool,
}

fn resolve_flush_reason(
    request_reason: Option<RegionFlushReason>,
    is_downgrading: bool,
) -> FlushReason {
    match request_reason {
        Some(reason) => FlushReason::from(reason),
        None if is_downgrading => FlushReason::Downgrading,
        None => FlushReason::Manual,
    }
}

impl<S: LogStore> RegionWorkerLoop<S> {
    /// On region flush job failed.
    pub(crate) async fn handle_flush_failed(&mut self, region_id: RegionId, request: FlushFailed) {
        let mut pending_ddls = self.flush_scheduler.on_flush_failed(region_id, request.err);
        if !pending_ddls.is_empty() {
            info!(
                "Flush terminated for region {}, handling {} pending lifecycle DDL requests",
                region_id,
                pending_ddls.len()
            );
            self.handle_ddl_requests(&mut pending_ddls).await;
        }
        // Maybe flush worker again.
        self.maybe_flush_worker();

        // Handle stalled requests.
        self.handle_stalled_requests().await;
        self.dispatch_pending_region_requests(region_id).await;
    }

    /// Checks whether the engine reaches flush threshold. If so, finds regions in this
    /// worker to flush.
    pub(crate) fn maybe_flush_worker(&mut self) {
        if !self.write_buffer_manager.should_flush_engine() {
            debug!("No need to flush worker");
            // No need to flush worker.
            return;
        }

        // If the engine needs flush, each worker will find some regions to flush. We might
        // flush more memory than expect but it should be acceptable.
        if let Err(e) = self.flush_regions_on_engine_full() {
            error!(e; "Failed to flush worker");
        }
    }

    /// Finds some regions to flush to reduce write buffer usage.
    fn flush_regions_on_engine_full(&mut self) -> Result<()> {
        let regions = self.regions.list_regions();
        let now = self.time_provider.current_time_millis();
        let mut pending_regions = vec![];

        for region in &regions {
            if self.flush_scheduler.is_flush_requested(region.region_id)
                || !region.is_writable()
                || region.is_staging()
            {
                // Already flushing, not writable, or in staging, whose rows must stay in the WAL.
                continue;
            }

            let version = region.version();
            let region_memtable_size = region_memtable_usage(&version);

            let auto_flush_interval = version
                .options
                .auto_flush_interval_or(self.config.auto_flush_interval);
            let min_last_flush_time = now.saturating_sub(
                auto_flush_interval
                    .as_millis()
                    .try_into()
                    .unwrap_or(i64::MAX),
            );

            if region.last_flush_millis() < min_last_flush_time {
                // If flush time of this region is earlier than `min_last_flush_time`, we can flush this region.
                let task =
                    self.new_flush_task(region, FlushReason::EngineFull, None, self.config.clone());
                self.flush_scheduler.schedule_flush(
                    region.region_id,
                    &region.version_control,
                    task,
                )?;
            } else if region_memtable_size > 0 {
                // We should only consider regions with memtable size > 0 to flush.
                pending_regions.push((region, region_memtable_size));
            }
        }
        pending_regions.sort_unstable_by_key(|(_, size)| std::cmp::Reverse(*size));
        // The flush target is the mutable memtable limit (half of the global buffer).
        // When memory is full, we aggressively flush regions until usage drops below this target,
        // not just below the full limit.
        let target_memory_usage = self.write_buffer_manager.flush_limit();
        let mut memory_usage = self.write_buffer_manager.memory_usage();

        #[cfg(test)]
        {
            debug!(
                "Flushing regions on engine full, target memory usage: {}, memory usage: {}, pending regions: {:?}",
                target_memory_usage,
                memory_usage,
                pending_regions
                    .iter()
                    .map(|(region, mem_size)| (region.region_id, mem_size))
                    .collect::<Vec<_>>()
            );
        }
        // Iterate over pending regions in descending order of their memory size and schedule flush tasks
        // for each region until the overall memory usage drops below the flush limit.
        for (region, region_mem_size) in pending_regions.into_iter() {
            // Make sure the first region is always flushed.
            if memory_usage < target_memory_usage {
                // Stop flushing regions if memory usage is already below the flush limit
                break;
            }
            let task =
                self.new_flush_task(region, FlushReason::EngineFull, None, self.config.clone());
            debug!("Scheduling flush task for region {}", region.region_id);
            // Schedule a flush task for the current region
            self.flush_scheduler
                .schedule_flush(region.region_id, &region.version_control, task)?;
            // Reduce memory usage by the region's size, ensuring it doesn't go negative
            memory_usage = memory_usage.saturating_sub(region_mem_size);
        }

        Ok(())
    }

    /// Flushes write regions that exceed their flush threshold and returns region pressure.
    pub(crate) fn maybe_flush_write_regions(
        &mut self,
        region_ids: HashSet<RegionId>,
    ) -> RegionWriteBufferPressure {
        let mut pressure = RegionWriteBufferPressure::default();

        for region_id in region_ids {
            let Some(region) = self.regions.get_region(region_id) else {
                continue;
            };
            if !region.is_writable() {
                continue;
            }

            let status = region_write_buffer_status(
                &region.version(),
                self.config.default_region_write_buffer_size,
                self.stalled_requests.estimated_size(&region_id),
            );

            if region.is_staging() {
                // Staging regions don't flush, so waiting would not release their memory. Writes
                // are rejected at twice the region limit or at half of the engine's flush limit,
                // which reserves memory headroom for other regions.
                let staging_usage = self.staging_regions.record_and_memtable_usage(&region);
                if status.should_reject
                    || staging_usage >= self.write_buffer_manager.flush_limit() / 2
                {
                    pressure.rejected_region_ids.insert(region_id);
                }
                continue;
            }

            if status.should_flush && !self.flush_scheduler.is_flush_requested(region.region_id) {
                let task = self.new_flush_task(
                    &region,
                    FlushReason::RegionFull,
                    None,
                    self.config.clone(),
                );
                if let Err(e) = self.flush_scheduler.schedule_flush(
                    region.region_id,
                    &region.version_control,
                    task,
                ) {
                    error!(e; "Failed to schedule flush task for region {}", region.region_id);
                }
            }

            if status.should_reject {
                pressure.rejected_region_ids.insert(region_id);
            } else if status.should_stall {
                pressure.stalled_region_ids.insert(region_id);
            }
        }

        pressure
    }

    /// Creates a flush task with specific `reason` for the `region`.
    pub(crate) fn new_flush_task(
        &self,
        region: &MitoRegionRef,
        reason: FlushReason,
        row_group_size: Option<usize>,
        engine_config: Arc<MitoConfig>,
    ) -> RegionFlushTask {
        RegionFlushTask {
            region_id: region.region_id,
            reason,
            senders: Vec::new(),
            request_sender: self.sender.clone(),
            access_layer: region.access_layer.clone(),
            listener: self.listener.clone(),
            engine_config,
            // The request-provided size (e.g. manual flush) takes precedence over the
            // region option.
            row_group_size: row_group_size.or(region.version().options.max_row_group_row_count),
            cache_manager: self.cache_manager.clone(),
            manifest_ctx: region.manifest_ctx.clone(),
            index_options: region.version().options.index_options.clone(),
            flush_semaphore: self.flush_semaphore.clone(),
            partition_expr: region.metadata().partition_expr.clone(),
            durability_barrier: self
                .wal
                .durability_barrier(region.region_id, &region.provider),
        }
    }
}

pub(crate) fn region_memtable_usage(version: &VersionRef) -> usize {
    version.memtables.mutable_usage() + version.memtables.immutables_usage()
}

pub(crate) fn region_write_buffer_size(
    version: &VersionRef,
    default_region_write_buffer_size: ReadableSize,
) -> Option<usize> {
    let size = version
        .options
        .write_buffer_size
        .unwrap_or(default_region_write_buffer_size);
    (size.as_bytes() > 0).then_some(size.as_bytes() as usize)
}

/// Returns the pressure state of a region's write buffer.
fn region_write_buffer_status(
    version: &VersionRef,
    default_region_write_buffer_size: ReadableSize,
    stalled_request_size: usize,
) -> RegionWriteBufferStatus {
    let Some(write_buffer_size) =
        region_write_buffer_size(version, default_region_write_buffer_size)
    else {
        return RegionWriteBufferStatus::default();
    };

    let mutable_usage = version.memtables.mutable_usage();
    let memory_usage = region_memtable_usage(version);
    region_write_buffer_status_from_usage(
        write_buffer_size,
        mutable_usage,
        memory_usage,
        stalled_request_size,
    )
}

fn region_write_buffer_status_from_usage(
    write_buffer_size: usize,
    mutable_usage: usize,
    memory_usage: usize,
    stalled_request_size: usize,
) -> RegionWriteBufferStatus {
    let mutable_limit = std::cmp::max(1, write_buffer_size / 2);
    let should_stall = memory_usage >= write_buffer_size;
    let reject_limit = write_buffer_size.saturating_mul(2);
    let should_reject = memory_usage.saturating_add(stalled_request_size) >= reject_limit;

    RegionWriteBufferStatus {
        should_flush: mutable_usage >= mutable_limit || should_stall,
        should_stall,
        should_reject,
    }
}

impl<S: LogStore> RegionWorkerLoop<S> {
    /// Handles manual flush request.
    pub(crate) fn handle_flush_request(
        &mut self,
        region_id: RegionId,
        request: RegionFlushRequest,
        sender: OptionOutputTx,
    ) {
        let region = match self.regions.flushable_region(region_id) {
            Ok(region) => region,
            Err(e) => {
                sender.send(Err(e));
                return;
            }
        };

        // `update_topic_latest_entry_id` updates `topic_latest_entry_id` when memtables are empty.
        // But the flush is skipped if memtables are empty. Thus should update the `topic_latest_entry_id`
        // when handling flush request instead of in `schedule_flush` or `flush_finished`.
        self.update_topic_latest_entry_id(&region);

        // A staging region must keep its WAL until staging ends, so it flushes nothing.
        if region.is_staging() && !region.version().memtables.is_empty() {
            sender.send(match request.reason {
                // These only ask to advance the flushed position, which can wait.
                Some(RegionFlushReason::RegionMigration | RegionFlushReason::RemoteWalPrune) => {
                    Ok(0)
                }
                _ => RegionBusySnafu { region_id }.fail(),
            });
            return;
        }

        let reason = resolve_flush_reason(request.reason, region.is_downgrading());
        let mut task =
            self.new_flush_task(&region, reason, request.row_group_size, self.config.clone());
        task.push_sender(sender);
        if let Err(e) =
            self.flush_scheduler
                .schedule_flush(region.region_id, &region.version_control, task)
        {
            error!(e; "Failed to schedule flush task for region {}", region.region_id);
        }
    }

    /// Flushes regions periodically.
    pub(crate) fn flush_periodically(&mut self) -> Result<()> {
        let regions = self.regions.list_regions();
        let now = self.time_provider.current_time_millis();

        for region in &regions {
            if self.flush_scheduler.is_flush_requested(region.region_id) || !region.is_writable() {
                // Already flushing or not writable.
                continue;
            }
            self.update_topic_latest_entry_id(region);
            // The rows of a staging region must stay in the WAL until staging ends, so a staging
            // region with rows pins its topic's prune point: the update above only keeps the id
            // fresh while the memtable is empty.
            if region.is_staging() {
                continue;
            }

            let auto_flush_interval = region
                .version()
                .options
                .auto_flush_interval_or(self.config.auto_flush_interval);
            let min_last_flush_time = now.saturating_sub(
                auto_flush_interval
                    .as_millis()
                    .try_into()
                    .unwrap_or(i64::MAX),
            );

            if region.last_flush_millis() < min_last_flush_time {
                // If flush time of this region is earlier than `min_last_flush_time`, we can flush this region.
                let task = self.new_flush_task(
                    region,
                    FlushReason::Periodically,
                    None,
                    self.config.clone(),
                );
                self.flush_scheduler.schedule_flush(
                    region.region_id,
                    &region.version_control,
                    task,
                )?;
            }
        }

        Ok(())
    }

    /// On region flush job finished.
    pub(crate) async fn handle_flush_finished(
        &mut self,
        region_id: RegionId,
        mut request: FlushFinished,
    ) {
        // Notifies other workers. Even the remaining steps of this method fail we still
        // wake up other workers as we have released some memory by flush.
        self.notify_group();

        let region = match self.regions.get_region(region_id) {
            Some(region) => region,
            None => {
                request.on_failure(RegionNotFoundSnafu { region_id }.build());
                return;
            }
        };

        region.version_control.apply_edit(
            Some(request.edit.clone()),
            &request.memtables_to_remove,
            region.file_purger.clone(),
        );

        // The flush has committed its manifest. Count only its Parquet output, before
        // WAL cleanup (which can fail after the commit).
        let ingested_parquet_bytes = request
            .edit
            .files_to_add
            .iter()
            .map(|file| file.file_size)
            .sum();
        region
            .region_stats
            .ingested_parquet_bytes
            .fetch_add(ingested_parquet_bytes, Ordering::Relaxed);

        region.update_flush_millis();
        // Update topic latest entry id as soon as possible after flush to make sure the prunable entry id is updated timely,
        // which is important for remote WAL pruning.
        self.update_topic_latest_entry_id(&region);

        // Delete wal.
        info!(
            "Region {} flush finished, tries to bump wal to {}",
            region_id, request.flushed_entry_id
        );
        if let Err(e) = self
            .wal
            .obsolete(region_id, request.flushed_entry_id, &region.provider)
            .await
        {
            error!(e; "Failed to write wal, region: {}", region_id);
            request.on_failure(e);
            return;
        }

        let flush_on_close = request.flush_reason == FlushReason::Closing;

        let index_build_file_metas = std::mem::take(&mut request.edit.files_to_add);

        // In async mode, create indexes after flush.
        if self.config.index.build_mode == IndexBuildMode::Async {
            self.handle_rebuild_index(
                BuildIndexRequest {
                    region_id,
                    build_type: IndexBuildType::Flush,
                    file_metas: index_build_file_metas,
                    result_senders: Vec::new(),
                },
                OptionOutputTx::new(None),
            )
            .await;
        }

        if flush_on_close {
            self.remove_region(region_id).await;
            info!("Region {} closed after flush", region_id);
            request.on_success();
            self.listener.on_flush_success(region_id);
            return;
        }

        // Notifies waiters and observes the flush timer.
        request.on_success();
        // Handle pending requests for the region.
        if let Some((mut ddl_requests, mut write_requests, mut bulk_writes)) =
            self.flush_scheduler.on_flush_success(region_id)
        {
            // Perform DDLs first because they require empty memtables.
            // These ALTERs already own the flush turn; bypass initial admission.
            for ddl in ddl_requests.drain(..) {
                if let crate::request::DdlRequest::Alter(request) = ddl.request {
                    self.execute_alter_request(ddl.region_id, request, ddl.sender)
                        .await;
                } else {
                    self.handle_ddl_requests(&mut vec![ddl]).await;
                }
            }
            if self.flush_scheduler.is_flush_requested(region_id) {
                // The DDL may schedule another flush, e.g. a close-time flush after writes
                // arrived in the mutable memtable during the previous flush. Keep pending
                // writes fenced until that flush reaches its terminal state instead of
                // accepting them while the DDL is still in progress.
                for write_request in write_requests {
                    self.flush_scheduler
                        .add_write_request_to_pending(write_request);
                }
                for bulk_write in bulk_writes {
                    self.flush_scheduler.add_bulk_request_to_pending(bulk_write);
                }
                self.listener.on_flush_success(region_id);
                return;
            }
            // A pending close DDL may have removed the region. Reject queued writes as
            // not found, then stop instead of scheduling compaction for a closed region.
            if !self.regions.is_region_exists(region_id) {
                self.handle_write_requests(&mut write_requests, &mut bulk_writes, false)
                    .await;
                self.listener.on_flush_success(region_id);
                return;
            }
            // Handle pending write requests, we don't stall these requests.
            self.handle_write_requests(&mut write_requests, &mut bulk_writes, false)
                .await;
        }
        // Maybe flush worker again.
        self.maybe_flush_worker();
        // Handle stalled requests.
        self.handle_stalled_requests().await;
        // Schedules compaction.
        self.schedule_compaction(&region).await;

        self.dispatch_pending_region_requests(region_id).await;
        self.listener.on_flush_success(region_id);
    }

    /// Updates the latest entry id since flush of the region.
    /// **This is only used for remote WAL pruning.**
    pub(crate) fn update_topic_latest_entry_id(&mut self, region: &MitoRegionRef) {
        if region.provider.is_remote_wal() && region.version().memtables.is_empty() {
            let latest_offset = self
                .wal
                .store()
                .latest_entry_id(&region.provider)
                .unwrap_or(0);
            let topic_last_entry_id = region.topic_latest_entry_id.load(Ordering::Relaxed);

            if latest_offset > topic_last_entry_id {
                region
                    .topic_latest_entry_id
                    .store(latest_offset, Ordering::Relaxed);
                debug!(
                    "Region {} latest entry id updated to {}",
                    region.region_id, latest_offset
                );
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_resolve_flush_reason_uses_request_reason() {
        assert_eq!(
            resolve_flush_reason(Some(RegionFlushReason::RegionMigration), true),
            FlushReason::RegionMigration
        );
        assert_eq!(
            resolve_flush_reason(Some(RegionFlushReason::Repartition), false),
            FlushReason::Repartition
        );
        assert_eq!(
            resolve_flush_reason(Some(RegionFlushReason::RemoteWalPrune), false),
            FlushReason::RemoteWalPrune
        );
        assert_eq!(
            resolve_flush_reason(Some(RegionFlushReason::Closing), false),
            FlushReason::Closing
        );
        assert_eq!(
            resolve_flush_reason(Some(RegionFlushReason::Downgrading), false),
            FlushReason::Downgrading
        );
    }

    #[test]
    fn test_resolve_flush_reason_fallback_unchanged() {
        assert_eq!(resolve_flush_reason(None, true), FlushReason::Downgrading);
        assert_eq!(resolve_flush_reason(None, false), FlushReason::Manual);
    }

    #[test]
    fn test_region_write_buffer_status_boundaries() {
        assert_eq!(
            RegionWriteBufferStatus::default(),
            region_write_buffer_status_from_usage(100, 49, 99, 0)
        );
        assert_eq!(
            RegionWriteBufferStatus {
                should_flush: true,
                should_stall: false,
                should_reject: false,
            },
            region_write_buffer_status_from_usage(100, 50, 99, 0)
        );
        assert_eq!(
            RegionWriteBufferStatus {
                should_flush: true,
                should_stall: true,
                should_reject: false,
            },
            region_write_buffer_status_from_usage(100, 50, 100, 99)
        );
        assert_eq!(
            RegionWriteBufferStatus {
                should_flush: true,
                should_stall: true,
                should_reject: true,
            },
            region_write_buffer_status_from_usage(100, 50, 100, 100)
        );
    }

    #[test]
    fn test_region_write_buffer_status_saturates_reject_limit() {
        assert_eq!(
            RegionWriteBufferStatus {
                should_flush: true,
                should_stall: true,
                should_reject: true,
            },
            region_write_buffer_status_from_usage(usize::MAX, usize::MAX, usize::MAX, usize::MAX,)
        );
    }
}