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

mod test_blob_store;
mod test_commit_ts;
mod test_engine_seek_region;
mod test_engine_snapshot;
mod test_remote_coprocessor;
mod test_stats;
mod test_store;

use std::{sync::Arc, thread, time::Duration};

use api_version::ApiV2;
use bytes::Bytes;
use futures::executor::block_on;
use kvengine::{Shard, SnapAccess};
use kvenginepb::ChangeSet;
use kvproto::kvrpcpb::Op;
use pd_client::PdClient;
use protobuf::Message;
use replication_worker::CdcApplyObserver;
use rfstore::store::ApplyContext;
use test_cloud_server::{
    ServerCluster,
    client::{ClusterClient, MutateOptions, TxnMutations},
    must_wait, try_wait,
    util::Mutation,
};
use tikv::storage::{Scanner, config::ExtraCfScanMode, mvcc::CloudReader, txn::CloudStoreScanner};
use tikv_util::{codec::bytes::encode_bytes, config::ReadableSize};
use txn_types::{Key, TsSet};

use crate::alloc_node_id;

#[test]
fn test_engine_auto_switch() {
    let node_id = alloc_node_id();
    let mut cluster = ServerCluster::new(vec![node_id], |_, conf| {
        conf.rocksdb.writecf.write_buffer_size = ReadableSize::kb(256);
    });
    let mut client = cluster.new_client();
    client.put_kv(0..100, i_to_key, i_to_val);
    client.put_kv(100..200, i_to_key, i_to_val);
    client.put_kv(200..300, i_to_key, i_to_val);
    let region_id = client.get_region_id(&[]);
    let engine = cluster.get_kvengine(node_id);
    let stats = engine.get_shard_stat(region_id);
    assert!(stats.mem_table_count + stats.l0_table_count > 1);
    cluster.stop();
}

fn i_to_key(i: usize) -> Vec<u8> {
    format!("x123key_{:05}", i).into_bytes()
}

fn i_to_val(i: usize) -> Vec<u8> {
    format!("val_{:05}", i).into_bytes().repeat(100)
}

#[test]
fn test_split_by_key() {
    let node_id = alloc_node_id();
    let mut cluster = ServerCluster::new(vec![node_id], |_, conf| {
        conf.coprocessor.region_split_keys = Some(3);
    });
    let mut client = cluster.new_client();
    client.put_kv(0..5, i_to_key, i_to_key);
    let engine = cluster.get_kvengine(node_id);
    let _ = try_wait(|| engine.get_all_shard_id_vers().len() == 2, 10);
    let shard_stats = engine.get_all_shard_stats();
    assert!(shard_stats.len() == 2, "{:?}", &shard_stats);
    client.put_kv(6..15, i_to_key, i_to_key);
    let _ = try_wait(|| engine.get_all_shard_id_vers().len() == 5, 10);
    let shard_stats = engine.get_all_shard_stats();
    assert!(shard_stats.len() == 5, "{:?}", &shard_stats);
    cluster.stop();
}

#[test]
fn test_split_whole_keyspace() {
    let node_id = alloc_node_id();
    let mut cluster = ServerCluster::new(vec![node_id], |_, _| {});
    let mut client = cluster.new_client();
    let default_keyspace_i_to_key =
        |i: usize| -> Vec<u8> { format!("t123key_{:05}", i).into_bytes() };
    client.put_kv(0..10, default_keyspace_i_to_key, i_to_val);
    client.split_keyspace(1);
    cluster.wait_pd_region_count(3);
    let pd_client = cluster.get_pd_client();
    let keyspace_prefix = ApiV2::get_txn_keyspace_prefix(1);
    let region = pd_client
        .get_region(&encode_bytes(&keyspace_prefix))
        .unwrap();
    let kv = cluster.get_kvengine(node_id);
    let shard = kv.get_shard(region.id).unwrap();
    let stats = shard.get_stats();
    assert_eq!(stats.kv_size, 0);
    cluster.stop();
}

#[test]
fn test_remove_and_add_peer() {
    let node_ids = vec![alloc_node_id(), alloc_node_id(), alloc_node_id()];
    let mut cluster = ServerCluster::new(node_ids.clone(), |_, _| {});
    cluster.wait_region_replicated(&[], 3);
    let mut client = cluster.new_client();
    let split_key = i_to_key(5);
    client.split(&split_key);
    // Wait for region heartbeat to update region epoch in PD.
    sleep();

    client.put_kv(0..10, i_to_key, i_to_key);
    let pd = cluster.get_pd_client();
    cluster.wait_pd_region_count(2);
    pd.disable_default_operator();
    let &first_node = node_ids.first().unwrap();
    cluster.remove_node_peers(first_node);
    // After one store has removed peer, the cluster is still available.
    let mut client = cluster.new_client();
    client.put_kv(0..10, i_to_key, i_to_key);
    cluster.stop_node(first_node);
    thread::sleep(Duration::from_millis(100));
    cluster.start_node(first_node, |_, _| {});
    pd.enable_default_operator();
    cluster.wait_region_replicated(&[], 3);
    cluster.wait_region_replicated(&split_key, 3);
    cluster.stop();
}

