metric_engine/

metadata_region.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::collections::hash_map::Entry;
use std::collections::{BTreeMap, HashMap};
use std::sync::{Arc, Mutex, Weak};
use std::time::Duration;

use api::v1::helper::row;
use api::v1::value::ValueData;
use api::v1::{ColumnDataType, ColumnSchema, Rows, SemanticType};
use async_stream::try_stream;
use base64::Engine;
use base64::engine::general_purpose::STANDARD_NO_PAD;
use common_base::readable_size::ReadableSize;
use common_recordbatch::{RecordBatch, SendableRecordBatchStream};
use common_telemetry::{debug, info, warn};
use datafusion::prelude::{col, lit};
use futures_util::TryStreamExt;
use futures_util::stream::BoxStream;
use mito2::engine::MitoEngine;
use moka::future::Cache;
use moka::policy::EvictionPolicy;
use snafu::{OptionExt, ResultExt};
use store_api::metadata::ColumnMetadata;
use store_api::metric_engine_consts::{
    METADATA_SCHEMA_KEY_COLUMN_INDEX, METADATA_SCHEMA_KEY_COLUMN_NAME,
    METADATA_SCHEMA_TIMESTAMP_COLUMN_NAME, METADATA_SCHEMA_VALUE_COLUMN_INDEX,
    METADATA_SCHEMA_VALUE_COLUMN_NAME,
};
use store_api::region_engine::RegionEngine;
use store_api::region_request::{RegionDeleteRequest, RegionPutRequest, RegionRequest};
use store_api::storage::{RegionId, ScanRequest};
use tokio::sync::{
    OwnedRwLockReadGuard, OwnedRwLockWriteGuard, RwLock, RwLockReadGuard, RwLockWriteGuard,
};

use crate::error::{
    CacheGetSnafu, CollectRecordBatchStreamSnafu, DecodeColumnValueSnafu,
    DeserializeColumnMetadataSnafu, LogicalRegionNotFoundSnafu, MitoReadOperationSnafu,
    MitoWriteOperationSnafu, ParseRegionIdSnafu, Result,
};
use crate::utils;

const REGION_PREFIX: &str = "__region_";
const COLUMN_PREFIX: &str = "__column_";

/// The other two fields key and value will be used as a k-v storage.
/// It contains two group of key:
/// - `__region_<LOGICAL_REGION_ID>` is used for marking table existence. It doesn't have value.
/// - `__column_<LOGICAL_REGION_ID>_<COLUMN_NAME>` is used for marking column existence,
///   the value is column's semantic type. To avoid the key conflict, this column key
///   will be encoded by base64([STANDARD_NO_PAD]).
///
/// This is a generic handler like [MetricEngine](crate::engine::MetricEngine). It
/// will handle all the metadata related operations across physical tables. Thus
/// every operation should be associated to a [RegionId], which is the physical
/// table id + region sequence. This handler will transform the region group by
/// itself.
pub struct MetadataRegion {
    pub(crate) mito: MitoEngine,
    /// The cache for contents(key-value pairs) of region metadata.
    ///
    /// The cache should be invalidated when any new values are put into the metadata region or any
    /// values are deleted from the metadata region.
    cache: Cache<RegionId, RegionMetadataCacheEntry>,
    /// Serializes cache fills with metadata writes and invalidation per metadata region.
    ///
    /// Holds weak references only; strong references live in [`CacheAccessLockLease`]s
    /// returned by [`Self::cache_access_lock`]. The last lease dropped for a region
    /// removes its entry, so the map self-cleans on success, error, or cancellation
    /// without coupling cleanup to region drop.
    cache_access_locks: CacheAccessLockRegistry,
    /// Logical lock for operations that need to be serialized. Like update & read region columns.
    ///
    /// Region entry will be registered on creating and opening logical region, and deregistered on
    /// removing logical region.
    logical_region_lock: RwLock<HashMap<RegionId, Arc<RwLock<()>>>>,
}

#[derive(Clone)]
struct RegionMetadataCacheEntry {
    key_values: Arc<BTreeMap<String, String>>,
    size: usize,
}

/// Weak index of per-region cache access locks.
///
/// The mutex is never held across `.await`.
type CacheAccessLockRegistry = Arc<Mutex<HashMap<RegionId, Weak<RwLock<()>>>>>;

/// Lease that keeps a per-region cache access lock alive.
///
/// Dropping the last lease for a region removes the matching registry entry.
struct CacheAccessLockLease {
    registry: CacheAccessLockRegistry,
    region_id: RegionId,
    /// Always `Some` while the lease is alive.
    ///
    /// `Option` lets `Drop` release the strong `Arc` before pruning the weak entry.
    lock: Option<Arc<RwLock<()>>>,
}

