greptimedb/tests-integration/src/standalone.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::Arc;
use std::time::Duration;

use cache::{
    build_datanode_layered_cache_registry, build_fundamental_cache_registry,
    with_default_composite_cache_registry,
};
use catalog::information_schema::NoopInformationExtension;
use catalog::kvbackend::KvBackendCatalogManagerBuilder;
use catalog::process_manager::ProcessManager;
use common_base::Plugins;
use common_catalog::consts::{MIN_USER_FLOW_ID, MIN_USER_TABLE_ID};
use common_config::KvBackendConfig;
use common_datasource::object_store::LocalFileAccess;
use common_event_recorder::EventRecorderOptions;
use common_meta::cache::LayeredCacheRegistryBuilder;
use common_meta::ddl::flow_meta::FlowMetadataAllocator;
use common_meta::ddl::table_meta::TableMetadataAllocator;
use common_meta::ddl::{DdlContext, NoopRegionFailureDetectorControl};
use common_meta::ddl_manager::DdlManager;
use common_meta::key::TableMetadataManager;
use common_meta::key::flow::FlowMetadataManager;
use common_meta::kv_backend::KvBackendRef;
use common_meta::procedure_executor::LocalProcedureExecutor;
use common_meta::region_keeper::MemoryRegionKeeper;
use common_meta::region_registry::LeaderRegionRegistry;
use common_meta::sequence::SequenceBuilder;
use common_procedure::ProcedureManagerRef;
use common_procedure::local::EventRecorderHandle;
use common_procedure::options::ProcedureConfig;
use common_telemetry::logging::SlowQueryOptions;
use common_test_util::find_workspace_path;
use common_wal::config::DatanodeWalConfig;
use datanode::datanode::{Datanode, DatanodeBuilder};
use flow::{FlownodeBuilder, FrontendClient, GrpcQueryHandlerWithBoxedError};
use frontend::frontend::Frontend;
use frontend::instance::Instance;
use frontend::instance::builder::FrontendBuilder;
use frontend::server::Services;
use frontend::service_config::{BatcherOptions, PendingRowsBatcherOptions};
use meta_srv::metasrv::{FLOW_ID_SEQ, TABLE_ID_SEQ};
use servers::grpc::GrpcOptions;
use standalone::options::StandaloneOptions;
use standalone::{StandaloneDatanodeManager, StandaloneRepartitionProcedureFactory};

use crate::test_util::{self, StorageType, TestGuard, create_tmp_dir_and_datanode_opts};

pub struct GreptimeDbStandalone {
    /// Storage engine for assertions across protocol and storage boundaries.
    pub mito_engine: mito2::engine::MitoEngine,
    /// Owns the region server and the log store, so a test can shut them down.
    pub datanode: Datanode,
    pub frontend: Arc<Frontend>,
    pub opts: StandaloneOptions,
    pub guard: TestGuard,
    // Used in rebuild.
    pub kv_backend: KvBackendRef,
    pub procedure_manager: ProcedureManagerRef,
    pub event_recorder_handle: EventRecorderHandle,
}

impl GreptimeDbStandalone {
    pub fn fe_instance(&self) -> &Arc<Instance> {
        &self.frontend.instance
    }
}

pub struct GreptimeDbStandaloneBuilder {
    instance_name: String,
    datanode_wal_config: DatanodeWalConfig,
    store_providers: Option<Vec<StorageType>>,
    default_store: Option<StorageType>,
    plugin: Option<Plugins>,
    slow_query_options: SlowQueryOptions,
    event_recorder_options: EventRecorderOptions,
    auto_create_table: bool,
    logical_batcher: Option<BatcherOptions>,
    table_batcher: BatcherOptions,
    mito_config: Option<mito2::config::MitoConfig>,
    pipeline_options: pipeline::PipelineOptions,
}

impl GreptimeDbStandaloneBuilder {
    pub fn new(instance_name: &str) -> Self {
        Self {
            instance_name: instance_name.to_string(),
            store_providers: None,
            plugin: None,
            default_store: None,
            datanode_wal_config: DatanodeWalConfig::default(),
            // Enable slow query log with 1s threshold by default for integration tests.
            slow_query_options: SlowQueryOptions {
                enable: true,
                threshold: Duration::from_secs(1),
                ..Default::default()
            },
            event_recorder_options: EventRecorderOptions::default(),
            auto_create_table: true,
            logical_batcher: None,
            table_batcher: BatcherOptions::default(),
            mito_config: None,
            pipeline_options: pipeline::PipelineOptions::default(),
        }
    }

    /// Overrides the Mito configuration for this test instance.
    #[must_use]
    pub fn with_mito_config(mut self, config: mito2::config::MitoConfig) -> Self {
        self.mito_config = Some(config);
        self
    }

    /// Configures pipeline compilation and caching for this test instance.
    #[must_use]
    pub fn with_pipeline_options(mut self, options: pipeline::PipelineOptions) -> Self {
        self.pipeline_options = options;
        self
    }

