// Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

#![feature(test)]

extern crate test;
#[macro_use]
extern crate slog_global;

pub mod encryption;
mod kv_generator;
mod logging;
mod macros;
mod runner;
mod security;

use std::{
    env,
    net::TcpListener,
    sync::atomic::{AtomicU16, Ordering},
    thread,
    time::Duration,
};

use rand::Rng;
use tikv_util::sys::thread::StdThreadBuildWrapper;

pub use crate::{
    encryption::*,
    kv_generator::*,
    logging::*,
    runner::{TestHook, clear_failpoints, run_failpoint_tests, run_test_with_hook, run_tests},
    security::*,
};

pub fn setup_for_ci() {
    // We use backtrace in tests to record suspicious problems. And loading
    // backtrace the first time can take several seconds. Spawning a thread and
    // load it ahead of time to avoid causing timeout.
    thread::Builder::new()
        .name(tikv_util::thd_name!("backtrace-loader"))
        .spawn_wrapper(::backtrace::Backtrace::new)
        .unwrap();

    if env::var("CI").is_ok() {
        // HACK! Use `epollex` as the polling engine for gRPC when running CI tests on
        // Linux and it hasn't been set before.
        // See more: https://github.com/grpc/grpc/blob/v1.17.2/src/core/lib/iomgr/ev_posix.cc#L124
        // See more: https://grpc.io/grpc/core/md_doc_core_grpc-polling-engines.html
        #[cfg(target_os = "linux")]
        {
            if env::var("GRPC_POLL_STRATEGY").is_err() {
                env::set_var("GRPC_POLL_STRATEGY", "epollex");
            }
        }

        logging::init_log_for_test();
    }

    if env::var("PANIC_ABORT").is_ok() {
        // Panics as aborts, it's helpful for debugging,
        // but also stops tests immediately.
        tikv_util::set_panic_hook(true, "./");
    }

    tikv_util::check_environment_variables();

    if let Err(e) = tikv_util::config::check_max_open_fds(4096) {
        panic!(
            "To run test, please make sure the maximum number of open file descriptors not \
             less than 4096: {:?}",
            e
        );
    }
}

static INITIAL_PORT: AtomicU16 = AtomicU16::new(0);
// Keep test ports above test_cloud_server's highest node port (28999) and
// below Linux's default ephemeral range (32768 and above).
const MIN_TEST_PORT: u16 = 30000;
const MAX_TEST_PORT: u16 = 32767;

/// Allocates a port for testing purpose.
pub fn alloc_port() -> u16 {
    let p = INITIAL_PORT.load(Ordering::Relaxed);
    if p == 0 {
        let _ = INITIAL_PORT.compare_exchange(
            0,
            rand::thread_rng().gen_range(MIN_TEST_PORT..=MAX_TEST_PORT),
            Ordering::SeqCst,
            Ordering::SeqCst,
        );
    }
    let mut p = INITIAL_PORT.load(Ordering::SeqCst);
    loop {
        let next = if p >= MAX_TEST_PORT {
            MIN_TEST_PORT
        } else {
            p + 1
        };
        match INITIAL_PORT.compare_exchange_weak(p, next, Ordering::SeqCst, Ordering::SeqCst) {
            Ok(_) if TcpListener::bind(("127.0.0.1", next)).is_ok() => return next,
            Ok(_) => p = next,
            Err(e) => p = e,
        }
    }
}

static MEM_DISK: &str = "TIKV_TEST_MEMORY_DISK_MOUNT_POINT";

/// Gets a temporary path. The directory will be removed when dropped.
///
/// The returned path will point to memory only when memory disk is available
/// and specified.
pub fn temp_dir(prefix: impl Into<Option<&'static str>>, prefer_mem: bool) -> tempfile::TempDir {
    let mut builder = tempfile::Builder::new();
    if let Some(prefix) = prefix.into() {
        builder.prefix(prefix);
    }
    match env::var(MEM_DISK) {
        Ok(dir) if prefer_mem => {
            debug!("using memory disk"; "path" => %dir);
            builder.tempdir_in(dir).unwrap()
        }
        _ => builder.tempdir().unwrap(),
    }
}

#[track_caller]
pub fn eventually(tick: Duration, total: Duration, mut check: impl FnMut() -> bool) {
    let start = std::time::Instant::now();
    loop {
        if check() {
            return;
        }
        if start.elapsed() < total {
            std::thread::sleep(tick);
            continue;
        }
        panic!(
            "failed to pass the check after {:?} elapsed",
            start.elapsed()
        );
    }
}