mito2/region/

utils.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::sync::Arc;
use std::time::Instant;

use common_base::readable_size::ReadableSize;
use common_telemetry::{debug, error, info};
use futures::future::try_join_all;
use object_store::manager::ObjectStoreManagerRef;
use snafu::{ResultExt, ensure};
use store_api::metadata::RegionMetadataRef;
use store_api::region_request::PathType;
use store_api::storage::{FileId, IndexVersion, RegionId};

use crate::access_layer::AccessLayerRef;
use crate::config::MitoConfig;
use crate::error::{self, InvalidSourceAndTargetRegionSnafu, Result};
use crate::manifest::action::RegionManifest;
use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions};
use crate::region::opener::get_object_store;
use crate::region::options::RegionOptions;
use crate::sst::file::{RegionFileId, RegionIndexId};
use crate::sst::location;

/// A loader for loading metadata from a region dir.
#[derive(Debug, Clone)]
pub struct RegionMetadataLoader {
    config: Arc<MitoConfig>,
    object_store_manager: ObjectStoreManagerRef,
    lease_enabled: bool,
}

impl RegionMetadataLoader {
    /// Creates a new `RegionMetadataLoader`.
    pub fn new(
        config: Arc<MitoConfig>,
        object_store_manager: ObjectStoreManagerRef,
        lease_enabled: bool,
    ) -> Self {
        Self {
            config,
            object_store_manager,
            lease_enabled,
        }
    }

    /// Loads the metadata of the region from the region dir.
    pub async fn load(
        &self,
        region_dir: &str,
        region_options: &RegionOptions,
    ) -> Result<Option<RegionMetadataRef>> {
        let manifest = self.load_manifest(region_dir, region_options).await?;
        Ok(manifest.map(|m| m.metadata.clone()))
    }

    /// Loads the manifest of the region from the region dir.
    pub async fn load_manifest(
        &self,
        region_dir: &str,
        region_options: &RegionOptions,
    ) -> Result<Option<Arc<RegionManifest>>> {
        let object_store = get_object_store(&region_options.storage, &self.object_store_manager)?;
        let region_manifest_options =
            RegionManifestOptions::new(&self.config, region_dir, &object_store)
                .with_conditional_publication(
                    region_options.requires_conditional_manifest_publication(self.lease_enabled),
                );
        let Some(manifest_manager) =
            RegionManifestManager::open(region_manifest_options, &Default::default()).await?
        else {
            return Ok(None);
        };

        let manifest = manifest_manager.manifest();
        Ok(Some(manifest))
    }
}

/// A copier for copying files from a region to another region.
#[derive(Debug, Clone)]
pub struct RegionFileCopier {
    access_layer: AccessLayerRef,
}

/// A descriptor for a file.
#[derive(Debug, Clone, Copy)]
pub enum FileDescriptor {
    /// An index file.
    Index {
        file_id: FileId,
        version: IndexVersion,
        size: u64,
    },
    /// A data file.
    Data { file_id: FileId, size: u64 },
}

impl FileDescriptor {
    pub fn size(&self) -> u64 {
        match self {
            FileDescriptor::Index { size, .. } => *size,
            FileDescriptor::Data { size, .. } => *size,
        }
    }

    fn file_path(&self, region_id: RegionId, table_dir: &str, path_type: PathType) -> String {
        match *self {
            FileDescriptor::Index {
                file_id, version, ..
            } => location::index_file_path(
                table_dir,
                RegionIndexId::new(RegionFileId::new(region_id, file_id), version),
                path_type,
            ),
            FileDescriptor::Data { file_id, .. } => {
                location::sst_file_path(table_dir, RegionFileId::new(region_id, file_id), path_type)
            }
        }
    }
}

/// Builds the source and target file paths for a given file descriptor.
///
/// # Arguments
///
/// * `source_region_id`: The ID of the source region.
/// * `target_region_id`: The ID of the target region.
/// * `file_id`: The ID of the file.
///
/// # Returns
///
/// A tuple containing the source and target file paths.
fn build_copy_file_paths(
    source_region_id: RegionId,
    target_region_id: RegionId,
    file_descriptor: FileDescriptor,
    table_dir: &str,
    path_type: PathType,
) -> (String, String) {
    (
        file_descriptor.file_path(source_region_id, table_dir, path_type),
        file_descriptor.file_path(target_region_id, table_dir, path_type),
    )
}