    /// Configures ordinary-table batching for protocol integration tests.
    #[must_use]
    pub fn with_table_batcher(mut self, options: BatcherOptions) -> Self {
        self.table_batcher = options;
        self
    }

    /// Configures logical-table batching for integration tests.
    #[must_use]
    pub fn with_logical_batcher(mut self, options: BatcherOptions) -> Self {
        self.logical_batcher = Some(options);
        self
    }

    #[must_use]
    pub fn with_auto_create_table(mut self, auto_create_table: bool) -> Self {
        self.auto_create_table = auto_create_table;
        self
    }

    /// Sets the event recorder options for the standalone instance.
    #[must_use]
    pub fn with_event_recorder_options(
        mut self,
        event_recorder_options: EventRecorderOptions,
    ) -> Self {
        self.event_recorder_options = event_recorder_options;
        self
    }

    #[must_use]
    pub fn with_default_store_type(self, store_type: StorageType) -> Self {
        Self {
            default_store: Some(store_type),
            ..self
        }
    }

    #[cfg(test)]
    #[must_use]
    pub fn with_store_providers(self, store_providers: Vec<StorageType>) -> Self {
        Self {
            store_providers: Some(store_providers),
            ..self
        }
    }

    #[must_use]
    pub fn with_plugin(self, plugin: Plugins) -> Self {
        Self {
            plugin: Some(plugin),
            ..self
        }
    }

    #[must_use]
    pub fn with_slow_query_threshold(mut self, threshold: Duration) -> Self {
        self.slow_query_options.threshold = threshold;
        self
    }

    #[must_use]
    pub fn with_datanode_wal_config(mut self, datanode_wal_config: DatanodeWalConfig) -> Self {
        self.datanode_wal_config = datanode_wal_config;
        self
    }

