greptimedb/src/mito2/src/engine/open_test.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::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

use api::v1::Rows;
use common_base::Plugins;
use common_error::ext::ErrorExt;
use common_error::status_code::StatusCode;
use common_recordbatch::RecordBatches;
use either::Either;
use store_api::logstore::LogStore;
use store_api::logstore::provider::Provider;
use store_api::region_engine::{RegionEngine, RegionRole, SettableRegionRoleState};
use store_api::region_request::{
    PathType, RegionCleanUpRequest, RegionCloseRequest, RegionOpenRequest, RegionPutRequest,
    RegionRequest,
};
use store_api::storage::{RegionId, ScanRequest};
use tokio::sync::{Notify, oneshot};

use crate::compaction::compactor::{CompactionRegionOpener, OpenCompactionRegionRequest};
use crate::config::MitoConfig;
use crate::engine::flush_test::MockRegionHook;
use crate::engine::listener::EventListener;
use crate::engine::region_hook::RegionHookRef;
use crate::error;
use crate::region::opener::{PartitionExprFetcher, PartitionExprFetcherRef};
use crate::region::options::RegionOptions;
use crate::sst::location::region_dir_from_table_dir;
use crate::test_util::{
    CreateRequestBuilder, LogStoreImpl, TestEnv, build_rows, flush_region, put_rows, reopen_region,
    rows_schema,
};

#[tokio::test]
async fn test_engine_open_empty() {
    test_engine_open_empty_with_format(false).await;
    test_engine_open_empty_with_format(true).await;
}

async fn test_engine_open_empty_with_format(flat_format: bool) {
    let mut env = TestEnv::with_prefix("open-empty").await;
    let engine = env
        .create_engine(MitoConfig {
            default_flat_format: flat_format,
            ..Default::default()
        })
        .await;

    let region_id = RegionId::new(1, 1);
    let err = engine
        .handle_request(
            region_id,
            RegionRequest::Open(RegionOpenRequest {
                engine: String::new(),
                table_dir: "empty".to_string(),
                path_type: PathType::Bare,
                options: HashMap::default(),
                skip_wal_replay: false,
                checkpoint: None,
                requirements: Default::default(),
            }),
        )
        .await
        .unwrap_err();
    assert_eq!(StatusCode::RegionNotFound, err.status_code());
    let err = engine
        .set_region_role(region_id, RegionRole::Leader)
        .unwrap_err();
    assert_eq!(StatusCode::RegionNotFound, err.status_code());
    let role = engine.role(region_id);
    assert_eq!(role, None);
}

#[tokio::test]
async fn test_engine_offline_cleanup_closed_region() {
    let mut env = TestEnv::with_prefix("offline-cleanup").await;
    let engine = env.create_engine(MitoConfig::default()).await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new().build();
    let table_dir = request.table_dir.clone();
    let path_type = request.path_type;
    let options = request.options.clone();
    let region_dir = region_dir_from_table_dir(&table_dir, region_id, path_type);
    let object_store = env.get_object_store().unwrap();

    engine
        .handle_request(region_id, RegionRequest::Create(request.clone()))
        .await
        .unwrap();
    assert!(object_store.exists(&region_dir).await.unwrap());

    put_rows(
        &engine,
        region_id,
        Rows {
            schema: rows_schema(&request),
            rows: build_rows(0, 3),
        },
    )
    .await;
    let Some(LogStoreImpl::RaftEngine(log_store)) = env.get_log_store() else {
        unreachable!()
    };
    let provider = Provider::raft_engine_provider(region_id.as_u64());
    assert!(log_store.latest_entry_id(&provider).unwrap() > 0);

    engine
        .handle_request(
            region_id,
            RegionRequest::Close(RegionCloseRequest::default()),
        )
        .await
        .unwrap();
    assert!(!engine.is_region_exists(region_id));

    let cleanup_request = RegionCleanUpRequest {
        engine: String::new(),
        table_dir,
        path_type,
        options,
    };
    for _ in 0..2 {
        engine
            .handle_request(region_id, RegionRequest::CleanUp(cleanup_request.clone()))
            .await
            .unwrap();
    }

    assert!(!engine.is_region_exists(region_id));
    assert!(!object_store.exists(&region_dir).await.unwrap());
    assert_eq!(0, log_store.latest_entry_id(&provider).unwrap());
}

#[tokio::test]
async fn test_engine_offline_cleanup_rejects_opened_region() {
    let mut env = TestEnv::with_prefix("offline-cleanup-opened").await;
    let engine = env.create_engine(MitoConfig::default()).await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new().build();
    let table_dir = request.table_dir.clone();
    let path_type = request.path_type;
    let options = request.options.clone();
    let region_dir = region_dir_from_table_dir(&table_dir, region_id, path_type);
    let object_store = env.get_object_store().unwrap();

    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();

    let err = engine
        .handle_request(
            region_id,
            RegionRequest::CleanUp(RegionCleanUpRequest {
                engine: String::new(),
                table_dir,
                path_type,
                options,
            }),
        )
        .await
        .unwrap_err();

    assert_eq!(StatusCode::RegionBusy, err.status_code());
    assert!(engine.is_region_exists(region_id));
    assert!(object_store.exists(&region_dir).await.unwrap());
}