fn build_delete_file_path(
    target_region_id: RegionId,
    file_descriptor: FileDescriptor,
    table_dir: &str,
    path_type: PathType,
) -> String {
    file_descriptor.file_path(target_region_id, table_dir, path_type)
}

impl RegionFileCopier {
    pub fn new(access_layer: AccessLayerRef) -> Self {
        Self { access_layer }
    }

    /// Copies files from a source region to a target region.
    ///
    /// # Arguments
    ///
    /// * `source_region_id`: The ID of the source region.
    /// * `target_region_id`: The ID of the target region.
    /// * `file_ids`: The IDs of the files to copy.
    pub async fn copy_files(
        &self,
        source_region_id: RegionId,
        target_region_id: RegionId,
        file_ids: Vec<FileDescriptor>,
        parallelism: usize,
    ) -> Result<()> {
        ensure!(
            source_region_id.table_id() == target_region_id.table_id(),
            InvalidSourceAndTargetRegionSnafu {
                source_region_id,
                target_region_id,
            },
        );
        let table_dir = self.access_layer.table_dir();
        let path_type = self.access_layer.path_type();
        let object_store = self.access_layer.object_store();

        info!(
            "Copying {} files from region {} to region {}",
            file_ids.len(),
            source_region_id,
            target_region_id
        );
        debug!(
            "Copying files: {:?} from region {} to region {}",
            file_ids, source_region_id, target_region_id
        );
        let mut tasks = Vec::with_capacity(parallelism);
        for skip in 0..parallelism {
            let target_file_ids = file_ids.iter().skip(skip).step_by(parallelism).copied();
            let object_store = object_store.clone();
            tasks.push(async move {
                for file_desc in target_file_ids {
                    let (source_path, target_path) = build_copy_file_paths(
                        source_region_id,
                        target_region_id,
                        file_desc,
                        table_dir,
                        path_type,
                    );
                    let now = Instant::now();
                    object_store
                        .copy(&source_path, &target_path)
                        .await
                        .inspect_err(
                            |e| error!(e; "Failed to copy file {} to {}", source_path, target_path),
                        )
                        .context(error::OpenDalSnafu)?;
                    let file_size = ReadableSize(file_desc.size());
                    info!(
                        "Copied file {} to {}, file size: {}, elapsed: {:?}",
                        source_path,
                        target_path,
                        file_size,
                        now.elapsed(),
                    );
                }

                Ok(())
            });
        }

        if let Err(err) = try_join_all(tasks).await {
            error!(err; "Failed to copy files from region {} to region {}", source_region_id, target_region_id);
            self.clean_target_region(target_region_id, file_ids).await;
            return Err(err);
        }

        Ok(())
    }

    /// Cleans the copied files from the target region.
    async fn clean_target_region(&self, target_region_id: RegionId, file_ids: Vec<FileDescriptor>) {
        let table_dir = self.access_layer.table_dir();
        let path_type = self.access_layer.path_type();
        let object_store = self.access_layer.object_store();
        let delete_file_path = file_ids
            .into_iter()
            .map(|file_descriptor| {
                build_delete_file_path(target_region_id, file_descriptor, table_dir, path_type)
            })
            .collect::<Vec<_>>();
        debug!(
            "Deleting files: {:?} after failed to copy files to target region {}",
            delete_file_path, target_region_id
        );
        if let Err(err) = object_store.delete_iter(delete_file_path).await {
            error!(err; "Failed to delete files from region {}", target_region_id);
        }
    }
}

#[cfg(test)]
mod tests {
    #[cfg(feature = "hdfs-object-store")]
    use object_store::ObjectStore;
    #[cfg(feature = "hdfs-object-store")]
    use object_store::layers::HdfsCompatibilityLayer;
    #[cfg(feature = "hdfs-object-store")]
    use object_store::services::Fs;

