// Copyright 2022 TiKV Project Authors. Licensed under Apache-2.0.

mod sql_util;
mod test_all;
mod test_columnar;
mod test_drop_table;
mod test_jepsen;
mod test_load_data;
mod test_native_br;
mod test_tidb;
mod test_txn_file;
mod test_unique;

use std::{
    collections::HashSet,
    str::FromStr,
    sync::{
        atomic::{AtomicU16, AtomicUsize, Ordering},
        Arc, RwLock,
    },
    thread::{sleep, JoinHandle},
    time::Duration,
};

use api_version::ApiV2;
use cloud_encryption::KeyspaceEncryptionConfig;
use codec::number::NumberEncoder;
use futures::executor::block_on;
use http::{Request, StatusCode, Uri};
use hyper::Body;
use kvengine::{
    dfs::{self, Dfs, FileType, S3Fs},
    table::columnar::build_schema_file,
};
use kvproto::metapb::Store;
use pd_client::PdClient;
use rand::prelude::*;
use security::SecurityConfig;
use test_cloud_server::{
    client::ClusterTxnClient,
    keyspace::{ClusterKeyspaceClient, KeyspaceManager},
    scheduler::Scheduler,
    tidb::TidbCluster,
    try_wait_result, ServerCluster,
};
use test_pd_client::TestPdClient;
use tidb_query_datatype::{
    codec::row::v2::encoder_for_test::{Column, RowEncoder},
    expr::EvalContext,
};
use tikv::config::TikvConfig;
use tikv_client::TimestampExt;
use tikv_util::{box_err, error, info, time::Instant, warn};
use tokio::sync::{OwnedRwLockWriteGuard, Semaphore};
use txn_types::Key;

pub(crate) type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
pub(crate) type Result<T> = std::result::Result<T, Error>;

lazy_static::lazy_static! {
    pub static ref WRITE_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref TXN_FILE_WRITE_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref MOVE_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref MERGE_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref TRANSFER_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref NODE_RESTART_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref BACKUP_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref BACKUP_TOLERATED_ERR_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref RESTORE_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref RESTORE_TOLERATED_ERR_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref LOAD_DATA_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref KEYSPACE_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref TABLE_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref DROP_TABLE_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref MANUAL_MAJOR_COMPACT_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref GC_ADVANCE_SAFE_POINT_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref TPCC_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref JEPSEN_BANK_TXN_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref JEPSEN_BANK_TXN_RETRY_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref UNIQUE_WORKLOAD_TXN_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref UNIQUE_WORKLOAD_CONFLICT_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref COLUMNAR_WRITE_COUNTER: AtomicUsize = AtomicUsize::new(0);
    pub static ref COLUMNAR_RETRY_COUNTER: AtomicUsize = AtomicUsize::new(0);
}

pub const TIMEOUT: Duration = Duration::from_secs(90);
pub const WRITE_CONCURRENCY: usize = 4;
pub const TXN_FILE_WRITE_CONCURRENCY: usize = 2;

const DEFAULT_INNER_KEY_OFFSET: usize = 4;
const REQUEST_MAJOR_COMPACT_ON_STORE_TIMEOUT: Duration = Duration::from_secs(20);

const KEYSPACE_CLEANUP_LOCKS_CONCURRENCY: usize = 4;
const KEYSPACE_CLEANUP_LOCKS_TIMEOUT: Duration = Duration::from_secs(30);
const GC_INTERVAL: Duration = Duration::from_secs(10);

const DROP_TABLE_CONCURRENCY: usize = 2; // More than 1 thread to reduce the chance of blocked by other mutual-exclusive workloads for a long time.

static NODE_ALLOCATOR: AtomicU16 = AtomicU16::new(1);

pub(crate) fn alloc_node_id() -> u16 {
    let node_id = NODE_ALLOCATOR.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
    info!("allocated node_id {}", node_id);
    node_id
}

pub(crate) fn alloc_node_id_vec(count: usize) -> Vec<u16> {
    let mut nodes = vec![];
    nodes.resize_with(count, || alloc_node_id());
    nodes
}