/// Offline cleanup (soft-drop PURGE) removes the region directory directly,
/// bypassing the drop GC worker, so it must notify extensions via
/// `on_region_gc` — otherwise sidecar state (e.g. an Iceberg export) leaks.
///
/// Covers the wiring: `handle_offline_cleanup_request` fires `on_region_gc`
/// with `is_region_dropped = true` and `full_file_listing = true` once the
/// region dir is physically removed, and does NOT fire `on_region_files_removed`
/// (the region is offline and carries no `RegionMetadataRef`).
#[tokio::test]
async fn test_engine_offline_cleanup_fires_on_region_gc() {
    common_telemetry::init_default_ut_logging();
    let mut env = TestEnv::with_prefix("offline-cleanup-gc-hook").await;

    let hook = Arc::new(MockRegionHook::new());
    let plugins = Plugins::new();
    plugins.insert(hook.clone() as RegionHookRef);
    let engine = env
        .create_engine_with_plugins(MitoConfig::default(), plugins)
        .await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new().build();
    let table_dir = request.table_dir.clone();
    let path_type = request.path_type;
    let options = request.options.clone();
    let region_dir = region_dir_from_table_dir(&table_dir, region_id, path_type);
    let object_store = env.get_object_store().unwrap();

    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();
    assert!(object_store.exists(&region_dir).await.unwrap());

    // Close the region first — offline cleanup refuses a still-registered region.
    engine
        .handle_request(
            region_id,
            RegionRequest::Close(RegionCloseRequest::default()),
        )
        .await
        .unwrap();
    assert!(!engine.is_region_exists(region_id));

    let gc_before = hook.gc_count.load(Ordering::Relaxed);
    let files_removed_before = hook.files_removed_count.load(Ordering::Relaxed);

    engine
        .handle_request(
            region_id,
            RegionRequest::CleanUp(RegionCleanUpRequest {
                engine: String::new(),
                table_dir,
                path_type,
                options,
            }),
        )
        .await
        .unwrap();

    // The region directory is gone …
    assert!(!engine.is_region_exists(region_id));
    assert!(!object_store.exists(&region_dir).await.unwrap());

    // … and the extension was notified exactly once via `on_region_gc`, with the
    // flags that authorize full sidecar reclamation. `on_region_files_removed`
    // must NOT fire for the offline path (no region metadata is available).
    assert_eq!(
        hook.gc_count.load(Ordering::Relaxed),
        gc_before + 1,
        "offline cleanup must fire on_region_gc"
    );
    assert!(
        hook.last_gc_is_region_dropped.load(Ordering::Relaxed),
        "on_region_gc during offline cleanup must report is_region_dropped = true"
    );
    assert!(
        hook.last_gc_full_file_listing.load(Ordering::Relaxed),
        "on_region_gc during offline cleanup must report full_file_listing = true"
    );
    assert_eq!(
        hook.files_removed_count.load(Ordering::Relaxed),
        files_removed_before,
        "offline cleanup must not fire on_region_files_removed"
    );
}

/// A failing `on_region_gc` during offline cleanup must surface as an error so
/// the caller (`PurgeDroppedTableProcedure`) retries the `CleanUp`, and the
/// retry must succeed once the hook stops failing — the handler is idempotent
/// (it re-runs the directory removal, WAL obsolete and the hook safely).
#[tokio::test]
async fn test_engine_offline_cleanup_propagates_on_region_gc_error() {
    common_telemetry::init_default_ut_logging();
    let mut env = TestEnv::with_prefix("offline-cleanup-gc-hook-retry").await;

    let hook = Arc::new(MockRegionHook::new());
    let plugins = Plugins::new();
    plugins.insert(hook.clone() as RegionHookRef);
    let engine = env
        .create_engine_with_plugins(MitoConfig::default(), plugins)
        .await;

    let region_id = RegionId::new(2, 1);
    let request = CreateRequestBuilder::new().build();
    let table_dir = request.table_dir.clone();
    let path_type = request.path_type;
    let options = request.options.clone();
    let region_dir = region_dir_from_table_dir(&table_dir, region_id, path_type);
    let object_store = env.get_object_store().unwrap();

    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();
    engine
        .handle_request(
            region_id,
            RegionRequest::Close(RegionCloseRequest::default()),
        )
        .await
        .unwrap();

    // Force the extension hook to fail: the offline cleanup must return the
    // error so the caller retries, instead of swallowing it.
    hook.gc_should_fail.store(true, Ordering::Relaxed);
    let gc_before = hook.gc_count.load(Ordering::Relaxed);
    let err = engine
        .handle_request(
            region_id,
            RegionRequest::CleanUp(RegionCleanUpRequest {
                engine: String::new(),
                table_dir: table_dir.clone(),
                path_type,
                options: options.clone(),
            }),
        )
        .await
        .unwrap_err();
    assert_eq!(
        err.status_code(),
        StatusCode::Unexpected,
        "offline cleanup must propagate the on_region_gc error"
    );
    assert_eq!(
        hook.gc_count.load(Ordering::Relaxed),
        gc_before + 1,
        "on_region_gc must fire even when it fails"
    );

    // Retry: clear the failure and the same CleanUp must now succeed. The
    // directory removal and WAL obsolete are idempotent, so the second attempt
    // re-runs them and the hook without complaint.
    hook.gc_should_fail.store(false, Ordering::Relaxed);
    engine
        .handle_request(
            region_id,
            RegionRequest::CleanUp(RegionCleanUpRequest {
                engine: String::new(),
                table_dir,
                path_type,
                options,
            }),
        )
        .await
        .unwrap();
    assert!(!engine.is_region_exists(region_id));
    assert!(!object_store.exists(&region_dir).await.unwrap());
    assert_eq!(
        hook.gc_count.load(Ordering::Relaxed),
        gc_before + 2,
        "on_region_gc must fire again on the successful retry"
    );
}

