relational_engine/

engine.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::any::Any;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::time::Instant;

use api::region::RegionResponse;
use async_trait::async_trait;
use common_catalog::consts::RELATIONAL_ENGINE;
use common_error::ext::BoxedError;
use common_telemetry::{error, info};
use object_store::ObjectStore;
use snafu::{OptionExt, ensure};
use store_api::metadata::RegionMetadataRef;
use store_api::region_engine::{
    RegionEngine, RegionRole, RegionScannerRef, RegionStatistic, RemapManifestsRequest,
    RemapManifestsResponse, SetRegionRoleStateResponse, SetRegionRoleStateSuccess,
    SettableRegionRoleState, SyncRegionFromRequest, SyncRegionFromResponse,
};
use store_api::region_request::{
    AffectedRows, RegionCloseRequest, RegionCreateRequest, RegionDropRequest, RegionOpenRequest,
    RegionRequest, RegionRequirements,
};
use store_api::storage::{RegionId, ScanRequest, SequenceNumber};
use tokio::sync::Mutex;

use crate::error::{
    RegionNotFoundSnafu, Result as EngineResult, UnexpectedEngineSnafu, UnsupportedSnafu,
};
use crate::manifest::RelationalRegionManifest;
use crate::region::{RelationalRegion, RelationalRegionRef};

pub struct RelationalEngine {
    inner: EngineInnerRef,
}

impl RelationalEngine {
    pub fn new(object_store: ObjectStore) -> Self {
        Self {
            inner: Arc::new(EngineInner::new(object_store)),
        }
    }
}

#[async_trait]
impl RegionEngine for RelationalEngine {
    fn name(&self) -> &str {
        RELATIONAL_ENGINE
    }

    async fn handle_request(
        &self,
        region_id: RegionId,
        request: RegionRequest,
    ) -> Result<RegionResponse, BoxedError> {
        self.inner
            .handle_request(region_id, request)
            .await
            .map_err(BoxedError::new)
    }

    async fn handle_query(
        &self,
        _region_id: RegionId,
        _request: ScanRequest,
    ) -> Result<RegionScannerRef, BoxedError> {
        Err(BoxedError::new(
            UnsupportedSnafu { operation: "scan" }.build(),
        ))
    }

    async fn get_metadata(&self, region_id: RegionId) -> Result<RegionMetadataRef, BoxedError> {
        self.inner
            .get_region(region_id)
            .await
            .map(|r| r.metadata())
            .context(RegionNotFoundSnafu { region_id })
            .map_err(BoxedError::new)
    }

    async fn stop(&self) -> Result<(), BoxedError> {
        self.inner.stop().await.map_err(BoxedError::new)
    }

    fn region_statistic(&self, _: RegionId) -> Option<RegionStatistic> {
        None
    }

    async fn get_committed_sequence(&self, _: RegionId) -> Result<SequenceNumber, BoxedError> {
        Ok(Default::default())
    }

    fn set_region_role(&self, region_id: RegionId, role: RegionRole) -> Result<(), BoxedError> {
        self.inner
            .set_region_role(region_id, role)
            .map_err(BoxedError::new)
    }

    /// Regions of this engine take no writes, so the lease is not needed.
    fn set_region_role_with_lease(
        &self,
        region_id: RegionId,
        role: RegionRole,
        _sent_at: Instant,
        _deadline: Instant,
    ) -> Result<bool, BoxedError> {
        self.set_region_role(region_id, role).map(|_| true)
    }

    async fn set_region_role_state_gracefully(
        &self,
        region_id: RegionId,
        _region_role_state: SettableRegionRoleState,
    ) -> Result<SetRegionRoleStateResponse, BoxedError> {
        let exists = self.inner.get_region(region_id).await.is_some();

        if exists {
            Ok(SetRegionRoleStateResponse::success(
                SetRegionRoleStateSuccess::relational(),
            ))
        } else {
            Ok(SetRegionRoleStateResponse::NotFound)
        }
    }

    async fn sync_region(
        &self,
        _region_id: RegionId,
        _request: SyncRegionFromRequest,
    ) -> Result<SyncRegionFromResponse, BoxedError> {
        // Relational engine doesn't need to sync region manifest.
        Ok(SyncRegionFromResponse::NotSupported)
    }

    async fn remap_manifests(
        &self,
        _request: RemapManifestsRequest,
    ) -> Result<RemapManifestsResponse, BoxedError> {
        Err(BoxedError::new(
            UnsupportedSnafu {
                operation: "remap_manifests",
            }
            .build(),
        ))
    }

    fn role(&self, region_id: RegionId) -> Option<RegionRole> {
        self.inner.state(region_id)
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}

struct EngineInner {
    /// All regions opened by the engine.
    ///
    /// Writing to `regions` should also hold the `region_mutex`.
    regions: RwLock<HashMap<RegionId, RelationalRegionRef>>,

