mito2/

access_layer.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.

use std::path::Path;
use std::sync::Arc;
use std::time::{Duration, Instant};

use async_stream::try_stream;
use common_base::readable_size::ReadableSize;
use common_runtime::runtime::RuntimeTrait;
use common_telemetry::warn;
use common_time::Timestamp;
use futures::{Stream, TryStreamExt};
use object_store::services::Fs;
use object_store::util::{join_dir, with_instrument_layers};
use object_store::{ATOMIC_WRITE_DIR, ErrorKind, OLD_ATOMIC_WRITE_DIR, ObjectStore};
use parquet::file::metadata::PageIndexPolicy;
use smallvec::SmallVec;
use snafu::ResultExt;
use store_api::metadata::RegionMetadataRef;
use store_api::region_request::PathType;
use store_api::sst_entry::StorageSstEntry;
use store_api::storage::{FileId, RegionId, SequenceNumber};

use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey};
use crate::cache::write_cache::SstUploadRequest;
use crate::cache::{CacheManagerRef, SstMetaPreparation, prepare_sst_meta_sync};
use crate::config::{BloomFilterConfig, FulltextIndexConfig, IndexConfig, InvertedIndexConfig};
use crate::error::{
    CleanDirSnafu, DeleteIndexSnafu, DeleteIndexesSnafu, DeleteSstsSnafu, OpenDalSnafu, Result,
};
use crate::metrics::{COMPACTION_STAGE_ELAPSED, FLUSH_ELAPSED};
use crate::read::FlatSource;
use crate::region::options::IndexOptions;
use crate::sst::file::{FileHandle, RegionFileId, RegionIndexId};
use crate::sst::index::intermediate::IntermediateManager;
use crate::sst::index::puffin_manager::{PuffinManagerFactory, SstPuffinManager};
use crate::sst::index::{IndexerBuilderImpl, build_index_from_parquet_bytes};
use crate::sst::location::{self, region_dir_from_table_dir};
use crate::sst::parquet::reader::ParquetReaderBuilder;
use crate::sst::parquet::writer::ParquetWriter;
use crate::sst::parquet::{SstInfo, WriteOptions};
use crate::sst::{DEFAULT_WRITE_CONCURRENCY, FormatType};

pub type AccessLayerRef = Arc<AccessLayer>;
/// SST write results.
pub type SstInfoArray = SmallVec<[SstInfo; 2]>;

/// Write operation type.
#[derive(Eq, PartialEq, Debug)]
pub enum WriteType {
    /// Writes from flush
    Flush,
    /// Writes from compaction.
    Compaction,
}

#[derive(Debug)]
pub struct Metrics {
    pub(crate) write_type: WriteType,
    pub(crate) iter_source: Duration,
    pub(crate) write_batch: Duration,
    pub(crate) update_index: Duration,
    pub(crate) upload_parquet: Duration,
    pub(crate) upload_puffin: Duration,
    pub(crate) compact_memtable: Duration,
}

impl Metrics {
    pub fn new(write_type: WriteType) -> Self {
        Self {
            write_type,
            iter_source: Default::default(),
            write_batch: Default::default(),
            update_index: Default::default(),
            upload_parquet: Default::default(),
            upload_puffin: Default::default(),
            compact_memtable: Default::default(),
        }
    }

    pub(crate) fn merge(mut self, other: Self) -> Self {
        assert_eq!(self.write_type, other.write_type);
        self.iter_source += other.iter_source;
        self.write_batch += other.write_batch;
        self.update_index += other.update_index;
        self.upload_parquet += other.upload_parquet;
        self.upload_puffin += other.upload_puffin;
        self.compact_memtable += other.compact_memtable;
        self
    }