pub(crate) fn spawn_move(scheduler: Scheduler, two_node_down: Arc<RwLock<()>>) -> JoinHandle<()> {
    std::thread::spawn(move || {
        let start_time = Instant::now();
        while start_time.saturating_elapsed() < TIMEOUT {
            sleep(Duration::from_millis(1000));
            let guard = two_node_down.read().unwrap();
            let ok = scheduler.move_random_region();
            drop(guard);
            if ok {
                MOVE_COUNTER.fetch_add(1, Ordering::SeqCst);
            }
        }
        info!("move thread exit");
    })
}

pub(crate) fn spawn_merge(scheduler: Scheduler, disallow_cross_keyspace: bool) -> JoinHandle<()> {
    std::thread::spawn(move || {
        let start_time = Instant::now();
        while start_time.saturating_elapsed() < TIMEOUT {
            if scheduler.merge_random_region(disallow_cross_keyspace) {
                MERGE_COUNTER.fetch_add(1, Ordering::SeqCst);
            }
        }
        info!("merge thread exit");
    })
}

pub(crate) fn spawn_transfer(scheduler: Scheduler) -> JoinHandle<()> {
    std::thread::spawn(move || {
        let start_time = Instant::now();
        while start_time.saturating_elapsed() < TIMEOUT {
            sleep(Duration::from_millis(100));
            if scheduler.transfer_random_leader() {
                TRANSFER_COUNTER.fetch_add(1, Ordering::SeqCst);
            }
        }
        info!("transfer thread exit");
    })
}

pub(crate) fn spawn_gc_worker(
    client: ClusterTxnClient,
    pd_client: Arc<TestPdClient>,
    keyspace_manager: KeyspaceManager,
    timeout: Duration,
) -> tokio::task::JoinHandle<()> {
    tokio::spawn(async move {
        let semaphore = Arc::new(Semaphore::new(KEYSPACE_CLEANUP_LOCKS_CONCURRENCY));
        let start_time = Instant::now();
        while start_time.saturating_elapsed() < timeout {
            let safepoint = client.current_timestamp().await.unwrap();
            let all_keyspaces = keyspace_manager.get_all_keyspaces();

            tokio::time::sleep(GC_INTERVAL).await;
            info!(
                "gc_worker: try advance safepoint to {}",
                safepoint.version()
            );

            let mut handles = Vec::with_capacity(all_keyspaces.len());
            for &keyspace_id in &all_keyspaces {
                let start = ApiV2::get_txn_keyspace_prefix(keyspace_id);
                let end = ApiV2::get_txn_keyspace_prefix(keyspace_id + 1);

                let semaphore = semaphore.clone();
                let inner_client = client.inner.clone();
                let safepoint = safepoint.clone();
                let h = tokio::spawn(async move {
                    let permit = semaphore.acquire().await.unwrap();
                    let res = tokio::time::timeout(
                        KEYSPACE_CLEANUP_LOCKS_TIMEOUT,
                        inner_client.cleanup_locks(start..end, &safepoint, Default::default()),
                    )
                    .await;
                    drop(permit);

                    match res {
                        Ok(Ok(res)) => {
                            // Should not have any error in `res`.
                            assert!(
                                res.key_error.is_none() && res.region_error.is_none(),
                                "unexpected CleanupLocksResult, check client rust: key_error {:?}, region_error {:?}, keyspace_id {}",
                                res.key_error,
                                res.region_error,
                                keyspace_id,
                            );
                            info!("gc_worker: cleanup_locks succeed"; "resolved_locks" => res.resolved_locks, "safepoint" => safepoint.version(), "keyspace_id" => keyspace_id);
                            true
                        }
                        Ok(Err(err)) => {
                            error!("gc_worker: cleanup_locks failed"; "err" => ?err, "safepoint" => safepoint.version(), "keyspace_id" => keyspace_id);
                            false
                        }
                        Err(_elapsed) => {
                            // Should not timeout for 30s.
                            // TODO: inspect the reason.
                            warn!("gc_worker: cleanup_locks timeout"; "safepoint" => safepoint.version(), "keyspace_id" => keyspace_id);
                            false
                        }
                    }
                });
                handles.push(h);
            }
            let results = futures::future::join_all(handles).await;
            if results.into_iter().all(|r| r.unwrap()) {
                let new_safepoint = pd_client.set_gc_safe_point(safepoint.version()).unwrap();
                assert!(new_safepoint <= safepoint.version());
                if new_safepoint == safepoint.version() {
                    info!("gc_worker: advance safepoint to {}", safepoint.version());
                    GC_ADVANCE_SAFE_POINT_COUNTER.fetch_add(1, Ordering::SeqCst);
                } else {
                    warn!("gc_worker: update safepoint failed"; "safepoint" => safepoint.version());
                }

                for keyspace_id in all_keyspaces {
                    test_drop_table::pick_and_run_pending_destroy_range(
                        keyspace_id,
                        new_safepoint,
                        &keyspace_manager,
                        &client,
                    )
                    .await
                    .unwrap();
                }
            }
        }
        info!("gc_worker thread exit");
    })
}