#[test]
fn test_increasing_put_and_split() {
    let node_id = alloc_node_id();
    let mut cluster = ServerCluster::new(vec![node_id], |_, _| {});
    let mut client = cluster.new_client();
    client.put_kv(0..50, i_to_key, i_to_val);
    for i in 1..5 {
        let split_idx = i * 10;
        let split_key = i_to_key(split_idx);
        client.split(&split_key);
        for _ in 0..10 {
            client.put_kv(split_idx..split_idx + 5, i_to_key, i_to_val);
        }
    }
    cluster.stop()
}

#[test]
fn test_cloud_reader_scan_write_for_key() {
    let node_id = alloc_node_id();
    let cluster = ServerCluster::new(vec![node_id], |_, _| {});
    let mut client = cluster.new_client();
    let keyspace_id = ApiV2::get_u32_keyspace_id_by_key("x123".as_bytes()).unwrap();
    client.split_keyspace(keyspace_id);
    cluster.wait_pd_region_count(3);

    let keys: Vec<_> = (1..=6).collect();
    let mut ts_list = vec![];

    // keyX has X versions
    // key1: [1]
    // key2: [1, 2]
    // key3: [1, 2, 3]
    // key4: [1, 2, 3, 4]
    // key5: [1, 2, 3, 4, 5]
    // key6: [1, 2, 3, 4, 5, 6]
    // list of start ts: [1, 2, 3, 4, 5, 6]

    for round in 0..=5 {
        // make sure each keys commit ts large than this round start ts
        let start_ts = client.get_ts();
        for key in keys.iter().skip(round) {
            client
                .try_put_kv(
                    std::iter::once(*key),
                    i_to_key,
                    i_to_val,
                    MutateOptions {
                        start_ts: Some(start_ts),
                        ..Default::default()
                    },
                )
                .unwrap();
        }
        ts_list.push(start_ts);
    }

    let region_id = client.get_region_id("x123".as_bytes());
    let engine = cluster.get_kvengine(node_id);
    let snapshot = engine.get_snap_access(region_id).unwrap();

    for (key_idx, key) in keys.iter().enumerate() {
        let key_raw = Key::from_raw(&i_to_key(*key));
        let expect_ts: Vec<u64> = ts_list
            .iter()
            .take(key_idx + 1)
            .cloned()
            .map(|x| x.into_inner())
            .collect();

        let mut reader = CloudReader::new(snapshot.clone(), false);
        let writes = reader
            .scan_write_for_key(&key_raw, txn_types::TimeStamp::max())
            .unwrap();
        let mut actual_ts: Vec<u64> = writes.iter().map(|w| w.1.start_ts.into_inner()).collect();
        // the largest version is on the top.
        actual_ts.reverse();
        assert_eq!(
            actual_ts, expect_ts,
            "key {} expect {:?}, got {:?}",
            key, expect_ts, actual_ts
        );

        for &cutoff in &expect_ts {
            let mut reader = CloudReader::new(snapshot.clone(), false);
            let writes = reader
                .scan_write_for_key(&key_raw, txn_types::TimeStamp::new(cutoff))
                .unwrap();
            let mut actual_ts: Vec<u64> =
                writes.iter().map(|w| w.1.start_ts.into_inner()).collect();
            actual_ts.reverse();
            let expect_filtered: Vec<u64> =
                expect_ts.iter().cloned().filter(|x| *x < cutoff).collect();
            assert_eq!(
                actual_ts, expect_filtered,
                "key {} scan ts {:?} expect {:?}, got {:?}",
                key, cutoff, expect_filtered, actual_ts
            );
        }

        let min_ts = *expect_ts.iter().min().unwrap();
        let mut reader = CloudReader::new(snapshot.clone(), false);
        let writes = reader
            .scan_write_for_key(&key_raw, txn_types::TimeStamp::new(min_ts - 1))
            .unwrap();
        assert!(
            writes.is_empty(),
            "key {} with scan ts={} should be empty, got {:?}",
            key,
            min_ts - 1,
            writes
        );
    }
}

#[test]
fn test_cloud_reader_find_extra_cf_conflict_record_returns_self_rollback() {
    test_util::init_log_for_test();
    let (mut cluster, mut client, node_id) = new_cloud_reader_test_cluster();
    let key = i_to_key(11);
    let start_ts = client.get_ts();
    let muts = single_key_txn_mutations(&key, &i_to_val(11), Op::Put);

    client
        .kv_prewrite_with_retry(muts.primary(), None, muts.clone(), start_ts)
        .unwrap();
    client.kv_rollback(muts, start_ts).unwrap();

    let snapshot = snapshot_for_key(&cluster, &mut client, node_id, &key);
    let mut reader = CloudReader::new_with_scan_mode(snapshot, false, ExtraCfScanMode::Pruned);
    let user_meta = reader
        .find_extra_cf_conflict_record(&Key::from_raw(&key), start_ts, txn_types::TimeStamp::max())
        .unwrap()
        .expect("self rollback should be found");

    assert_eq!(user_meta.start_ts, start_ts.into_inner());
    assert_eq!(user_meta.commit_ts, 0);
    assert!(user_meta.is_rollback());
    cluster.stop();
}