#[tokio::test]
async fn test_engine_open_existing() {
    test_engine_open_existing_with_format(false).await;
    test_engine_open_existing_with_format(true).await;
}

async fn test_engine_open_existing_with_format(flat_format: bool) {
    let mut env = TestEnv::with_prefix("open-exiting").await;
    let engine = env
        .create_engine(MitoConfig {
            default_flat_format: flat_format,
            ..Default::default()
        })
        .await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new().build();
    let table_dir = request.table_dir.clone();
    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();

    engine
        .handle_request(
            region_id,
            RegionRequest::Open(RegionOpenRequest {
                engine: String::new(),
                table_dir,
                path_type: PathType::Bare,
                options: HashMap::default(),
                skip_wal_replay: false,
                checkpoint: None,
                requirements: Default::default(),
            }),
        )
        .await
        .unwrap();
}

#[tokio::test]
async fn test_engine_reopen_region() {
    test_engine_reopen_region_with_format(false).await;
    test_engine_reopen_region_with_format(true).await;
}

async fn test_engine_reopen_region_with_format(flat_format: bool) {
    let mut env = TestEnv::with_prefix("reopen-region").await;
    let engine = env
        .create_engine(MitoConfig {
            default_flat_format: flat_format,
            ..Default::default()
        })
        .await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new().build();
    let table_dir = request.table_dir.clone();
    let column_schemas = rows_schema(&request);
    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();
    put_rows(
        &engine,
        region_id,
        Rows {
            schema: column_schemas,
            rows: build_rows(0, 2),
        },
    )
    .await;

    reopen_region(&engine, region_id, table_dir, false, Default::default()).await;
    assert!(engine.is_region_exists(region_id));

    let scanner = engine
        .scanner(region_id, ScanRequest::default())
        .await
        .unwrap();
    let stream = scanner.scan().await.unwrap();
    let batches = RecordBatches::try_collect(stream).await.unwrap();
    assert_eq!(2, batches.iter().map(|b| b.num_rows()).sum::<usize>());
}

#[tokio::test]
async fn test_engine_open_readonly() {
    test_engine_open_readonly_with_format(false).await;
    test_engine_open_readonly_with_format(true).await;
}

async fn test_engine_open_readonly_with_format(flat_format: bool) {
    let mut env = TestEnv::new().await;
    let engine = env
        .create_engine(MitoConfig {
            default_flat_format: flat_format,
            ..Default::default()
        })
        .await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new().build();
    let table_dir = request.table_dir.clone();
    let column_schemas = rows_schema(&request);
    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();

    reopen_region(&engine, region_id, table_dir, false, Default::default()).await;

    // Region is readonly.
    let rows = Rows {
        schema: column_schemas.clone(),
        rows: build_rows(0, 2),
    };
    let err = engine
        .handle_request(
            region_id,
            RegionRequest::Put(RegionPutRequest {
                skip_wal: false,
                rows: rows.clone(),
                hint: None,
                partition_expr_version: None,
            }),
        )
        .await
        .unwrap_err();
    assert_eq!(StatusCode::RegionNotReady, err.status_code());

    assert_eq!(Some(RegionRole::Follower), engine.role(region_id));
    // Converts region to leader.
    engine
        .set_region_role(region_id, RegionRole::Leader)
        .unwrap();
    assert_eq!(Some(RegionRole::Leader), engine.role(region_id));

    put_rows(&engine, region_id, rows).await;
}

#[tokio::test]
async fn test_engine_region_open_with_options() {
    test_engine_region_open_with_options_with_format(false).await;
    test_engine_region_open_with_options_with_format(true).await;
}

async fn test_engine_region_open_with_options_with_format(flat_format: bool) {
    let mut env = TestEnv::new().await;
    let engine = env
        .create_engine(MitoConfig {
            default_flat_format: flat_format,
            ..Default::default()
        })
        .await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new().build();
    let table_dir = request.table_dir.clone();
    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();

    // Close the region.
    engine
        .handle_request(
            region_id,
            RegionRequest::Close(RegionCloseRequest::default()),
        )
        .await
        .unwrap();

    // Open the region again with options.
    engine
        .handle_request(
            region_id,
            RegionRequest::Open(RegionOpenRequest {
                engine: String::new(),
                table_dir,
                path_type: PathType::Bare,
                options: HashMap::from([("ttl".to_string(), "4d".to_string())]),
                skip_wal_replay: false,
                checkpoint: None,
                requirements: Default::default(),
            }),
        )
        .await
        .unwrap();

    let region = engine.get_region(region_id).unwrap();
    assert_eq!(
        region.version().options.ttl,
        Some(Duration::from_secs(3600 * 24 * 4).into())
    );
}