    pub(crate) fn observe(self) {
        match self.write_type {
            WriteType::Flush => {
                FLUSH_ELAPSED
                    .with_label_values(&["iter_source"])
                    .observe(self.iter_source.as_secs_f64());
                FLUSH_ELAPSED
                    .with_label_values(&["write_batch"])
                    .observe(self.write_batch.as_secs_f64());
                FLUSH_ELAPSED
                    .with_label_values(&["update_index"])
                    .observe(self.update_index.as_secs_f64());
                FLUSH_ELAPSED
                    .with_label_values(&["upload_parquet"])
                    .observe(self.upload_parquet.as_secs_f64());
                FLUSH_ELAPSED
                    .with_label_values(&["upload_puffin"])
                    .observe(self.upload_puffin.as_secs_f64());
                if !self.compact_memtable.is_zero() {
                    FLUSH_ELAPSED
                        .with_label_values(&["compact_memtable"])
                        .observe(self.upload_puffin.as_secs_f64());
                }
            }
            WriteType::Compaction => {
                COMPACTION_STAGE_ELAPSED
                    .with_label_values(&["iter_source"])
                    .observe(self.iter_source.as_secs_f64());
                COMPACTION_STAGE_ELAPSED
                    .with_label_values(&["write_batch"])
                    .observe(self.write_batch.as_secs_f64());
                COMPACTION_STAGE_ELAPSED
                    .with_label_values(&["update_index"])
                    .observe(self.update_index.as_secs_f64());
                COMPACTION_STAGE_ELAPSED
                    .with_label_values(&["upload_parquet"])
                    .observe(self.upload_parquet.as_secs_f64());
                COMPACTION_STAGE_ELAPSED
                    .with_label_values(&["upload_puffin"])
                    .observe(self.upload_puffin.as_secs_f64());
            }
        };
    }
}

/// A layer to access SST files under the same directory.
pub struct AccessLayer {
    table_dir: String,
    /// Path type for generating file paths.
    path_type: PathType,
    /// Target object store.
    object_store: ObjectStore,
    /// Puffin manager factory for index.
    puffin_manager_factory: PuffinManagerFactory,
    /// Intermediate manager for inverted index.
    intermediate_manager: IntermediateManager,
}

impl std::fmt::Debug for AccessLayer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AccessLayer")
            .field("table_dir", &self.table_dir)
            .finish()
    }
}

impl AccessLayer {
    /// Returns a new [AccessLayer] for specific `table_dir`.
    pub fn new(
        table_dir: impl Into<String>,
        path_type: PathType,
        object_store: ObjectStore,
        puffin_manager_factory: PuffinManagerFactory,
        intermediate_manager: IntermediateManager,
    ) -> AccessLayer {
        AccessLayer {
            table_dir: table_dir.into(),
            path_type,
            object_store,
            puffin_manager_factory,
            intermediate_manager,
        }
    }

    /// Returns the directory of the table.
    pub fn table_dir(&self) -> &str {
        &self.table_dir
    }

    /// Returns the object store of the layer.
    pub fn object_store(&self) -> &ObjectStore {
        &self.object_store
    }

    /// Returns the path type of the layer.
    pub fn path_type(&self) -> PathType {
        self.path_type
    }

    /// Returns the puffin manager factory.
    pub fn puffin_manager_factory(&self) -> &PuffinManagerFactory {
        &self.puffin_manager_factory
    }

    /// Returns the intermediate manager.
    pub fn intermediate_manager(&self) -> &IntermediateManager {
        &self.intermediate_manager
    }

    /// Build the puffin manager.
    pub(crate) fn build_puffin_manager(&self) -> SstPuffinManager {
        let store = self.object_store.clone();
        let path_provider =
            RegionFilePathFactory::new(self.table_dir().to_string(), self.path_type());
        self.puffin_manager_factory.build(store, path_provider)
    }

    pub(crate) async fn delete_index(
        &self,
        index_file_id: RegionIndexId,
    ) -> Result<(), crate::error::Error> {
        let path = location::index_file_path(
            &self.table_dir,
            RegionIndexId::new(index_file_id.file_id, index_file_id.version),
            self.path_type,
        );
        self.object_store
            .delete(&path)
            .await
            .context(DeleteIndexSnafu {
                file_id: index_file_id.file_id(),
            })?;
        Ok(())
    }

    pub(crate) async fn delete_ssts(
        &self,
        region_id: RegionId,
        file_ids: &[FileId],
    ) -> Result<(), crate::error::Error> {
        if file_ids.is_empty() {
            return Ok(());
        }

        let attempted_files = file_ids.to_vec();
        // Deleter does not normalize leading slashes like Operator::delete does.
        let paths: Vec<_> = file_ids
            .iter()
            .map(|file_id| {
                location::sst_file_path(
                    &self.table_dir,
                    RegionFileId::new(region_id, *file_id),
                    self.path_type,
                )
                .trim_start_matches('/')
                .to_string()
            })
            .collect();

        let mut deleter = self
            .object_store
            .deleter()
            .await
            .with_context(|_| DeleteSstsSnafu {
                region_id,
                file_ids: attempted_files.clone(),
            })?;
        deleter
            .delete_iter(paths)
            .await
            .with_context(|_| DeleteSstsSnafu {
                region_id,
                file_ids: attempted_files.clone(),
            })?;
        deleter.close().await.with_context(|_| DeleteSstsSnafu {
            region_id,
            file_ids: attempted_files,
        })?;

        Ok(())
    }