impl CacheAccessLockLease {
    async fn read(&self) -> RwLockReadGuard<'_, ()> {
        self.lock().read().await
    }

    async fn write(&self) -> RwLockWriteGuard<'_, ()> {
        self.lock().write().await
    }

    fn lock(&self) -> &RwLock<()> {
        self.lock
            .as_deref()
            // Safety: `lock` is initialized when the lease is created and is taken only
            // by `Drop`. `read` and `write` borrow `self`, so `Drop` cannot run while
            // this reference is in use.
            .expect("cache access lock lease must hold a lock")
    }
}

impl Drop for CacheAccessLockLease {
    fn drop(&mut self) {
        let Some(lock) = self.lock.take() else {
            return;
        };
        let weak = Arc::downgrade(&lock);
        // Release this lease before pruning; concurrent drops must not observe
        // each other's strong refs.
        drop(lock);

        let mut registry = self
            .registry
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        let Some(current) = registry.get(&self.region_id) else {
            return;
        };
        // `ptr_eq` protects a newer lock for the same region; `upgrade` ensures it is dead.
        if current.ptr_eq(&weak) && current.upgrade().is_none() {
            registry.remove(&self.region_id);
        }
    }
}

/// The max size of the region metadata cache.
const MAX_CACHE_SIZE: u64 = ReadableSize::mb(128).as_bytes();
/// The TTL of the region metadata cache.
const CACHE_TTL: Duration = Duration::from_secs(5 * 60);

impl MetadataRegion {
    pub fn new(mito: MitoEngine) -> Self {
        let cache = Cache::builder()
            .max_capacity(MAX_CACHE_SIZE)
            // Use the LRU eviction policy to minimize frequent mito scans.
            // Recently accessed items are retained longer in the cache.
            .eviction_policy(EvictionPolicy::lru())
            .time_to_live(CACHE_TTL)
            .weigher(|_, v: &RegionMetadataCacheEntry| v.size as u32)
            .build();
        Self {
            mito,
            cache,
            cache_access_locks: CacheAccessLockRegistry::default(),
            logical_region_lock: RwLock::new(HashMap::new()),
        }
    }

    /// Open a logical region.
    ///
    /// Returns true if the logical region is opened for the first time.
    pub async fn open_logical_region(&self, logical_region_id: RegionId) -> bool {
        match self
            .logical_region_lock
            .write()
            .await
            .entry(logical_region_id)
        {
            Entry::Occupied(_) => false,
            Entry::Vacant(vacant_entry) => {
                vacant_entry.insert(Arc::new(RwLock::new(())));
                true
            }
        }
    }

    /// Retrieve a read lock guard of given logical region id.
    pub async fn read_lock_logical_region(
        &self,
        logical_region_id: RegionId,
    ) -> Result<OwnedRwLockReadGuard<()>> {
        let lock = self
            .logical_region_lock
            .read()
            .await
            .get(&logical_region_id)
            .context(LogicalRegionNotFoundSnafu {
                region_id: logical_region_id,
            })?
            .clone();
        Ok(RwLock::read_owned(lock).await)
    }

    /// Retrieve a write lock guard of given logical region id.
    pub async fn write_lock_logical_region(
        &self,
        logical_region_id: RegionId,
    ) -> Result<OwnedRwLockWriteGuard<()>> {
        let lock = self
            .logical_region_lock
            .read()
            .await
            .get(&logical_region_id)
            .context(LogicalRegionNotFoundSnafu {
                region_id: logical_region_id,
            })?
            .clone();
        Ok(RwLock::write_owned(lock).await)
    }

    /// Remove a registered logical region from metadata.
    ///
    /// This method doesn't check if the previous key exists.
    pub async fn remove_logical_region(
        &self,
        physical_region_id: RegionId,
        logical_region_id: RegionId,
    ) -> Result<()> {
        // concat region key
        let region_id = utils::to_metadata_region_id(physical_region_id);
        let region_key = Self::concat_region_key(logical_region_id);

        // concat column keys
        let logical_columns = self
            .logical_columns(physical_region_id, logical_region_id)
            .await?;
        let mut column_keys = logical_columns
            .into_iter()
            .map(|(col, _)| Self::concat_column_key(logical_region_id, &col))
            .collect::<Vec<_>>();

        // remove region key and column keys
        column_keys.push(region_key);
        self.delete(region_id, &column_keys).await?;

        self.logical_region_lock
            .write()
            .await
            .remove(&logical_region_id);

        Ok(())
    }