#[tokio::test]
async fn test_engine_region_open_with_custom_store() {
    test_engine_region_open_with_custom_store_with_format(false).await;
    test_engine_region_open_with_custom_store_with_format(true).await;
}

async fn test_engine_region_open_with_custom_store_with_format(flat_format: bool) {
    let mut env = TestEnv::new().await;
    let engine = env
        .create_engine_with_multiple_object_stores(
            MitoConfig {
                default_flat_format: flat_format,
                ..Default::default()
            },
            None,
            None,
            &["Gcs"],
        )
        .await;
    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new()
        .insert_option("storage", "Gcs")
        .build();
    let table_dir = request.table_dir.clone();

    // Create a custom region.
    engine
        .handle_request(region_id, RegionRequest::Create(request.clone()))
        .await
        .unwrap();

    // Close the custom region.
    engine
        .handle_request(
            region_id,
            RegionRequest::Close(RegionCloseRequest::default()),
        )
        .await
        .unwrap();

    // Open the custom region.
    engine
        .handle_request(
            region_id,
            RegionRequest::Open(RegionOpenRequest {
                engine: String::new(),
                table_dir,
                path_type: PathType::Bare,
                options: HashMap::from([("storage".to_string(), "Gcs".to_string())]),
                skip_wal_replay: false,
                checkpoint: None,
                requirements: Default::default(),
            }),
        )
        .await
        .unwrap();

    // The region should not be opened with the default object store.
    let region = engine.get_region(region_id).unwrap();
    let object_store_manager = env.get_object_store_manager().unwrap();
    assert!(
        !object_store_manager
            .default_object_store()
            .exists(&region.access_layer.build_region_dir(region_id))
            .await
            .unwrap()
    );
    assert!(
        object_store_manager
            .find("Gcs")
            .unwrap()
            .exists(&region.access_layer.build_region_dir(region_id))
            .await
            .unwrap()
    );
}

#[tokio::test]
async fn test_open_region_skip_wal_replay() {
    test_open_region_skip_wal_replay_with_format(false).await;
    test_open_region_skip_wal_replay_with_format(true).await;
}

async fn test_open_region_skip_wal_replay_with_format(flat_format: bool) {
    let mut env = TestEnv::new().await;
    let engine = env
        .create_engine(MitoConfig {
            default_flat_format: flat_format,
            ..Default::default()
        })
        .await;

    let region_id = RegionId::new(1, 1);
    env.get_schema_metadata_manager()
        .register_region_table_info(
            region_id.table_id(),
            "test_table",
            "test_catalog",
            "test_schema",
            None,
            env.get_kv_backend(),
        )
        .await;

    let request = CreateRequestBuilder::new().build();
    let table_dir = request.table_dir.clone();

    let column_schemas = rows_schema(&request);
    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();

    let rows = Rows {
        schema: column_schemas.clone(),
        rows: build_rows(0, 3),
    };
    put_rows(&engine, region_id, rows).await;

    flush_region(&engine, region_id, None).await;

    let rows = Rows {
        schema: column_schemas.clone(),
        rows: build_rows(3, 5),
    };
    put_rows(&engine, region_id, rows).await;

    let engine = env
        .reopen_engine(
            engine,
            MitoConfig {
                default_flat_format: flat_format,
                ..Default::default()
            },
        )
        .await;
    // Skip the WAL replay .
    engine
        .handle_request(
            region_id,
            RegionRequest::Open(RegionOpenRequest {
                engine: String::new(),
                table_dir: table_dir.clone(),
                path_type: PathType::Bare,
                options: Default::default(),
                skip_wal_replay: true,
                checkpoint: None,
                requirements: Default::default(),
            }),
        )
        .await
        .unwrap();

    let request = ScanRequest::default();
    let stream = engine.scan_to_stream(region_id, request).await.unwrap();
    let batches = RecordBatches::try_collect(stream).await.unwrap();
    let expected = "\
+-------+---------+---------------------+
| tag_0 | field_0 | ts                  |
+-------+---------+---------------------+
| 0     | 0.0     | 1970-01-01T00:00:00 |
| 1     | 1.0     | 1970-01-01T00:00:01 |
| 2     | 2.0     | 1970-01-01T00:00:02 |
+-------+---------+---------------------+";
    assert_eq!(expected, batches.pretty_print().unwrap());

    // Replay the WAL.
    let engine = env
        .reopen_engine(
            engine,
            MitoConfig {
                default_flat_format: flat_format,
                ..Default::default()
            },
        )
        .await;
    // Open the region again with options.
    engine
        .handle_request(
            region_id,
            RegionRequest::Open(RegionOpenRequest {
                engine: String::new(),
                table_dir,
                path_type: PathType::Bare,
                options: Default::default(),
                skip_wal_replay: false,
                checkpoint: None,
                requirements: Default::default(),
            }),
        )
        .await
        .unwrap();

    let request = ScanRequest::default();
    let stream = engine.scan_to_stream(region_id, request).await.unwrap();
    let batches = RecordBatches::try_collect(stream).await.unwrap();
    let expected = "\
+-------+---------+---------------------+
| tag_0 | field_0 | ts                  |
+-------+---------+---------------------+
| 0     | 0.0     | 1970-01-01T00:00:00 |
| 1     | 1.0     | 1970-01-01T00:00:01 |
| 2     | 2.0     | 1970-01-01T00:00:02 |
| 3     | 3.0     | 1970-01-01T00:00:03 |
| 4     | 4.0     | 1970-01-01T00:00:04 |
+-------+---------+---------------------+";
    assert_eq!(expected, batches.pretty_print().unwrap());
}