#[test]
fn test_cloud_reader_find_extra_cf_conflict_record_returns_newer_lock_conflict() {
    test_util::init_log_for_test();
    let (mut cluster, mut client, node_id) = new_cloud_reader_test_cluster();
    let key = i_to_key(12);
    let start_ts = client.get_ts();
    let muts = single_key_txn_mutations(&key, &[], Op::Lock);

    client
        .kv_prewrite_with_retry(muts.primary(), None, muts.clone(), start_ts)
        .unwrap();
    let commit_ts = client.get_ts();
    client.kv_commit(muts, start_ts, commit_ts, false).unwrap();

    let snapshot = snapshot_for_key(&cluster, &mut client, node_id, &key);
    let mut reader = CloudReader::new_with_scan_mode(snapshot, false, ExtraCfScanMode::Pruned);
    let probe_ts = client.get_ts();
    let user_meta = reader
        .find_extra_cf_conflict_record(
            &Key::from_raw(&key),
            probe_ts,
            (commit_ts.into_inner() - 1).into(),
        )
        .unwrap()
        .expect("newer lock conflict should be found");

    assert_eq!(user_meta.start_ts, start_ts.into_inner());
    assert_eq!(user_meta.commit_ts, commit_ts.into_inner());
    assert!(!user_meta.is_rollback());
    cluster.stop();
}

#[test]
fn test_cloud_reader_find_extra_cf_conflict_record_returns_none_when_all_candidates_pruned() {
    test_util::init_log_for_test();
    let (mut cluster, mut client, node_id) = new_cloud_reader_test_cluster();
    let key = i_to_key(13);
    let start_ts = client.get_ts();
    let muts = single_key_txn_mutations(&key, &[], Op::Lock);

    client
        .kv_prewrite_with_retry(muts.primary(), None, muts.clone(), start_ts)
        .unwrap();
    let commit_ts = client.get_ts();
    client.kv_commit(muts, start_ts, commit_ts, false).unwrap();

    let snapshot = snapshot_for_key(&cluster, &mut client, node_id, &key);
    let mut reader = CloudReader::new_with_scan_mode(snapshot, false, ExtraCfScanMode::Pruned);
    let probe_ts = client.get_ts();
    let result = reader
        .find_extra_cf_conflict_record(&Key::from_raw(&key), probe_ts, commit_ts)
        .unwrap();

    assert!(result.is_none());
    cluster.stop();
}

#[test]
fn test_cloud_reader_find_extra_cf_conflict_record_returns_none_for_exact_non_rollback_on_self_rollback_check()
 {
    test_util::init_log_for_test();
    let (mut cluster, mut client, node_id) = new_cloud_reader_test_cluster();
    let key = i_to_key(15);
    let start_ts = commit_single_extra_cf_record(&mut client, &key, Op::Lock);

    let snapshot = snapshot_for_key(&cluster, &mut client, node_id, &key);
    let mut reader = CloudReader::new_with_scan_mode(snapshot, false, ExtraCfScanMode::Pruned);
    let result = reader
        .find_extra_cf_conflict_record(&Key::from_raw(&key), start_ts, txn_types::TimeStamp::max())
        .unwrap();

    assert!(
        result.is_none(),
        "exact non-rollback record must not be treated as self rollback"
    );
    cluster.stop();
}

#[test]
fn test_cloud_reader_find_extra_cf_conflict_record_returns_none_when_surviving_source_has_no_conflict()
 {
    test_util::init_log_for_test();
    let (mut cluster, mut client, node_id) = new_cloud_reader_test_cluster();
    let target_key = i_to_key(16);
    let other_key = i_to_key(17);

    let target_start_ts = commit_single_extra_cf_record(&mut client, &target_key, Op::Lock);
    let target_commit_ts = must_get_extra_cf_commit_ts(&mut client, &target_key, target_start_ts);

    let other_start_ts = commit_single_extra_cf_record(&mut client, &other_key, Op::Lock);
    let other_commit_ts = must_get_extra_cf_commit_ts(&mut client, &other_key, other_start_ts);
    assert!(other_commit_ts > target_commit_ts);

    let snapshot = snapshot_for_key(&cluster, &mut client, node_id, &target_key);
    let mut reader = CloudReader::new_with_scan_mode(snapshot, false, ExtraCfScanMode::Pruned);
    let probe_ts = client.get_ts();
    let result = reader
        .find_extra_cf_conflict_record(
            &Key::from_raw(&target_key),
            probe_ts,
            (target_commit_ts.into_inner() + 1).into(),
        )
        .unwrap();

    assert!(
        result.is_none(),
        "source should survive because of another key, but target key should still have no conflict"
    );
    cluster.stop();
}