pub(crate) fn check_gc() {
    // TODO: check times of GC after address the cleanup_locks timeout issue.
    // let counter = GC_ADVANCE_SAFE_POINT_COUNTER.load(Ordering::SeqCst);
    // assert!(
    //     counter > 0,
    //     "gc_worker: times of safe point advanced less than expected, actual
    // {}",     counter,
    // );
}

pub(crate) fn spawn_major_compact(
    pd_client: Arc<TestPdClient>,
    keyspace_manager: KeyspaceManager,
    timeout: Duration,
) -> JoinHandle<()> {
    std::thread::spawn(move || {
        let mut rng = rand::thread_rng();
        let runtime = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .enable_all()
            .thread_name("major-compact")
            .build()
            .unwrap();

        let start_time = Instant::now();
        while start_time.saturating_elapsed() < timeout {
            // Manual major compact should not conform the zipf distribution.
            let keyspace_id = keyspace_manager.get_uniform_random_keyspace(&mut rng);

            let stores = pd_client.get_all_stores(true).unwrap();
            {
                // To be mutual-exclusive with load_data.
                let lock = keyspace_manager.get_keyspace_lock(keyspace_id);
                let _guard = lock.extra_lock();

                let mut handles = Vec::with_capacity(stores.len());
                for store in stores {
                    handles.push(
                        runtime.spawn(request_major_compact_on_store(store.clone(), keyspace_id)),
                    );
                }
                for handle in handles {
                    runtime.block_on(handle).unwrap().unwrap();
                }
            }

            MANUAL_MAJOR_COMPACT_COUNTER.fetch_add(1, Ordering::SeqCst);
            sleep(Duration::from_secs(10));
        }
        info!("major compact worker thread exit");
    })
}

pub(crate) async fn request_major_compact_on_store(store: Store, keyspace_id: u32) -> Result<()> {
    let uri = Uri::from_str(&format!(
        "http://{}/major-compact?major_compact=true&keyspace_id={}",
        &store.status_address, keyspace_id
    ))
    .unwrap();
    let client = hyper::Client::new();
    let mut last_err: Option<Error> = None;
    let mut retry = 0;
    let start_time = Instant::now();
    while start_time.saturating_elapsed() < REQUEST_MAJOR_COMPACT_ON_STORE_TIMEOUT {
        retry += 1;
        let req = Request::post(&uri).body(Body::empty()).unwrap();
        match client.request(req).await {
            // Treat 404 as success.
            Ok(resp) if (resp.status().is_success() || resp.status() == StatusCode::NOT_FOUND) => {
                let body = hyper::body::to_bytes(resp.into_body()).await.unwrap();
                let msg = String::from_utf8_lossy(&body);
                info!(
                    "request_major_compact_on_store success, keyspace {}: {}",
                    keyspace_id,
                    msg.as_ref()
                );
                return Ok(());
            }
            Ok(resp) => {
                let body = hyper::body::to_bytes(resp.into_body()).await.unwrap();
                let msg = String::from_utf8_lossy(&body);
                // Return error only when there is bad request argument.
                panic!(
                    "request_major_compact_on_store failed: {}, retry {}",
                    msg.as_ref(),
                    retry
                );
            }
            Err(err) => {
                last_err = Some(box_err!(
                    "request_major_compact_on_store failed: {:?}, retry {}",
                    err,
                    retry
                ));
                warn!("{:?}", last_err.as_ref().unwrap());
                tokio::time::sleep(Duration::from_millis(500)).await;
                continue;
            }
        }
    }
    Err(last_err.unwrap())
}