    // TODO(ruihang): avoid using `get_all`
    /// Get all the columns of a given logical region.
    /// Return a list of (column_name, column_metadata).
    pub async fn logical_columns(
        &self,
        physical_region_id: RegionId,
        logical_region_id: RegionId,
    ) -> Result<Vec<(String, ColumnMetadata)>> {
        let metadata_region_id = utils::to_metadata_region_id(physical_region_id);
        let region_column_prefix = Self::concat_column_key_prefix(logical_region_id);

        let mut columns = vec![];
        for (k, v) in self
            .get_all_with_prefix(metadata_region_id, &region_column_prefix)
            .await?
        {
            if !k.starts_with(&region_column_prefix) {
                continue;
            }
            // Safety: we have checked the prefix
            let (_, column_name) = Self::parse_column_key(&k)?.unwrap();
            let column_metadata = Self::deserialize_column_metadata(&v)?;
            columns.push((column_name, column_metadata));
        }

        Ok(columns)
    }

    /// Return all logical regions associated with the physical region.
    pub async fn logical_regions(&self, physical_region_id: RegionId) -> Result<Vec<RegionId>> {
        let metadata_region_id = utils::to_metadata_region_id(physical_region_id);

        let mut regions = vec![];
        for k in self
            .get_all_key_with_prefix(metadata_region_id, REGION_PREFIX)
            .await?
        {
            if !k.starts_with(REGION_PREFIX) {
                continue;
            }
            // Safety: we have checked the prefix
            let region_id = Self::parse_region_key(&k).unwrap();
            let region_id = region_id.parse::<u64>().unwrap().into();
            regions.push(region_id);
        }

        Ok(regions)
    }
}

// utils to concat and parse key/value
impl MetadataRegion {
    pub fn concat_region_key(region_id: RegionId) -> String {
        format!("{REGION_PREFIX}{}", region_id.as_u64())
    }

    /// Column name will be encoded by base64([STANDARD_NO_PAD])
    pub fn concat_column_key(region_id: RegionId, column_name: &str) -> String {
        let encoded_column_name = STANDARD_NO_PAD.encode(column_name);
        format!(
            "{COLUMN_PREFIX}{}_{}",
            region_id.as_u64(),
            encoded_column_name
        )
    }

    /// Concat a column key prefix without column name
    pub fn concat_column_key_prefix(region_id: RegionId) -> String {
        format!("{COLUMN_PREFIX}{}_", region_id.as_u64())
    }

    pub fn parse_region_key(key: &str) -> Option<&str> {
        key.strip_prefix(REGION_PREFIX)
    }

    /// Parse column key to (logical_region_id, column_name)
    pub fn parse_column_key(key: &str) -> Result<Option<(RegionId, String)>> {
        if let Some(stripped) = key.strip_prefix(COLUMN_PREFIX) {
            let mut iter = stripped.split('_');

            let region_id_raw = iter.next().unwrap();
            let region_id = region_id_raw
                .parse::<u64>()
                .with_context(|_| ParseRegionIdSnafu { raw: region_id_raw })?
                .into();

            let encoded_column_name = iter.next().unwrap();
            let column_name = STANDARD_NO_PAD
                .decode(encoded_column_name)
                .context(DecodeColumnValueSnafu)?;

            Ok(Some((region_id, String::from_utf8(column_name).unwrap())))
        } else {
            Ok(None)
        }
    }

    pub fn serialize_column_metadata(column_metadata: &ColumnMetadata) -> String {
        serde_json::to_string(column_metadata).unwrap()
    }

    pub fn deserialize_column_metadata(column_metadata: &str) -> Result<ColumnMetadata> {
        serde_json::from_str(column_metadata).with_context(|_| DeserializeColumnMetadataSnafu {
            raw: column_metadata,
        })
    }
}

/// Decode a record batch stream to a stream of items.
pub fn decode_batch_stream<T: Send + 'static>(
    mut record_batch_stream: SendableRecordBatchStream,
    decode: fn(RecordBatch) -> Vec<T>,
) -> BoxStream<'static, Result<T>> {
    let stream = try_stream! {
        while let Some(batch) = record_batch_stream.try_next().await.context(CollectRecordBatchStreamSnafu)? {
            for item in decode(batch) {
                yield item;
            }
        }
    };
    Box::pin(stream)
}

/// Decode a record batch to a list of key and value.
fn decode_record_batch_to_key_and_value(batch: RecordBatch) -> Vec<(String, String)> {
    let keys = batch.iter_column_as_string(0);
    let values = batch.iter_column_as_string(1);
    keys.zip(values)
        .filter_map(|(k, v)| match (k, v) {
            (Some(k), Some(v)) => Some((k, v)),
            (Some(k), None) => Some((k, "".to_string())),
            (None, _) => None,
        })
        .collect::<Vec<_>>()
}

/// Decode a record batch to a list of key.
fn decode_record_batch_to_key(batch: RecordBatch) -> Vec<String> {
    batch.iter_column_as_string(0).flatten().collect::<Vec<_>>()
}