#[test]
fn test_cloud_reader_find_extra_cf_conflict_record_ignores_commit_ts_equal_to_threshold() {
    test_util::init_log_for_test();
    let (mut cluster, mut client, node_id) = new_cloud_reader_test_cluster();
    let target_key = i_to_key(18);
    let other_key = i_to_key(19);

    let target_start_ts = commit_single_extra_cf_record(&mut client, &target_key, Op::Lock);
    let target_commit_ts = must_get_extra_cf_commit_ts(&mut client, &target_key, target_start_ts);

    let other_start_ts = commit_single_extra_cf_record(&mut client, &other_key, Op::Lock);
    let other_commit_ts = must_get_extra_cf_commit_ts(&mut client, &other_key, other_start_ts);
    assert!(other_commit_ts > target_commit_ts);

    let snapshot = snapshot_for_key(&cluster, &mut client, node_id, &target_key);
    let mut reader = CloudReader::new_with_scan_mode(snapshot, false, ExtraCfScanMode::Pruned);
    let probe_ts = client.get_ts();
    let result = reader
        .find_extra_cf_conflict_record(&Key::from_raw(&target_key), probe_ts, target_commit_ts)
        .unwrap();

    assert!(
        result.is_none(),
        "commit_ts equal to threshold must not be reported as conflict"
    );
    cluster.stop();
}

#[test]
fn test_cloud_reader_find_extra_cf_conflict_record_legacy_mode_returns_newer_lock_conflict() {
    test_util::init_log_for_test();
    let (mut cluster, mut client, node_id) = new_cloud_reader_test_cluster();
    let key = i_to_key(20);
    let start_ts = client.get_ts();
    let muts = single_key_txn_mutations(&key, &[], Op::Lock);

    client
        .kv_prewrite_with_retry(muts.primary(), None, muts.clone(), start_ts)
        .unwrap();
    let commit_ts = client.get_ts();
    client.kv_commit(muts, start_ts, commit_ts, false).unwrap();

    let snapshot = snapshot_for_key(&cluster, &mut client, node_id, &key);
    let mut reader = CloudReader::new_with_scan_mode(snapshot, false, ExtraCfScanMode::Legacy);
    let probe_ts = client.get_ts();
    let user_meta = reader
        .find_extra_cf_conflict_record(
            &Key::from_raw(&key),
            probe_ts,
            (commit_ts.into_inner() - 1).into(),
        )
        .unwrap()
        .expect("legacy mode should still find newer lock conflict");

    assert_eq!(user_meta.start_ts, start_ts.into_inner());
    assert_eq!(user_meta.commit_ts, commit_ts.into_inner());
    assert!(!user_meta.is_rollback());
    cluster.stop();
}

#[test]
fn test_cloud_reader_legacy_restores_ordered_conflict_selection_across_strategy_boundary() {
    test_util::init_log_for_test();
    let (mut cluster, mut client, node_id) = new_cloud_reader_test_cluster();
    let key = i_to_key(22);
    let rollback_start_ts = client.get_ts();
    let threshold_at_boundary = txn_types::TimeStamp::compose(
        rollback_start_ts.physical() + 1_000,
        rollback_start_ts.logical(),
    );
    let threshold_above_boundary = txn_types::TimeStamp::compose(
        rollback_start_ts.physical() + 1_001,
        rollback_start_ts.logical(),
    );

    // Write an exact rollback record for the transaction being checked.
    let rollback_muts = single_key_txn_mutations(&key, &i_to_val(22), Op::Put);
    client
        .kv_prewrite_with_retry(
            rollback_muts.primary(),
            None,
            rollback_muts.clone(),
            rollback_start_ts,
        )
        .unwrap();
    client
        .kv_rollback(rollback_muts, rollback_start_ts)
        .unwrap();

    // Commit a newer LOCK record that precedes the rollback in EXTRA_CF order.
    // Move the test PD clock forward so its commit timestamp is beyond both
    // conflict thresholds without relying on wall-clock sleeps.
    let current_ts = client.get_ts();
    cluster
        .get_pd_client()
        .set_tso(txn_types::TimeStamp::compose(
            current_ts.physical() + 2_000,
            current_ts.logical(),
        ));
    let newer_lock_start_ts = client.get_ts();
    let newer_lock_muts = single_key_txn_mutations(&key, &[], Op::Lock);
    client
        .kv_prewrite_with_retry(
            newer_lock_muts.primary(),
            None,
            newer_lock_muts.clone(),
            newer_lock_start_ts,
        )
        .unwrap();
    let newer_lock_commit_ts = client.get_ts();
    client
        .kv_commit(
            newer_lock_muts,
            newer_lock_start_ts,
            newer_lock_commit_ts,
            false,
        )
        .unwrap();

    let snapshot = snapshot_for_key(&cluster, &mut client, node_id, &key);
    let find_conflict = |mode, max_allowed_write| {
        let mut reader = CloudReader::new_with_scan_mode(snapshot.clone(), false, mode);
        reader
            .find_extra_cf_conflict_record(
                &Key::from_raw(&key),
                rollback_start_ts,
                max_allowed_write,
            )
            .unwrap()
            .expect("a rollback or newer LOCK conflict should be found")
    };

    // Legacy preserves the pre-PR ordered full scan on both sides of the
    // strategy boundary, so the newer LOCK record is selected first.
    for threshold in [threshold_at_boundary, threshold_above_boundary] {
        let conflict = find_conflict(ExtraCfScanMode::Legacy, threshold);
        assert!(!conflict.is_rollback());
        assert_eq!(conflict.start_ts, newer_lock_start_ts.into_inner());
        assert_eq!(conflict.commit_ts, newer_lock_commit_ts.into_inner());
    }

    // Pruned mode still uses the heuristic: unified at the boundary, then
    // point-GET-first two-step selection above it.
    let conflict = find_conflict(ExtraCfScanMode::Pruned, threshold_at_boundary);
    assert!(!conflict.is_rollback());
    assert_eq!(conflict.start_ts, newer_lock_start_ts.into_inner());

    let conflict = find_conflict(ExtraCfScanMode::Pruned, threshold_above_boundary);
    assert!(conflict.is_rollback());
    assert_eq!(conflict.start_ts, rollback_start_ts.into_inner());

    cluster.stop();
}