pub(crate) fn spawn_keyspace_write(
    idx: usize,
    client: ClusterKeyspaceClient,
    timeout: Duration,
) -> tokio::task::JoinHandle<()> {
    tokio::spawn(async move {
        // Make sure each write thread don't conflict with others.
        let begin = idx * 2000;
        let end = begin + 2000 - 10;
        let start_time = Instant::now();
        while start_time.saturating_elapsed() < timeout {
            let random_keyspace = || {
                let mut rng = rand::thread_rng();
                client.keyspace_manager().get_zipf_random_keyspace(&mut rng)
            };
            let keyspace_id = random_keyspace();
            {
                let lock = client.keyspace_manager().get_keyspace_lock(keyspace_id);
                let guard = lock.try_shared_lock();
                if guard.is_none() {
                    continue;
                }
                let _guard = guard.unwrap();

                let random = || {
                    let mut rng = rand::thread_rng();
                    let table_meta = client.keyspace_manager().get_random_available_table(
                        keyspace_id,
                        &mut rng,
                        true,
                    );
                    let i = rng.gen_range(begin..end);
                    let put_kv = rng.gen_ratio(2, 3);
                    (table_meta, i, put_kv)
                };
                let (table_meta, i, put_kv) = random();
                let table_id = match table_meta.as_ref() {
                    Some(table_meta) => table_meta.id(),
                    None => continue,
                };
                let is_schema_enabled = table_meta.unwrap().is_schema_enabled();

                let ver = client.current_timestamp().await.unwrap().version();
                info!(
                    "[{}] thread write on keyspace {} table {} schema_enabled {}, ver {}",
                    idx, keyspace_id, table_id, is_schema_enabled, ver
                );

                if put_kv {
                    client
                        .keyspace_put_kv(
                            keyspace_id,
                            table_id,
                            i..(i + 10),
                            gen_row_key_suffix,
                            |i| gen_row_val(&mut EvalContext::default(), i, ver),
                        )
                        .await
                        .unwrap();
                } else {
                    client
                        .keyspace_del_kv(keyspace_id, table_id, i..(i + 10), gen_row_key_suffix)
                        .await
                        .unwrap();
                }
            }
            WRITE_COUNTER.fetch_add(10, Ordering::SeqCst);
        }
        info!("keyspace write thread {} exit", idx);
    })
}

pub fn spawn_create_keyspace(
    pd_client: Arc<TestPdClient>,
    keyspace_manager: KeyspaceManager,
    fs: &S3Fs,
    initial_table_count: usize,
    schema_enable_ratio: f64,
    timeout: Duration,
) -> JoinHandle<()> {
    let fs = fs.clone();
    std::thread::spawn(move || {
        let mut rng = rand::thread_rng();
        let start_time = Instant::now();
        while start_time.saturating_elapsed() < timeout {
            sleep(Duration::from_secs(rng.gen_range(0..5)));

            let _ = fs.get_runtime().block_on(create_new_keyspace(
                &pd_client,
                &keyspace_manager,
                &fs,
                initial_table_count,
                schema_enable_ratio,
                false,
            ));
        }
        info!("create keyspace thread exit");
    })
}

