greptimedb/tests/runner/src/cmd/bare.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::path::PathBuf;
use std::sync::Arc;

use clap::{Parser, ValueEnum};
use sqlness::interceptor::Registry;
use sqlness::{ConfigBuilder, Runner};

use crate::cmd::SqlnessConfig;
use crate::env::bare::{Env, ServiceProvider, StoreConfig, WalConfig};
use crate::{protocol_interceptor, util};

#[derive(ValueEnum, Debug, Clone)]
#[clap(rename_all = "snake_case")]
enum Wal {
    RaftEngine,
    Kafka,
}

// add a group to ensure that all server addresses are set together
#[derive(clap::Args, Debug, Clone, Default)]
pub(crate) struct ServerAddr {
    /// Address of the grpc server.
    #[clap(short, long)]
    pub(crate) server_addr: Option<String>,

    /// Address of the postgres server. Must be set if server_addr is set.
    #[clap(short, long, requires = "server_addr")]
    pub(crate) pg_server_addr: Option<String>,

    /// Address of the mysql server. Must be set if server_addr is set.
    #[clap(short, long, requires = "server_addr")]
    pub(crate) mysql_server_addr: Option<String>,
}

#[derive(Debug, Parser)]
/// Run sqlness tests in bare mode.
pub struct BareCommand {
    #[clap(flatten)]
    config: SqlnessConfig,

    /// Addresses of the server.
    #[command(flatten)]
    server_addr: ServerAddr,

    /// The type of Wal.
    #[clap(short, long, default_value = "raft_engine")]
    wal: Wal,

    /// The kafka wal broker endpoints. This config will suppress sqlness runner
    /// from starting a kafka cluster, and use the given endpoint as kafka backend.
    #[clap(short, long)]
    kafka_wal_broker_endpoints: Option<String>,

    /// The path to the directory where GreptimeDB's binaries resides.
    /// If not set, sqlness will build GreptimeDB on the fly.
    #[clap(long)]
    bins_dir: Option<PathBuf>,

    /// Preserve persistent state in the temporary directory.
    /// This may affect future test runs.
    #[clap(long)]
    preserve_state: bool,

    /// Pull Different versions of GreptimeDB on need.
    #[clap(long, default_value = "true")]
    pull_version_on_need: bool,

    /// The store addresses for metadata, if empty, will use memory store.
    #[clap(long)]
    store_addrs: Vec<String>,

    /// Whether to setup etcd, by default it is false.
    #[clap(long, default_value = "false")]
    setup_etcd: bool,

    /// Whether to setup pg, by default it is false.
    #[clap(long, default_missing_value = "", num_args(0..=1))]
    setup_pg: Option<ServiceProvider>,

    /// Whether to setup mysql, by default it is false.
    #[clap(long, default_missing_value = "", num_args(0..=1))]
    setup_mysql: Option<ServiceProvider>,

    /// The number of jobs to run in parallel. Default to half of the cores.
    #[clap(short, long, default_value = "0")]
    jobs: usize,

    /// Extra command line arguments when starting GreptimeDB binaries.
    #[clap(long)]
    extra_args: Vec<String>,

    /// Enable flat format for storage engine (sets default_flat_format = true).
    #[clap(long, default_value = "false")]
    enable_flat_format: bool,

    /// Enable garbage collection in metasrv and datanodes.
    #[clap(long, default_value = "false")]
    enable_gc: bool,
}

#[derive(Debug, PartialEq, Eq)]
struct Parallelism {
    /// How many environments (standalone, distributed) run at the same time.
    env: usize,
    /// How many database instances run at the same time within one environment.
    instance: usize,
}

impl BareCommand {
    fn parallelism(&self) -> Parallelism {
        let self_managed_kafka =
            matches!(self.wal, Wal::Kafka) && self.kafka_wal_broker_endpoints.is_none();
        // External servers are shared by both environments, a filtered run is for
        // debugging a single case, and each environment starts and tears down the
        // self-managed Kafka cluster on its own.
        if self.server_addr.server_addr.is_some()
            || self.config.test_filter != ".*"
            || self_managed_kafka
        {
            return Parallelism {
                env: 1,
                instance: 1,
            };
        }

        // Instances of one environment would share the same metadata table and
        // Kafka topics. Only distributed uses the kv backend, and the two
        // environments use different Kafka topic prefixes, so the environments
        // themselves can still run concurrently.
        if self.setup_etcd
            || !self.store_addrs.is_empty()
            || self.setup_pg.is_some()
            || self.setup_mysql.is_some()
            || self.kafka_wal_broker_endpoints.is_some()
        {
            return Parallelism {
                env: self.jobs,
                instance: 1,
            };
        }

        Parallelism {
            env: self.jobs,
            instance: self.jobs,
        }
    }