    /// Region mutex is used to protect the operations such as creating/opening/closing
    /// a region, to avoid things like opening the same region simultaneously.
    region_mutex: Mutex<()>,

    object_store: ObjectStore,
}

type EngineInnerRef = Arc<EngineInner>;

fn ensure_region_requirements(
    requirements: RegionRequirements,
    object_store: &ObjectStore,
) -> EngineResult<()> {
    if !requirements.object_storage {
        return Ok(());
    }

    ensure!(
        object_store::util::is_object_storage(object_store),
        UnsupportedSnafu {
            operation: "create or open region with object storage requirement on non-object storage"
        }
    );

    Ok(())
}

impl EngineInner {
    fn new(object_store: ObjectStore) -> Self {
        Self {
            regions: RwLock::new(HashMap::new()),
            region_mutex: Mutex::new(()),
            object_store,
        }
    }

    async fn handle_request(
        &self,
        region_id: RegionId,
        request: RegionRequest,
    ) -> EngineResult<RegionResponse> {
        let result = match request {
            RegionRequest::Create(req) => self.handle_create(region_id, req).await,
            RegionRequest::Drop(req) => self.handle_drop(region_id, req).await,
            RegionRequest::Open(req) => self.handle_open(region_id, req).await,
            RegionRequest::Close(req) => self.handle_close(region_id, req).await,
            _ => UnsupportedSnafu {
                operation: request.to_string(),
            }
            .fail(),
        };
        result.map(RegionResponse::new)
    }

    async fn stop(&self) -> EngineResult<()> {
        let _lock = self.region_mutex.lock().await;
        self.regions.write().unwrap().clear();
        Ok(())
    }

    fn set_region_role(&self, _region_id: RegionId, _region_role: RegionRole) -> EngineResult<()> {
        // TODO(zhongzc): Improve the semantics and implementation of this API.
        Ok(())
    }

    fn state(&self, region_id: RegionId) -> Option<RegionRole> {
        if self.regions.read().unwrap().get(&region_id).is_some() {
            Some(RegionRole::Leader)
        } else {
            None
        }
    }
}

impl EngineInner {
    async fn handle_create(
        &self,
        region_id: RegionId,
        request: RegionCreateRequest,
    ) -> EngineResult<AffectedRows> {
        ensure!(
            request.engine == RELATIONAL_ENGINE,
            UnexpectedEngineSnafu {
                engine: request.engine
            }
        );

        // If the region is already open, validate the definition before treating the
        // create as idempotent: a mismatched create must fail, not silently succeed.
        if let Some(region) = self.get_region(region_id).await {
            region.check_create_request(region_id, &request)?;
            return Ok(0);
        }

        info!("Try to create region, region_id: {}", region_id);

        let _lock = self.region_mutex.lock().await;
        // Check again after acquiring the lock
        if let Some(region) = self.get_region(region_id).await {
            region.check_create_request(region_id, &request)?;
            return Ok(0);
        }

        ensure_region_requirements(request.requirements, &self.object_store)?;

        // The manifest is the source of truth for whether a region already
        // exists. The in-memory `regions` map only tracks regions opened in
        // this process, so after a restart it is empty even though the
        // manifest was already persisted. Recover the region from the manifest
        // instead of failing with `ManifestExists`, keeping `CREATE EXTERNAL
        // TABLE` idempotent.
        let region_dir = RelationalRegion::region_dir(&request.table_dir, region_id);
        if RelationalRegionManifest::exists(&region_dir, &self.object_store).await? {
            let manifest =
                RelationalRegionManifest::load(region_id, &region_dir, &self.object_store).await?;
            let region = RelationalRegion::from_manifest(&region_dir, manifest)?;
            // The persisted definition must match the incoming request, otherwise a
            // retried CREATE with a different schema would silently reuse the old
            // manifest. Reject the mismatch instead.
            region.check_create_request(region_id, &request)?;
            self.regions.write().unwrap().insert(region_id, region);
            info!(
                "Region recovered from existing manifest, region_id: {}",
                region_id
            );
            return Ok(0);
        }

        let res = RelationalRegion::create(region_id, request, &self.object_store).await;
        let region = res.inspect_err(|err| {
            error!(
                err;
                "Failed to create region, region_id: {}",
                region_id
            );
        })?;
        self.regions.write().unwrap().insert(region_id, region);

        info!("A new region is created, region_id: {}", region_id);
        Ok(0)
    }