    pub(crate) async fn delete_indexes(
        &self,
        index_ids: &[RegionIndexId],
    ) -> Result<(), crate::error::Error> {
        if index_ids.is_empty() {
            return Ok(());
        }

        let file_ids: Vec<_> = index_ids
            .iter()
            .map(|index_id| index_id.file_id())
            .collect();
        let paths: Vec<_> = index_ids
            .iter()
            .map(|index_id| {
                location::index_file_path(&self.table_dir, *index_id, self.path_type)
                    .trim_start_matches('/')
                    .to_string()
            })
            .collect();

        let mut deleter = self
            .object_store
            .deleter()
            .await
            .context(DeleteIndexesSnafu {
                file_ids: file_ids.clone(),
            })?;
        deleter
            .delete_iter(paths)
            .await
            .context(DeleteIndexesSnafu {
                file_ids: file_ids.clone(),
            })?;
        deleter
            .close()
            .await
            .context(DeleteIndexesSnafu { file_ids })?;

        Ok(())
    }

    /// Returns the directory of the region in the table.
    pub fn build_region_dir(&self, region_id: RegionId) -> String {
        region_dir_from_table_dir(&self.table_dir, region_id, self.path_type)
    }

    /// Returns a reader builder for specific `file`.
    pub(crate) fn read_sst(&self, file: FileHandle) -> ParquetReaderBuilder {
        ParquetReaderBuilder::new(
            self.table_dir.clone(),
            self.path_type,
            file,
            self.object_store.clone(),
        )
    }

    /// Writes a SST with specific `file_id` and `metadata` to the layer.
    ///
    /// Returns the info of the SST. If no data written, returns None.
    pub async fn write_sst(
        &self,
        request: SstWriteRequest,
        write_opts: &WriteOptions,
        metrics: &mut Metrics,
    ) -> Result<SstInfoArray> {
        self.write_sst_inner(request, write_opts, metrics, None)
            .await
    }

    /// Writes an SST whose output is registered with `context.guard`.
    ///
    /// The write cache is bypassed, upload concurrency is fixed to one, and
    /// the guard runs before either the SST or its index becomes visible.
    /// `context.limits` bounds the encoder and footer workspace.
    #[cfg(feature = "enterprise")]
    pub(crate) async fn write_sst_with_guard(
        &self,
        request: SstWriteRequest,
        write_opts: &WriteOptions,
        metrics: &mut Metrics,
        context: crate::sst::parquet::writer::SstWriteContext,
    ) -> Result<SstInfoArray> {
        self.write_sst_inner(request, write_opts, metrics, Some(context))
            .await
    }