async fn create_new_keyspace(
    pd_client: &Arc<TestPdClient>,
    keyspace_manager: &KeyspaceManager,
    fs: &S3Fs,
    initial_table_count: usize,
    schema_enable_ratio: f64,
    locked: bool,
) -> (u32, Option<OwnedRwLockWriteGuard<()>>) {
    let random = || -> (u32, bool) {
        let mut rng = thread_rng();
        let delta = rng.gen_range(1..=3);
        let enable_encryption = rng.gen();
        (delta, enable_encryption)
    };
    let (delta, enable_encryption) = random();

    let new_keyspace = keyspace_manager.new_keyspace_id(delta);
    must_split_region_for_keyspace(pd_client, new_keyspace, enable_encryption).await;
    // `create_single_keyspace` do not shuffle keyspaces. Then we will have a few
    // big keyspaces.
    let lock = keyspace_manager
        .create_single_keyspace(
            new_keyspace,
            TidbCluster::keyspace_name(new_keyspace as u16),
            DEFAULT_INNER_KEY_OFFSET,
            initial_table_count,
            schema_enable_ratio,
            locked,
        )
        .await;
    // Setup schema file to stores
    // TODO: support tables with schema enabled created dynamically.
    let keyspace_meta = keyspace_manager.get_keyspace_meta(new_keyspace).unwrap();
    let stores = pd_client.get_all_stores(true).unwrap();
    if !keyspace_meta.schemas().is_empty() {
        let schema_version = 10;
        let schema_data =
            build_schema_file(new_keyspace, schema_version, keyspace_meta.schemas(), 0);
        let schema_file_id = pd_client.alloc_id().unwrap();
        fs.create(
            schema_file_id,
            schema_data.into(),
            dfs::Options::default().with_type(FileType::Schema),
        )
        .await
        .unwrap();
        // Setup schema file to stores.
        broadcast_schema_file_request(&stores, new_keyspace, schema_file_id).await;
    }
    KEYSPACE_COUNTER.fetch_add(1, Ordering::Relaxed);
    TABLE_COUNTER.fetch_add(initial_table_count, Ordering::Relaxed);

    (new_keyspace, lock)
}

pub async fn broadcast_schema_file_request(
    stores: &[Store],
    keyspace_id: u32,
    schema_file_id: u64,
) {
    for store in stores {
        let status_addr = store.get_status_address();
        let http_client = hyper::client::Client::new();
        let start_time = Instant::now_coarse();
        while start_time.saturating_elapsed() < TIMEOUT {
            let request = hyper::http::Request::builder()
                .method(http::method::Method::POST)
                .uri(format!(
                    "http://{}/schema_file?keyspace_id={}&file_id={}",
                    status_addr, keyspace_id, schema_file_id
                ))
                .body(Body::empty())
                .unwrap();
            if let Ok(resp) = http_client.request(request).await {
                if resp.status().is_success() {
                    break;
                }
            }
            tokio::time::sleep(Duration::from_millis(500)).await;
        }
    }
}

pub fn gen_row_key_suffix(i: usize) -> Vec<u8> {
    let mut suffix = vec![];
    suffix.write_i64(i as i64).unwrap();
    suffix
}

pub fn gen_row_val(ctx: &mut EvalContext, i: usize, ver: u64) -> Vec<u8> {
    let mut row_val = vec![];
    let str_val = generate_random_string(format!("write-{}-", ver));
    let cols = vec![
        Column::new(1, Some(i as i64)),
        Column::new(2, Some(str_val(10))),
    ];
    row_val.write_row(ctx, cols).unwrap();
    row_val
}

async fn must_split_region_for_keyspace(
    pd_client: &TestPdClient,
    keyspace_id: u32,
    enable_encryption: bool,
) {
    let keys = vec![
        ApiV2::get_txn_keyspace_prefix(keyspace_id),
        ApiV2::get_txn_keyspace_prefix(keyspace_id + 1),
    ];
    if enable_encryption {
        let cfg = KeyspaceEncryptionConfig { enabled: true };
        pd_client.set_keyspace_encryption(keyspace_id, cfg).unwrap();
    }
    let split_keys = keys
        .iter()
        .map(|k| Key::from_raw(k).into_encoded())
        .collect::<Vec<_>>();
    pd_client
        .split_regions_with_retry(split_keys, Duration::from_secs(60))
        .await
        .unwrap();
    let new_region = pd_client.get_region(&keys[0]).unwrap();
    info!(
        "split region for keyspace {}: {:?}",
        keyspace_id, new_region
    );
}