/// Holds the first background opener after registration without blocking the worker.
#[derive(Default)]
struct OpenRegistrationGate {
    registrations: AtomicUsize,
    entered: Notify,
    resume: Notify,
}

#[async_trait::async_trait]
impl EventListener for OpenRegistrationGate {
    async fn on_region_open_registered(&self, _region_id: RegionId) {
        if self.registrations.fetch_add(1, Ordering::Relaxed) == 0 {
            self.entered.notify_one();
            self.resume.notified().await;
        }
    }
}

/// Releases the opener even if a regression assertion fails.
struct OpenRegistrationGateGuard(Arc<OpenRegistrationGate>);

impl Drop for OpenRegistrationGateGuard {
    fn drop(&mut self) {
        self.0.resume.notify_one();
    }
}

#[tokio::test]
async fn test_reopen_does_not_wait_for_previous_open_task() {
    let mut env = TestEnv::with_prefix("reopen-before-open-task-exits").await;
    let gate = Arc::new(OpenRegistrationGate::default());
    let engine = env
        .create_engine_with(
            MitoConfig {
                num_workers: 1,
                ..Default::default()
            },
            None,
            Some(gate.clone()),
            None,
        )
        .await;
    let region_id = RegionId::new(1, 1);
    let create = CreateRequestBuilder::new().build();
    let open = RegionOpenRequest {
        engine: String::new(),
        table_dir: create.table_dir.clone(),
        path_type: PathType::Bare,
        options: Default::default(),
        skip_wal_replay: false,
        checkpoint: None,
        requirements: Default::default(),
    };
    engine
        .handle_request(region_id, RegionRequest::Create(create))
        .await
        .unwrap();
    engine
        .handle_request(
            region_id,
            RegionRequest::Close(RegionCloseRequest::default()),
        )
        .await
        .unwrap();

    let guard = OpenRegistrationGateGuard(gate.clone());
    let first_engine = engine.clone();
    let first_request = open.clone();
    let first_open = tokio::spawn(async move {
        first_engine
            .handle_request(region_id, RegionRequest::Open(first_request))
            .await
    });
    tokio::time::timeout(Duration::from_secs(5), gate.entered.notified())
        .await
        .expect("the first opener did not receive its registration acknowledgement");
    assert!(engine.is_region_exists(region_id));

    // A later open must load the region again, not join the first opener's waiters.
    tokio::time::timeout(Duration::from_secs(5), async {
        engine
            .handle_request(
                region_id,
                RegionRequest::Close(RegionCloseRequest::default()),
            )
            .await
            .unwrap();
        assert!(!engine.is_region_exists(region_id));
        engine
            .handle_request(region_id, RegionRequest::Open(open))
            .await
            .unwrap();
    })
    .await
    .expect("reopening must not wait for the previous background opener to exit");
    assert!(engine.is_region_exists(region_id));
    assert!(!engine.is_region_opening(region_id));

    drop(guard);
    tokio::time::timeout(Duration::from_secs(5), first_open)
        .await
        .unwrap()
        .unwrap()
        .unwrap();
}

#[tokio::test]
async fn test_open_region_wait_for_opening_region_ok() {
    test_open_region_wait_for_opening_region_ok_with_format(false).await;
    test_open_region_wait_for_opening_region_ok_with_format(true).await;
}

async fn test_open_region_wait_for_opening_region_ok_with_format(flat_format: bool) {
    let mut env = TestEnv::with_prefix("wait-for-opening-region-ok").await;
    let engine = env
        .create_engine(MitoConfig {
            default_flat_format: flat_format,
            ..Default::default()
        })
        .await;
    let region_id = RegionId::new(1, 1);
    let worker = engine.inner.workers.worker(region_id);
    let (tx, rx) = oneshot::channel();
    let opening_regions = worker.opening_regions().clone();
    opening_regions.insert_sender(region_id, tx.into());
    assert!(engine.is_region_opening(region_id));

    let handle_open = tokio::spawn(async move {
        engine
            .handle_request(
                region_id,
                RegionRequest::Open(RegionOpenRequest {
                    engine: String::new(),
                    table_dir: "empty".to_string(),
                    path_type: PathType::Bare,
                    options: HashMap::default(),
                    skip_wal_replay: false,
                    checkpoint: None,
                    requirements: Default::default(),
                }),
            )
            .await
    });

    // Wait for conditions
    while opening_regions.sender_len(region_id) != 2 {
        tokio::time::sleep(Duration::from_millis(100)).await;
    }

    let senders = opening_regions.remove_sender(region_id);
    for sender in senders {
        sender.send(Ok(0));
    }

    assert_eq!(handle_open.await.unwrap().unwrap().affected_rows, 0);
    assert_eq!(rx.await.unwrap().unwrap(), 0);
}