    async fn write_sst_inner(
        &self,
        request: SstWriteRequest,
        write_opts: &WriteOptions,
        metrics: &mut Metrics,
        context: Option<crate::sst::parquet::writer::SstWriteContext>,
    ) -> Result<SstInfoArray> {
        let limits = context.as_ref().map(|context| context.limits);
        let guard = context.map(|context| context.guard);
        let op_type = request.op_type;
        let region_id = request.metadata.region_id;
        let region_metadata = request.metadata.clone();
        let cache_manager = request.cache_manager.clone();
        let override_sequence = if request.preserve_row_sequence {
            None
        } else {
            request.max_sequence
        };

        let sst_info = if let Some(write_cache) = cache_manager.write_cache()
            && guard.is_none()
        {
            // Write to the write cache.
            write_cache
                .write_and_upload_sst(
                    request,
                    SstUploadRequest {
                        dest_path_provider: RegionFilePathFactory::new(
                            self.table_dir.clone(),
                            self.path_type,
                        ),
                        remote_store: self.object_store.clone(),
                    },
                    write_opts,
                    metrics,
                )
                .await?
        } else {
            // Write cache is disabled, or bypassed by an SST write guard.
            let store = self.object_store.clone();
            let path_provider = RegionFilePathFactory::new(self.table_dir.clone(), self.path_type);
            let puffin_factory = if limits.is_some() {
                self.puffin_manager_factory
                    .clone()
                    .with_write_buffer_size(write_opts.write_buffer_size.as_bytes() as usize)
            } else {
                self.puffin_manager_factory.clone()
            };
            let indexer_builder = IndexerBuilderImpl {
                build_type: request.op_type.into(),
                metadata: request.metadata.clone(),
                puffin_manager: puffin_factory.build(store, path_provider.clone()),
                write_cache_enabled: false,
                intermediate_manager: self.intermediate_manager.clone(),
                index_options: request.index_options,
                inverted_index_config: request.inverted_index_config,
                fulltext_index_config: request.fulltext_index_config,
                bloom_filter_index_config: request.bloom_filter_index_config,
            };
            // We disable write cache on file system but we still use atomic write.
            // TODO(yingwen): If we support other non-fs stores without the write cache, then
            // we may have find a way to check whether we need the cleaner.
            let cleaner = TempFileCleaner::new(region_id, self.object_store.clone());
            let mut writer = ParquetWriter::new_with_object_store(
                self.object_store.clone(),
                request.metadata,
                request.index_config,
                indexer_builder,
                path_provider,
                metrics,
            )
            .await
            .with_file_cleaner(cleaner)
            .with_write_concurrency(if guard.is_some() {
                1
            } else {
                crate::sst::DEFAULT_WRITE_CONCURRENCY
            })
            .with_write_guard(guard)
            .with_memory_limits(limits);
            match request.sst_write_format {
                FormatType::PrimaryKey => {
                    writer
                        .write_all_flat_as_primary_key(
                            request.source,
                            override_sequence,
                            write_opts,
                        )
                        .await?
                }
                FormatType::Flat => {
                    writer
                        .write_all_flat(request.source, override_sequence, write_opts)
                        .await?
                }
            }
        };

        // Put parquet metadata to cache manager.
        if !sst_info.is_empty() && cache_manager.sst_meta_cache_on_write_enabled() {
            let runtime = match op_type {
                OperationType::Compact => common_runtime::compact_runtime(),
                OperationType::Flush => common_runtime::global_runtime(),
            };
            for sst in &sst_info {
                if let Some(parquet_metadata) = &sst.file_metadata {
                    let file_id = RegionFileId::new(region_id, sst.file_id);
                    let file_path = format!(
                        "region_id={}, file_id={}",
                        file_id.region_id(),
                        file_id.file_id()
                    );
                    let page_index_policy = if parquet_metadata.offset_index().is_some() {
                        PageIndexPolicy::Optional
                    } else {
                        PageIndexPolicy::Skip
                    };
                    let parquet_metadata = parquet_metadata.clone();
                    let region_metadata = region_metadata.clone();
                    let cache_manager = cache_manager.clone();
                    // Prepare and cache metadata in the background without delaying the SST write.
                    let task_runtime = runtime.clone();
                    runtime.spawn(async move {
                        let preparation = task_runtime.spawn_blocking(move || {
                            prepare_sst_meta_sync(
                                &file_path,
                                Arc::unwrap_or_clone(parquet_metadata),
                                Some(region_metadata),
                                page_index_policy,
                            )
                        }).await.context(crate::error::JoinSnafu).and_then(|result| result);
                        match preparation {
                            Ok(SstMetaPreparation::Prepared(metadata)) => {
                                cache_manager.put_compact_sst_meta(file_id, metadata.compact);
                                cache_manager.put_sst_meta_data(file_id, metadata.decoded).await;
                            }
                            Ok(SstMetaPreparation::DecodedOnly { encoding_error, .. }) => warn!(
                                encoding_error;
                                "Failed to encode parquet metadata for cache, file: {}",
                                file_id.file_id()
                            ),
                            Err(err) => {
                                warn!(err; "Failed to cache parquet metadata for {}", file_id.file_id());
                            }
                        }
                    });
                }
            }
        }

        Ok(sst_info)
    }