pub(crate) fn random_node_restart(cluster: &mut ServerCluster) {
    let mut rng = rand::thread_rng();

    // Some regions would loss majority for a wile when the sleep duration is small.
    // Data corruption should not happen, and workloads should tolerate this.
    sleep(Duration::from_secs(rng.gen_range(3..17)));

    let nodes = cluster.get_nodes();
    let node_id = *nodes.choose(&mut rng).unwrap();
    let sleep_sec = rng.gen_range(0..3);
    let force_stop = rng.gen();
    cluster.restart_node(node_id, Duration::from_secs(sleep_sec), force_stop);
    NODE_RESTART_COUNTER.fetch_add(1, Ordering::Relaxed);
}

pub(crate) fn i_to_key(i: usize) -> Vec<u8> {
    format!("tkey{:08}", i).into_bytes()
}

pub(crate) fn generate_keyspace_key(keyspace_id: u32) -> impl Fn(usize) -> Vec<u8> {
    move |i: usize| -> Vec<u8> {
        let mut key = ApiV2::get_txn_keyspace_prefix(keyspace_id);
        key.extend(i_to_key(i));
        key
    }
}

pub(crate) fn generate_random_string(prefix: String) -> impl Fn(usize) -> Vec<u8> {
    move |i: usize| -> Vec<u8> {
        let mut val = prefix.as_bytes().to_vec();
        let n = i % 512 + 1;
        val.reserve(n);
        val.extend(
            rand::thread_rng()
                .sample_iter(&rand::distributions::Alphanumeric)
                .take(n),
        );
        val
    }
}

pub(crate) fn new_security_config() -> SecurityConfig {
    let mut conf = SecurityConfig::default();
    conf.master_key.vendor = "test".to_string();
    conf.master_key.key_id = "random".to_string();
    conf
}

// Switch acquired from env variable, `0` for false, `>0` for true.
pub(crate) fn env_switch(env_key: &str) -> bool {
    env_switch_opt(env_key, 1)
}

pub(crate) fn env_switch_opt(env_key: &str, default: i32) -> bool {
    std::env::var(env_key)
        .map(|s| s.parse().unwrap())
        .unwrap_or(default)
        > 0
}

pub(crate) fn env_param<T>(env_key: &str, default: T) -> T
where
    T: FromStr,
    <T as FromStr>::Err: std::fmt::Debug,
{
    std::env::var(env_key)
        .map(|s| s.parse::<T>().unwrap())
        .unwrap_or(default)
}

pub(crate) fn verify_cluster_stats(cluster: &ServerCluster, bucket_size: u64, timeout: Duration) {
    let mut success_shards = HashSet::new();
    try_wait_result(
        || {
            let stats = cluster.get_data_stats_ext(Some(&success_shards));
            stats.check_data().map_err(|(errs, this_success_shards)| {
                success_shards.extend(this_success_shards);
                (errs, stats)
            })
        },
        timeout.as_secs() as usize,
    )
    .unwrap_or_else(|(err, stats)| {
        stats.log_all();
        panic!("check_data failed: {:?}", err);
    });

    let mut starts_from_encoded_key = vec![];
    try_wait_result(
        || {
            let data_stats = cluster.get_data_stats();
            let res = data_stats.check_buckets(
                cluster.get_pd_client_ext().as_ref(),
                bucket_size,
                &starts_from_encoded_key,
            );
            if let Err((reason, region)) = &res {
                warn!(
                    "check_buckets failed, region {:?}, reason {:?}",
                    region, reason
                );
                if let Some(region) = region {
                    starts_from_encoded_key = region.start_key.clone();
                }
            }
            res.map_err(|err| (err, data_stats))
        },
        timeout.as_secs() as usize,
    )
    .unwrap_or_else(|(err, stats)| {
        stats.log_all();
        panic!("check_buckets failed: {:?}", err);
    });
}