    async fn handle_open(
        &self,
        region_id: RegionId,
        request: RegionOpenRequest,
    ) -> EngineResult<AffectedRows> {
        if self.exists(region_id).await {
            return Ok(0);
        }

        info!("Try to open region, region_id: {}", region_id);

        let _lock = self.region_mutex.lock().await;
        // Check again after acquiring the lock
        if self.exists(region_id).await {
            return Ok(0);
        }

        ensure_region_requirements(request.requirements, &self.object_store)?;

        let res = RelationalRegion::open(region_id, request, &self.object_store).await;
        let region = res.inspect_err(|err| {
            error!(
                err;
                "Failed to open region, region_id: {}",
                region_id
            );
        })?;
        self.regions.write().unwrap().insert(region_id, region);

        info!("Region opened, region_id: {}", region_id);
        Ok(0)
    }

    async fn handle_close(
        &self,
        region_id: RegionId,
        _request: RegionCloseRequest,
    ) -> EngineResult<AffectedRows> {
        let _lock = self.region_mutex.lock().await;

        let mut regions = self.regions.write().unwrap();
        if regions.remove(&region_id).is_some() {
            info!("Region closed, region_id: {}", region_id);
        }

        Ok(0)
    }

    async fn handle_drop(
        &self,
        region_id: RegionId,
        _request: RegionDropRequest,
    ) -> EngineResult<AffectedRows> {
        // Serialize against create/open/close so the region can't be closed
        // between the existence check and the manifest deletion. A region that
        // is no longer in memory (e.g. closed concurrently) is reported as not
        // found instead of silently returning success with a leaked manifest.
        let _lock = self.region_mutex.lock().await;

        let region = self.get_region(region_id).await;
        let Some(region) = region else {
            return RegionNotFoundSnafu { region_id }.fail();
        };

        info!("Try to drop region, region_id: {}", region_id);

        let res = RelationalRegion::drop(&region, &self.object_store).await;
        res.inspect_err(|err| {
            error!(
                err;
                "Failed to drop region, region_id: {}",
                region_id
            );
        })?;

        let _ = self.regions.write().unwrap().remove(&region_id);

        info!("Region dropped, region_id: {}", region_id);
        Ok(0)
    }

    async fn get_region(&self, region_id: RegionId) -> Option<RelationalRegionRef> {
        self.regions.read().unwrap().get(&region_id).cloned()
    }

    async fn exists(&self, region_id: RegionId) -> bool {
        self.regions.read().unwrap().contains_key(&region_id)
    }
}

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

    use object_store::ObjectStore;
    use object_store::services::{Fs, S3};
    use store_api::region_request::{
        PathType, RegionCreateRequest, RegionDropRequest, RegionRequirements,
    };
    use store_api::storage::RegionId;

    use super::{EngineInner, ensure_region_requirements};
    use crate::error::Error;
    use crate::test_util::{new_test_column_metadata, new_test_object_store, new_test_options};

    fn build_fs_object_store() -> ObjectStore {
        ObjectStore::new(Fs::default().root("/tmp")).unwrap()
    }

    fn build_s3_object_store() -> ObjectStore {
        ObjectStore::new(
            S3::default()
                .bucket("test-bucket")
                .region("us-east-1")
                .disable_ec2_metadata(),
        )
        .unwrap()
    }

    #[test]
    fn test_empty_region_requirements_are_supported() {
        ensure_region_requirements(RegionRequirements::empty(), &build_fs_object_store()).unwrap();
    }

    #[test]
    fn test_object_storage_region_requirement_rejects_fs_object_store() {
        let err = ensure_region_requirements(
            RegionRequirements::object_storage(),
            &build_fs_object_store(),
        )
        .unwrap_err();

        assert!(matches!(err, Error::Unsupported { .. }));
    }

    #[test]
    fn test_object_storage_region_requirement_accepts_s3_object_store() {
        ensure_region_requirements(
            RegionRequirements::object_storage(),
            &build_s3_object_store(),
        )
        .unwrap();
    }

    #[tokio::test]
    async fn test_drop_returns_not_found_when_region_absent() {
        let (_dir, object_store) = new_test_object_store("test_drop_absent");
        let inner = EngineInner::new(object_store);

        let err = inner
            .handle_drop(
                RegionId::new(1, 0),
                RegionDropRequest {
                    fast_path: false,
                    force: false,
                    partial_drop: false,
                },
            )
            .await
            .unwrap_err();

        assert_matches!(err, Error::RegionNotFound { .. });
    }

    #[tokio::test]
    async fn test_drop_removes_manifest() {
        let (_dir, object_store) = new_test_object_store("test_drop_removes_manifest");
        let inner = EngineInner::new(object_store.clone());

        let region_id = RegionId::new(1, 0);
        let request = RegionCreateRequest {
            engine: "relational".to_string(),
            column_metadatas: new_test_column_metadata(),
            primary_key: vec![1],
            options: new_test_options(),
            table_dir: "drop_dir/".to_string(),
            path_type: PathType::Bare,
            partition_expr_json: Some("".to_string()),
            requirements: Default::default(),
        };

        inner.handle_create(region_id, request).await.unwrap();

        let manifest_path = "drop_dir/1_0000000000/manifest/_relational_manifest";
        assert!(object_store.exists(manifest_path).await.unwrap());

        inner
            .handle_drop(
                region_id,
                RegionDropRequest {
                    fast_path: false,
                    force: false,
                    partial_drop: false,
                },
            )
            .await
            .unwrap();

        assert!(!object_store.exists(manifest_path).await.unwrap());
        assert!(!inner.exists(region_id).await);
    }