    /// Puts encoded SST bytes to the write cache (if enabled) and uploads it to the object store.
    ///
    /// Builds the SST's indexes from `data` into `sst_info.index_metadata` when
    /// `index_build_params` is set.
    pub(crate) async fn put_sst(
        &self,
        data: &bytes::Bytes,
        region_id: RegionId,
        sst_info: &mut SstInfo,
        cache_manager: &CacheManagerRef,
        write_buffer_size: ReadableSize,
        index_build_params: Option<EncodedSstIndexBuildParams>,
    ) -> Result<Metrics> {
        let region_file_id = RegionFileId::new(region_id, sst_info.file_id);
        if let Some(write_cache) = cache_manager.write_cache() {
            // Write to cache and upload to remote store
            let upload_request = SstUploadRequest {
                dest_path_provider: RegionFilePathFactory::new(
                    self.table_dir.clone(),
                    self.path_type,
                ),
                remote_store: self.object_store.clone(),
            };
            write_cache
                .put_and_upload_sst(
                    data,
                    region_id,
                    sst_info,
                    upload_request,
                    write_buffer_size,
                    index_build_params,
                )
                .await
        } else {
            let start = Instant::now();
            let cleaner =
                TempFileCleaner::new(region_file_id.region_id(), self.object_store.clone());
            let path_provider = RegionFilePathFactory::new(self.table_dir.clone(), self.path_type);
            let sst_file_path = path_provider.build_sst_file_path(region_file_id);
            let mut writer = self
                .object_store
                .writer_with(&sst_file_path)
                .chunk(write_buffer_size.as_bytes() as usize)
                .concurrent(DEFAULT_WRITE_CONCURRENCY)
                .await
                .context(OpenDalSnafu)?;
            if let Err(err) = writer.write(data.clone()).await.context(OpenDalSnafu) {
                cleaner.clean_by_file_id(region_file_id.file_id()).await;
                return Err(err);
            }
            if let Err(err) = writer.close().await.context(OpenDalSnafu) {
                cleaner.clean_by_file_id(region_file_id.file_id()).await;
                return Err(err);
            }
            let mut metrics = Metrics::new(WriteType::Flush);
            metrics.write_batch = start.elapsed();

            if let Some(params) = index_build_params {
                let start = Instant::now();
                let builder = params.indexer_builder(
                    self.puffin_manager_factory
                        .build(self.object_store.clone(), path_provider),
                    false,
                    self.intermediate_manager.clone(),
                );
                sst_info.index_metadata =
                    build_index_from_parquet_bytes(&builder, region_file_id, data).await;
                metrics.update_index = start.elapsed();
            }
            Ok(metrics)
        }
    }

    /// Lists the SST entries from the storage layer in the table directory.
    pub fn storage_sst_entries(&self) -> impl Stream<Item = Result<StorageSstEntry>> + use<> {
        let object_store = self.object_store.clone();
        let table_dir = self.table_dir.clone();

        try_stream! {
            let mut lister = object_store
                .lister_with(table_dir.as_str())
                .recursive(true)
                .await
                .context(OpenDalSnafu)?;

            while let Some(entry) = lister.try_next().await.context(OpenDalSnafu)? {
                let metadata = entry.metadata();
                if metadata.is_dir() {
                    continue;
                }

                let path = entry.path();
                if !path.ends_with(".parquet") && !path.ends_with(".puffin") {
                    continue;
                }

                let file_size = metadata.content_length();
                let file_size = if file_size == 0 { None } else { Some(file_size) };
                let last_modified_ms = metadata
                    .last_modified()
                    .map(|ts| Timestamp::new_millisecond(ts.into_inner().as_millisecond()));

                let entry = StorageSstEntry {
                    file_path: path.to_string(),
                    file_size,
                    last_modified_ms,
                    node_id: None,
                };

                yield entry;
            }
        }
    }
}

/// `OperationType` represents the origin of the `SstWriteRequest`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum OperationType {
    Flush,
    Compact,
}

/// Contents to build a SST.
pub struct SstWriteRequest {
    pub op_type: OperationType,
    pub metadata: RegionMetadataRef,
    pub source: FlatSource,
    pub cache_manager: CacheManagerRef,
    /// Optional uniform row sequence for writes that do not preserve sequences.
    /// Compaction and LastNonNull flushes pass `None` to retain input sequences.
    pub max_sequence: Option<SequenceNumber>,
    pub sst_write_format: FormatType,

    pub preserve_row_sequence: bool,

    /// Configs for index
    pub index_options: IndexOptions,
    pub index_config: IndexConfig,
    pub inverted_index_config: InvertedIndexConfig,
    pub fulltext_index_config: FulltextIndexConfig,
    pub bloom_filter_index_config: BloomFilterConfig,
}

/// Index inputs for SST bytes that were encoded outside [ParquetWriter].
pub(crate) struct EncodedSstIndexBuildParams {
    pub(crate) metadata: RegionMetadataRef,
    pub(crate) index_options: IndexOptions,
    pub(crate) inverted_index_config: InvertedIndexConfig,
    pub(crate) fulltext_index_config: FulltextIndexConfig,
    pub(crate) bloom_filter_index_config: BloomFilterConfig,
}

impl EncodedSstIndexBuildParams {
    pub(crate) fn indexer_builder(
        self,
        puffin_manager: SstPuffinManager,
        write_cache_enabled: bool,
        intermediate_manager: IntermediateManager,
    ) -> IndexerBuilderImpl {
        IndexerBuilderImpl {
            build_type: OperationType::Flush.into(),
            metadata: self.metadata,
            puffin_manager,
            write_cache_enabled,
            intermediate_manager,
            index_options: self.index_options,
            inverted_index_config: self.inverted_index_config,
            fulltext_index_config: self.fulltext_index_config,
            bloom_filter_index_config: self.bloom_filter_index_config,
        }
    }
}