#[tokio::test]
async fn test_open_region_wait_for_opening_region_err() {
    test_open_region_wait_for_opening_region_err_with_format(false).await;
    test_open_region_wait_for_opening_region_err_with_format(true).await;
}

async fn test_open_region_wait_for_opening_region_err_with_format(flat_format: bool) {
    let mut env = TestEnv::with_prefix("wait-for-opening-region-err").await;
    let engine = env
        .create_engine(MitoConfig {
            default_flat_format: flat_format,
            ..Default::default()
        })
        .await;
    let region_id = RegionId::new(1, 1);
    let worker = engine.inner.workers.worker(region_id);
    let (tx, rx) = oneshot::channel();
    let opening_regions = worker.opening_regions().clone();
    opening_regions.insert_sender(region_id, tx.into());
    assert!(engine.is_region_opening(region_id));

    let handle_open = tokio::spawn(async move {
        engine
            .handle_request(
                region_id,
                RegionRequest::Open(RegionOpenRequest {
                    engine: String::new(),
                    table_dir: "empty".to_string(),
                    path_type: PathType::Bare,
                    options: HashMap::default(),
                    skip_wal_replay: false,
                    checkpoint: None,
                    requirements: Default::default(),
                }),
            )
            .await
    });

    // Wait for conditions
    while opening_regions.sender_len(region_id) != 2 {
        tokio::time::sleep(Duration::from_millis(100)).await;
    }

    let senders = opening_regions.remove_sender(region_id);
    for sender in senders {
        sender.send(Err(error::RegionNotFoundSnafu { region_id }.build()));
    }

    assert_eq!(
        handle_open.await.unwrap().unwrap_err().status_code(),
        StatusCode::RegionNotFound
    );
    assert_eq!(
        rx.await.unwrap().unwrap_err().status_code(),
        StatusCode::RegionNotFound
    );
}

#[tokio::test]
async fn test_open_compaction_region() {
    test_open_compaction_region_with_format(false).await;
    test_open_compaction_region_with_format(true).await;
}

async fn test_open_compaction_region_with_format(flat_format: bool) {
    let mut env = TestEnv::new().await;
    let mut mito_config = MitoConfig {
        default_flat_format: flat_format,
        ..Default::default()
    };
    mito_config
        .sanitize(&env.data_home().display().to_string())
        .unwrap();

    let engine = env.create_engine(mito_config.clone()).await;

    let region_id = RegionId::new(1, 1);
    let schema_metadata_manager = env.get_schema_metadata_manager();
    schema_metadata_manager
        .register_region_table_info(
            region_id.table_id(),
            "test_table",
            "test_catalog",
            "test_schema",
            None,
            env.get_kv_backend(),
        )
        .await;
    let request = CreateRequestBuilder::new().build();
    let table_dir = request.table_dir.clone();
    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();

    // Close the region.
    engine
        .handle_request(
            region_id,
            RegionRequest::Close(RegionCloseRequest::default()),
        )
        .await
        .unwrap();

    let object_store_manager = env.get_object_store_manager().unwrap();

    let req = OpenCompactionRegionRequest {
        region_id,
        table_dir: table_dir.clone(),
        path_type: PathType::Bare,
        region_options: RegionOptions::default(),
        max_parallelism: 1,
        plugins: Plugins::new(),
    };

    mito_config.index.aux_path = env
        .data_home()
        .join("compactor-index")
        .display()
        .to_string();
    let opener = CompactionRegionOpener::new(mito_config, object_store_manager)
        .await
        .unwrap();
    let compaction_region = opener
        .open(&req, Either::Right(schema_metadata_manager))
        .await
        .unwrap();

    assert_eq!(region_id, compaction_region.region_id);
    compaction_region.stop_purger_scheduler().await.unwrap();
}