// simulate to `KvBackend`
//
// methods in this block assume the given region id is transformed.
impl MetadataRegion {
    fn build_prefix_read_request(prefix: &str, key_only: bool) -> ScanRequest {
        let filter_expr = col(METADATA_SCHEMA_KEY_COLUMN_NAME).like(lit(prefix));

        let projection = if key_only {
            vec![METADATA_SCHEMA_KEY_COLUMN_INDEX]
        } else {
            vec![
                METADATA_SCHEMA_KEY_COLUMN_INDEX,
                METADATA_SCHEMA_VALUE_COLUMN_INDEX,
            ]
        };
        ScanRequest {
            projection: Some(projection),
            filters: vec![filter_expr],
            ..Default::default()
        }
    }

    fn build_read_request() -> ScanRequest {
        let projection = vec![
            METADATA_SCHEMA_KEY_COLUMN_INDEX,
            METADATA_SCHEMA_VALUE_COLUMN_INDEX,
        ];
        ScanRequest {
            projection: Some(projection),
            ..Default::default()
        }
    }

    async fn load_all(&self, metadata_region_id: RegionId) -> Result<RegionMetadataCacheEntry> {
        let scan_req = MetadataRegion::build_read_request();
        let record_batch_stream = self
            .mito
            .scan_to_stream(metadata_region_id, scan_req)
            .await
            .context(MitoReadOperationSnafu)?;

        let kv = decode_batch_stream(record_batch_stream, decode_record_batch_to_key_and_value)
            .try_collect::<BTreeMap<_, _>>()
            .await?;
        let mut size = 0;
        for (k, v) in kv.iter() {
            size += k.len();
            size += v.len();
        }
        let kv = Arc::new(kv);
        Ok(RegionMetadataCacheEntry {
            key_values: kv,
            size,
        })
    }

    /// Acquires the cache access lock lease for `metadata_region_id`.
    ///
    /// The lease must be kept alive for as long as any guard taken from it.
    fn cache_access_lock(&self, metadata_region_id: RegionId) -> CacheAccessLockLease {
        let mut registry = self
            .cache_access_locks
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        let lock = match registry.get(&metadata_region_id).and_then(Weak::upgrade) {
            Some(lock) => lock,
            None => {
                let lock = Arc::new(RwLock::new(()));
                registry.insert(metadata_region_id, Arc::downgrade(&lock));
                lock
            }
        };

        CacheAccessLockLease {
            registry: Arc::clone(&self.cache_access_locks),
            region_id: metadata_region_id,
            lock: Some(lock),
        }
    }

    async fn get_all_with_prefix(
        &self,
        metadata_region_id: RegionId,
        prefix: &str,
    ) -> Result<HashMap<String, String>> {
        let cache_access_lock = self.cache_access_lock(metadata_region_id);
        let _cache_guard = cache_access_lock.read().await;
        let region_metadata = self
            .cache
            .try_get_with(metadata_region_id, self.load_all(metadata_region_id))
            .await
            .context(CacheGetSnafu)?;

        let mut result = HashMap::new();
        get_all_with_prefix(&region_metadata, prefix, |k, v| {
            result.insert(k.to_string(), v.to_string());
            Ok(())
        })?;
        Ok(result)
    }

    pub async fn get_all_key_with_prefix(
        &self,
        region_id: RegionId,
        prefix: &str,
    ) -> Result<Vec<String>> {
        let scan_req = MetadataRegion::build_prefix_read_request(prefix, true);
        let record_batch_stream = self
            .mito
            .scan_to_stream(region_id, scan_req)
            .await
            .context(MitoReadOperationSnafu)?;

        decode_batch_stream(record_batch_stream, decode_record_batch_to_key)
            .try_collect::<Vec<_>>()
            .await
    }

    /// Delete the given keys. For performance consideration, this method
    /// doesn't check if those keys exist or not.
    async fn delete(&self, metadata_region_id: RegionId, keys: &[String]) -> Result<()> {
        let delete_request = Self::build_delete_request(keys);
        self.write_metadata(metadata_region_id, RegionRequest::Delete(delete_request))
            .await
    }

    /// Writes metadata and invalidates the corresponding cache entry.
    async fn write_metadata(
        &self,
        metadata_region_id: RegionId,
        request: RegionRequest,
    ) -> Result<()> {
        let cache_access_lock = self.cache_access_lock(metadata_region_id);
        let _cache_guard = cache_access_lock.write().await;
        self.mito
            .handle_request(metadata_region_id, request)
            .await
            .context(MitoWriteOperationSnafu)?;
        self.cache.invalidate(&metadata_region_id).await;

        Ok(())
    }