#[test]
fn test_cloud_reader_find_extra_cf_conflict_record_uses_test_build_default_mode() {
    test_util::init_log_for_test();
    let (mut cluster, mut client, node_id) = new_cloud_reader_test_cluster();
    let key = i_to_key(21);
    let start_ts = client.get_ts();
    let muts = single_key_txn_mutations(&key, &[], Op::Lock);

    client
        .kv_prewrite_with_retry(muts.primary(), None, muts.clone(), start_ts)
        .unwrap();
    let commit_ts = client.get_ts();
    client.kv_commit(muts, start_ts, commit_ts, false).unwrap();

    let snapshot = snapshot_for_key(&cluster, &mut client, node_id, &key);
    let mut reader = CloudReader::new(snapshot, false);
    let probe_ts = client.get_ts();
    let user_meta = reader
        .find_extra_cf_conflict_record(
            &Key::from_raw(&key),
            probe_ts,
            (commit_ts.into_inner() - 1).into(),
        )
        .unwrap()
        .expect("default CloudReader mode should use the test build default");

    assert_eq!(user_meta.start_ts, start_ts.into_inner());
    assert_eq!(user_meta.commit_ts, commit_ts.into_inner());
    assert!(!user_meta.is_rollback());
    cluster.stop();
}

fn new_cloud_reader_test_cluster() -> (ServerCluster, ClusterClient, u16) {
    let node_id = alloc_node_id();
    let cluster = ServerCluster::new(vec![node_id], |_, _| {});
    cluster.wait_region_replicated(&[], 1);
    let mut client = cluster.new_client();
    let keyspace_id = ApiV2::get_u32_keyspace_id_by_key(&i_to_key(0)).unwrap();
    client.split_keyspace(keyspace_id);
    cluster.wait_pd_region_count(3);
    (cluster, client, node_id)
}

fn snapshot_for_key(
    cluster: &ServerCluster,
    client: &mut ClusterClient,
    node_id: u16,
    key: &[u8],
) -> SnapAccess {
    let region_id = client.get_region_id(key);
    cluster
        .get_kvengine(node_id)
        .get_snap_access(region_id)
        .unwrap()
}

fn single_key_txn_mutations(key: &[u8], value: &[u8], op: Op) -> TxnMutations {
    TxnMutations::from_normal(vec![Mutation {
        key: key.to_vec().into(),
        value: value.to_vec().into(),
        op,
        assertion: kvproto::kvrpcpb::Assertion::None,
    }])
}

fn commit_single_extra_cf_record(
    client: &mut ClusterClient,
    key: &[u8],
    op: Op,
) -> txn_types::TimeStamp {
    let value = match op {
        Op::Put => i_to_val(100),
        _ => Vec::new(),
    };
    let start_ts = client.get_ts();
    let muts = single_key_txn_mutations(key, &value, op);
    client
        .kv_prewrite_with_retry(muts.primary(), None, muts.clone(), start_ts)
        .unwrap();
    let commit_ts = client.get_ts();
    client.kv_commit(muts, start_ts, commit_ts, false).unwrap();
    start_ts
}