// Regression: per-job initialization used to delete another compaction's live
// atomic-write, external-sort, and staging files under the shared aux_path.
#[tokio::test]
async fn test_compaction_region_opens_preserve_live_index_files() {
    use common_time::TimeToLive;
    use futures::{AsyncReadExt, AsyncWriteExt};
    use index::external_provider::ExternalTempFileProvider;
    use store_api::storage::FileId;

    use crate::sst::index::intermediate::{IntermediateLocation, TempFileProvider};

    let mut env = TestEnv::new().await;
    let engine = env.create_engine(MitoConfig::default()).await;
    let mut requests = Vec::new();
    for (index, path_type) in [PathType::Bare, PathType::Data].into_iter().enumerate() {
        let region_id = RegionId::new(100 + index as u32, 0);
        let mut create = CreateRequestBuilder::new()
            .table_dir(&format!("data/{}/", region_id.table_id()))
            .build();
        create.path_type = path_type;
        engine
            .handle_request(region_id, RegionRequest::Create(create.clone()))
            .await
            .unwrap();
        put_rows(
            &engine,
            region_id,
            Rows {
                schema: rows_schema(&create),
                rows: build_rows(0, 3),
            },
        )
        .await;
        flush_region(&engine, region_id, None).await;
        requests.push(OpenCompactionRegionRequest {
            region_id,
            table_dir: create.table_dir,
            path_type,
            region_options: RegionOptions::default(),
            max_parallelism: 1,
            plugins: Plugins::new(),
        });
    }

    let mut config = MitoConfig::default();
    config.sanitize(env.data_home().to_str().unwrap()).unwrap();
    let aux_path = env.data_home().join("compactor-index");
    config.index.aux_path = aux_path.display().to_string();
    let opener = CompactionRegionOpener::new(config, env.get_object_store_manager().unwrap())
        .await
        .unwrap();
    let first = opener
        .open(&requests[0], Either::Left(TimeToLive::default()))
        .await
        .unwrap();
    let provider = TempFileProvider::new(
        IntermediateLocation::new(&first.region_id, &FileId::random()),
        first.access_layer.intermediate_manager().clone(),
    );
    let mut committed = provider.create("sort", "committed").await.unwrap();
    committed.write_all(b"committed sort input").await.unwrap();
    committed.close().await.unwrap();
    let mut pending = provider.create("sort", "pending").await.unwrap();
    pending.write_all(b"pending sort input").await.unwrap();
    pending.flush().await.unwrap();
    let staging_file = aux_path.join("staging/inflight.tmp");
    tokio::fs::write(&staging_file, b"staging input")
        .await
        .unwrap();

    let cloned_opener = opener.clone();
    let (second, reopened) = tokio::try_join!(
        opener.open(&requests[1], Either::Left(TimeToLive::default())),
        cloned_opener.open(&requests[0], Either::Left(TimeToLive::default())),
    )
    .unwrap();
    assert_eq!(second.access_layer.path_type(), PathType::Data);
    assert_eq!(reopened.access_layer.path_type(), PathType::Bare);
    let missing = OpenCompactionRegionRequest {
        region_id: RegionId::new(999, 0),
        table_dir: "missing/".into(),
        ..requests[0].clone()
    };
    assert!(matches!(
        opener
            .open(&missing, Either::Left(TimeToLive::default()))
            .await,
        Err(error::Error::EmptyRegionDir { .. })
    ));
    drop(cloned_opener);
    drop(opener);

    // Open regions retain the resources even after the opener itself is dropped.
    pending.close().await.unwrap();
    let mut contents = Vec::new();
    for (_, mut reader) in provider.read_all("sort").await.unwrap() {
        let mut bytes = Vec::new();
        reader.read_to_end(&mut bytes).await.unwrap();
        contents.push(bytes);
    }
    contents.sort();
    assert_eq!(
        contents,
        vec![
            b"committed sort input".to_vec(),
            b"pending sort input".to_vec()
        ]
    );
    assert_eq!(
        tokio::fs::read(&staging_file).await.unwrap(),
        b"staging input"
    );
    provider.cleanup().await.unwrap();
    for region in [first, second, reopened] {
        region.stop_purger_scheduler().await.unwrap();
    }
}

/// Returns a dummy fetcher that returns the expr_json only once.
fn new_dummy_fetcher(region_id: RegionId, expr_json: String) -> PartitionExprFetcherRef {
    struct DummyFetcher {
        region_id: RegionId,
        expr_json: String,
        fetch_count: AtomicUsize,
    }
    #[async_trait::async_trait]
    impl PartitionExprFetcher for DummyFetcher {
        async fn fetch_expr(&self, region_id: RegionId) -> Option<String> {
            if region_id == self.region_id {
                let count = self.fetch_count.fetch_add(1, Ordering::Relaxed);
                (count == 0).then(|| self.expr_json.clone())
            } else {
                None
            }
        }
    }

    Arc::new(DummyFetcher {
        region_id,
        expr_json,
        fetch_count: AtomicUsize::new(0),
    })
}

#[tokio::test]
async fn test_open_backfills_partition_expr_with_fetcher() {
    let mut env = TestEnv::with_prefix("open-backfill-fetcher").await;

    // prepare plugins with dummy fetcher
    let region_id = RegionId::new(100, 1);
    let expr_json = r#"{\"Expr\":{\"lhs\":{\"Column\":\"a\"},\"op\":\"GtEq\",\"rhs\":{\"Value\":{\"UInt32\":10}}}}"#.to_string();
    let fetcher = new_dummy_fetcher(region_id, expr_json.clone());
    let engine = env
        .create_engine_with(MitoConfig::default(), None, None, Some(fetcher))
        .await;

    // create region with partition_expr=None (legacy style)
    let request = CreateRequestBuilder::new()
        .table_dir("t")
        .partition_expr_json(None)
        .build();
    engine
        .handle_request(region_id, RegionRequest::Create(request.clone()))
        .await
        .unwrap();

    let meta = engine.get_region(region_id).unwrap().metadata();
    assert!(meta.partition_expr.is_none());

    // close and reopen to trigger backfill in opener
    engine
        .handle_request(
            region_id,
            RegionRequest::Close(RegionCloseRequest::default()),
        )
        .await
        .unwrap();
    engine
        .handle_request(
            region_id,
            RegionRequest::Open(RegionOpenRequest {
                engine: String::new(),
                table_dir: request.table_dir.clone(),
                path_type: PathType::Bare,
                options: HashMap::default(),
                skip_wal_replay: false,
                checkpoint: None,
                requirements: Default::default(),
            }),
        )
        .await
        .unwrap();

    // verify partition_expr is backfilled and persisted
    let meta = engine.get_region(region_id).unwrap().metadata();
    assert_eq!(meta.partition_expr.as_deref(), Some(expr_json.as_str()));

    // Set leader and trigger backfill
    engine
        .set_region_role(region_id, RegionRole::Leader)
        .unwrap();
    engine
        .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Leader)
        .await
        .unwrap();

    // reopen again to ensure no further changes and still Some
    engine
        .handle_request(
            region_id,
            RegionRequest::Close(RegionCloseRequest::default()),
        )
        .await
        .unwrap();
    let engine = env.reopen_engine(engine, MitoConfig::default()).await;
    engine
        .handle_request(
            region_id,
            RegionRequest::Open(RegionOpenRequest {
                engine: String::new(),
                table_dir: request.table_dir.clone(),
                path_type: PathType::Bare,
                options: HashMap::default(),
                skip_wal_replay: false,
                checkpoint: None,
                requirements: Default::default(),
            }),
        )
        .await
        .unwrap();
    let meta = engine.get_region(region_id).unwrap().metadata();
    assert_eq!(meta.partition_expr.as_deref(), Some(expr_json.as_str()));
}