/// Cleaner to remove temp files on the atomic write dir.
pub(crate) struct TempFileCleaner {
    region_id: RegionId,
    object_store: ObjectStore,
}

impl TempFileCleaner {
    /// Constructs the cleaner for the region and store.
    pub(crate) fn new(region_id: RegionId, object_store: ObjectStore) -> Self {
        Self {
            region_id,
            object_store,
        }
    }

    /// Removes the SST and index file from the local atomic dir by the file id.
    /// This only removes the initial index, since the index version is always 0 for a new SST, this method should be safe to pass 0.
    pub(crate) async fn clean_by_file_id(&self, file_id: FileId) {
        let sst_key = IndexKey::new(self.region_id, file_id, FileType::Parquet).to_string();
        let index_key = IndexKey::new(self.region_id, file_id, FileType::Puffin(0)).to_string();

        Self::clean_atomic_dir_files(&self.object_store, &[&sst_key, &index_key]).await;
    }

    /// Removes the files from the local atomic dir by their names.
    pub(crate) async fn clean_atomic_dir_files(
        local_store: &ObjectStore,
        names_to_remove: &[&str],
    ) {
        // We don't know the actual suffix of the file under atomic dir, so we have
        // to list the dir. The cost should be acceptable as there won't be to many files.
        let Ok(entries) = local_store.list(ATOMIC_WRITE_DIR).await.inspect_err(|e| {
            if e.kind() != ErrorKind::NotFound {
                common_telemetry::error!(e; "Failed to list tmp files for {:?}", names_to_remove)
            }
        }) else {
            return;
        };

        // In our case, we can ensure the file id is unique so it is safe to remove all files
        // with the same file id under the atomic write dir.
        let actual_files: Vec<_> = entries
            .into_iter()
            .filter_map(|entry| {
                if entry.metadata().is_dir() {
                    return None;
                }

                // Remove name that matches files_to_remove.
                let should_remove = names_to_remove
                    .iter()
                    .any(|file| entry.name().starts_with(file));
                if should_remove {
                    Some(entry.path().to_string())
                } else {
                    None
                }
            })
            .collect();

        common_telemetry::warn!(
            "Clean files {:?} under atomic write dir for {:?}",
            actual_files,
            names_to_remove
        );

        if let Err(e) = local_store.delete_iter(actual_files).await {
            common_telemetry::error!(e; "Failed to delete tmp file for {:?}", names_to_remove);
        }
    }
}

pub(crate) async fn new_fs_cache_store(root: &str) -> Result<ObjectStore> {
    // Preserve native filesystem prefixes such as Windows UNC shares.
    let atomic_write_dir = Path::new(root).join(ATOMIC_WRITE_DIR);
    clean_dir(&atomic_write_dir.to_string_lossy()).await?;

    // Compatible code. Remove this after a major release.
    let old_atomic_temp_dir = join_dir(root, OLD_ATOMIC_WRITE_DIR);
    clean_dir(&old_atomic_temp_dir).await?;

    let builder = Fs::default()
        .root(root)
        .atomic_write_dir(&atomic_write_dir.to_string_lossy());
    let store = ObjectStore::new(builder).context(OpenDalSnafu)?;

    Ok(with_instrument_layers(store, false))
}

/// Clean the directory.
async fn clean_dir(dir: &str) -> Result<()> {
    if tokio::fs::try_exists(dir)
        .await
        .context(CleanDirSnafu { dir })?
    {
        tokio::fs::remove_dir_all(dir)
            .await
            .context(CleanDirSnafu { dir })?;
    }

    Ok(())
}

/// Path provider for SST file and index file.
pub trait FilePathProvider: Send + Sync {
    /// Creates index file path of given file id. Version default to 0, and not shown in the path.
    fn build_index_file_path(&self, file_id: RegionFileId) -> String;

    /// Creates index file path of given index id (with version support).
    fn build_index_file_path_with_version(&self, index_id: RegionIndexId) -> String;

    /// Creates SST file path of given file id.
    fn build_sst_file_path(&self, file_id: RegionFileId) -> String;
}

/// Path provider that builds paths in local write cache.
#[derive(Clone)]
pub(crate) struct WriteCachePathProvider {
    file_cache: FileCacheRef,
}

impl WriteCachePathProvider {
    /// Creates a new `WriteCachePathProvider` instance.
    pub fn new(file_cache: FileCacheRef) -> Self {
        Self { file_cache }
    }
}