fn must_get_extra_cf_commit_ts(
    client: &mut ClusterClient,
    key: &[u8],
    start_ts: txn_types::TimeStamp,
) -> txn_types::TimeStamp {
    let info = client
        .kv_get_mvcc_by_key(key)
        .unwrap()
        .expect("mvcc info should exist");
    let write = info
        .get_writes()
        .iter()
        .find(|write| write.get_start_ts() == start_ts.into_inner())
        .expect("write record should exist");
    write.get_commit_ts().into()
}

#[test]
fn test_cloud_store_reverse_scan() {
    let node_id = alloc_node_id();
    let cluster = ServerCluster::new(vec![node_id], |_, _| {});
    let mut client = cluster.new_client();
    let keyspace_id = ApiV2::get_u32_keyspace_id_by_key("x123".as_bytes()).unwrap();
    client.split_keyspace(keyspace_id);
    cluster.wait_pd_region_count(3);
    client.put_kv(1..6, i_to_key, i_to_val);
    let region_id = client.get_region_id("x123".as_bytes());
    let engine = cluster.get_kvengine(node_id);
    let snapshot = engine.get_snap_access(region_id).unwrap();
    for rev in [true, false] {
        let lower_bound = Some(Key::from_raw(&i_to_key(2)));
        let upper_bound = Some(Key::from_raw(&i_to_key(5)));
        let mut scanner = CloudStoreScanner::new(
            snapshot.clone(),
            rev,
            true,
            TsSet::Empty,
            u64::MAX,
            lower_bound,
            upper_bound,
            false,
            false,
            false,
        )
        .unwrap();
        let mut indices = vec![2, 3, 4];
        if rev {
            indices.reverse();
        }
        for i in indices {
            let (key, _) = scanner.next().unwrap().unwrap();
            assert_eq!(key.into_raw().unwrap(), i_to_key(i));
        }
        assert!(scanner.next().unwrap().is_none());
    }
}

#[test]
fn test_cloud_store_scan_processed_keys_counts_returned_keys_only() {
    let node_id = alloc_node_id();
    let mut cluster = ServerCluster::new(vec![node_id], |_, _| {});
    let mut client = cluster.new_client();
    let keyspace_id = ApiV2::get_u32_keyspace_id_by_key("x123".as_bytes()).unwrap();
    client.split_keyspace(keyspace_id);
    cluster.wait_pd_region_count(3);

    client.put_kv(1..50, i_to_key, i_to_val);
    client.del_kv(2..49, i_to_key);

    let region_id = client.get_region_id("x123".as_bytes());
    let engine = cluster.get_kvengine(node_id);

    // No mvcc versions, should return key1 only, processed_keys should be 1.
    // total_size should be approximate to processed_size.
    {
        let snapshot = engine.get_snap_access(region_id).unwrap();
        let mut scanner = CloudStoreScanner::new(
            snapshot,
            false,
            true,
            TsSet::Empty,
            u64::MAX,
            None,
            Some(Key::from_raw(&i_to_key(2))),
            false,
            false,
            false,
        )
        .unwrap();
        let mut keys = Vec::new();
        while let Some((key, _)) = scanner.next().unwrap() {
            keys.push(key.into_raw().unwrap());
        }

        assert_eq!(keys, vec![i_to_key(1)]);
        let stats = scanner.take_statistics();
        assert_eq!(stats.write.next, 2);
        assert_eq!(stats.write.processed_keys, 1);
        assert_eq!(stats.write.next_tombstone, 0);
        assert_eq!(stats.processed_size, 913);
        assert_eq!(stats.total_size, 936);
    }

    // Lots of MVCC versions, should return key1 and key49, processed_keys should be
    // 2. total_size should be large than processed_size.
    {
        let snapshot = engine.get_snap_access(region_id).unwrap();
        let mut scanner = CloudStoreScanner::new(
            snapshot,
            false,
            true,
            TsSet::Empty,
            u64::MAX,
            None,
            None,
            false,
            false,
            false,
        )
        .unwrap();

        let mut keys = Vec::new();
        while let Some((key, _)) = scanner.next().unwrap() {
            keys.push(key.into_raw().unwrap());
        }

        assert_eq!(keys, vec![i_to_key(1), i_to_key(49)]);
        let stats = scanner.take_statistics();
        assert_eq!(stats.write.next, 49);
        assert_eq!(stats.write.processed_keys, 2);
        assert_eq!(stats.write.next_tombstone, 47);
        assert_eq!(stats.processed_size, 1826);
        assert_eq!(stats.total_size, 2682);
    }

    cluster.stop();
}