    pub(crate) fn build_put_request_from_iter(
        kv: impl Iterator<Item = (String, String)>,
    ) -> RegionPutRequest {
        let cols = vec![
            ColumnSchema {
                column_name: METADATA_SCHEMA_TIMESTAMP_COLUMN_NAME.to_string(),
                datatype: ColumnDataType::TimestampMillisecond as _,
                semantic_type: SemanticType::Timestamp as _,
                ..Default::default()
            },
            ColumnSchema {
                column_name: METADATA_SCHEMA_KEY_COLUMN_NAME.to_string(),
                datatype: ColumnDataType::String as _,
                semantic_type: SemanticType::Tag as _,
                ..Default::default()
            },
            ColumnSchema {
                column_name: METADATA_SCHEMA_VALUE_COLUMN_NAME.to_string(),
                datatype: ColumnDataType::String as _,
                semantic_type: SemanticType::Field as _,
                ..Default::default()
            },
        ];
        let rows = Rows {
            schema: cols,
            rows: kv
                .into_iter()
                .map(|(key, value)| {
                    row(vec![
                        ValueData::TimestampMillisecondValue(0),
                        ValueData::StringValue(key),
                        ValueData::StringValue(value),
                    ])
                })
                .collect(),
        };

        RegionPutRequest {
            rows,
            hint: None,
            partition_expr_version: None,
        }
    }

    fn build_delete_request(keys: &[String]) -> RegionDeleteRequest {
        let cols = vec![
            ColumnSchema {
                column_name: METADATA_SCHEMA_TIMESTAMP_COLUMN_NAME.to_string(),
                datatype: ColumnDataType::TimestampMillisecond as _,
                semantic_type: SemanticType::Timestamp as _,
                ..Default::default()
            },
            ColumnSchema {
                column_name: METADATA_SCHEMA_KEY_COLUMN_NAME.to_string(),
                datatype: ColumnDataType::String as _,
                semantic_type: SemanticType::Tag as _,
                ..Default::default()
            },
        ];
        let rows = keys
            .iter()
            .map(|key| {
                row(vec![
                    ValueData::TimestampMillisecondValue(0),
                    ValueData::StringValue(key.clone()),
                ])
            })
            .collect();
        let rows = Rows { schema: cols, rows };

        RegionDeleteRequest {
            rows,
            hint: None,
            partition_expr_version: None,
        }
    }

    /// Add logical regions to the metadata region.
    pub async fn add_logical_regions(
        &self,
        physical_region_id: RegionId,
        write_region_id: bool,
        logical_regions: impl Iterator<Item = (RegionId, HashMap<&str, &ColumnMetadata>)>,
    ) -> Result<()> {
        let metadata_region_id = utils::to_metadata_region_id(physical_region_id);
        let iter = logical_regions
            .into_iter()
            .flat_map(|(logical_region_id, column_metadatas)| {
                if write_region_id {
                    Some((
                        MetadataRegion::concat_region_key(logical_region_id),
                        String::new(),
                    ))
                } else {
                    None
                }
                .into_iter()
                .chain(column_metadatas.into_iter().map(
                    move |(name, column_metadata)| {
                        (
                            MetadataRegion::concat_column_key(logical_region_id, name),
                            MetadataRegion::serialize_column_metadata(column_metadata),
                        )
                    },
                ))
            })
            .collect::<Vec<_>>();

        let put_request = MetadataRegion::build_put_request_from_iter(iter.into_iter());
        self.write_metadata(metadata_region_id, RegionRequest::Put(put_request))
            .await
    }