impl FilePathProvider for WriteCachePathProvider {
    fn build_index_file_path(&self, file_id: RegionFileId) -> String {
        let puffin_key = IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Puffin(0));
        self.file_cache.cache_file_path(puffin_key)
    }

    fn build_index_file_path_with_version(&self, index_id: RegionIndexId) -> String {
        let puffin_key = IndexKey::new(
            index_id.region_id(),
            index_id.file_id(),
            FileType::Puffin(index_id.version),
        );
        self.file_cache.cache_file_path(puffin_key)
    }

    fn build_sst_file_path(&self, file_id: RegionFileId) -> String {
        let parquet_file_key =
            IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Parquet);
        self.file_cache.cache_file_path(parquet_file_key)
    }
}

/// Path provider that builds paths in region storage path.
#[derive(Clone, Debug)]
pub(crate) struct RegionFilePathFactory {
    pub(crate) table_dir: String,
    pub(crate) path_type: PathType,
}

impl RegionFilePathFactory {
    /// Creates a new `RegionFilePathFactory` instance.
    pub fn new(table_dir: String, path_type: PathType) -> Self {
        Self {
            table_dir,
            path_type,
        }
    }
}

impl FilePathProvider for RegionFilePathFactory {
    fn build_index_file_path(&self, file_id: RegionFileId) -> String {
        location::index_file_path_legacy(&self.table_dir, file_id, self.path_type)
    }

    fn build_index_file_path_with_version(&self, index_id: RegionIndexId) -> String {
        location::index_file_path(&self.table_dir, index_id, self.path_type)
    }

    fn build_sst_file_path(&self, file_id: RegionFileId) -> String {
        location::sst_file_path(&self.table_dir, file_id, self.path_type)
    }
}

#[cfg(test)]
mod tests {
    use std::path::PathBuf;

    use bytes::Bytes;
    use common_test_util::temp_dir::create_temp_dir;

    use super::*;
    use crate::cache::CacheManager;
    use crate::cache::test_util::new_fs_store;
    use crate::test_util::TestEnv;
    use crate::test_util::sst_util::WriteChunkRecorder;