#[test]
fn test_snap_marshal() {
    let node_id = alloc_node_id();
    let mut cluster = ServerCluster::new(vec![node_id], |_, _| {});
    let mut client = cluster.new_client();
    let keyspace_id = ApiV2::get_u32_keyspace_id_by_key("x123".as_bytes()).unwrap_or_default();
    client.split_keyspace(keyspace_id);
    client.put_kv(0..100, i_to_key, i_to_val);
    client.put_kv(100..200, i_to_key, i_to_val);
    let region_id = client.get_region_id(&i_to_key(0));
    let kvengine = cluster.get_kvengine(node_id);
    let check_table_count = |start: Bytes, end: Bytes, check_l0: bool| -> bool {
        let snap = kvengine.get_snap_access(region_id).unwrap();
        let (_, cs_bin) = snap.marshal(&[(start, end)], false, false);
        let mut cs = ChangeSet::default();
        cs.merge_from_bytes(&cs_bin).unwrap();
        let cs_snap = cs.get_snapshot();
        if check_l0 {
            !cs_snap.get_l0_creates().is_empty()
        } else {
            !cs_snap.get_table_creates().is_empty()
        }
    };
    must_wait(
        || check_table_count(Bytes::from(i_to_key(0)), Bytes::from(i_to_key(500)), true),
        10,
        || "snapshot table count is zero".to_string(),
    );
    client.put_kv(200..300, i_to_key, i_to_val);
    client.put_kv(300..400, i_to_key, i_to_val);
    client.put_kv(400..500, i_to_key, i_to_val);
    must_wait(
        || check_table_count(Bytes::from("x123"), Bytes::from("x124"), false),
        10,
        || "snapshot table count is zero".to_string(),
    );
    cluster.stop();
}

#[test]
fn test_apply_observer() {
    let node_id = alloc_node_id();
    let cluster = ServerCluster::new(vec![node_id], |_, conf| {
        conf.rocksdb.writecf.write_buffer_size = ReadableSize::mb(1);
    });
    let mut client = cluster.new_client();
    let region_id = client.get_region_id(&i_to_key(0));
    let rfengine = cluster.get_rfengine(node_id);
    let store_id = rfengine.get_engine_id();
    let region_peer_map = rfengine.get_region_peer_map();
    let &peer_id = region_peer_map.get(&region_id).unwrap();
    let kvengine = cluster.get_kvengine(node_id);
    let shard_meta = rfstore::store::load_engine_meta(&rfengine, store_id, peer_id).unwrap();
    let shard_cs = shard_meta.to_change_set();
    let shard = Arc::new(Shard::new_for_ingest(
        kvengine.get_engine_id(),
        &shard_cs,
        kvengine.opts.clone(),
        &kvengine.get_master_key(),
        kvengine.get_encryption_key_manager(),
    ));
    for i in 0..10 {
        let start = i * 10;
        client.put_kv(start..start + 10, i_to_key, i_to_val);
    }
    let recoverer = rfstore::store::RecoverHandler::new(rfengine.clone());
    let mut apply_ctx = ApplyContext::new(kvengine.clone(), None, None, None);

    let (tx, rx) = tikv_util::mpsc::unbounded();
    let runtime = cluster.get_dfs().unwrap().get_runtime().handle().clone();
    let observer = CdcApplyObserver::new(kvengine.clone(), tx.clone(), runtime);
    apply_ctx.set_apply_observer(Box::new(observer));
    recoverer
        .recover_with_apply_ctx(&mut apply_ctx, &shard, &shard_meta, false)
        .unwrap();
    let mut observer = apply_ctx.take_apply_observer().unwrap();
    observer.flush();
    let msg = rx.recv().unwrap();
    match msg {
        replication_worker::CdcMsg::Applied { region_events, .. } => {
            let mut i = 0;
            for event in &region_events.events {
                for row in event.get_entries().get_entries() {
                    if row.get_type() == kvproto::cdcpb::EventLogType::Committed {
                        let trim_keyspace_key = i_to_key(i)[4..].to_vec();
                        assert_eq!(row.get_key(), &trim_keyspace_key);
                        assert_eq!(row.get_value(), i_to_val(i));
                        i += 1;
                    }
                }
            }
            assert_eq!(i, 100);
        }
        _ => panic!("unexpected message"),
    }
}