    /// Updates logical region metadata so that any entries previously referencing
    /// `source_region_id` are modified to reference the data region of `physical_region_id`.
    ///
    /// This method should be called after copying files from `source_region_id`
    /// into the target region. It scans the metadata for the target physical
    /// region, finds logical regions with the same region number as the source,
    /// and reinserts region and column entries updated to use the target's
    /// region number.
    pub async fn transform_logical_region_metadata(
        &self,
        physical_region_id: RegionId,
        source_region_id: RegionId,
    ) -> Result<()> {
        let metadata_region_id = utils::to_metadata_region_id(physical_region_id);
        let data_region_id = utils::to_data_region_id(physical_region_id);
        let logical_regions = self
            .logical_regions(data_region_id)
            .await?
            .into_iter()
            .filter(|r| r.region_number() == source_region_id.region_number())
            .collect::<Vec<_>>();
        if logical_regions.is_empty() {
            info!(
                "No logical regions found from source region {}, physical region id: {}",
                source_region_id, physical_region_id,
            );
            return Ok(());
        }

        let metadata = self.load_all(metadata_region_id).await?;
        let mut output = Vec::new();
        for logical_region_id in &logical_regions {
            let prefix = MetadataRegion::concat_column_key_prefix(*logical_region_id);
            get_all_with_prefix(&metadata, &prefix, |k, v| {
                // Safety: we have checked the prefix
                let (src_logical_region_id, column_name) = Self::parse_column_key(k)?.unwrap();
                // Change the region number to the data region number.
                let new_key = MetadataRegion::concat_column_key(
                    RegionId::new(
                        src_logical_region_id.table_id(),
                        data_region_id.region_number(),
                    ),
                    &column_name,
                );
                output.push((new_key, v.to_string()));
                Ok(())
            })?;

            let new_key = MetadataRegion::concat_region_key(RegionId::new(
                logical_region_id.table_id(),
                data_region_id.region_number(),
            ));
            output.push((new_key, String::new()));
        }

        if output.is_empty() {
            warn!(
                "No logical regions metadata found from source region {}, physical region id: {}",
                source_region_id, physical_region_id
            );
            return Ok(());
        }

        debug!(
            "Transform logical regions metadata to physical region {}, source region: {}, transformed metadata: {}",
            data_region_id,
            source_region_id,
            output.len(),
        );

        let put_request = MetadataRegion::build_put_request_from_iter(output.into_iter());
        self.write_metadata(metadata_region_id, RegionRequest::Put(put_request))
            .await?;
        info!(
            "Transformed {} logical regions metadata to physical region {}, source region: {}",
            logical_regions.len(),
            data_region_id,
            source_region_id
        );
        Ok(())
    }
}

fn get_all_with_prefix(
    region_metadata: &RegionMetadataCacheEntry,
    prefix: &str,
    mut callback: impl FnMut(&str, &str) -> Result<()>,
) -> Result<()> {
    let range = region_metadata.key_values.range(prefix.to_string()..);
    for (k, v) in range {
        if !k.starts_with(prefix) {
            break;
        }
        callback(k, v)?;
    }
    Ok(())
}

#[cfg(test)]
impl MetadataRegion {
    /// Retrieves the value associated with the given key in the specified region.
    /// Returns `Ok(None)` if the key is not found.
    pub async fn get(&self, region_id: RegionId, key: &str) -> Result<Option<String>> {
        use datatypes::arrow::array::{Array, AsArray};

        let filter_expr = datafusion::prelude::col(METADATA_SCHEMA_KEY_COLUMN_NAME)
            .eq(datafusion::prelude::lit(key));

        let projection = Some(vec![METADATA_SCHEMA_VALUE_COLUMN_INDEX]);
        let scan_req = ScanRequest {
            projection,
            filters: vec![filter_expr],
            ..Default::default()
        };
        let record_batch_stream = self
            .mito
            .scan_to_stream(region_id, scan_req)
            .await
            .context(MitoReadOperationSnafu)?;
        let scan_result = common_recordbatch::util::collect(record_batch_stream)
            .await
            .context(CollectRecordBatchStreamSnafu)?;

        let Some(first_batch) = scan_result.first() else {
            return Ok(None);
        };

        let column = first_batch.column(0);
        let column = column.as_string::<i32>();
        let val = column.is_valid(0).then(|| column.value(0).to_string());

        Ok(val)
    }

    /// Check if the given column exists. Return the semantic type if exists.
    pub async fn column_semantic_type(
        &self,
        physical_region_id: RegionId,
        logical_region_id: RegionId,
        column_name: &str,
    ) -> Result<Option<SemanticType>> {
        let region_id = utils::to_metadata_region_id(physical_region_id);
        let column_key = Self::concat_column_key(logical_region_id, column_name);
        let semantic_type = self.get(region_id, &column_key).await?;
        semantic_type
            .map(|s| Self::deserialize_column_metadata(&s).map(|c| c.semantic_type))
            .transpose()
    }
}

#[cfg(test)]
mod test {
    use datatypes::data_type::ConcreteDataType;
    use datatypes::schema::ColumnSchema;

    use super::*;
    use crate::test_util::TestEnv;
    use crate::utils::to_metadata_region_id;

    #[test]
    fn test_concat_table_key() {
        let region_id = RegionId::new(1234, 7844);
        let expected = "__region_5299989651108".to_string();
        assert_eq!(MetadataRegion::concat_region_key(region_id), expected);
    }

    #[test]
    fn test_concat_column_key() {
        let region_id = RegionId::new(8489, 9184);
        let column_name = "my_column";
        let expected = "__column_36459977384928_bXlfY29sdW1u".to_string();
        assert_eq!(
            MetadataRegion::concat_column_key(region_id, column_name),
            expected
        );
    }