#[tokio::test]
async fn test_open_keeps_none_without_fetcher() {
    let mut env = TestEnv::with_prefix("open-no-fetcher").await;
    // engine without fetcher
    let engine = env.create_engine(MitoConfig::default()).await;

    let region_id = RegionId::new(200, 1);
    let request = CreateRequestBuilder::new()
        .table_dir("t2")
        .partition_expr_json(None)
        .build();
    engine
        .handle_request(region_id, RegionRequest::Create(request.clone()))
        .await
        .unwrap();
    let meta = engine.get_region(region_id).unwrap().metadata();
    assert!(meta.partition_expr.is_none());

    engine
        .handle_request(
            region_id,
            RegionRequest::Close(RegionCloseRequest::default()),
        )
        .await
        .unwrap();
    let engine = env.reopen_engine(engine, MitoConfig::default()).await;
    engine
        .handle_request(
            region_id,
            RegionRequest::Open(RegionOpenRequest {
                engine: String::new(),
                table_dir: request.table_dir.clone(),
                path_type: PathType::Bare,
                options: HashMap::default(),
                skip_wal_replay: false,
                checkpoint: None,
                requirements: Default::default(),
            }),
        )
        .await
        .unwrap();

    let meta = engine.get_region(region_id).unwrap().metadata();
    assert!(meta.partition_expr.is_none());
}

#[tokio::test]
async fn test_region_hook_on_open() {
    common_telemetry::init_default_ut_logging();
    let mut env = TestEnv::with_prefix("open_hook").await;

    let hook = Arc::new(MockRegionHook::new());
    let plugins = Plugins::new();
    plugins.insert(hook.clone() as RegionHookRef);

    let engine = env
        .create_engine_with_plugins(MitoConfig::default(), plugins)
        .await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new().build();
    let table_dir = request.table_dir.clone();
    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();

    // Creating a region fires on_region_opened exactly once.
    assert_eq!(hook.opened_count.load(Ordering::Relaxed), 1);
    assert_eq!(hook.closed_count.load(Ordering::Relaxed), 0);

    // `reopen_region` closes then opens; the open must fire the hook again.
    reopen_region(&engine, region_id, table_dir, false, Default::default()).await;

    assert_eq!(hook.opened_count.load(Ordering::Relaxed), 2);
    assert_eq!(hook.closed_count.load(Ordering::Relaxed), 1);
    assert_eq!(hook.dropped_count.load(Ordering::Relaxed), 0);
}

#[tokio::test]
async fn test_reopen_region_with_compressed_manifest_cache() {
    let mut env = TestEnv::new().await;
    let engine = env
        .create_engine(MitoConfig {
            enable_write_cache: true,
            compress_manifest: true,
            manifest_checkpoint_distance: 1,
            ..Default::default()
        })
        .await;
    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new().build();
    let table_dir = request.table_dir.clone();
    let schema = rows_schema(&request);
    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();
    for i in 0..3 {
        put_rows(
            &engine,
            region_id,
            Rows {
                schema: schema.clone(),
                rows: build_rows(i * 3, (i + 1) * 3),
            },
        )
        .await;
        flush_region(&engine, region_id, None).await;
    }
    let region = engine.get_region(region_id).unwrap();
    let mut manager = region.manifest_ctx.manifest_manager.write().await;
    tokio::time::timeout(
        Duration::from_secs(10),
        manager.wait_for_pending_checkpoint(),
    )
    .await
    .unwrap();
    assert!(manager.checkpointer().last_checkpoint_version() > 0);
    drop(manager);
    let stream = engine
        .scan_to_stream(region_id, ScanRequest::default())
        .await
        .unwrap();
    let expected = RecordBatches::try_collect(stream)
        .await
        .unwrap()
        .pretty_print()
        .unwrap();
    reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await;
    let stream = engine
        .scan_to_stream(region_id, ScanRequest::default())
        .await
        .unwrap();
    let actual = RecordBatches::try_collect(stream)
        .await
        .unwrap()
        .pretty_print()
        .unwrap();
    assert_eq!(expected, actual);
}