    #[rstest::rstest]
    #[tokio::test]
    async fn test_put_sst_uses_configured_buffer_size(
        #[values(false, true)] enable_write_cache: bool,
        #[values(1024, 4096)] chunk_size: usize,
    ) {
        let mut env = TestEnv::new().await;
        let remote_chunks = WriteChunkRecorder::default();
        let remote_store = env.init_object_store_manager().layer(remote_chunks.layer());
        let local_dir = create_temp_dir("encoded-sst-cache");
        let local_chunks = WriteChunkRecorder::default();
        let local_store =
            new_fs_store(local_dir.path().to_str().unwrap()).layer(local_chunks.layer());
        let write_cache = if enable_write_cache {
            Some(
                env.create_write_cache(local_store.clone(), ReadableSize::mb(10))
                    .await,
            )
        } else {
            None
        };
        let cache_manager = Arc::new(
            CacheManager::builder()
                .write_cache(write_cache.clone())
                .build(),
        );
        let access_layer = AccessLayer::new(
            "test",
            PathType::Bare,
            remote_store.clone(),
            env.get_puffin_manager(),
            env.get_intermediate_manager(),
        );
        let file_id = RegionFileId::new(RegionId::new(1024, 1), FileId::random());
        // Two full chunks and a partial tail distinguish the configured size from the default.
        let encoded = Bytes::from(
            (0..2 * chunk_size + 17)
                .map(|i| (i % 251) as u8)
                .collect::<Vec<_>>(),
        );
        access_layer
            .put_sst(
                &encoded,
                file_id.region_id(),
                &mut SstInfo {
                    file_id: file_id.file_id(),
                    ..Default::default()
                },
                &cache_manager,
                ReadableSize(chunk_size as u64),
                None,
            )
            .await
            .unwrap();

        let remote_path = location::sst_file_path("test", file_id, PathType::Bare);
        remote_chunks.assert_chunks(&remote_path, chunk_size, encoded.len());
        assert_eq!(
            remote_store.read(&remote_path).await.unwrap().to_bytes(),
            encoded
        );
        if let Some(write_cache) = write_cache {
            let key = IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Parquet);
            let cache_path = write_cache.file_cache().cache_file_path(key);
            local_chunks.assert_chunks(&cache_path, chunk_size, encoded.len());
            assert_eq!(
                local_store.read(&cache_path).await.unwrap().to_bytes(),
                encoded
            );
        }
    }

    #[rstest::rstest]
    #[case::flush(OperationType::Flush, false)]
    #[case::local_compaction(OperationType::Compact, false)]
    #[case::input_only_compaction(OperationType::Compact, true)]
    #[tokio::test]
    async fn test_written_sst_metadata_cache_policy(
        #[case] op_type: OperationType,
        #[case] input_only: bool,
    ) {
        use crate::cache::CacheStrategy;
        use crate::sst::file::FileMeta;
        use crate::sst::file_purger::NoopFilePurger;
        use crate::sst::parquet::reader::MetadataCacheMetrics;
        use crate::test_util::sst_util::{
            new_flat_source_from_record_batches, new_record_batch_by_range, sst_region_metadata,
        };

        let mut env = TestEnv::new().await;
        let store = env.init_object_store_manager();
        let access_layer = AccessLayer::new(
            "test",
            PathType::Bare,
            store,
            env.get_puffin_manager(),
            env.get_intermediate_manager(),
        );
        let builder = CacheManager::builder().sst_meta_cache_size(1024 * 1024);
        let cache = Arc::new(if input_only {
            builder.disable_sst_meta_cache_on_write().build()
        } else {
            builder.build()
        });
        let metadata = Arc::new(sst_region_metadata());
        let region_id = metadata.region_id;
        let infos =
            access_layer
                .write_sst(
                    SstWriteRequest {
                        op_type,
                        metadata,
                        source: new_flat_source_from_record_batches(vec![
                            new_record_batch_by_range(&["a", "b"], 0, 16),
                        ]),
                        cache_manager: cache.clone(),
                        max_sequence: None,
                        sst_write_format: FormatType::Flat,
                        preserve_row_sequence: false,
                        index_options: Default::default(),
                        index_config: Default::default(),
                        inverted_index_config: Default::default(),
                        fulltext_index_config: Default::default(),
                        bloom_filter_index_config: Default::default(),
                    },
                    &WriteOptions::default(),
                    &mut Metrics::new(WriteType::Compaction),
                )
                .await
                .unwrap();
        assert_eq!(1, infos.len());
        // A prewarming task owns a cache reference until both tiers are populated.
        // Drain it before checking absence so the negative assertion cannot race it.
        tokio::time::timeout(Duration::from_secs(10), async {
            while Arc::strong_count(&cache) > 1 {
                tokio::task::yield_now().await;
            }
        })
        .await
        .unwrap();
        let info = &infos[0];
        let file_id = RegionFileId::new(region_id, info.file_id);
        let cached = cache
            .get_sst_meta_data(
                file_id,
                &mut MetadataCacheMetrics::default(),
                PageIndexPolicy::Skip,
                &common_runtime::compact_runtime(),
            )
            .await;
        assert_eq!(!input_only, cached.is_some());

        let file = FileHandle::new(
            FileMeta {
                region_id,
                file_id: info.file_id,
                file_size: info.file_size,
                ..Default::default()
            },
            Arc::new(NoopFilePurger),
        );
        let path = file.file_path(access_layer.table_dir(), access_layer.path_type());
        let reader = access_layer
            .read_sst(file)
            .cache(CacheStrategy::Compaction(cache));
        let mut cold = MetadataCacheMetrics::default();
        let (first, _) = reader
            .read_parquet_metadata(&path, info.file_size, &mut cold, PageIndexPolicy::Skip)
            .await
            .unwrap();
        assert_eq!(input_only, cold.num_reads > 0);
        let mut warm = MetadataCacheMetrics::default();
        let (second, _) = reader
            .read_parquet_metadata(&path, info.file_size, &mut warm, PageIndexPolicy::Skip)
            .await
            .unwrap();
        assert_eq!(16, second.parquet_metadata().file_metadata().num_rows());
        assert!(Arc::ptr_eq(&first, &second));
        assert_eq!(0, warm.num_reads);
        assert_eq!(1, warm.mem_cache_hit);
    }

    #[tokio::test]
    async fn test_new_fs_cache_store() {
        let root = common_test_util::temp_dir::create_temp_dir("fs-cache-store");
        let dirs = [
            root.path().join(ATOMIC_WRITE_DIR),
            PathBuf::from(join_dir(
                root.path().to_str().unwrap(),
                OLD_ATOMIC_WRITE_DIR,
            )),
        ];
        for dir in &dirs {
            tokio::fs::create_dir_all(dir).await.unwrap();
            tokio::fs::write(dir.join("stale"), b"stale").await.unwrap();
        }
        let store = new_fs_cache_store(root.path().to_str().unwrap())
            .await
            .unwrap();
        for dir in &dirs {
            assert!(!dir.join("stale").exists());
        }
        store.write("index", "contents").await.unwrap();
        assert_eq!(
            tokio::fs::read(root.path().join("index")).await.unwrap(),
            b"contents"
        );
    }
}