    #[test]
    fn test_parse_table_key() {
        let region_id = RegionId::new(87474, 10607);
        let encoded = MetadataRegion::concat_column_key(region_id, "my_column");
        assert_eq!(encoded, "__column_375697969260911_bXlfY29sdW1u");

        let decoded = MetadataRegion::parse_column_key(&encoded).unwrap();
        assert_eq!(decoded, Some((region_id, "my_column".to_string())));
    }

    #[test]
    fn test_parse_valid_column_key() {
        let region_id = RegionId::new(176, 910);
        let encoded = MetadataRegion::concat_column_key(region_id, "my_column");
        assert_eq!(encoded, "__column_755914245006_bXlfY29sdW1u");

        let decoded = MetadataRegion::parse_column_key(&encoded).unwrap();
        assert_eq!(decoded, Some((region_id, "my_column".to_string())));
    }

    #[test]
    fn test_parse_invalid_column_key() {
        let key = "__column_asdfasd_????";
        let result = MetadataRegion::parse_column_key(key);
        assert!(result.is_err());
    }

    #[test]
    fn test_serialize_column_metadata() {
        let semantic_type = SemanticType::Tag;
        let column_metadata = ColumnMetadata {
            column_schema: ColumnSchema::new("blabla", ConcreteDataType::string_datatype(), false),
            semantic_type,
            column_id: 5,
        };
        let old_fmt = "{\"column_schema\":{\"name\":\"blabla\",\"data_type\":{\"String\":null},\"is_nullable\":false,\"is_time_index\":false,\"default_constraint\":null,\"metadata\":{}},\"semantic_type\":\"Tag\",\"column_id\":5}".to_string();
        let new_fmt = "{\"column_schema\":{\"name\":\"blabla\",\"data_type\":{\"String\":{\"size_type\":\"Utf8\"}},\"is_nullable\":false,\"is_time_index\":false,\"default_constraint\":null,\"metadata\":{}},\"semantic_type\":\"Tag\",\"column_id\":5}".to_string();
        assert_eq!(
            MetadataRegion::serialize_column_metadata(&column_metadata),
            new_fmt
        );
        // Ensure both old and new formats can be deserialized.
        assert_eq!(
            MetadataRegion::deserialize_column_metadata(&old_fmt).unwrap(),
            column_metadata
        );
        assert_eq!(
            MetadataRegion::deserialize_column_metadata(&new_fmt).unwrap(),
            column_metadata
        );

        let semantic_type = "\"Invalid Column Metadata\"";
        assert!(MetadataRegion::deserialize_column_metadata(semantic_type).is_err());
    }

    fn test_column_metadatas() -> HashMap<String, ColumnMetadata> {
        HashMap::from([
            (
                "label1".to_string(),
                ColumnMetadata {
                    column_schema: ColumnSchema::new(
                        "label1".to_string(),
                        ConcreteDataType::string_datatype(),
                        false,
                    ),
                    semantic_type: SemanticType::Tag,
                    column_id: 5,
                },
            ),
            (
                "label2".to_string(),
                ColumnMetadata {
                    column_schema: ColumnSchema::new(
                        "label2".to_string(),
                        ConcreteDataType::string_datatype(),
                        false,
                    ),
                    semantic_type: SemanticType::Tag,
                    column_id: 5,
                },
            ),
        ])
    }

    async fn add_test_logical_region(
        metadata_region: &MetadataRegion,
        physical_region_id: RegionId,
        logical_region_id: RegionId,
    ) -> Result<()> {
        let column_metadatas = test_column_metadatas();
        let logical_regions = std::iter::once((
            logical_region_id,
            column_metadatas
                .iter()
                .map(|(name, metadata)| (name.as_str(), metadata))
                .collect::<HashMap<_, _>>(),
        ));
        metadata_region
            .add_logical_regions(physical_region_id, true, logical_regions)
            .await
    }

    #[tokio::test]
    async fn add_logical_regions_to_meta_region() {
        let env = TestEnv::new().await;
        env.init_metric_region().await;
        let metadata_region = env.metadata_region();
        let physical_region_id = to_metadata_region_id(env.default_physical_region_id());
        let column_metadatas = test_column_metadatas();
        let logical_region_id = RegionId::new(1024, 1);

        let iter = vec![(
            logical_region_id,
            column_metadatas
                .iter()
                .map(|(k, v)| (k.as_str(), v))
                .collect::<HashMap<_, _>>(),
        )];
        metadata_region
            .add_logical_regions(physical_region_id, true, iter.into_iter())
            .await
            .unwrap();
        // Add logical region again.
        let iter = vec![(
            logical_region_id,
            column_metadatas
                .iter()
                .map(|(k, v)| (k.as_str(), v))
                .collect::<HashMap<_, _>>(),
        )];
        metadata_region
            .add_logical_regions(physical_region_id, true, iter.into_iter())
            .await
            .unwrap();

        // Check if the logical region is added.
        let logical_regions = metadata_region
            .logical_regions(physical_region_id)
            .await
            .unwrap();
        assert_eq!(logical_regions.len(), 2);

        // Check if the logical region columns are added.
        let logical_columns = metadata_region
            .logical_columns(physical_region_id, logical_region_id)
            .await
            .unwrap()
            .into_iter()
            .collect::<HashMap<_, _>>();
        assert_eq!(logical_columns.len(), 2);
        assert_eq!(column_metadatas, logical_columns);
    }