    pub async fn build_with(
        &self,
        kv_backend: KvBackendRef,
        guard: TestGuard,
        opts: StandaloneOptions,
        procedure_manager: ProcedureManagerRef,
        event_recorder_handle: EventRecorderHandle,
        register_procedure_loaders: bool,
    ) -> GreptimeDbStandalone {
        let plugins = self.plugin.clone().unwrap_or_default();

        let layered_cache_registry =
            Arc::new(build_datanode_layered_cache_registry(kv_backend.clone()));

        let mut builder =
            DatanodeBuilder::new(opts.datanode_options(), plugins.clone(), kv_backend.clone());
        let local_file_access = LocalFileAccess::sandboxed(find_workspace_path(".")).unwrap();
        builder.with_cache_registry(layered_cache_registry);
        builder.with_local_file_access(local_file_access.clone());
        let datanode = builder.build().await.unwrap();

        let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend.clone()));
        table_metadata_manager.init().await.unwrap();

        let flow_metadata_manager = Arc::new(FlowMetadataManager::new(kv_backend.clone()));

        let layered_cache_builder = LayeredCacheRegistryBuilder::default();
        let fundamental_cache_registry = build_fundamental_cache_registry(kv_backend.clone());
        let cache_registry = Arc::new(
            with_default_composite_cache_registry(
                layered_cache_builder.add_cache_registry(fundamental_cache_registry),
            )
            .unwrap()
            .build(),
        );

        let catalog_manager = KvBackendCatalogManagerBuilder::new(
            Arc::new(NoopInformationExtension),
            kv_backend.clone(),
            cache_registry.clone(),
        )
        .with_procedure_manager(procedure_manager.clone())
        .build();

        let (frontend_client, frontend_instance_handler) =
            FrontendClient::from_empty_grpc_handler(opts.query.clone());
        let flow_builder = FlownodeBuilder::new(
            Default::default(),
            plugins.clone(),
            table_metadata_manager.clone(),
            catalog_manager.clone(),
            flow_metadata_manager.clone(),
            Arc::new(frontend_client),
        );
        let flownode = Arc::new(flow_builder.build().await.unwrap());

        let node_manager = Arc::new(StandaloneDatanodeManager {
            region_server: datanode.region_server(),
            flow_server: flownode.flow_engine(),
        });

        let table_id_allocator = Arc::new(
            SequenceBuilder::new(TABLE_ID_SEQ, kv_backend.clone())
                .initial(MIN_USER_TABLE_ID as u64)
                .step(10)
                .build(),
        );
        let flow_id_sequence = Arc::new(
            SequenceBuilder::new(FLOW_ID_SEQ, kv_backend.clone())
                .initial(MIN_USER_FLOW_ID as u64)
                .step(10)
                .build(),
        );
        let wal_provider = cmd::standalone::build_standalone_wal_provider(
            &opts.wal,
            opts.datanode_options().node_id,
            kv_backend.clone(),
        )
        .await
        .unwrap();
        let wal_provider = Arc::new(wal_provider);
        let table_metadata_allocator = Arc::new(TableMetadataAllocator::new(
            table_id_allocator,
            wal_provider.clone(),
        ));
        let flow_metadata_allocator = Arc::new(FlowMetadataAllocator::with_noop_peer_allocator(
            flow_id_sequence,
        ));

        let repartition_procedure_factory = Arc::new(StandaloneRepartitionProcedureFactory);
        let ddl_manager = Arc::new(DdlManager::new(
            DdlContext {
                affinity_resolver: None,
                affinity_fallback: Default::default(),
                peer_alloc_context_provider: None,
                node_manager: node_manager.clone(),
                cache_invalidator: cache_registry.clone(),
                memory_region_keeper: Arc::new(MemoryRegionKeeper::default()),
                leader_region_registry: Arc::new(LeaderRegionRegistry::default()),
                table_metadata_manager,
                table_metadata_allocator,
                flow_metadata_manager,
                flow_metadata_allocator,
                region_failure_detector_controller: Arc::new(NoopRegionFailureDetectorControl),
                soft_drop_enabled: false,
                soft_drop_retention: None,
                create_database_metadata_committer: None,
            },
            procedure_manager.clone(),
            repartition_procedure_factory,
        ));
        if register_procedure_loaders {
            ddl_manager.register_loaders().unwrap();
        }
        let procedure_executor = Arc::new(LocalProcedureExecutor::new(
            ddl_manager,
            procedure_manager.clone(),
        ));

        let server_addr = opts.frontend_options().grpc.server_addr.clone();

        let instance = FrontendBuilder::new(
            opts.frontend_options(),
            kv_backend.clone(),
            cache_registry.clone(),
            catalog_manager.clone(),
            node_manager.clone(),
            procedure_executor.clone(),
            Arc::new(ProcessManager::new(server_addr, None)),
        )
        .with_local_file_access(local_file_access)
        .with_plugin(plugins.clone())
        .try_build()
        .await
        .unwrap();
        let instance = Arc::new(instance);

        event_recorder_handle.install(instance.event_recorder());

        // set the frontend client for flownode
        let grpc_handler = instance.clone() as Arc<dyn GrpcQueryHandlerWithBoxedError>;
        let weak_grpc_handler = Arc::downgrade(&grpc_handler);
        frontend_instance_handler
            .set_handler(weak_grpc_handler)
            .await;
        procedure_manager.start().await.unwrap();
        wal_provider.start().await.unwrap();

        test_util::prepare_another_catalog_and_schema(&instance).await;

        let servers = Services::new(opts.clone(), instance.clone(), plugins)
            .build()
            .unwrap();
        let frontend = Frontend {
            instance,
            servers,
            heartbeat_task: None,
        };

        GreptimeDbStandalone {
            mito_engine: datanode.region_server().mito_engine().unwrap(),
            datanode,
            frontend: Arc::new(frontend),
            opts,
            guard,
            kv_backend,
            procedure_manager,
            event_recorder_handle,
        }
    }

    pub async fn build(&self) -> GreptimeDbStandalone {
        let default_store_type = self.default_store.unwrap_or(StorageType::File);
        let store_types = self.store_providers.clone().unwrap_or_default();

        let (opts, guard) = create_tmp_dir_and_datanode_opts(
            default_store_type,
            store_types,
            &self.instance_name,
            self.datanode_wal_config.clone(),
            Default::default(),
        );

        let kv_backend_config = KvBackendConfig::default();
        let procedure_config = ProcedureConfig::default();

        let kv_backend = standalone::build_metadata_kvbackend(
            format!("{}/kv", opts.storage.data_home),
            kv_backend_config,
        )
        .unwrap();
        let (procedure_manager, event_recorder_handle) =
            standalone::build_procedure_manager(kv_backend.clone(), procedure_config);

        let mut standalone_opts = StandaloneOptions {
            storage: opts.storage,
            procedure: procedure_config,
            metadata_store: kv_backend_config,
            wal: opts.wal,
            grpc: GrpcOptions::default().with_server_addr("127.0.0.1:4001"),
            slow_query: self.slow_query_options.clone(),
            event_recorder: self.event_recorder_options.clone(),
            auto_create_table: self.auto_create_table,
            pipeline: self.pipeline_options.clone(),
            pending_rows_batcher: PendingRowsBatcherOptions {
                logical_table: self.logical_batcher.clone(),
                table: self.table_batcher.clone(),
            },
            // Tests cover the descriptor, so they run with it enabled.
            otlp: frontend::service_config::OtlpOptions {
                experimental_enable_resource_info: true,
                ..Default::default()
            },
            ..StandaloneOptions::default()
        };

        if let Some(config) = &self.mito_config {
            for engine in &mut standalone_opts.region_engine {
                if let datanode::config::RegionEngineConfig::Mito(mito) = engine {
                    *mito = config.clone();
                }
            }
        }

        self.build_with(
            kv_backend,
            guard,
            standalone_opts,
            procedure_manager,
            event_recorder_handle,
            true,
        )
        .await
    }
}