    use super::*;
    #[cfg(feature = "hdfs-object-store")]
    use crate::access_layer::AccessLayer;
    #[cfg(feature = "hdfs-object-store")]
    use crate::sst::index::intermediate::IntermediateManager;
    #[cfg(feature = "hdfs-object-store")]
    use crate::sst::index::puffin_manager::PuffinManagerFactory;

    #[test]
    fn test_build_copy_file_paths() {
        common_telemetry::init_default_ut_logging();
        let file_id = FileId::random();
        let source_region_id = RegionId::new(1, 1);
        let target_region_id = RegionId::new(1, 2);
        let file_descriptor = FileDescriptor::Data { file_id, size: 100 };
        let table_dir = "/table_dir";
        let path_type = PathType::Bare;
        let (source_path, target_path) = build_copy_file_paths(
            source_region_id,
            target_region_id,
            file_descriptor,
            table_dir,
            path_type,
        );
        assert_eq!(
            source_path,
            format!("/table_dir/1_0000000001/{}.parquet", file_id)
        );
        assert_eq!(
            target_path,
            format!("/table_dir/1_0000000002/{}.parquet", file_id)
        );

        let version = 1;
        let file_descriptor = FileDescriptor::Index {
            file_id,
            version,
            size: 100,
        };
        let (source_path, target_path) = build_copy_file_paths(
            source_region_id,
            target_region_id,
            file_descriptor,
            table_dir,
            path_type,
        );
        assert_eq!(
            source_path,
            format!(
                "/table_dir/1_0000000001/index/{}.{}.puffin",
                file_id, version
            )
        );
        assert_eq!(
            target_path,
            format!(
                "/table_dir/1_0000000002/index/{}.{}.puffin",
                file_id, version
            )
        );
    }

    #[test]
    fn test_build_delete_file_path() {
        common_telemetry::init_default_ut_logging();
        let file_id = FileId::random();
        let target_region_id = RegionId::new(1, 2);
        let table_dir = "/table_dir";
        let path_type = PathType::Bare;

        let file_descriptor = FileDescriptor::Data { file_id, size: 100 };
        let path = build_delete_file_path(target_region_id, file_descriptor, table_dir, path_type);
        assert_eq!(path, format!("/table_dir/1_0000000002/{}.parquet", file_id));

        let file_descriptor = FileDescriptor::Index {
            file_id,
            version: 1,
            size: 100,
        };
        let path = build_delete_file_path(target_region_id, file_descriptor, table_dir, path_type);
        assert_eq!(
            path,
            format!("/table_dir/1_0000000002/index/{}.1.puffin", file_id)
        );
    }

    #[cfg(feature = "hdfs-object-store")]
    #[tokio::test]
    async fn test_copy_region_files_with_hdfs_fallback() {
        let (temp_dir, puffin_manager) =
            PuffinManagerFactory::new_for_test_async("hdfs-copy-region").await;
        let intermediate_manager = IntermediateManager::init_fs(temp_dir.path().to_string_lossy())
            .await
            .unwrap();
        let storage_dir = temp_dir.path().join("storage");
        std::fs::create_dir(&storage_dir).unwrap();
        let object_store = ObjectStore::new(Fs::default().root(storage_dir.to_str().unwrap()))
            .unwrap()
            .layer(HdfsCompatibilityLayer::new_for_test());
        let access_layer = Arc::new(AccessLayer::new(
            "table_dir",
            PathType::Bare,
            object_store.clone(),
            puffin_manager,
            intermediate_manager,
        ));
        let copier = RegionFileCopier::new(access_layer);
        let source_region_id = RegionId::new(1, 1);
        let target_region_id = RegionId::new(1, 2);
        let file_id = FileId::random();
        let descriptor = FileDescriptor::Data { file_id, size: 8 };
        let (source_path, target_path) = build_copy_file_paths(
            source_region_id,
            target_region_id,
            descriptor,
            "table_dir",
            PathType::Bare,
        );
        object_store.write(&source_path, "contents").await.unwrap();

        copier
            .copy_files(source_region_id, target_region_id, vec![descriptor], 1)
            .await
            .unwrap();

        assert_eq!(
            b"contents",
            object_store
                .read(&target_path)
                .await
                .unwrap()
                .to_bytes()
                .as_ref()
        );
    }
}