    #[tokio::test]
    async fn metadata_writes_are_synchronized_per_region() {
        let env = TestEnv::new().await;
        env.init_metric_region().await;
        let other_physical_region_id = RegionId::new(2, 2);
        env.create_physical_region(other_physical_region_id, "/test_dir2", vec![])
            .await;
        let metadata_region = Arc::new(env.metadata_region());
        let physical_region_id = env.default_physical_region_id();
        let metadata_region_id = to_metadata_region_id(physical_region_id);

        let (snapshot_loaded_tx, snapshot_loaded_rx) = tokio::sync::oneshot::channel();
        let (release_snapshot_tx, release_snapshot_rx) = tokio::sync::oneshot::channel();
        let cache_fill = {
            let metadata_region = Arc::clone(&metadata_region);
            tokio::spawn(async move {
                let cache_access_lock = metadata_region.cache_access_lock(metadata_region_id);
                let _cache_guard = cache_access_lock.read().await;
                metadata_region
                    .cache
                    .try_get_with(metadata_region_id, async {
                        let snapshot = metadata_region.load_all(metadata_region_id).await?;
                        snapshot_loaded_tx.send(()).unwrap();
                        release_snapshot_rx.await.unwrap();
                        Ok::<_, crate::error::Error>(snapshot)
                    })
                    .await
                    .unwrap();
            })
        };
        snapshot_loaded_rx.await.unwrap();

        let logical_region_id = RegionId::new(1024, 1);
        let metadata_write = {
            let metadata_region = Arc::clone(&metadata_region);
            tokio::spawn(async move {
                add_test_logical_region(&metadata_region, physical_region_id, logical_region_id)
                    .await
            })
        };
        let mut metadata_write = metadata_write;
        assert!(
            tokio::time::timeout(std::time::Duration::from_millis(50), &mut metadata_write)
                .await
                .is_err()
        );

        let other_logical_region_id = RegionId::new(2048, 2);
        tokio::time::timeout(
            std::time::Duration::from_secs(5),
            add_test_logical_region(
                &metadata_region,
                other_physical_region_id,
                other_logical_region_id,
            ),
        )
        .await
        .unwrap()
        .unwrap();

        release_snapshot_tx.send(()).unwrap();
        cache_fill.await.unwrap();
        metadata_write.await.unwrap().unwrap();

        let logical_columns = metadata_region
            .logical_columns(physical_region_id, logical_region_id)
            .await
            .unwrap();
        assert_eq!(logical_columns.len(), 2);
        assert!(
            metadata_region
                .cache_access_locks
                .lock()
                .unwrap()
                .is_empty()
        );
    }

    #[tokio::test]
    async fn cache_access_locks_self_clean() {
        let env = TestEnv::new().await;
        env.init_metric_region().await;
        let metadata_region = env.metadata_region();
        let physical_region_id = env.default_physical_region_id();
        let metadata_region_id = to_metadata_region_id(physical_region_id);
        let registry_len = || metadata_region.cache_access_locks.lock().unwrap().len();

        // Metadata reads and writes leave no lock entries behind.
        let logical_region_id = RegionId::new(1024, 1);
        add_test_logical_region(&metadata_region, physical_region_id, logical_region_id)
            .await
            .unwrap();
        metadata_region
            .logical_columns(physical_region_id, logical_region_id)
            .await
            .unwrap();
        assert_eq!(registry_len(), 0);

        // Concurrent leases share one lock; the entry lives until the last lease drops.
        let first = metadata_region.cache_access_lock(metadata_region_id);
        let second = metadata_region.cache_access_lock(metadata_region_id);
        assert!(Arc::ptr_eq(
            first.lock.as_ref().unwrap(),
            second.lock.as_ref().unwrap()
        ));
        drop(first);
        assert_eq!(registry_len(), 1);
        drop(second);
        assert_eq!(registry_len(), 0);

        // A new acquisition after cleanup mints a fresh lock and entry.
        let third = metadata_region.cache_access_lock(metadata_region_id);
        assert_eq!(registry_len(), 1);
        drop(third);
        assert_eq!(registry_len(), 0);
    }
}