    pub async fn run(mut self) {
        let temp_dir = tempfile::Builder::new()
            .prefix("sqlness")
            .tempdir()
            .unwrap();
        let sqlness_home = temp_dir.keep();

        let mut interceptor_registry: Registry = Default::default();
        interceptor_registry.register(
            protocol_interceptor::PREFIX,
            Arc::new(protocol_interceptor::ProtocolInterceptorFactory),
        );

        if let Some(d) = &self.config.case_dir
            && !d.is_dir()
        {
            panic!("{} is not a directory", d.display());
        }
        if self.jobs == 0 {
            self.jobs = num_cpus::get() / 2;
        }
        let parallelism = self.parallelism();
        println!(
            "Environment parallelism: {}, instance parallelism per environment: {}",
            parallelism.env, parallelism.instance
        );

        let config = ConfigBuilder::default()
            .case_dir(util::get_case_dir(self.config.case_dir))
            .fail_fast(self.config.fail_fast)
            .test_filter(self.config.test_filter)
            .follow_links(true)
            .env_config_file(self.config.env_config_file)
            .interceptor_registry(interceptor_registry)
            .parallelism(parallelism.instance)
            .env_parallelism(parallelism.env)
            .build()
            .unwrap();

        let wal = match self.wal {
            Wal::RaftEngine => WalConfig::RaftEngine,
            Wal::Kafka => WalConfig::Kafka {
                needs_kafka_cluster: self.kafka_wal_broker_endpoints.is_none(),
                broker_endpoints: self
                    .kafka_wal_broker_endpoints
                    .map(|s| s.split(',').map(|s| s.to_string()).collect())
                    // otherwise default to the same port in `kafka-cluster.yml`
                    .unwrap_or(vec!["127.0.0.1:9092".to_string()]),
            },
        };

        let store = StoreConfig {
            store_addrs: self.store_addrs.clone(),
            setup_etcd: self.setup_etcd,
            setup_pg: self.setup_pg,
            setup_mysql: self.setup_mysql,
            enable_flat_format: self.enable_flat_format,
            enable_gc: self.enable_gc,
        };

        let runner = Runner::new(
            config,
            Env::new(
                sqlness_home.clone(),
                self.server_addr,
                wal,
                self.pull_version_on_need,
                self.bins_dir,
                store,
                self.extra_args,
            ),
        );
        match runner.run().await {
            Ok(_) => println!("\x1b[32mAll sqlness tests passed!\x1b[0m"),
            Err(e) => {
                println!("\x1b[31mTest failed: {}\x1b[0m", e);
                std::process::exit(1);
            }
        }

        // clean up and exit
        if !self.preserve_state {
            if self.setup_etcd {
                println!("Stopping etcd");
                util::stop_rm_etcd();
            }
            // TODO(weny): remove postgre and mysql containers
            println!("Removing state in {:?}", sqlness_home);
            tokio::fs::remove_dir_all(sqlness_home).await.unwrap();
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn parallelism(jobs: &str, args: &[&str]) -> Parallelism {
        let cmd = BareCommand::try_parse_from(["bare", "-j", jobs].iter().chain(args)).unwrap();
        cmd.parallelism()
    }

    #[test]
    fn test_parallelism() {
        let serial = Parallelism {
            env: 1,
            instance: 1,
        };
        let env_only = Parallelism {
            env: 4,
            instance: 1,
        };

        assert_eq!(
            parallelism("4", &[]),
            Parallelism {
                env: 4,
                instance: 4
            }
        );
        assert_eq!(parallelism("4", &["--setup-pg"]), env_only);
        assert_eq!(parallelism("4", &["--setup-mysql"]), env_only);
        assert_eq!(parallelism("4", &["--setup-etcd"]), env_only);
        assert_eq!(
            parallelism("4", &["--store-addrs", "127.0.0.1:2379"]),
            env_only
        );
        assert_eq!(
            parallelism("4", &["-w", "kafka", "-k", "127.0.0.1:9092"]),
            env_only
        );
        assert_eq!(parallelism("4", &["-w", "kafka"]), serial);
        assert_eq!(parallelism("4", &["-t", "basic"]), serial);
        assert_eq!(
            parallelism(
                "4",
                &[
                    "-s",
                    "127.0.0.1:4001",
                    "-p",
                    "127.0.0.1:4003",
                    "-m",
                    "127.0.0.1:4002"
                ]
            ),
            serial
        );
        assert_eq!(parallelism("1", &["--setup-pg"]), serial);
    }
}