#[test]
fn test_cloud_store_reset_range() {
    let node_id = alloc_node_id();
    let cluster = ServerCluster::new(vec![node_id], |_, _| {});
    let mut client = cluster.new_client();

    let keyspace_id =
        api_version::ApiV2::get_u32_keyspace_id_by_key("x123".as_bytes()).unwrap_or_default();
    client.split_keyspace(keyspace_id);

    client.put_kv(3..10, i_to_key, i_to_val);
    client.put_kv(30..40, i_to_key, i_to_val);

    let region_id = client.get_region_id("x123".as_bytes());
    let engine = cluster.get_kvengine(node_id);
    let snap = engine.get_snap_access(region_id).unwrap();
    let read_ts = block_on(cluster.get_pd_client().get_tso())
        .unwrap()
        .into_inner();
    let mut scanner = new_scanner(&snap, false, read_ts);
    let cases = vec![
        ((9, 10), Some(9), true), // old range is not set should seek.
        ((10, 19), None, true),   // inner is smaller, should seek.
        ((19, 25), None, false),
        ((23, 27), None, true), // none monotonic reset range do not skip seek.
        ((28, 35), Some(30), false),
        ((37, 39), Some(37), true),
        ((40, 70), None, true),  // inner is smaller, should seek.
        ((70, 80), None, false), // inner is invalid, should not seek.
    ];
    for ((start, end), key, seek) in cases {
        check_scanner(&mut scanner, false, start, end, key, seek);
    }
    let mut desc_scanner = new_scanner(&snap, true, read_ts);
    let desc_cases = vec![
        ((40, 70), None, true),      // old range is not set should seek.
        ((37, 40), Some(39), false), // inner is equal to upper, skip seek.
        ((36, 38), Some(37), true),  // none monotonic reset range do not skip seek.
        ((23, 27), None, true),      // inner is greater should seek.
        ((15, 20), None, false),     // inner is at 9, skip seek.
        ((5, 9), Some(8), false),
        ((2, 3), None, true),  // inner is greater, should seek.
        ((0, 2), None, false), // inner is invalid, should not seek.
    ];
    for ((start, end), key, seek) in desc_cases {
        check_scanner(&mut desc_scanner, true, start, end, key, seek);
    }

    // test lock.
    let keys_nums = [21, 31, 51];
    let mutations: Vec<Mutation> = keys_nums
        .iter()
        .map(|&k_num| {
            let mut m = Mutation::default();
            m.set_key(i_to_key(k_num));
            m.set_value(i_to_val(k_num));
            m.set_op(Op::Put);
            m
        })
        .collect();
    let start_ts = block_on(cluster.get_pd_client().get_tso()).unwrap();
    let txn_muts = TxnMutations::from_normal(mutations);
    client
        .kv_prewrite(txn_muts.primary(), None, txn_muts, start_ts)
        .expect("kv_prewrite");

    let snap = engine.get_snap_access(region_id).unwrap();
    let (locked, seeked) = (true, true);
    let read_ts = block_on(cluster.get_pd_client().get_tso())
        .unwrap()
        .into_inner();
    let mut scanner = new_scanner(&snap, false, read_ts);
    let lock_cases = vec![
        ((0, 10), false, seeked),
        ((10, 20), false, false),
        ((20, 25), locked, false),
        ((25, 30), false, seeked),
        ((30, 40), locked, false),
        ((55, 60), false, seeked),
        ((60, 65), false, false),
        ((70, 85), false, false),
    ];
    for ((lower, upper), is_err, seeked) in lock_cases {
        reset_range(&mut scanner, false, lower, upper);
        if is_err {
            scanner.next().unwrap_err();
        } else {
            scanner.next().unwrap();
        }
        assert_eq!(scanner.take_statistics().lock.seek, seeked as usize);
    }

    let mut desc_scanner = new_scanner(&snap, true, read_ts);
    let desc_lock_cases = vec![
        ((70, 85), false, seeked),
        ((60, 65), false, false),
        ((55, 60), false, false),
        ((30, 40), locked, seeked), // locked(31)
        ((25, 30), false, seeked),
        ((20, 25), locked, false), // locked(21)
        ((10, 20), false, seeked),
        ((0, 10), false, false),
    ];
    for ((lower, upper), is_err, seeked) in desc_lock_cases {
        reset_range(&mut desc_scanner, true, lower, upper);
        if is_err {
            desc_scanner.next().unwrap_err();
        } else {
            desc_scanner.next().unwrap();
        }
        assert_eq!(desc_scanner.take_statistics().lock.seek, seeked as usize);
    }
}

fn new_scanner(snap: &SnapAccess, desc: bool, read_ts: u64) -> CloudStoreScanner {
    CloudStoreScanner::new(
        snap.clone(),
        desc,
        true,
        TsSet::Empty,
        read_ts,
        None,
        None,
        false,
        false,
        false,
    )
    .unwrap()
}

fn check_scanner(
    scanner: &mut CloudStoreScanner,
    desc: bool,
    start: usize,
    end: usize,
    key: Option<usize>,
    seek: bool,
) {
    reset_range(scanner, desc, start, end);
    if let Some(key) = key {
        assert_eq!(next_key(scanner), i_to_key(key));
    } else {
        assert!(scanner.next().unwrap().is_none());
    }
    let stats = scanner.take_statistics();
    if seek {
        assert_eq!(stats.write.seek, 1);
    } else {
        assert_eq!(stats.write.seek, 0);
        assert_eq!(stats.lock.seek, 0);
    }
}

fn reset_range(scanner: &mut CloudStoreScanner, desc: bool, lower: usize, upper: usize) {
    scanner
        .reset_range(
            desc,
            false,
            Some(Key::from_raw(&i_to_key(lower))),
            Some(Key::from_raw(&i_to_key(upper))),
        )
        .unwrap();
}

fn next_key(scanner: &mut CloudStoreScanner) -> Vec<u8> {
    scanner
        .next()
        .unwrap()
        .map(|(k, _)| k.into_raw().unwrap())
        .unwrap()
}

fn sleep() {
    std::thread::sleep(Duration::from_millis(100));
}