    #[tokio::test]
    async fn test_create_is_idempotent_across_engine_rebuild() {
        let (_dir, object_store) = new_test_object_store("test_create_idempotent");
        let region_id = RegionId::new(1, 0);
        let request = RegionCreateRequest {
            engine: "relational".to_string(),
            column_metadatas: new_test_column_metadata(),
            primary_key: vec![1],
            options: new_test_options(),
            table_dir: "idempotent_dir/".to_string(),
            path_type: PathType::Bare,
            partition_expr_json: Some("".to_string()),
            requirements: Default::default(),
        };

        // First create persists the manifest to object storage.
        {
            let inner = EngineInner::new(object_store.clone());
            inner
                .handle_create(region_id, request.clone())
                .await
                .unwrap();
            assert!(inner.exists(region_id).await);
        }

        // Rebuild the engine on the same object store: the in-memory region
        // map is empty, but the manifest is still persisted. Retrying the
        // same create must recover from the manifest instead of failing.
        let rebuilt = EngineInner::new(object_store.clone());
        let affected = rebuilt.handle_create(region_id, request).await.unwrap();
        assert_eq!(affected, 0);
        assert!(rebuilt.exists(region_id).await);

        let manifest_path = "idempotent_dir/1_0000000000/manifest/_relational_manifest";
        assert!(object_store.exists(manifest_path).await.unwrap());
    }

    #[tokio::test]
    async fn test_create_rejects_mismatched_definition() {
        let (_dir, object_store) = new_test_object_store("test_create_mismatch");
        let region_id = RegionId::new(1, 0);

        let request = RegionCreateRequest {
            engine: "relational".to_string(),
            column_metadatas: new_test_column_metadata(),
            primary_key: vec![1],
            options: new_test_options(),
            table_dir: "mismatch_dir/".to_string(),
            path_type: PathType::Bare,
            partition_expr_json: Some("".to_string()),
            requirements: Default::default(),
        };

        // First create persists the manifest.
        let inner = EngineInner::new(object_store.clone());
        inner.handle_create(region_id, request).await.unwrap();

        // Rebuild the engine and retry the create with a different primary key:
        // the persisted definition no longer matches, so this must fail instead of
        // silently reusing the old manifest.
        let rebuilt = EngineInner::new(object_store.clone());
        let mismatched = RegionCreateRequest {
            engine: "relational".to_string(),
            column_metadatas: new_test_column_metadata(),
            primary_key: vec![2],
            options: new_test_options(),
            table_dir: "mismatch_dir/".to_string(),
            path_type: PathType::Bare,
            partition_expr_json: Some("".to_string()),
            requirements: Default::default(),
        };
        let err = rebuilt
            .handle_create(region_id, mismatched)
            .await
            .unwrap_err();
        assert_matches!(err, Error::RegionExists { .. });
    }

    #[tokio::test]
    async fn test_create_rejects_mismatched_definition_in_memory() {
        let (_dir, object_store) = new_test_object_store("test_create_mismatch_in_memory");
        let inner = EngineInner::new(object_store);
        let region_id = RegionId::new(1, 0);

        let request = RegionCreateRequest {
            engine: "relational".to_string(),
            column_metadatas: new_test_column_metadata(),
            primary_key: vec![1],
            options: new_test_options(),
            table_dir: "mismatch_in_memory_dir/".to_string(),
            path_type: PathType::Bare,
            partition_expr_json: Some("".to_string()),
            requirements: Default::default(),
        };

        inner
            .handle_create(region_id, request.clone())
            .await
            .unwrap();

        // An identical create is idempotent while the region stays in memory.
        assert_eq!(inner.handle_create(region_id, request).await.unwrap(), 0);

        // A mismatched create must fail, not silently return success.
        let mismatched = RegionCreateRequest {
            engine: "relational".to_string(),
            column_metadatas: new_test_column_metadata(),
            primary_key: vec![2],
            options: new_test_options(),
            table_dir: "mismatch_in_memory_dir/".to_string(),
            path_type: PathType::Bare,
            partition_expr_json: Some("".to_string()),
            requirements: Default::default(),
        };
        let err = inner
            .handle_create(region_id, mismatched)
            .await
            .unwrap_err();
        assert_matches!(
            err,
            Error::RegionExists { reason, .. } if reason.contains("primary key")
        );
    }
}