// Copyright 2024 PingCAP, Inc.

use std::{
    collections::HashMap,
    ops::Range,
    sync::{Arc, Barrier},
    thread,
    time::Duration,
};

use anyhow::{self, bail};
use bytes::Bytes;
use futures::executor::block_on;
use hyper::Method;
use kvengine::{
    dfs::{self, Dfs, FileType},
    table::{
        InnerKey, PessimisticAction, TXN_FILE_FORMAT_V2,
        txn_file::{OP_PUT, TxnChunkBuilder},
    },
};
use kvproto::{
    kvrpcpb,
    kvrpcpb::{
        BatchRollbackRequest, CheckTxnStatusRequest, CommitRequest, PrewriteRequest,
        PrewriteRequestPessimisticAction, ResolveLockRequest, TxnHeartBeatRequest,
    },
    tikvpb::TikvClient,
};
use matches::assert_matches;
use pd_client::PdClient;
use rstest::rstest;
use security::{RestfulClient, SecurityConfig, SecurityManager};
use test_cloud_server::{
    MajorCompactionTarget, ServerCluster, ServerClusterBuilder, TikvWorkerOptions, client,
    client::{
        ClusterClient, ClusterClientOptions, CommitAction, Error as ClientError, MutateOptions,
        PrewriteExt, RequestOptions, TxnMutations, TxnWriteMethod,
    },
    keyspace::CreateKeyspaceOptions,
    must_wait,
    oss::prepare_dfs,
    try_wait,
    txn::txn_file::TxnFileHelper,
    util::Mutation,
};
use test_pd_client::PdWrapper;
use tikv_util::{config::ReadableDuration, info, time::Instant, warn};

use crate::{
    TestEncryptionMode, alloc_node_id_vec, generate_keyspace_key, i_to_tidb_key, i_to_val,
    i_to_val_opt,
};

const NODES_COUNT: usize = 3;
const KEYSPACE_ID: u32 = 10;

fn update_active_resolve_settings(
    cluster: &ServerCluster,
    node_ids: &[u16],
    threshold: &str,
    maintenance_tick_interval: &str,
) {
    let mut change = HashMap::new();
    change.insert(
        "raftstore.txn-file-lock-active-resolve.threshold".to_owned(),
        threshold.to_owned(),
    );
    change.insert(
        "raftstore.maintenance-tick-interval".to_owned(),
        maintenance_tick_interval.to_owned(),
    );
    let body = serde_json::to_vec(&change).unwrap();
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();

    for &node_id in node_ids {
        let (status, body) = rt.block_on(async {
            let client = hyper::Client::new();
            let req = hyper::Request::builder()
                .method(hyper::Method::POST)
                .uri(format!("http://{}/config", cluster.status_addr(node_id)))
                .body(hyper::Body::from(body.clone()))
                .unwrap();
            let resp = client.request(req).await.unwrap();
            let status = resp.status();
            let body = String::from_utf8(
                hyper::body::to_bytes(resp.into_body())
                    .await
                    .unwrap()
                    .to_vec(),
            )
            .unwrap_or_default();
            (status, body)
        });
        assert!(
            status.is_success(),
            "update config on node {} failed: status {:?}, body {}",
            node_id,
            status,
            body
        );
    }

    thread::sleep(Duration::from_millis(200));
}

#[test]
fn test_txn_file_commands() {
    let mut cluster = ServerCluster::new(alloc_node_id_vec(3), |_, _| {});
    cluster.wait_region_replicated(&[], 3);
    let dfs = cluster.get_dfs().unwrap();

    let mut client = cluster.new_client();
    let mut prepare_ctx = |key: &[u8]| -> (u64, kvrpcpb::Context, TikvClient) {
        let region_id = client.get_region_id(key);
        let ctx = client.new_rpc_ctx(region_id, key).unwrap();
        let kv_client = client.get_kv_client(ctx.get_peer().get_store_id());
        (region_id, ctx, kv_client)
    };

    let mut client = cluster.new_client();
    client.split_keyspace(KEYSPACE_ID);

    let gen_key = generate_keyspace_key(KEYSPACE_ID);
    let primary_lock = gen_key(0);
    cluster.wait_region_replicated(&primary_lock, 3);
    let (mut region_id, mut ctx, mut kv_client) = prepare_ctx(&primary_lock);

    // test rollback lock.
    let start_ts = client.get_ts().into_inner();
    let chunk_ids = build_txn_files(&dfs, start_ts, 0, 300);

    let mut req = PrewriteRequest::new();
    req.set_context(ctx.clone());
    req.set_txn_file_chunks(chunk_ids);
    req.set_primary_lock(primary_lock.clone());
    req.set_lock_ttl(6000);
    req.set_start_version(start_ts);
    req.set_min_commit_ts(start_ts + 1);
    let ok = try_wait(
        || {
            let resp = kv_client.kv_prewrite(&req).unwrap();
            assert!(resp.get_errors().is_empty(), "resp {:?}", resp);
            if resp.has_region_error() {
                warn!("prewrite got region error {:?}", resp.get_region_error());
                (region_id, ctx, kv_client) = prepare_ctx(&primary_lock);
                req.set_context(ctx.clone());
                false
            } else {
                true
            }
        },
        10,
    );
    assert!(ok);

    let mut req = CheckTxnStatusRequest::new();
    req.set_context(ctx.clone());
    req.set_primary_key(primary_lock.clone());
    req.set_lock_ts(start_ts);
    req.set_caller_start_ts(start_ts + 5);
    req.set_is_txn_file(true);
    let resp = kv_client.kv_check_txn_status(&req).unwrap();
    assert!(
        !resp.has_error() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );
    let lock_info = resp.get_lock_info();
    // verify that min commit ts is pushed.
    assert_eq!(
        lock_info.min_commit_ts,
        start_ts + 6,
        "lock info {:?}",
        lock_info
    );

    let mut req = TxnHeartBeatRequest::new();
    req.set_context(ctx.clone());
    req.set_start_version(start_ts);
    req.set_primary_lock(primary_lock);
    req.set_advise_lock_ttl(10000);
    req.set_is_txn_file(true);
    let resp = kv_client.kv_txn_heart_beat(&req).unwrap();
    assert!(
        !resp.has_error() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );
    // verify that lock ttl is pushed.
    assert_eq!(resp.lock_ttl, 10000);

    let mut req = BatchRollbackRequest::new();
    req.set_context(ctx.clone());
    req.set_is_txn_file(true);
    req.set_start_version(start_ts);
    let resp = kv_client.kv_batch_rollback(&req).unwrap();
    assert!(
        !resp.has_error() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );

    // test rollback lock.
    let start_ts = client.get_ts().into_inner();
    let chunk_ids = build_txn_files(&dfs, start_ts, 0, 300);
    let primary_lock = gen_key(0);

    let mut req = PrewriteRequest::new();
    req.set_context(ctx.clone());
    req.set_txn_file_chunks(chunk_ids);
    req.set_primary_lock(primary_lock);
    req.set_lock_ttl(6000);
    req.set_start_version(start_ts);
    req.set_min_commit_ts(start_ts + 1);
    let resp = kv_client.kv_prewrite(&req).unwrap();
    assert!(
        resp.get_errors().is_empty() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );

    let mut req = BatchRollbackRequest::new();
    req.set_context(ctx.clone());
    req.set_is_txn_file(true);
    req.set_start_version(start_ts);
    let resp = kv_client.kv_batch_rollback(&req).unwrap();
    assert!(
        !resp.has_error() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );

    // test resolve lock.
    let start_ts = client.get_ts().into_inner();
    let chunk_ids = build_txn_files(&dfs, start_ts, 0, 300);
    let primary_lock = gen_key(0);

    let mut req = PrewriteRequest::new();
    req.set_context(ctx.clone());
    req.set_txn_file_chunks(chunk_ids);
    req.set_primary_lock(primary_lock);
    req.set_lock_ttl(6000);
    req.set_start_version(start_ts);
    req.set_min_commit_ts(start_ts + 1);
    let resp = kv_client.kv_prewrite(&req).unwrap();
    assert!(
        resp.get_errors().is_empty() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );

    let mut req = ResolveLockRequest::new();
    req.set_context(ctx.clone());
    req.set_start_version(start_ts);
    req.set_is_txn_file(true);
    let resp = kv_client.kv_resolve_lock(&req).unwrap();
    assert!(
        !resp.has_error() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );

    // test resolve lock (batch).
    let start_ts = client.get_ts().into_inner();
    let chunk_ids = build_txn_files(&dfs, start_ts, 0, 300);
    let primary_lock = gen_key(0);

    let mut req = PrewriteRequest::new();
    req.set_context(ctx.clone());
    req.set_txn_file_chunks(chunk_ids);
    req.set_primary_lock(primary_lock);
    req.set_lock_ttl(6000);
    req.set_start_version(start_ts);
    req.set_min_commit_ts(start_ts + 1);
    let resp = kv_client.kv_prewrite(&req).unwrap();
    assert!(
        resp.get_errors().is_empty() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );
    let txn_info = kvrpcpb::TxnInfo {
        txn: start_ts,
        status: 0,
        is_txn_file: true,
        ..Default::default()
    };
    let mut req = ResolveLockRequest::new();
    req.set_context(ctx.clone());
    req.set_txn_infos(vec![txn_info].into());
    let resp = kv_client.kv_resolve_lock(&req).unwrap();
    assert!(
        !resp.has_error() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );

    // test success 2pc.
    let start_ts = client.get_ts().into_inner();
    let chunk_ids = build_txn_files(&dfs, start_ts, 0, 300);
    let primary_lock = gen_key(0);

    let mut req = PrewriteRequest::new();
    req.set_context(ctx.clone());
    req.set_txn_file_chunks(chunk_ids);
    req.set_primary_lock(primary_lock);
    req.set_lock_ttl(6000);
    req.set_start_version(start_ts);
    req.set_min_commit_ts(start_ts + 1);
    let resp = kv_client.kv_prewrite(&req).unwrap();
    assert!(
        resp.get_errors().is_empty() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );

    let mut req = CommitRequest::new();
    req.set_context(ctx);
    req.set_start_version(start_ts);
    req.set_is_txn_file(true);
    let commit_ts = client.get_ts().into_inner();
    req.set_commit_version(commit_ts);
    let resp = kv_client.kv_commit(&req).unwrap();
    assert!(
        !resp.has_error() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );

    verify_range(&mut client, 0, 300, commit_ts);
    cluster.stop();
}

#[test]
fn test_txn_file_pessimistic_commit_cleans_in_memory_locks() {
    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test");
    let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None);
    let mut cluster = ServerClusterBuilder::new(alloc_node_id_vec(1), |_, conf| {
        conf.dfs = dfs_config.clone();
        conf.pessimistic_txn.pipelined = true;
        conf.pessimistic_txn.in_memory = true;
    })
    .pd(pd_wrapper)
    .build();
    cluster.start_tikv_workers(alloc_node_id_vec(1), TikvWorkerOptions::default());
    cluster.wait_region_replicated(&[], 1);

    let gen_key = generate_keyspace_key(KEYSPACE_ID);
    let key = Bytes::from(gen_key(0));
    let value = Bytes::from_static(b"v1");

    let mut client = cluster.new_client_opt(ClusterClientOptions {
        txn_file_max_chunk_size: Some(1024),
        ..Default::default()
    });
    client.split_keyspace(KEYSPACE_ID);

    let start_ts = client.get_ts();
    let for_update_ts = client.get_ts();

    let mut txn = client.begin_transaction(Some(start_ts));
    let (_, key_errors) = client
        .kv_pessimistic_lock(
            key.clone(),
            vec![key.clone()],
            start_ts,
            3000,
            for_update_ts,
            None,
        )
        .unwrap();
    assert!(key_errors.is_empty(), "key_errors {:?}", key_errors);

    let mutation = Mutation {
        op: kvrpcpb::Op::Put,
        key: key.clone(),
        value: value.clone(),
        ..Default::default()
    };
    let mut options = MutateOptions::pessimistic(start_ts, for_update_ts);
    options.write_method = TxnWriteMethod::FileBased;
    let commit_ts = client
        .prewrite_pessimistic_and_commit(
            vec![mutation],
            vec![PrewriteRequestPessimisticAction::DoPessimisticCheck],
            options,
            PrewriteExt::with_for_update_ts(for_update_ts),
            &mut txn,
        )
        .unwrap();
    assert!(commit_ts > start_ts);

    let new_start_ts = client.get_ts();
    let new_for_update_ts = client.get_ts();
    let (_, key_errors) = client
        .kv_pessimistic_lock(
            key.clone(),
            vec![key.clone()],
            new_start_ts,
            3000,
            new_for_update_ts,
            None,
        )
        .unwrap();
    assert!(
        key_errors.is_empty(),
        "stale in-memory pessimistic locks should be cleaned up: {:?}",
        key_errors
    );

    cluster.stop();
    oss.shutdown();
}

#[rstest]
#[case::file_based_large(
    KEYSPACE_ID + 1,
    10,
    TxnWriteMethod::FileBased,
    TestEncryptionMode::Disabled,
    false
)]
#[case::file_based_large_metadata(
    KEYSPACE_ID + 2,
    10,
    TxnWriteMethod::FileBased,
    TestEncryptionMode::Metadata,
    false
)]
#[case::file_based_large_legacy_pessimistic(
    KEYSPACE_ID + 3,
    10,
    TxnWriteMethod::FileBased,
    TestEncryptionMode::Legacy,
    true
)]
#[case::file_based_small(
    KEYSPACE_ID + 4,
    1,
    TxnWriteMethod::FileBased,
    TestEncryptionMode::Disabled,
    false
)]
fn test_txn_file_basic(
    #[case] keyspace_id: u32,
    #[case] size_factor: usize,
    #[case] write_method: TxnWriteMethod,
    #[case] encryption_mode: TestEncryptionMode,
    #[case] is_pessimistic: bool,
) {
    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test");

    let node_ids = alloc_node_id_vec(1);
    let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None);
    let cluster_id = pd_wrapper.client().get_cluster_id().unwrap();
    info!("test_txn_file_basic";
        "keyspace_id" => keyspace_id,
        "size_factor" => size_factor,
        "encryption_mode" => ?encryption_mode,
        "write_method" => ?write_method,
        "cluster_id" => cluster_id,
        "is_pessimistic" => is_pessimistic,
    );
    let mut cluster = ServerClusterBuilder::new(node_ids, |_, conf| {
        conf.dfs = dfs_config.clone();
    })
    .pd(pd_wrapper)
    .build();
    cluster.start_tikv_workers(alloc_node_id_vec(1), TikvWorkerOptions::default());
    cluster.wait_region_replicated(&[], 1);

    let gen_key = generate_keyspace_key(keyspace_id);

    let mut verify_data = {
        let rt = tokio::runtime::Runtime::new().unwrap();
        let mut client = cluster.new_client();
        let mut txn_client = rt.block_on(cluster.new_txn_client());

        move |range: Option<(&[u8], &[u8])>, expected: usize| -> anyhow::Result<()> {
            let guard = client.ref_store();
            let ref_store = guard.lock().unwrap();
            let cnt = rt
                .block_on(txn_client.verify_data_by_scan(&ref_store, range))
                .unwrap();
            if cnt != expected {
                bail!("txn client verify data failed, expect {expected}, got {cnt}");
            }
            let (cnt, _) = client
                .verify_data_with_given_ref_store(&ref_store, range, &RequestOptions::default())
                .unwrap();
            if cnt != expected {
                bail!("client verify data failed, expect {expected}, got {cnt}");
            }
            Ok(())
        }
    };

    let rt = tokio::runtime::Runtime::new().unwrap();
    let _enter = rt.enter();

    if encryption_mode.is_enabled() {
        cluster
            .get_pd_client()
            .set_keyspace_encryption(keyspace_id, encryption_mode.keyspace_encryption_config())
            .unwrap();
    }

    let mut client = cluster.new_client_opt(ClusterClientOptions {
        txn_file_max_chunk_size: Some(1024),
        ..Default::default()
    });

    client.split_keyspace(keyspace_id);

    {
        // To work around that data not existed in ref store will not be checked.
        let guard = client.ref_store();
        let mut ref_store = guard.lock().unwrap();
        for k in 0..20 * size_factor {
            ref_store.del_kv(gen_key(k));
        }
    }

    {
        client
            .try_put_kv_with_txn(
                0..10 * size_factor,
                &gen_key,
                i_to_val_opt("value0_", 3),
                MutateOptions {
                    commit_action: CommitAction::AsyncCommitSecondaryKeys(Duration::ZERO),
                    write_method,
                    ..Default::default()
                },
                is_pessimistic,
            )
            .unwrap();
        verify_data(None, 10 * size_factor).unwrap();
    }

    {
        // Verify locks of committed txn.
        client
            .try_put_kv_with_txn(
                size_factor..11 * size_factor,
                &gen_key,
                i_to_val_opt("value1_", 3),
                MutateOptions {
                    commit_action: CommitAction::AsyncCommitSecondaryKeys(Duration::MAX),
                    write_method,
                    ..Default::default()
                },
                is_pessimistic,
            )
            .unwrap();
        // Verify on non-primary txn file when there are more than one.
        verify_data(
            Some((&gen_key(6 * size_factor), &gen_key(11 * size_factor))),
            5 * size_factor,
        )
        .unwrap();
    }

    {
        // Verify rollback & resolve locks.
        // Pessimistic prewrite has slightly different TTL bookkeeping, so keep
        // enough margin for the scan to observe and resolve its expired lock.
        let lock_ttl = if is_pessimistic {
            Duration::from_secs(1)
        } else {
            Duration::from_millis(300)
        };
        client
            .try_put_kv_with_txn(
                2 * size_factor..12 * size_factor,
                &gen_key,
                i_to_val_opt("value2_", 3),
                MutateOptions {
                    commit_action: CommitAction::NoCommit,
                    write_method,
                    lock_ttl,
                    ..Default::default()
                },
                is_pessimistic,
            )
            .unwrap();
        let start = Instant::now_coarse();
        // Verify on non-primary txn file when there are more than one.
        verify_data(
            Some((&gen_key(7 * size_factor), &gen_key(12 * size_factor))),
            4 * size_factor,
        )
        .unwrap();
        // Wait until locks expired.
        // Note that txn client will not wait lock timeout during scan because of min
        // commit ts pushed.
        let elapsed = start.saturating_elapsed();
        let minimum_wait = if is_pessimistic {
            Duration::from_millis(500)
        } else {
            Duration::from_millis(200)
        };
        assert!(elapsed > minimum_wait, "elapsed {:?}", elapsed);
    }

    verify_data(None, 11 * size_factor).unwrap();
    cluster.stop();
    oss.shutdown();
}

#[test]
fn test_txn_file_split_merge() {
    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test");

    let node_ids = alloc_node_id_vec(NODES_COUNT);
    let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None);
    let mut cluster = ServerClusterBuilder::new(node_ids, |_, conf| {
        conf.dfs = dfs_config.clone();
    })
    .pd(pd_wrapper)
    .build();
    cluster.start_tikv_workers(alloc_node_id_vec(1), TikvWorkerOptions::default());
    cluster.wait_region_replicated(&[], NODES_COUNT);

    let gen_key = generate_keyspace_key(KEYSPACE_ID);

    let rt = tokio::runtime::Runtime::new().unwrap();
    let _enter = rt.enter();

    let mut client = cluster.new_client_opt(ClusterClientOptions {
        txn_file_max_chunk_size: Some(1024),
        ..Default::default()
    });
    client.split_keyspace(KEYSPACE_ID);
    client.try_split(&gen_key(100), 5).unwrap();
    {
        // To work around that data not existed in ref store will not be checked.
        let guard = client.ref_store();
        let mut ref_store = guard.lock().unwrap();
        for k in 0..200 {
            ref_store.del_kv(gen_key(k));
        }
    }

    // Put the locks:
    client
        .try_put_kv(
            100..200,
            &gen_key,
            i_to_val_opt("value0_", 3),
            MutateOptions {
                commit_action: CommitAction::NoCommit,
                write_method: TxnWriteMethod::FileBased,
                ..Default::default()
            },
        )
        .unwrap();
    client
        .try_put_kv(
            0..100,
            &gen_key,
            i_to_val_opt("value1_", 3),
            MutateOptions {
                commit_action: CommitAction::NoCommit,
                write_method: TxnWriteMethod::FileBased,
                ..Default::default()
            },
        )
        .unwrap();

    // Fail to split:
    let err = client.try_split(&gen_key(20), 1).unwrap_err();
    match err {
        client::Error::KeyErrors(key_errs) => {
            assert_eq!(key_errs.len(), 1);
            let key_err = &key_errs[0];
            assert!(key_err.has_locked());
            let lock_info = key_err.get_locked();
            assert_eq!(lock_info.key, gen_key(20));
            assert_eq!(lock_info.primary_lock, gen_key(0));
        }
        _ => {
            panic!("unexpected error: {:?}", err);
        }
    }
    // Another split path. Note this split request would succeed after locks are
    // cleanup.
    client.try_split_by_pd(&gen_key(50), 1).unwrap_err();

    // Clean up the lock:
    info!("cleanup the locks (0..100");
    client
        .try_put_kv(
            0..100,
            &gen_key,
            i_to_val_opt("value2_", 3),
            MutateOptions {
                commit_action: CommitAction::SyncCommit,
                write_method: TxnWriteMethod::FileBased,
                ..Default::default()
            },
        )
        .unwrap();
    // Can split:
    let ok = try_wait(
        || {
            info!("split on {:?}", gen_key(20));
            match client.try_split(&gen_key(20), 5) {
                Ok(_) => true,
                Err(client::Error::KeyErrors(key_errs)) => {
                    info!("try split encounter locks, try again"; "key_errs" => ?key_errs);
                    false
                }
                Err(err) => panic!("split failed: {:?}", err),
            }
        },
        5,
    );
    assert!(ok, "split key 20 failed");
    client
        .try_split(&gen_key(80), 5)
        .expect("split key 80 failed");

    // Clean up the lock:
    info!("cleanup the locks (100..200");
    // Note: verify data will not clean up locks when locks are not expired.
    client.verify_data_with_ref_store();
    // Can split:
    let ok = try_wait(
        || {
            info!("split on {:?}", gen_key(150));
            match client.try_split(&gen_key(150), 5) {
                Ok(_) => true,
                Err(client::Error::KeyErrors(key_errs)) => {
                    info!("try split encounter locks, try again"; "key_errs" => ?key_errs);
                    false
                }
                Err(err) => panic!("split failed: {:?}", err),
            }
        },
        5,
    );
    assert!(ok, "split key 150 failed");

    // Current regions in [, 100):
    // [, 20), [20, 80), [80, 100), or [,20), [20,50), [50,80), [80, 100)

    // Write the lock to the middle region(s):
    client
        .try_put_kv(
            40..60,
            &gen_key,
            i_to_val_opt("value3_", 3),
            MutateOptions {
                commit_action: CommitAction::NoCommit,
                write_method: TxnWriteMethod::FileBased,
                ..Default::default()
            },
        )
        .unwrap();

    // Fail to merge source region has locks.
    let ok = client.try_merge_and_wait(&gen_key(50), &gen_key(100), 1);
    assert!(!ok, "source region having locks merged");

    // Succeed to merge when target region has locks
    // [, 80), [80, 100), or [,50), [50,80), [80, 100)
    let ok = client.try_merge_and_wait(&gen_key(0), &gen_key(45), 5);
    assert!(ok, "source region having no lock not merged");

    // Clean up locks:
    client
        .try_put_kv(
            40..60,
            &gen_key,
            i_to_val_opt("value4_", 3),
            MutateOptions {
                commit_action: CommitAction::SyncCommit,
                write_method: TxnWriteMethod::FileBased,
                ..Default::default()
            },
        )
        .unwrap();

    // Succeed to merge:
    let ok = client.try_merge_and_wait(&gen_key(50), &gen_key(80), 5);
    assert!(ok, "region having no lock not merged");

    // Verify:
    client.verify_data_with_ref_store();
    cluster.stop();
    oss.shutdown();
}

#[rstest]
#[case::tikv_rpc_split(ActiveResolveSplitMethod::TikvRpc)]
#[case::pd_split(ActiveResolveSplitMethod::Pd)]
fn test_txn_file_split_active_resolve_old_lock(#[case] split_method: ActiveResolveSplitMethod) {
    test_util::init_log_for_test();
    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test");

    let node_ids = alloc_node_id_vec(1);
    let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None);
    let mut cluster = ServerClusterBuilder::new(node_ids, |_, conf| {
        conf.dfs = dfs_config.clone();
        conf.raft_store.txn_file_lock_active_resolve.threshold = ReadableDuration::secs(0);
    })
    .pd(pd_wrapper)
    .build();
    cluster.wait_region_replicated(&[], 1);
    let dfs = cluster.get_dfs().unwrap();

    let gen_key = generate_keyspace_key(KEYSPACE_ID);
    let mut client = cluster.new_client_opt(ClusterClientOptions {
        txn_file_max_chunk_size: Some(1024),
        ..Default::default()
    });
    client.split_keyspace(KEYSPACE_ID);
    let primary_lock = gen_key(0);
    cluster.wait_region_replicated(&primary_lock, 1);

    let prepare_ctx = |client: &mut ClusterClient, key: &[u8]| -> (kvrpcpb::Context, TikvClient) {
        let region_id = client.get_region_id(key);
        let ctx = client.new_rpc_ctx(region_id, key).unwrap();
        let kv_client = client.get_kv_client(ctx.get_peer().get_store_id());
        (ctx, kv_client)
    };
    let (mut ctx, mut kv_client) = prepare_ctx(&mut client, &primary_lock);
    let start_ts = client.get_ts().into_inner();
    let chunk_ids = build_txn_files_with_key_fn(&dfs, start_ts, 0, 100, |i| gen_key(i));

    let mut req = PrewriteRequest::new();
    req.set_context(ctx.clone());
    req.set_txn_file_chunks(chunk_ids);
    req.set_primary_lock(primary_lock);
    req.set_lock_ttl(200);
    req.set_start_version(start_ts);
    req.set_min_commit_ts(start_ts + 1);
    let ok = try_wait(
        || {
            let resp = kv_client.kv_prewrite(&req).unwrap();
            assert!(resp.get_errors().is_empty(), "resp {:?}", resp);
            if resp.has_region_error() {
                warn!("prewrite got region error {:?}", resp.get_region_error());
                let (new_ctx, new_kv_client) = prepare_ctx(&mut client, &gen_key(0));
                ctx = new_ctx;
                kv_client = new_kv_client;
                req.set_context(ctx.clone());
                false
            } else {
                true
            }
        },
        10,
    );
    assert!(ok);
    thread::sleep(Duration::from_millis(500));

    let split_key = gen_key(20);
    let initial_region_count = assert_txn_file_lock_blocks_split(
        split_method,
        &mut client,
        &split_key,
        "active resolve disabled",
    );

    update_active_resolve_settings(&cluster, &cluster.get_nodes(), "100ms", "100ms");
    assert_txn_file_lock_allows_split(
        split_method,
        &cluster,
        &mut client,
        &split_key,
        initial_region_count,
        "active resolve",
    );

    cluster.stop();
    oss.shutdown();
}

#[test]
fn test_txn_file_split_active_resolve_uses_heartbeat_updated_ttl() {
    test_util::init_log_for_test();
    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test");

    let node_ids = alloc_node_id_vec(NODES_COUNT);
    let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None);
    let mut cluster = ServerClusterBuilder::new(node_ids, |_, conf| {
        conf.dfs = dfs_config.clone();
        conf.raft_store.txn_file_lock_active_resolve.threshold = ReadableDuration::millis(100);
    })
    .pd(pd_wrapper)
    .build();
    cluster.wait_region_replicated(&[], 3);
    let dfs = cluster.get_dfs().unwrap();

    let gen_key = generate_keyspace_key(KEYSPACE_ID);
    let mut client = cluster.new_client_opt(ClusterClientOptions {
        txn_file_max_chunk_size: Some(1024),
        ..Default::default()
    });
    client.split_keyspace(KEYSPACE_ID);
    let primary_lock = gen_key(0);
    cluster.wait_region_replicated(&primary_lock, 3);

    let prepare_ctx = |client: &mut ClusterClient, key: &[u8]| -> (kvrpcpb::Context, TikvClient) {
        let region_id = client.get_region_id(key);
        let ctx = client.new_rpc_ctx(region_id, key).unwrap();
        let kv_client = client.get_kv_client(ctx.get_peer().get_store_id());
        (ctx, kv_client)
    };
    let (mut ctx, mut kv_client) = prepare_ctx(&mut client, &primary_lock);
    let start_ts = client.get_ts().into_inner();
    let chunk_ids = build_txn_files_with_key_fn(&dfs, start_ts, 0, 100, |i| gen_key(i));

    let mut req = PrewriteRequest::new();
    req.set_context(ctx.clone());
    req.set_txn_file_chunks(chunk_ids);
    req.set_primary_lock(primary_lock.clone());
    req.set_lock_ttl(200);
    req.set_start_version(start_ts);
    req.set_min_commit_ts(start_ts + 1);
    let ok = try_wait(
        || {
            let resp = kv_client.kv_prewrite(&req).unwrap();
            assert!(resp.get_errors().is_empty(), "resp {:?}", resp);
            if resp.has_region_error() {
                warn!("prewrite got region error {:?}", resp.get_region_error());
                let (new_ctx, new_kv_client) = prepare_ctx(&mut client, &primary_lock);
                ctx = new_ctx;
                kv_client = new_kv_client;
                req.set_context(ctx.clone());
                false
            } else {
                true
            }
        },
        10,
    );
    assert!(ok);

    let mut heartbeat_req = TxnHeartBeatRequest::new();
    heartbeat_req.set_context(ctx.clone());
    heartbeat_req.set_start_version(start_ts);
    heartbeat_req.set_primary_lock(primary_lock.clone());
    heartbeat_req.set_advise_lock_ttl(20_000);
    heartbeat_req.set_is_txn_file(true);
    let heartbeat_resp = kv_client.kv_txn_heart_beat(&heartbeat_req).unwrap();
    assert!(
        !heartbeat_resp.has_error() && !heartbeat_resp.has_region_error(),
        "resp {:?}",
        heartbeat_resp
    );
    assert_eq!(heartbeat_resp.lock_ttl, 20_000);

    // Wait until the original prewrite TTL would have expired. If active
    // resolve still reconstructs the old lock prefix from ShardMeta, it may
    // roll back this live txn and allow the split to proceed.
    thread::sleep(Duration::from_millis(500));

    let split_key = gen_key(20);
    let initial_region_count = assert_txn_file_lock_blocks_split(
        ActiveResolveSplitMethod::TikvRpc,
        &mut client,
        &split_key,
        "active resolve ignored the heartbeat-refreshed TTL",
    );

    let (ctx, kv_client) = prepare_ctx(&mut client, &primary_lock);
    let mut commit_req = CommitRequest::new();
    commit_req.set_context(ctx);
    commit_req.set_start_version(start_ts);
    commit_req.set_is_txn_file(true);
    let commit_ts = client.get_ts().into_inner();
    commit_req.set_commit_version(commit_ts);
    let commit_resp = kv_client.kv_commit(&commit_req).unwrap();
    assert!(
        !commit_resp.has_error() && !commit_resp.has_region_error(),
        "resp {:?}",
        commit_resp
    );

    assert_txn_file_lock_allows_split(
        ActiveResolveSplitMethod::TikvRpc,
        &cluster,
        &mut client,
        &split_key,
        initial_region_count,
        "committing live lock",
    );

    cluster.stop();
    oss.shutdown();
}

// Tests for abnormal processes.
#[rstest]
#[case::txn_file_large(30, true)]
#[case::normal(10, false)]
fn test_txn_file_abnormal(#[case] data_count: usize, #[case] use_txn_file: bool) {
    let write_method = if use_txn_file {
        TxnWriteMethod::FileBased
    } else {
        TxnWriteMethod::Normal
    };
    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test");

    let node_ids = alloc_node_id_vec(NODES_COUNT);
    let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None);
    let cluster_id = pd_wrapper.client().get_cluster_id().unwrap();
    info!("test_txn_file_abnormal_process"; "data_count" => data_count, "write_method" => ?write_method, "cluster_id" => cluster_id);
    let mut cluster = ServerClusterBuilder::new(node_ids, |_, conf| {
        conf.dfs = dfs_config.clone();
    })
    .pd(pd_wrapper)
    .build();
    cluster.start_tikv_workers(alloc_node_id_vec(1), TikvWorkerOptions::default());
    cluster.wait_region_replicated(&[], NODES_COUNT);

    let gen_key = generate_keyspace_key(KEYSPACE_ID);

    let rt = tokio::runtime::Builder::new_multi_thread()
        .thread_name("test-abnormal-process")
        .worker_threads(4)
        .enable_all()
        .build()
        .unwrap();
    let _enter = rt.enter();

    let mut client = cluster.new_client_opt(ClusterClientOptions {
        txn_file_max_chunk_size: Some(1024),
        ..Default::default()
    });
    client.split_keyspace(KEYSPACE_ID);
    client.try_split(&gen_key(data_count / 2), 5).unwrap();
    {
        // To work around that data not existed in ref store will not be checked.
        let guard = client.ref_store();
        let mut ref_store = guard.lock().unwrap();
        for k in 0..data_count {
            ref_store.del_kv(gen_key(k));
        }
    }

    let txn_file_helper = client.txn_file_helper();
    let make_mutations_impl = |value, insert: bool| {
        let op_fn = if insert {
            |_: usize| kvrpcpb::Op::Insert
        } else {
            |_: usize| kvrpcpb::Op::Put
        };
        make_txn_mutations(
            value,
            0,
            data_count,
            &gen_key,
            op_fn,
            write_method,
            txn_file_helper.clone(),
        )
    };
    let make_mutations = |value| make_mutations_impl(value, false);
    let make_insert_mutations = |value| make_mutations_impl(value, true);

    // Already exist.
    {
        let (muts0, txn_muts0) = make_insert_mutations("already_exist");
        let start_ts0 = client.get_ts();
        client
            .kv_prewrite(txn_muts0.primary(), None, txn_muts0.clone(), start_ts0)
            .unwrap();
        let commit_ts0 = client.get_ts();
        client
            .kv_commit(txn_muts0.clone(), start_ts0, commit_ts0, false)
            .unwrap();
        client.put_kv_in_ref_store(muts0);
        client.verify_data_with_ref_store();

        let (_, txn_muts1) = make_insert_mutations("already_exist_1");
        let start_ts1 = client.get_ts();
        let err = client
            .kv_prewrite(txn_muts1.primary(), None, txn_muts1, start_ts1)
            .unwrap_err();
        assert_matches!(err, ClientError::AlreadyExist(_));
        client.verify_data_with_ref_store();
    }

    // Commit without prewrite:
    {
        let (_, txn_muts) = make_mutations("with_prewrite");
        let start_ts = client.get_ts();
        let commit_ts = client.get_ts();
        let err = client
            .kv_commit(txn_muts, start_ts, commit_ts, false)
            .unwrap_err();
        expect_err_msg(&err, "TxnLockNotFound");
        client.verify_data_with_ref_store();
    }

    // Commit after rollback.
    {
        let (_, txn_muts0) = make_mutations("after_rollback");
        let start_ts0 = client.get_ts();
        client
            .kv_prewrite(txn_muts0.primary(), None, txn_muts0.clone(), start_ts0)
            .unwrap();

        let (muts1, txn_muts1) = make_mutations("after_rollback_1");
        let start_ts1 = client.get_ts();
        // Will rollback txn0.
        client
            .kv_prewrite(txn_muts1.primary(), None, txn_muts1.clone(), start_ts1)
            .unwrap();

        let commit_ts0 = client.get_ts();
        let err = client
            .kv_commit(txn_muts0, start_ts0, commit_ts0, false)
            .unwrap_err();
        expect_err_msg(&err, "TxnLockNotFound");

        let commit_ts1 = client.get_ts();
        client
            .kv_commit(txn_muts1, start_ts1, commit_ts1, false)
            .unwrap();
        client.put_kv_in_ref_store(muts1);
        client.verify_data_with_ref_store();
    }

    // Write conflict & duplicated commit.
    {
        let (muts0, txn_muts0) = make_mutations("write_conflict");
        let start_ts0 = client.get_ts();
        client
            .kv_prewrite(txn_muts0.primary(), None, txn_muts0.clone(), start_ts0)
            .unwrap();

        let commit_ts0 = client.get_ts();
        let txn_muts0_clone = txn_muts0.clone();
        let mut client_clone = client.clone();
        let txn0 = rt.spawn_blocking(move || {
            thread::sleep(Duration::from_secs(1));
            let new_commit_ts = client_clone
                .kv_commit(txn_muts0_clone.clone(), start_ts0, commit_ts0, false)
                .unwrap();
            assert!(
                new_commit_ts > commit_ts0,
                "commit_ts0: {}, new_commit_ts: {}",
                commit_ts0,
                new_commit_ts
            );
            new_commit_ts
        });

        let (_, txn_muts1) = make_mutations("write_conflict_1");
        let start_ts1 = client.get_ts();
        let err = client
            .kv_prewrite(txn_muts1.primary(), None, txn_muts1, start_ts1)
            .unwrap_err();
        assert_matches!(err, ClientError::WriteConflict(_));

        let _new_commit_ts0 = rt.block_on(txn0).unwrap();
        client.put_kv_in_ref_store(muts0);
        client.verify_data_with_ref_store();

        // Duplicated commit.
        // Note: commit for secondary keys would fail for TxnLockNotFound error, when
        // there is no key in the region.
        let dup_commit_ts = client
            .kv_commit_ext(
                txn_muts0,
                start_ts0,
                commit_ts0,
                CommitAction::AsyncCommitSecondaryKeys(Duration::ZERO),
            )
            .unwrap();
        // Will return the commit_ts in request other than of committed record.
        assert_eq!(dup_commit_ts, commit_ts0);
        client.verify_data_with_ref_store();
    }

    // Rollback
    {
        let (muts, txn_muts) = make_mutations("rollback");

        // Rollback a not existed txn.
        let mut start_ts = client.get_ts();
        client.kv_rollback(txn_muts.clone(), start_ts).unwrap();
        if !use_txn_file {
            // Normal txn will write rollback record even the txn is not exist.
            // But file based txn can not do so.
            let err = client
                .kv_prewrite(txn_muts.primary(), None, txn_muts.clone(), start_ts)
                .unwrap_err();
            assert_matches!(err, ClientError::WriteConflict(_));
            start_ts = client.get_ts();
        }
        client
            .kv_prewrite(txn_muts.primary(), None, txn_muts.clone(), start_ts)
            .unwrap();
        client.kv_rollback(txn_muts.clone(), start_ts).unwrap();

        // Duplicated rollback.
        client.kv_rollback(txn_muts.clone(), start_ts).unwrap();

        // Prewrite after rollback.
        let err = client
            .kv_prewrite(txn_muts.primary(), None, txn_muts.clone(), start_ts)
            .unwrap_err();
        assert_matches!(err, ClientError::WriteConflict(_));

        // Commit after rollback.
        let commit_ts = client.get_ts();
        let err = client
            .kv_commit(txn_muts.clone(), start_ts, commit_ts, false)
            .unwrap_err();
        expect_err_msg(&err, "TxnLockNotFound");

        client.verify_data_with_ref_store();

        // Committed.
        let start_ts = client.get_ts();
        let commit_ts = client.get_ts();
        client
            .kv_prewrite(txn_muts.primary(), None, txn_muts.clone(), start_ts)
            .unwrap();
        client
            .kv_commit(txn_muts.clone(), start_ts, commit_ts, false)
            .unwrap();
        // Rollback after committed.
        let err = client.kv_rollback(txn_muts, start_ts).unwrap_err();
        expect_err_msg(&err, "Committed");

        client.put_kv_in_ref_store(muts);
        client.verify_data_with_ref_store();
    }

    // Resolve lock
    {
        let (muts, txn_muts) = make_mutations("resolve_lock");

        // Resolve a not existed txn.
        let mut start_ts = client.get_ts();
        client
            .kv_resolve_lock_for_mutations(start_ts.into_inner(), None, muts.clone(), use_txn_file)
            .unwrap();
        if !use_txn_file {
            // Normal txn will write rollback record even the txn is not exist.
            // But file based txn can not do so.
            let err = client
                .kv_prewrite(txn_muts.primary(), None, txn_muts.clone(), start_ts)
                .unwrap_err();
            assert_matches!(err, ClientError::WriteConflict(_));
            start_ts = client.get_ts();
        }
        client
            .kv_prewrite(txn_muts.primary(), None, txn_muts.clone(), start_ts)
            .unwrap();
        // Rollback the not committed txn.
        client
            .kv_resolve_lock_for_mutations(start_ts.into_inner(), None, muts.clone(), use_txn_file)
            .unwrap();

        // Duplicated resolve lock.
        client
            .kv_resolve_lock_for_mutations(start_ts.into_inner(), None, muts.clone(), use_txn_file)
            .unwrap();

        // Prewrite after rollback.
        let err = client
            .kv_prewrite(txn_muts.primary(), None, txn_muts.clone(), start_ts)
            .unwrap_err();
        assert_matches!(err, ClientError::WriteConflict(_));

        // Commit after rollback.
        let commit_ts = client.get_ts();
        let err = client
            .kv_commit(txn_muts.clone(), start_ts, commit_ts, false)
            .unwrap_err();
        expect_err_msg(&err, "TxnLockNotFound");

        client.verify_data_with_ref_store();

        // Committed.
        let start_ts = client.get_ts();
        let commit_ts = client.get_ts();
        client
            .kv_prewrite(txn_muts.primary(), None, txn_muts.clone(), start_ts)
            .unwrap();
        client
            .kv_commit(txn_muts, start_ts, commit_ts, false)
            .unwrap();
        // Rollback after committed.
        let err = client
            .kv_resolve_lock_for_mutations(start_ts.into_inner(), None, muts.clone(), use_txn_file)
            .unwrap_err();
        expect_err_msg(&err, "Committed");

        // Duplicated commit.
        client
            .kv_resolve_lock_for_mutations(
                start_ts.into_inner(),
                Some(commit_ts.into_inner()),
                muts.clone(),
                use_txn_file,
            )
            .unwrap();
        let commit_ts = client.get_ts();
        client
            .kv_resolve_lock_for_mutations(
                start_ts.into_inner(),
                Some(commit_ts.into_inner()),
                muts.clone(),
                use_txn_file,
            )
            .unwrap();

        client.put_kv_in_ref_store(muts);
        client.verify_data_with_ref_store();
    }

    // Check txn status
    {
        let (muts, txn_muts) = make_mutations("check_txn_status");
        let start_ts = client.get_ts();
        client
            .kv_prewrite(txn_muts.primary(), None, txn_muts.clone(), start_ts)
            .unwrap();
        {
            // primary mismatch.
            let second_key = muts[1].key.clone();
            let resp = client
                .kv_check_txn_status(
                    &second_key,
                    start_ts,
                    start_ts,
                    start_ts,
                    false,
                    false,
                    false,
                    use_txn_file,
                )
                .unwrap();
            assert!(resp.has_error());
            assert!(resp.get_error().has_primary_mismatch());
        }

        {
            let resp = client
                .kv_check_txn_status(
                    &txn_muts.primary(),
                    start_ts,
                    start_ts,
                    start_ts,
                    false,
                    false,
                    false,
                    use_txn_file,
                )
                .unwrap();
            assert!(!resp.has_error());
            assert_eq!(resp.commit_version, 0);
            assert_eq!(resp.get_lock_info().lock_version, start_ts.into_inner());
            assert_eq!(
                resp.get_lock_info().get_primary_lock(),
                txn_muts.primary().as_ref()
            );
        }
        // The status checks above intentionally leave the transaction active.
        // Clean it up explicitly instead of making the final verification wait
        // for the lock TTL.
        client.kv_rollback(txn_muts, start_ts).unwrap();
    }

    client.verify_data_with_ref_store();
    cluster.stop();
    oss.shutdown();
}

// Tests for abnormal processes.
#[rstest]
#[case::txn_file_normal(4, TxnWriteMethod::Normal)]
#[case::txn_file_pessimistic(4, TxnWriteMethod::FileBased)]
fn test_txn_file_abnormal_pessimistic(
    #[case] data_count: usize,
    #[case] write_method: TxnWriteMethod,
) {
    let is_txn_file = write_method == TxnWriteMethod::FileBased;
    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test");

    let node_ids = alloc_node_id_vec(1);
    let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None);
    let cluster_id = pd_wrapper.client().get_cluster_id().unwrap();
    info!("test_txn_file_abnormal_pessimistic"; "data_count" => data_count, "cluster_id" => cluster_id);
    let mut cluster = ServerClusterBuilder::new(node_ids, |_, conf| {
        conf.dfs = dfs_config.clone();
    })
    .pd(pd_wrapper)
    .build();
    cluster.start_tikv_workers(alloc_node_id_vec(1), TikvWorkerOptions::default());
    cluster.wait_region_replicated(&[], 1);

    let gen_key = generate_keyspace_key(KEYSPACE_ID);

    let rt = tokio::runtime::Builder::new_multi_thread()
        .thread_name("test-abnormal-pessimistic-process")
        .worker_threads(4)
        .enable_all()
        .build()
        .unwrap();
    let _enter = rt.enter();

    let mut client = cluster.new_client_opt(ClusterClientOptions {
        txn_file_max_chunk_size: Some(1024),
        ..Default::default()
    });
    client.split_keyspace(KEYSPACE_ID);
    client.try_split(&gen_key(data_count / 2), 5).unwrap();
    {
        // To work around that data not existed in ref store will not be checked.
        let guard = client.ref_store();
        let mut ref_store = guard.lock().unwrap();
        for k in 0..data_count {
            ref_store.del_kv(gen_key(k));
        }
    }

    let txn_file_helper = client.txn_file_helper();
    let make_mutations_impl = |value, insert: bool| {
        let op_fn = if insert {
            |_: usize| kvrpcpb::Op::Insert
        } else {
            |_: usize| kvrpcpb::Op::Put
        };
        make_pessimistic_txn_mutations(
            value,
            0,
            data_count,
            &gen_key,
            op_fn,
            write_method,
            txn_file_helper.clone(),
        )
    };
    let make_mutations = |value| make_mutations_impl(value, false);
    let make_insert_mutations = |value| make_mutations_impl(value, true);
    let lock_keys = |muts: &[Mutation]| {
        muts.iter()
            .filter_map(|m| {
                if m.op == kvrpcpb::Op::Put {
                    Some(m.key.clone())
                } else {
                    None
                }
            })
            .collect::<Vec<_>>()
    };
    let lock_ttl = Duration::from_secs(20);
    let pessimistic_lock = |client: &mut ClusterClient,
                            primary: Bytes,
                            keys: Vec<Bytes>,
                            start_ts,
                            for_update_ts|
     -> client::Result<Vec<kvrpcpb::KeyError>> {
        let mut key_errors = vec![];
        for key in keys {
            let (_results, errs) = client.kv_pessimistic_lock(
                primary.clone(),
                vec![key],
                start_ts,
                lock_ttl.as_millis() as u64,
                for_update_ts,
                None,
            )?;
            key_errors.extend(errs);
        }

        Ok(key_errors)
    };
    let pessimistic_prewrite =
        |client: &mut ClusterClient, muts: TxnMutations, start_ts, for_update_ts| {
            let mut txn = client.begin_transaction(Some(start_ts));
            let mut ext = PrewriteExt::with_for_update_ts(for_update_ts);
            ext.lock_ttl = lock_ttl;
            client.kv_prewrite_ext(muts.primary(), None, muts, &mut txn, ext)
        };

    // Already exist.
    {
        let (muts0, txn_muts0) = make_insert_mutations("already_exist");
        let start_ts0 = client.get_ts();
        let for_update_ts0 = client.get_ts();
        let key_errors = pessimistic_lock(
            &mut client,
            txn_muts0.primary(),
            lock_keys(&muts0),
            start_ts0,
            for_update_ts0,
        )
        .unwrap();
        assert!(key_errors.is_empty(), "key_errors {:?}", key_errors);
        pessimistic_prewrite(&mut client, txn_muts0.clone(), start_ts0, for_update_ts0).unwrap();
        let commit_ts0 = client.get_ts();
        client
            .kv_commit(txn_muts0.clone(), start_ts0, commit_ts0, false)
            .unwrap();
        client.put_kv_in_ref_store(muts0);
        client.verify_data_with_ref_store();

        let (muts1, txn_muts1) = make_insert_mutations("already_exist_1");
        let start_ts1 = client.get_ts();
        let for_update_ts1 = client.get_ts();
        let key_errors = pessimistic_lock(
            &mut client,
            txn_muts1.primary(),
            lock_keys(&muts1),
            start_ts1,
            for_update_ts1,
        )
        .unwrap();
        assert!(key_errors.is_empty(), "key_errors {:?}", key_errors);
        let err =
            pessimistic_prewrite(&mut client, txn_muts1, start_ts1, for_update_ts1).unwrap_err();
        assert_matches!(err, ClientError::AlreadyExist(_));
        client.verify_data_with_ref_store();
    }

    // Commit without prewrite:
    {
        let (_, txn_muts) = make_mutations("with_prewrite");
        let start_ts = client.get_ts();
        let commit_ts = client.get_ts();
        let err = client
            .kv_commit(txn_muts, start_ts, commit_ts, false)
            .unwrap_err();
        expect_err_msg(&err, "TxnLockNotFound");
        client.verify_data_with_ref_store();
    }

    // Commit after rollback.
    {
        let (muts0, txn_muts0) = make_mutations("after_rollback");
        let start_ts0 = client.get_ts();
        let for_update_ts0 = client.get_ts();
        let key_errors = pessimistic_lock(
            &mut client,
            txn_muts0.primary(),
            lock_keys(&muts0),
            start_ts0,
            for_update_ts0,
        )
        .unwrap();
        assert!(key_errors.is_empty(), "key_errors {:?}", key_errors);
        pessimistic_prewrite(&mut client, txn_muts0.clone(), start_ts0, for_update_ts0).unwrap();
        client.kv_rollback(txn_muts0.clone(), start_ts0).unwrap();

        let commit_ts0 = client.get_ts();
        let err = client
            .kv_commit(txn_muts0.clone(), start_ts0, commit_ts0, false)
            .unwrap_err();
        expect_err_msg(&err, "TxnLockNotFound");

        let (muts1, txn_muts1) = make_mutations("after_rollback_1");
        let start_ts1 = client.get_ts();
        let for_update_ts1 = client.get_ts();
        let key_errors = pessimistic_lock(
            &mut client,
            txn_muts1.primary(),
            lock_keys(&muts1),
            start_ts1,
            for_update_ts1,
        )
        .unwrap();
        assert!(key_errors.is_empty(), "key_errors {:?}", key_errors);
        pessimistic_prewrite(&mut client, txn_muts1.clone(), start_ts1, for_update_ts1).unwrap();

        let commit_ts1 = client.get_ts();
        client
            .kv_commit(txn_muts1.clone(), start_ts1, commit_ts1, false)
            .unwrap();
        client.put_kv_in_ref_store(muts1);
        client.verify_data_with_ref_store();
    }

    // Pessimistic lock conflict & duplicated commit.
    {
        let (muts0, txn_muts0) = make_mutations("write_conflict");
        let start_ts0 = client.get_ts();
        let for_update_ts0 = client.get_ts();
        let key_errors = pessimistic_lock(
            &mut client,
            txn_muts0.primary(),
            lock_keys(&muts0),
            start_ts0,
            for_update_ts0,
        )
        .unwrap();
        assert!(key_errors.is_empty(), "key_errors {:?}", key_errors);
        pessimistic_prewrite(&mut client, txn_muts0.clone(), start_ts0, for_update_ts0).unwrap();

        let commit_ts0 = client.get_ts();
        let txn_muts0_clone = txn_muts0.clone();
        let mut client_clone = client.clone();
        // let (tx, rx) = std::sync::mpsc::channel();
        let txn0 = rt.spawn_blocking(move || {
            thread::sleep(Duration::from_secs(1));
            let new_commit_ts = client_clone
                .kv_commit(txn_muts0_clone.clone(), start_ts0, commit_ts0, false)
                .unwrap();
            assert!(
                new_commit_ts > commit_ts0,
                "commit_ts0: {}, new_commit_ts: {}",
                commit_ts0,
                new_commit_ts
            );
            new_commit_ts
        });

        // do another parallel conflict prewrite to renew the commit ts.
        let (_, txn_muts1) = make_mutations("write_conflict_1");
        let start_ts1 = client.get_ts();
        let err = client
            .kv_prewrite(txn_muts1.primary(), None, txn_muts1, start_ts1)
            .unwrap_err();
        assert_matches!(err, ClientError::WriteConflict(_));

        let _new_commit_ts0 = rt.block_on(txn0).unwrap();
        client.put_kv_in_ref_store(muts0);
        client.verify_data_with_ref_store();

        // Duplicated commit.
        // Note: commit for secondary keys would fail for TxnLockNotFound error, when
        // there is no key in the region.
        let dup_commit_ts = client
            .kv_commit_ext(
                txn_muts0,
                start_ts0,
                commit_ts0,
                CommitAction::AsyncCommitSecondaryKeys(Duration::ZERO),
            )
            .unwrap();
        // Will return the commit_ts in request other than of committed record.
        assert_eq!(dup_commit_ts, commit_ts0);
        client.verify_data_with_ref_store();
    }

    // Rollback
    {
        let (muts, txn_muts) = make_mutations("rollback");
        // Rollback a not existed txn.
        let start_ts = client.get_ts();
        client.kv_rollback(txn_muts.clone(), start_ts).unwrap();

        let start_ts = client.get_ts();
        let for_update_ts = client.get_ts();
        let key_errors = pessimistic_lock(
            &mut client,
            txn_muts.primary(),
            lock_keys(&muts),
            start_ts,
            for_update_ts,
        )
        .unwrap();
        assert!(key_errors.is_empty(), "key_errors {:?}", key_errors);
        pessimistic_prewrite(&mut client, txn_muts.clone(), start_ts, for_update_ts).unwrap();
        client.kv_rollback(txn_muts.clone(), start_ts).unwrap();

        // Duplicated rollback.
        client.kv_rollback(txn_muts.clone(), start_ts).unwrap();

        // Commit after rollback.
        let commit_ts = client.get_ts();
        let err = client
            .kv_commit(txn_muts.clone(), start_ts, commit_ts, false)
            .unwrap_err();
        expect_err_msg(&err, "TxnLockNotFound");

        client.verify_data_with_ref_store();

        // Committed.
        let start_ts = client.get_ts();
        let for_update_ts = client.get_ts();
        let key_errors = pessimistic_lock(
            &mut client,
            txn_muts.primary(),
            lock_keys(&muts),
            start_ts,
            for_update_ts,
        )
        .unwrap();
        assert!(key_errors.is_empty(), "key_errors {:?}", key_errors);
        pessimistic_prewrite(&mut client, txn_muts.clone(), start_ts, for_update_ts).unwrap();
        let commit_ts = client.get_ts();
        client
            .kv_commit(txn_muts.clone(), start_ts, commit_ts, false)
            .unwrap();
        // Rollback after committed.
        let err = client.kv_rollback(txn_muts, start_ts).unwrap_err();
        expect_err_msg(&err, "Committed");

        client.put_kv_in_ref_store(muts);
        client.verify_data_with_ref_store();
    }

    // Resolve lock
    {
        let (muts, txn_muts) = make_mutations("resolve_lock");

        // Resolve a not existed txn.
        let start_ts = client.get_ts();
        client
            .kv_resolve_lock_for_mutations(start_ts.into_inner(), None, muts.clone(), true)
            .unwrap();

        let start_ts = client.get_ts();
        let for_update_ts = client.get_ts();
        let key_errors = pessimistic_lock(
            &mut client,
            txn_muts.primary(),
            lock_keys(&muts),
            start_ts,
            for_update_ts,
        )
        .unwrap();
        assert!(key_errors.is_empty(), "key_errors {:?}", key_errors);
        pessimistic_prewrite(&mut client, txn_muts.clone(), start_ts, for_update_ts).unwrap();
        // Rollback the not committed txn.
        client
            .kv_resolve_lock_for_mutations(start_ts.into_inner(), None, muts.clone(), is_txn_file)
            .unwrap();

        // Duplicated resolve lock.
        client
            .kv_resolve_lock_for_mutations(start_ts.into_inner(), None, muts.clone(), is_txn_file)
            .unwrap();

        // Commit after rollback.
        let commit_ts = client.get_ts();
        let err = client
            .kv_commit(txn_muts.clone(), start_ts, commit_ts, false)
            .unwrap_err();
        expect_err_msg(&err, "TxnLockNotFound");

        client.verify_data_with_ref_store();

        // Committed.
        let start_ts = client.get_ts();
        let for_update_ts = client.get_ts();
        let key_errors = pessimistic_lock(
            &mut client,
            txn_muts.primary(),
            lock_keys(&muts),
            start_ts,
            for_update_ts,
        )
        .unwrap();
        assert!(key_errors.is_empty(), "key_errors {:?}", key_errors);
        pessimistic_prewrite(&mut client, txn_muts.clone(), start_ts, for_update_ts).unwrap();
        let commit_ts = client.get_ts();
        client
            .kv_commit(txn_muts.clone(), start_ts, commit_ts, false)
            .unwrap();
        // Rollback after committed.
        let err = client
            .kv_resolve_lock_for_mutations(start_ts.into_inner(), None, muts.clone(), is_txn_file)
            .unwrap_err();
        expect_err_msg(&err, "Committed");

        // Duplicated commit.
        client
            .kv_resolve_lock_for_mutations(
                start_ts.into_inner(),
                Some(commit_ts.into_inner()),
                muts.clone(),
                is_txn_file,
            )
            .unwrap();
        let commit_ts = client.get_ts();
        client
            .kv_resolve_lock_for_mutations(
                start_ts.into_inner(),
                Some(commit_ts.into_inner()),
                muts.clone(),
                is_txn_file,
            )
            .unwrap();

        client.put_kv_in_ref_store(muts);
        client.verify_data_with_ref_store();
    }

    // Check txn status
    {
        let (muts, txn_muts) = make_mutations("check_txn_status");
        let start_ts = client.get_ts();
        let for_update_ts = client.get_ts();
        let key_errors = pessimistic_lock(
            &mut client,
            txn_muts.primary(),
            lock_keys(&muts),
            start_ts,
            for_update_ts,
        )
        .unwrap();
        assert!(key_errors.is_empty(), "key_errors {:?}", key_errors);
        pessimistic_prewrite(&mut client, txn_muts.clone(), start_ts, for_update_ts).unwrap();
        {
            // primary mismatch.
            let second_key = muts[1].key.clone();
            let resp = client
                .kv_check_txn_status(
                    &second_key,
                    start_ts,
                    start_ts,
                    start_ts,
                    false,
                    false,
                    false,
                    true,
                )
                .unwrap();
            assert!(resp.has_error());
            assert!(resp.get_error().has_primary_mismatch());
        }

        {
            let resp = client
                .kv_check_txn_status(
                    &txn_muts.primary(),
                    start_ts,
                    start_ts,
                    start_ts,
                    false,
                    false,
                    false,
                    true,
                )
                .unwrap();
            assert!(!resp.has_error());
            assert_eq!(resp.commit_version, 0);
            assert_eq!(resp.get_lock_info().lock_version, start_ts.into_inner());
            assert_eq!(
                resp.get_lock_info().get_primary_lock(),
                txn_muts.primary().as_ref()
            );
        }
        // The status checks above intentionally leave the transaction active.
        // Clean it up explicitly instead of making the final verification wait
        // for the 20-second lock TTL.
        client.kv_rollback(txn_muts, start_ts).unwrap();
    }

    client.verify_data_with_ref_store();
    cluster.stop();
    oss.shutdown();
}

#[test]
fn test_txn_file_move_down() {
    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test");
    let node_ids = alloc_node_id_vec(NODES_COUNT);
    let first_node = node_ids[0];
    let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None);
    let cluster_id = pd_wrapper.client().get_cluster_id().unwrap();
    info!("test_txn_file_move_down"; "cluster_id" => cluster_id);
    let mut cluster = ServerClusterBuilder::new(node_ids, |_, conf| {
        conf.dfs = dfs_config.clone();
        conf.kvengine.flush_split_l0 = true;
    })
    .pd(pd_wrapper)
    .build();
    cluster.start_tikv_workers(alloc_node_id_vec(1), TikvWorkerOptions::default());
    cluster.wait_region_replicated(&[], 3);

    let gen_key = generate_keyspace_key(KEYSPACE_ID);

    let rt = tokio::runtime::Runtime::new().unwrap();
    let _enter = rt.enter();

    let mut client = cluster.new_client_opt(ClusterClientOptions {
        txn_file_max_chunk_size: Some(256 * 1024),
        ..Default::default()
    });

    client.split_keyspace(KEYSPACE_ID);
    let num_keys = 10000;
    {
        client
            .try_put_kv(
                0..num_keys,
                &gen_key,
                i_to_val_opt("value0_", 3),
                MutateOptions {
                    commit_action: CommitAction::AsyncCommitSecondaryKeys(Duration::ZERO),
                    write_method: TxnWriteMethod::FileBased,
                    ..Default::default()
                },
            )
            .unwrap();
    }
    must_wait(
        || {
            let kvengine = cluster.get_kvengine(first_node);
            let stats = kvengine.get_all_shard_stats();
            stats
                .iter()
                .all(|s| s.l0_table_count == 0 && s.mem_table_size == 0)
        },
        5,
        || "wait for compaction".to_string(),
    );
    client.verify_data_with_ref_store();
    cluster.stop();
    oss.shutdown();
}

#[test]
fn test_txn_file_lock_limit_rejects_new_prewrite() {
    let mut cluster = ServerCluster::new(alloc_node_id_vec(1), |_, conf| {
        conf.storage.max_txn_file_locks_per_region = 1;
    });
    let dfs = cluster.get_dfs().unwrap();
    cluster.wait_region_replicated(&[], 1);

    let mut client = cluster.new_client();
    client.split_keyspace(KEYSPACE_ID);

    let gen_key = generate_keyspace_key(KEYSPACE_ID);
    let primary_lock0 = gen_key(0);
    let primary_lock1 = gen_key(100);

    let start_ts0 = client.get_ts().into_inner();
    let chunk_ids0 = build_txn_files(&dfs, start_ts0, 0, 100);
    let (_, ctx0, kv_client0) = prepare_region_ctx(&mut client, &primary_lock0);
    let resp0 = kv_client0
        .kv_prewrite(&build_txn_file_prewrite_req(
            ctx0.clone(),
            primary_lock0.clone(),
            start_ts0,
            chunk_ids0,
        ))
        .unwrap();
    assert!(
        resp0.get_errors().is_empty() && !resp0.has_region_error(),
        "resp {:?}",
        resp0
    );

    let mut check_req = CheckTxnStatusRequest::new();
    check_req.set_context(ctx0.clone());
    check_req.set_primary_key(primary_lock0.clone());
    check_req.set_lock_ts(start_ts0);
    check_req.set_caller_start_ts(start_ts0 + 1);
    check_req.set_is_txn_file(true);
    let check_resp = kv_client0.kv_check_txn_status(&check_req).unwrap();
    assert!(
        !check_resp.has_error() && !check_resp.has_region_error(),
        "resp {:?}",
        check_resp
    );

    let mut heartbeat_req = TxnHeartBeatRequest::new();
    heartbeat_req.set_context(ctx0.clone());
    heartbeat_req.set_start_version(start_ts0);
    heartbeat_req.set_primary_lock(primary_lock0.clone());
    heartbeat_req.set_advise_lock_ttl(10_000);
    heartbeat_req.set_is_txn_file(true);
    let heartbeat_resp = kv_client0.kv_txn_heart_beat(&heartbeat_req).unwrap();
    assert!(
        !heartbeat_resp.has_error() && !heartbeat_resp.has_region_error(),
        "resp {:?}",
        heartbeat_resp
    );

    let start_ts1 = client.get_ts().into_inner();
    let chunk_ids1 = build_txn_files_with_first_chunk_id(&dfs, start_ts1 + 1, 100, 200);
    let (_, ctx1, kv_client1) = prepare_region_ctx(&mut client, &primary_lock1);
    let resp1 = kv_client1
        .kv_prewrite(&build_txn_file_prewrite_req(
            ctx1.clone(),
            primary_lock1.clone(),
            start_ts1,
            chunk_ids1.clone(),
        ))
        .unwrap();
    assert!(resp1.get_errors().is_empty(), "resp {:?}", resp1);
    assert!(resp1.has_region_error(), "resp {:?}", resp1);
    assert!(
        resp1.get_region_error().has_server_is_busy(),
        "resp {:?}",
        resp1
    );
    assert_eq!(
        resp1.get_region_error().get_server_is_busy().reason,
        "txn file lock limit reached"
    );

    let mut commit_req = CommitRequest::new();
    commit_req.set_context(ctx0);
    commit_req.set_start_version(start_ts0);
    commit_req.set_is_txn_file(true);
    commit_req.set_commit_version(client.get_ts().into_inner());
    let commit_resp = kv_client0.kv_commit(&commit_req).unwrap();
    assert!(
        !commit_resp.has_error() && !commit_resp.has_region_error(),
        "resp {:?}",
        commit_resp
    );

    let retry_resp = kv_client1
        .kv_prewrite(&build_txn_file_prewrite_req(
            ctx1.clone(),
            primary_lock1.clone(),
            start_ts1,
            chunk_ids1,
        ))
        .unwrap();
    assert!(
        retry_resp.get_errors().is_empty() && !retry_resp.has_region_error(),
        "resp {:?}",
        retry_resp
    );

    let mut rollback_req = BatchRollbackRequest::new();
    rollback_req.set_context(ctx1);
    rollback_req.set_is_txn_file(true);
    rollback_req.set_start_version(start_ts1);
    let rollback_resp = kv_client1.kv_batch_rollback(&rollback_req).unwrap();
    assert!(
        !rollback_resp.has_error() && !rollback_resp.has_region_error(),
        "resp {:?}",
        rollback_resp
    );

    cluster.stop();
}

#[test]
fn test_txn_file_lock_limit_online_config() {
    let node_ids = alloc_node_id_vec(1);
    let node_id = node_ids[0];
    let mut cluster = ServerCluster::new(node_ids, |_, conf| {
        conf.storage.max_txn_file_locks_per_region = 1;
    });
    let dfs = cluster.get_dfs().unwrap();
    cluster.wait_region_replicated(&[], 1);

    let mut client = cluster.new_client();
    client.split_keyspace(KEYSPACE_ID);

    let gen_key = generate_keyspace_key(KEYSPACE_ID);
    let primary_lock0 = gen_key(0);
    let primary_lock1 = gen_key(100);
    let primary_lock2 = gen_key(400);

    let start_ts0 = client.get_ts().into_inner();
    let chunk_ids0 = build_txn_files(&dfs, start_ts0, 0, 100);
    let (_, ctx0, kv_client0) = prepare_region_ctx(&mut client, &primary_lock0);
    let resp0 = kv_client0
        .kv_prewrite(&build_txn_file_prewrite_req(
            ctx0.clone(),
            primary_lock0.clone(),
            start_ts0,
            chunk_ids0,
        ))
        .unwrap();
    assert!(
        resp0.get_errors().is_empty() && !resp0.has_region_error(),
        "resp {:?}",
        resp0
    );

    update_txn_file_lock_limit(cluster.status_addr(node_id), 2);

    let start_ts1 = client.get_ts().into_inner();
    let chunk_ids1 = build_txn_files(&dfs, start_ts1, 100, 200);
    let (_, ctx1, kv_client1) = prepare_region_ctx(&mut client, &primary_lock1);
    let resp1 = kv_client1
        .kv_prewrite(&build_txn_file_prewrite_req(
            ctx1.clone(),
            primary_lock1.clone(),
            start_ts1,
            chunk_ids1,
        ))
        .unwrap();
    assert!(
        resp1.get_errors().is_empty() && !resp1.has_region_error(),
        "resp {:?}",
        resp1
    );

    update_txn_file_lock_limit(cluster.status_addr(node_id), 1);

    let start_ts2 = client.get_ts().into_inner();
    let chunk_ids2 = build_txn_files(&dfs, start_ts2, 400, 500);
    let (_, ctx2, kv_client2) = prepare_region_ctx(&mut client, &primary_lock2);
    let resp2 = kv_client2
        .kv_prewrite(&build_txn_file_prewrite_req(
            ctx2,
            primary_lock2,
            start_ts2,
            chunk_ids2,
        ))
        .unwrap();
    assert!(resp2.get_errors().is_empty(), "resp {:?}", resp2);
    assert!(resp2.has_region_error(), "resp {:?}", resp2);
    assert!(
        resp2.get_region_error().has_server_is_busy(),
        "resp {:?}",
        resp2
    );
    assert_eq!(
        resp2.get_region_error().get_server_is_busy().reason,
        "txn file lock limit reached"
    );

    let merge_chunk_ids = build_txn_files_with_first_chunk_id(&dfs, start_ts1 + 1001, 200, 400);
    let merge_resp = kv_client1
        .kv_prewrite(&build_txn_file_prewrite_req(
            ctx1.clone(),
            primary_lock1.clone(),
            start_ts1,
            merge_chunk_ids,
        ))
        .unwrap();
    assert!(
        merge_resp.get_errors().is_empty() && !merge_resp.has_region_error(),
        "resp {:?}",
        merge_resp
    );

    let mut commit_req = CommitRequest::new();
    commit_req.set_context(ctx1.clone());
    commit_req.set_start_version(start_ts1);
    commit_req.set_is_txn_file(true);
    let commit_ts1 = client.get_ts().into_inner();
    commit_req.set_commit_version(commit_ts1);
    let commit_resp = kv_client1.kv_commit(&commit_req).unwrap();
    assert!(
        !commit_resp.has_error() && !commit_resp.has_region_error(),
        "resp {:?}",
        commit_resp
    );
    verify_range(&mut client, 100, 400, commit_ts1);

    let mut rollback_req = BatchRollbackRequest::new();
    rollback_req.set_context(ctx0);
    rollback_req.set_is_txn_file(true);
    rollback_req.set_start_version(start_ts0);
    let rollback_resp = kv_client0.kv_batch_rollback(&rollback_req).unwrap();
    assert!(
        !rollback_resp.has_error() && !rollback_resp.has_region_error(),
        "resp {:?}",
        rollback_resp
    );

    cluster.stop();
}

#[test]
fn test_txn_file_lock_limit_concurrent_prewrite() {
    let mut cluster = ServerCluster::new(alloc_node_id_vec(1), |_, conf| {
        conf.storage.max_txn_file_locks_per_region = 1;
        conf.storage.scheduler_worker_pool_size = 4;
        conf.kvengine.txn_file_worker_pool_size = Some(4);
    });
    let dfs = cluster.get_dfs().unwrap();
    cluster.wait_region_replicated(&[], 1);

    let mut client = cluster.new_client();
    client.split_keyspace(KEYSPACE_ID);

    let gen_key = generate_keyspace_key(KEYSPACE_ID);
    let primary_lock0 = gen_key(0);
    let primary_lock1 = gen_key(500);
    let start_ts0 = client.get_ts().into_inner();
    let start_ts1 = client.get_ts().into_inner();
    let chunk_ids0 = build_txn_files(&dfs, start_ts0, 0, 500);
    let chunk_ids1 = build_txn_files(&dfs, start_ts1, 500, 1000);

    let barrier = Arc::new(Barrier::new(2));
    let mut client0 = cluster.new_client();
    let mut client1 = cluster.new_client();
    let handle0 = {
        let barrier = barrier.clone();
        let primary_lock = primary_lock0.clone();
        std::thread::spawn(move || {
            let (_, ctx, kv_client) = prepare_region_ctx(&mut client0, &primary_lock);
            barrier.wait();
            let mut resp = kv_client
                .kv_prewrite(&build_txn_file_prewrite_req(
                    ctx,
                    primary_lock.clone(),
                    start_ts0,
                    chunk_ids0,
                ))
                .unwrap();
            assert!(resp.get_errors().is_empty(), "resp {:?}", resp);
            if resp.has_region_error() {
                Err(resp.take_region_error())
            } else {
                Ok((primary_lock, start_ts0))
            }
        })
    };
    let handle1 = {
        let barrier = barrier.clone();
        let primary_lock = primary_lock1.clone();
        std::thread::spawn(move || {
            let (_, ctx, kv_client) = prepare_region_ctx(&mut client1, &primary_lock);
            barrier.wait();
            let mut resp = kv_client
                .kv_prewrite(&build_txn_file_prewrite_req(
                    ctx,
                    primary_lock.clone(),
                    start_ts1,
                    chunk_ids1,
                ))
                .unwrap();
            assert!(resp.get_errors().is_empty(), "resp {:?}", resp);
            if resp.has_region_error() {
                Err(resp.take_region_error())
            } else {
                Ok((primary_lock, start_ts1))
            }
        })
    };

    let mut success_txn = None;
    let mut server_is_busy_count = 0;
    for result in [handle0.join().unwrap(), handle1.join().unwrap()] {
        match result {
            Ok(txn) => {
                assert!(
                    success_txn.is_none(),
                    "unexpected second success: {:?}",
                    txn
                );
                success_txn = Some(txn);
            }
            Err(region_err) => {
                assert!(
                    region_err.has_server_is_busy(),
                    "unexpected region error: {:?}",
                    region_err
                );
                assert_eq!(
                    region_err.get_server_is_busy().reason,
                    "txn file lock limit reached"
                );
                server_is_busy_count += 1;
            }
        }
    }
    assert!(success_txn.is_some(), "expected one successful prewrite");
    assert_eq!(server_is_busy_count, 1);

    let (primary_lock, start_ts) = success_txn.unwrap();
    let (_, ctx, kv_client) = prepare_region_ctx(&mut client, &primary_lock);
    let mut rollback_req = BatchRollbackRequest::new();
    rollback_req.set_context(ctx);
    rollback_req.set_is_txn_file(true);
    rollback_req.set_start_version(start_ts);
    let rollback_resp = kv_client.kv_batch_rollback(&rollback_req).unwrap();
    assert!(
        !rollback_resp.has_error() && !rollback_resp.has_region_error(),
        "resp {:?}",
        rollback_resp
    );

    cluster.stop();
}

#[test]
fn test_txn_file_merge() {
    let cases = vec![
        vec![1..2, 3..5, 0..1, 2..3],
        vec![1..3, 2..4, 3..5, 0..2],
        vec![0..1, 1..2, 2..3, 3..4, 4..5],
        vec![0..2, 2..5],
        vec![0..5; 3],
    ];

    for ranges in cases {
        test_txn_file_merge_impl(ranges, None);
    }
}

#[test]
fn test_txn_file_lock_limit_merge_same_start_ts() {
    test_txn_file_merge_impl(vec![0..1, 1..2, 2..3, 3..4, 4..5], Some(1));
}

fn test_txn_file_merge_impl(
    ranges: Vec<Range<usize>>,
    max_txn_file_locks_per_region: Option<usize>,
) {
    let mut cluster = ServerCluster::new(alloc_node_id_vec(3), |_, conf| {
        conf.storage.scheduler_worker_pool_size = 8;
        conf.kvengine.txn_file_worker_pool_size = Some(16);
        if let Some(limit) = max_txn_file_locks_per_region {
            conf.storage.max_txn_file_locks_per_region = limit;
        }
    });
    let dfs = cluster.get_dfs().unwrap();
    cluster.wait_region_replicated(&[], 3);

    let prepare_ctx =
        |client: &mut ClusterClient, key: &[u8]| -> (u64, kvrpcpb::Context, TikvClient) {
            let region_id = client.get_region_id(key);
            let ctx = client.new_rpc_ctx(region_id, key).unwrap();
            let kv_client = client.get_kv_client(ctx.get_peer().get_store_id());
            (region_id, ctx, kv_client)
        };

    let mut client = cluster.new_client();
    client.split_keyspace(KEYSPACE_ID);

    let start_ts = client.get_ts().into_inner();
    let chunk_ids = build_txn_files(&dfs, start_ts, 0, 500);
    assert_eq!(chunk_ids.len(), 5);

    let gen_key = generate_keyspace_key(KEYSPACE_ID);
    let primary_lock = gen_key(0);

    // Note: the primary key should be prewrite first. But it's OK here, as long as
    // there is no other command on primary.
    let mut handles = Vec::with_capacity(ranges.len());
    for range in ranges {
        let mut client = cluster.new_client();
        let sub_chunk_ids = chunk_ids[range].to_vec();
        let primary_lock = primary_lock.clone();
        let handle = std::thread::spawn(move || {
            let (mut region_id, mut ctx, mut kv_client) = prepare_ctx(&mut client, &primary_lock);

            let mut req = PrewriteRequest::new();
            req.set_context(ctx.clone());
            req.set_txn_file_chunks(sub_chunk_ids);
            req.set_primary_lock(primary_lock.clone());
            req.set_lock_ttl(6000);
            req.set_start_version(start_ts);
            req.set_min_commit_ts(start_ts + 1);

            let ok = try_wait(
                || {
                    let resp = kv_client.kv_prewrite(&req).unwrap();
                    assert!(resp.get_errors().is_empty(), "resp {:?}", resp);
                    if resp.has_region_error() {
                        warn!("prewrite got region error {:?}", resp.get_region_error());
                        (region_id, ctx, kv_client) = prepare_ctx(&mut client, &primary_lock);
                        req.set_context(ctx.clone());
                        false
                    } else {
                        true
                    }
                },
                10,
            );
            assert!(ok);
        });
        handles.push(handle);
    }
    for h in handles {
        h.join().unwrap();
    }

    let (_, ctx, kv_client) = prepare_ctx(&mut client, &primary_lock);
    let mut req = CommitRequest::new();
    req.set_context(ctx);
    req.set_start_version(start_ts);
    req.set_is_txn_file(true);
    let commit_ts = client.get_ts().into_inner();
    req.set_commit_version(commit_ts);
    let resp = kv_client.kv_commit(&req).unwrap();
    assert!(
        !resp.has_error() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );

    verify_range(&mut client, 0, 500, commit_ts);
    cluster.stop();
}

// Test for commit "not primary" lock to primary region.
// See https://github.com/tidbcloud/cloud-storage-engine/issues/1800.
#[test]
fn test_commit_primary_region() {
    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test");

    let node_ids = alloc_node_id_vec(NODES_COUNT);
    let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None);
    let mut cluster = ServerClusterBuilder::new(node_ids, |_, conf| {
        conf.dfs = dfs_config.clone();
    })
    .pd(pd_wrapper)
    .build();
    cluster.start_tikv_workers(alloc_node_id_vec(1), TikvWorkerOptions::default());
    cluster.wait_region_replicated(&[], 3);

    let gen_key = generate_keyspace_key(KEYSPACE_ID);

    let rt = tokio::runtime::Runtime::new().unwrap();
    let _enter = rt.enter();

    let mut client = cluster.new_client_opt(ClusterClientOptions {
        txn_file_max_chunk_size: Some(1024),
        ..Default::default()
    });
    let txn_file_helper = client.txn_file_helper();
    let make_mutations = |value, start, end| {
        make_txn_mutations(
            value,
            start,
            end,
            &gen_key,
            |_| kvrpcpb::Op::Put,
            TxnWriteMethod::FileBased,
            txn_file_helper.clone(),
        )
    };
    client.split_keyspace(KEYSPACE_ID);

    // Prewrite to two regions.
    client.try_split(&gen_key(100), 5).unwrap();
    let (_, txn_muts) = make_mutations("value_", 0, 200);
    let start_ts = client.get_ts();
    client
        .kv_prewrite(txn_muts.primary(), None, txn_muts.clone(), start_ts)
        .unwrap();

    // Rollback the primary region.
    let (_, txn_muts0) = make_mutations("value0_", 0, 1);
    client.kv_rollback(txn_muts0, start_ts).unwrap();
    let ok = client.try_merge_and_wait(&gen_key(0), &gen_key(100), 5);
    assert!(ok);

    // Commit must fail.
    let commit_ts = client.get_ts();
    let err = client
        .kv_commit(txn_muts, start_ts, commit_ts, false)
        .unwrap_err();
    expect_err_msg(&err, "TxnLockNotFound");

    {
        // To work around that data not existed in ref store will not be checked.
        let guard = client.ref_store();
        let mut ref_store = guard.lock().unwrap();
        for k in 0..200 {
            ref_store.del_kv(gen_key(k));
        }
    }
    client.verify_data_with_ref_store();
    cluster.stop();
    oss.shutdown();
}

#[test]
fn test_commit_primary_region_after_merge_reprewrite() {
    test_util::init_log_for_test();
    let (_temp_dir, _oss, dfs_config) = prepare_dfs("test");

    let node_ids = alloc_node_id_vec(NODES_COUNT);
    let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None);
    let cluster_id = pd_wrapper.client().get_cluster_id().unwrap();
    let mut cluster = ServerClusterBuilder::new(node_ids.clone(), |_, conf| {
        conf.dfs = dfs_config.clone();
        conf.rfengine.lightweight_backup = true;
        conf.storage.check_backup_ts = true;
    })
    .pd(pd_wrapper)
    .build();
    cluster.start_tikv_workers(alloc_node_id_vec(1), TikvWorkerOptions::default());
    cluster.wait_region_replicated(&[], 3);

    let gen_key = generate_keyspace_key(KEYSPACE_ID);
    block_on(cluster.keyspace_manager().create_single_keyspace(
        KEYSPACE_ID,
        format!("ks{KEYSPACE_ID}"),
        &CreateKeyspaceOptions::default(),
        false,
    ));

    let mut client = cluster.new_client_opt(ClusterClientOptions {
        txn_file_max_chunk_size: Some(1024),
        ..Default::default()
    });
    client.split_keyspace(KEYSPACE_ID);
    client.try_split(&gen_key(100), 5).unwrap();

    let dfs = cluster.get_dfs().unwrap();
    let primary_lock = gen_key(0);
    let secondary_key = gen_key(100);
    let start_ts = client.get_ts().into_inner();
    let primary_chunk_ids = build_txn_files(&dfs, start_ts, 0, 100);
    let secondary_chunk_ids = build_txn_files_with_first_chunk_id(&dfs, start_ts + 1001, 100, 200);
    assert_eq!(primary_chunk_ids.len(), 1);
    assert_eq!(secondary_chunk_ids.len(), 1);

    // Prewrite the primary-side chunk on the source region.
    prewrite_txn_file_chunk(
        &mut client,
        &primary_lock,
        &primary_lock,
        start_ts,
        &primary_chunk_ids,
        6000,
    );

    // Prewrite the secondary-side chunk on the target region.
    prewrite_txn_file_chunk(
        &mut client,
        &secondary_key,
        &primary_lock,
        start_ts,
        &secondary_chunk_ids,
        6000,
    );

    // Commit the source region first and keep the commit_ts for the later
    // secondary commit on the merged region.
    let requested_commit_ts = client.get_ts().into_inner();
    let (source_region_id, ctx, kv_client) = prepare_region_ctx(&mut client, &primary_lock);
    let mut commit_req = CommitRequest::new();
    commit_req.set_context(ctx);
    commit_req.set_start_version(start_ts);
    commit_req.set_is_txn_file(true);
    commit_req.set_commit_version(requested_commit_ts);
    let commit_resp = kv_client.kv_commit(&commit_req).unwrap();
    assert!(
        !commit_resp.has_error() && !commit_resp.has_region_error(),
        "resp {:?}",
        commit_resp
    );
    let commit_ts = if commit_resp.commit_version == 0 {
        requested_commit_ts
    } else {
        commit_resp.commit_version
    };

    // Overwrite the committed primary key so the source-region write becomes
    // GC-able before the merge.
    let overwrite_commit_ts = client.put_kv(0..1, |_| primary_lock.clone(), |_| i_to_val(1000));
    assert!(
        overwrite_commit_ts.into_inner() > commit_ts,
        "overwrite should commit after the source-region commit"
    );

    let pd_client = cluster.get_pd_client();
    let gc_safe_point = client.get_ts().into_inner();
    assert!(
        gc_safe_point > start_ts,
        "gc safe point should advance past start_ts"
    );
    block_on(pd_client.advance_txn_safe_point(KEYSPACE_ID, gc_safe_point.into())).unwrap();
    block_on(pd_client.advance_gc_safe_point(KEYSPACE_ID, gc_safe_point.into())).unwrap();
    block_on(cluster.update_gc_states_immediately()).unwrap();

    cluster.flush_memtable(source_region_id).unwrap();
    assert!(
        cluster.wait_for_memtable_flushed(source_region_id, Duration::from_secs(10)),
        "memtable not flushed for source region {}",
        source_region_id
    );
    cluster.request_major_compaction(
        MajorCompactionTarget::Region(source_region_id),
        false,
        Duration::from_secs(10),
    );
    assert!(
        try_wait(
            || {
                node_ids.iter().all(|&node_id| {
                    cluster
                        .get_kvengine(node_id)
                        .get_shard_stat(source_region_id)
                        .is_major_compacted()
                })
            },
            10,
        ),
        "source region {} was not major compacted",
        source_region_id
    );

    // Merge the source region into the secondary-side region after the primary
    // chunk is committed and compacted away.
    let ok = client.try_merge_and_wait(&primary_lock, &secondary_key, 5);
    assert!(ok, "primary region should merge into the secondary region");
    client.clear_region_cache();
    assert_eq!(
        client.get_region_id(&primary_lock),
        client.get_region_id(&secondary_key),
        "primary and secondary keys should be in the same merged region"
    );

    // Re-prewrite the secondary chunk on the merged region after the source-side
    // write has been compacted away.
    // In real scenario, TiDB should not send the re-prewrite request.
    let ok = try_wait(
        || {
            let (_, ctx, kv_client) = prepare_region_ctx(&mut client, &secondary_key);
            let req = build_txn_file_prewrite_req(
                ctx,
                primary_lock.clone(),
                start_ts,
                secondary_chunk_ids.clone(),
            );
            let resp = kv_client.kv_prewrite(&req).unwrap();
            if resp.has_region_error() {
                warn!("re-prewrite got region error {:?}", resp.get_region_error());
                client.clear_region_cache();
                return false;
            }
            assert!(resp.get_errors().is_empty(), "resp {:?}", resp);
            true
        },
        10,
    );
    assert!(ok, "re-prewrite failed for key {:?}", secondary_key);

    let backup_ts = client.get_ts().into_inner();
    assert!(
        backup_ts > commit_ts,
        "backup_ts should be newer than commit_ts"
    );
    request_lightweight_backup_to_all_stores(&cluster, &node_ids, cluster_id, backup_ts);

    // Commit on the merged region with the old commit_ts.
    // Before issue/4809 is fixed, this commit will be rejected with
    // CommitTsExpired.
    let (_, ctx, kv_client) = prepare_region_ctx(&mut client, &secondary_key);
    let mut expired_commit_req = CommitRequest::new();
    expired_commit_req.set_context(ctx);
    expired_commit_req.set_start_version(start_ts);
    expired_commit_req.set_is_txn_file(true);
    expired_commit_req.set_commit_version(commit_ts);
    let expired_commit_resp = kv_client.kv_commit(&expired_commit_req).unwrap();
    assert!(
        !expired_commit_resp.has_region_error(),
        "resp {:?}",
        expired_commit_resp
    );
    assert!(
        expired_commit_resp.has_error(),
        "resp {:?}",
        expired_commit_resp
    );
    assert!(
        expired_commit_resp
            .get_error()
            .get_retryable()
            .contains("TxnLockNotFound"),
        "resp {:?}",
        expired_commit_resp
    );

    // Before issue/4809 is fixed, this retry will panic for BadFormatLock.
    let retry_commit_ts = client.get_ts().into_inner();
    assert!(
        retry_commit_ts > backup_ts,
        "retry commit_ts should be newer than backup_ts"
    );
    let (_, ctx, kv_client) = prepare_region_ctx(&mut client, &secondary_key);
    let mut retry_commit_req = CommitRequest::new();
    retry_commit_req.set_context(ctx);
    retry_commit_req.set_start_version(start_ts);
    retry_commit_req.set_is_txn_file(true);
    retry_commit_req.set_commit_version(retry_commit_ts);
    let resp = kv_client.kv_commit(&retry_commit_req).unwrap();
    assert!(
        resp.get_error().get_retryable().contains("TxnLockNotFound"),
        "resp {:?}",
        resp
    );

    cluster.stop();
}

#[test]
fn test_commit_primary_region_after_merge_reprewrite_with_source_rollback() {
    test_util::init_log_for_test();
    let (_temp_dir, _oss, dfs_config) = prepare_dfs("test");

    let node_ids = alloc_node_id_vec(NODES_COUNT);
    let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None);
    let cluster_id = pd_wrapper.client().get_cluster_id().unwrap();
    let mut cluster = ServerClusterBuilder::new(node_ids.clone(), |_, conf| {
        conf.dfs = dfs_config.clone();
        conf.rfengine.lightweight_backup = true;
        conf.storage.check_backup_ts = true;
    })
    .pd(pd_wrapper)
    .build();
    cluster.start_tikv_workers(alloc_node_id_vec(1), TikvWorkerOptions::default());
    cluster.wait_region_replicated(&[], 3);

    let gen_key = generate_keyspace_key(KEYSPACE_ID);
    block_on(cluster.keyspace_manager().create_single_keyspace(
        KEYSPACE_ID,
        format!("ks{KEYSPACE_ID}"),
        &CreateKeyspaceOptions::default(),
        false,
    ));

    let mut client = cluster.new_client_opt(ClusterClientOptions {
        txn_file_max_chunk_size: Some(1024),
        ..Default::default()
    });
    client.split_keyspace(KEYSPACE_ID);
    client.try_split(&gen_key(100), 5).unwrap();

    let dfs = cluster.get_dfs().unwrap();
    let primary_lock = gen_key(0);
    let secondary_key = gen_key(100);
    let start_ts = client.get_ts().into_inner();
    let primary_chunk_ids = build_txn_files(&dfs, start_ts, 0, 100);
    let secondary_chunk_ids = build_txn_files_with_first_chunk_id(&dfs, start_ts + 1001, 100, 200);
    assert_eq!(primary_chunk_ids.len(), 1);
    assert_eq!(secondary_chunk_ids.len(), 1);

    // Prewrite the primary-side chunk on the source region.
    prewrite_txn_file_chunk(
        &mut client,
        &primary_lock,
        &primary_lock,
        start_ts,
        &primary_chunk_ids,
        6000,
    );

    // Prewrite the secondary-side chunk on the target region.
    prewrite_txn_file_chunk(
        &mut client,
        &secondary_key,
        &primary_lock,
        start_ts,
        &secondary_chunk_ids,
        6000,
    );

    // Roll back the source region so the primary lock is removed before merge.
    let (source_region_id, ctx, kv_client) = prepare_region_ctx(&mut client, &primary_lock);
    let mut rollback_req = BatchRollbackRequest::new();
    rollback_req.set_context(ctx);
    rollback_req.set_is_txn_file(true);
    rollback_req.set_start_version(start_ts);
    let rollback_resp = kv_client.kv_batch_rollback(&rollback_req).unwrap();
    assert!(
        !rollback_resp.has_error() && !rollback_resp.has_region_error(),
        "resp {:?}",
        rollback_resp
    );

    // Make the rollback record in EXTRA_CF GCed.
    let pd_client = cluster.get_pd_client();
    let gc_safe_point = client.get_ts().into_inner();
    assert!(
        gc_safe_point > start_ts,
        "gc safe point should advance past start_ts"
    );
    block_on(pd_client.advance_txn_safe_point(KEYSPACE_ID, gc_safe_point.into())).unwrap();
    block_on(pd_client.advance_gc_safe_point(KEYSPACE_ID, gc_safe_point.into())).unwrap();
    block_on(cluster.update_gc_states_immediately()).unwrap();

    cluster.flush_memtable(source_region_id).unwrap();
    assert!(
        cluster.wait_for_memtable_flushed(source_region_id, Duration::from_secs(10)),
        "memtable not flushed for source region {}",
        source_region_id
    );
    cluster.request_major_compaction(
        MajorCompactionTarget::Region(source_region_id),
        false,
        Duration::from_secs(10),
    );
    assert!(
        try_wait(
            || {
                node_ids.iter().all(|&node_id| {
                    cluster
                        .get_kvengine(node_id)
                        .get_shard_stat(source_region_id)
                        .is_extra_cf_major_compacted()
                })
            },
            10,
        ),
        "source region {} extra cf was not major compacted away",
        source_region_id
    );

    // Merge the source region into the secondary-side region after the
    // source-side rollback has been persisted and compacted away.
    let ok = client.try_merge_and_wait(&primary_lock, &secondary_key, 5);
    assert!(ok, "primary region should merge into the secondary region");
    client.clear_region_cache();
    assert_eq!(
        client.get_region_id(&primary_lock),
        client.get_region_id(&secondary_key),
        "primary and secondary keys should be in the same merged region"
    );

    // Re-prewrite the secondary chunk on the merged region after the source-side
    // rollback has been compacted away.
    let ok = try_wait(
        || {
            let (_, ctx, kv_client) = prepare_region_ctx(&mut client, &secondary_key);
            let req = build_txn_file_prewrite_req(
                ctx,
                primary_lock.clone(),
                start_ts,
                secondary_chunk_ids.clone(),
            );
            let resp = kv_client.kv_prewrite(&req).unwrap();
            if resp.has_region_error() {
                warn!("re-prewrite got region error {:?}", resp.get_region_error());
                client.clear_region_cache();
                return false;
            }
            assert!(resp.get_errors().is_empty(), "resp {:?}", resp);
            true
        },
        10,
    );
    assert!(ok, "re-prewrite failed for key {:?}", secondary_key);

    let commit_ts = client.get_ts().into_inner();
    assert!(
        commit_ts > start_ts,
        "commit_ts should be newer than start_ts"
    );
    let backup_ts = client.get_ts().into_inner();
    assert!(
        backup_ts > commit_ts,
        "backup_ts should be newer than commit_ts"
    );
    request_lightweight_backup_to_all_stores(&cluster, &node_ids, cluster_id, backup_ts);

    // Commit on the merged region with the old commit_ts.
    // Before issue/4809 is fixed, this commit will be rejected with
    // CommitTsExpired.
    let (_, ctx, kv_client) = prepare_region_ctx(&mut client, &secondary_key);
    let mut expired_commit_req = CommitRequest::new();
    expired_commit_req.set_context(ctx);
    expired_commit_req.set_start_version(start_ts);
    expired_commit_req.set_is_txn_file(true);
    expired_commit_req.set_commit_version(commit_ts);
    let expired_commit_resp = kv_client.kv_commit(&expired_commit_req).unwrap();
    assert!(
        !expired_commit_resp.has_region_error(),
        "resp {:?}",
        expired_commit_resp
    );
    assert!(
        expired_commit_resp.has_error(),
        "resp {:?}",
        expired_commit_resp
    );
    assert!(
        expired_commit_resp
            .get_error()
            .get_retryable()
            .contains("TxnLockNotFound"),
        "resp {:?}",
        expired_commit_resp
    );

    // Before issue/4809 is fixed, this retry will panic for BadFormatLock.
    let retry_commit_ts = client.get_ts().into_inner();
    assert!(
        retry_commit_ts > backup_ts,
        "retry commit_ts should be newer than backup_ts"
    );
    let (_, ctx, kv_client) = prepare_region_ctx(&mut client, &secondary_key);
    let mut retry_commit_req = CommitRequest::new();
    retry_commit_req.set_context(ctx);
    retry_commit_req.set_start_version(start_ts);
    retry_commit_req.set_is_txn_file(true);
    retry_commit_req.set_commit_version(retry_commit_ts);
    let resp = kv_client.kv_commit(&retry_commit_req).unwrap();
    assert!(
        resp.get_error().get_retryable().contains("TxnLockNotFound"),
        "resp {:?}",
        resp
    );

    cluster.stop();
}

#[test]
fn test_txn_file_lock_rejects_normal_rpc_then_commits() {
    // Prepare a keyspace and route all raw KV RPCs to the primary lock's region.
    let mut cluster = ServerCluster::new(alloc_node_id_vec(3), |_, _| {});
    cluster.wait_region_replicated(&[], 3);
    cluster.get_pd_client().disable_default_operator();
    let dfs = cluster.get_dfs().unwrap();

    let mut client = cluster.new_client();
    client.split_keyspace(KEYSPACE_ID);
    let gen_key = generate_keyspace_key(KEYSPACE_ID);
    let primary_lock = gen_key(0);
    cluster.wait_region_replicated(&primary_lock, 3);

    let start_ts = client.get_ts().into_inner();
    let chunk_ids = build_txn_files(&dfs, start_ts, 0, 300);

    // Prewrite txn-file locks backed by the generated txn-file chunks.
    prewrite_txn_file_chunk(
        &mut client,
        &primary_lock,
        &primary_lock,
        start_ts,
        &chunk_ids,
        6000,
    );
    let (_, ctx, kv_client) = prepare_region_ctx(&mut client, &primary_lock);

    // CheckTxnStatus without the txn-file marker must reject the live txn-file
    // lock.
    let mut req = CheckTxnStatusRequest::new();
    req.set_context(ctx.clone());
    req.set_primary_key(primary_lock.clone());
    req.set_lock_ts(start_ts);
    req.set_caller_start_ts(start_ts + 5);
    let resp = kv_client.kv_check_txn_status(&req).unwrap();
    assert!(
        resp.has_error() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );
    assert!(resp.get_error().get_abort().contains("LockTypeNotMatch"));

    // Commit without the txn-file marker must reject the live txn-file lock.
    let mut req = CommitRequest::new();
    req.set_context(ctx.clone());
    req.set_keys(vec![primary_lock.clone()].into());
    req.set_start_version(start_ts);
    req.set_commit_version(start_ts + 10);
    let resp = kv_client.kv_commit(&req).unwrap();
    assert!(
        resp.has_error() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );
    assert!(resp.get_error().get_abort().contains("LockTypeNotMatch"));

    // ResolveLock without the txn-file marker must also reject the txn-file lock.
    let mut req = ResolveLockRequest::new();
    req.set_context(ctx.clone());
    req.set_keys(vec![primary_lock.clone()].into());
    req.set_start_version(start_ts);
    req.set_commit_version(start_ts + 10);
    let resp = kv_client.kv_resolve_lock(&req).unwrap();
    assert!(
        resp.has_error() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );
    assert!(resp.get_error().get_abort().contains("LockTypeNotMatch"));

    // BatchRollback without the txn-file marker must fail closed through the
    // normal cleanup guard and leave the live txn-file lock intact.
    let mut req = BatchRollbackRequest::new();
    req.set_context(ctx.clone());
    req.set_keys(vec![primary_lock.clone()].into());
    req.set_start_version(start_ts);
    let resp = kv_client.kv_batch_rollback(&req).unwrap();
    assert!(
        resp.has_error() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );
    assert!(resp.get_error().has_locked(), "resp {:?}", resp);
    assert_eq!(resp.get_error().get_locked().lock_version, start_ts);

    // ResolveLock rollback without the txn-file marker takes the same cleanup
    // path as BatchRollback and must not write a normal rollback record.
    let mut req = ResolveLockRequest::new();
    req.set_context(ctx.clone());
    req.set_keys(vec![primary_lock.clone()].into());
    req.set_start_version(start_ts);
    let resp = kv_client.kv_resolve_lock(&req).unwrap();
    assert!(
        resp.has_error() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );
    assert!(resp.get_error().has_locked(), "resp {:?}", resp);
    assert_eq!(resp.get_error().get_locked().lock_version, start_ts);

    // Commit through the txn-file path and verify all txn-file values become
    // visible.
    let mut req = CommitRequest::new();
    req.set_context(ctx);
    req.set_start_version(start_ts);
    req.set_is_txn_file(true);
    let commit_ts = client.get_ts().into_inner();
    req.set_commit_version(commit_ts);
    let resp = kv_client.kv_commit(&req).unwrap();
    assert!(
        !resp.has_error() && !resp.has_region_error(),
        "resp {:?}",
        resp
    );

    verify_range(&mut client, 0, 300, commit_ts);
    cluster.stop();
}

fn make_txn_mutations<KeyF, OpF>(
    value: &str,
    start: usize,
    end: usize,
    gen_key: KeyF,
    gen_op: OpF,
    write_method: TxnWriteMethod,
    txn_file_helper: Option<Arc<TxnFileHelper>>,
) -> (Vec<Mutation>, TxnMutations)
where
    KeyF: Fn(usize) -> Vec<u8>,
    OpF: Fn(usize) -> kvrpcpb::Op,
{
    let gen_val = i_to_val_opt(value, 3);
    let mut mutations = vec![];
    for i in start..end {
        let mut m = Mutation::default();
        m.set_op(gen_op(i));
        m.set_key(gen_key(i));
        m.set_value(gen_val(i));
        mutations.push(m);
    }
    let txn_muts = block_on(TxnMutations::build(
        mutations.clone(),
        vec![],
        write_method,
        txn_file_helper,
    ))
    .unwrap();
    (mutations, txn_muts)
}

fn make_pessimistic_txn_mutations<KeyF, OpF>(
    value: &str,
    start: usize,
    end: usize,
    gen_key: KeyF,
    gen_op: OpF,
    write_method: TxnWriteMethod,
    txn_file_helper: Option<Arc<TxnFileHelper>>,
) -> (Vec<Mutation>, TxnMutations)
where
    KeyF: Fn(usize) -> Vec<u8>,
    OpF: Fn(usize) -> kvrpcpb::Op,
{
    let gen_val = i_to_val_opt(value, 3);
    let mut mutations = vec![];
    let mut pessimistic_actions = vec![];
    for i in start..end {
        let mut m = Mutation::default();
        let op = gen_op(i);
        m.set_op(op);
        m.set_key(gen_key(i));
        m.set_value(gen_val(i));
        mutations.push(m);
        let action = if op == kvrpcpb::Op::Insert {
            PrewriteRequestPessimisticAction::DoConstraintCheck
        } else {
            PrewriteRequestPessimisticAction::DoPessimisticCheck
        };
        pessimistic_actions.push(action);
    }
    let txn_muts = block_on(TxnMutations::build(
        mutations.clone(),
        pessimistic_actions,
        write_method,
        txn_file_helper,
    ))
    .unwrap();
    (mutations, txn_muts)
}

fn prepare_region_ctx(
    client: &mut ClusterClient,
    key: &[u8],
) -> (u64, kvrpcpb::Context, TikvClient) {
    let region_id = client.get_region_id(key);
    let ctx = client.new_rpc_ctx(region_id, key).unwrap();
    let kv_client = client.get_kv_client(ctx.get_peer().get_store_id());
    (region_id, ctx, kv_client)
}

fn build_txn_file_prewrite_req(
    ctx: kvrpcpb::Context,
    primary_lock: Vec<u8>,
    start_ts: u64,
    chunk_ids: Vec<u64>,
) -> PrewriteRequest {
    let mut req = PrewriteRequest::new();
    req.set_context(ctx);
    req.set_txn_file_chunks(chunk_ids);
    req.set_primary_lock(primary_lock);
    req.set_lock_ttl(6000);
    req.set_start_version(start_ts);
    req.set_min_commit_ts(start_ts + 1);
    req
}

fn prewrite_txn_file_chunk(
    client: &mut ClusterClient,
    request_key: &[u8],
    primary_lock: &[u8],
    start_ts: u64,
    chunk_ids: &[u64],
    lock_ttl: u64,
) {
    let ok = try_wait(
        || {
            let (_, ctx, kv_client) = prepare_region_ctx(client, request_key);
            let mut req = build_txn_file_prewrite_req(
                ctx,
                primary_lock.to_vec(),
                start_ts,
                chunk_ids.to_vec(),
            );
            req.set_lock_ttl(lock_ttl);
            let resp = kv_client.kv_prewrite(&req).unwrap();
            assert!(resp.get_errors().is_empty(), "resp {:?}", resp);
            if resp.has_region_error() {
                warn!("prewrite got region error {:?}", resp.get_region_error());
                client.clear_region_cache();
                false
            } else {
                true
            }
        },
        10,
    );
    assert!(ok, "prewrite failed for key {:?}", request_key);
}

fn update_txn_file_lock_limit(status_addr: String, limit: usize) {
    let security_mgr = Arc::new(SecurityManager::new(&SecurityConfig::default()).unwrap());
    let rt = tokio::runtime::Runtime::new().unwrap();
    let change = std::collections::HashMap::from([(
        "storage.max-txn-file-locks-per-region".to_owned(),
        limit.to_string(),
    )]);
    rt.block_on(async move {
        RestfulClient::new("txn-file-config", vec![status_addr], security_mgr)
            .unwrap()
            .request(
                "config",
                Method::POST,
                Some(Bytes::from(serde_json::to_vec(&change).unwrap())),
            )
            .await
            .unwrap();
    });
}

fn request_lightweight_backup_to_all_stores(
    cluster: &ServerCluster,
    node_ids: &[u16],
    cluster_id: u64,
    backup_ts: u64,
) {
    let body = serde_json::to_vec(&rfengine::BackupConfig {
        cluster_id,
        lightweight: true,
        backup_ts: Some(backup_ts),
        ..Default::default()
    })
    .unwrap();
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    let client = hyper::Client::new();

    for &node_id in node_ids {
        let client = client.clone();
        let request_body = body.clone();
        let (status, resp_body) = rt.block_on(async move {
            let req = hyper::Request::builder()
                .method(Method::POST)
                .uri(format!(
                    "http://{}/rfengine/backup",
                    cluster.status_addr(node_id)
                ))
                .body(hyper::Body::from(request_body))
                .unwrap();
            let resp = client.request(req).await.unwrap();
            let status = resp.status();
            let body = hyper::body::to_bytes(resp.into_body()).await.unwrap();
            (status, body)
        });
        assert!(
            status.is_success(),
            "backup request failed on node {}: status {:?}, body {}",
            node_id,
            status,
            String::from_utf8_lossy(&resp_body)
        );
    }
}

fn build_txn_files(dfs: &Arc<dyn Dfs>, start_ts: u64, start: usize, end: usize) -> Vec<u64> {
    build_txn_files_with_first_chunk_id(dfs, start_ts + 1, start, end)
}

fn build_txn_files_with_key_fn<F>(
    dfs: &Arc<dyn Dfs>,
    start_ts: u64,
    start: usize,
    end: usize,
    key_fn: F,
) -> Vec<u64>
where
    F: Fn(usize) -> Vec<u8>,
{
    build_txn_files_with_chunk_id_and_key_fn(dfs, start_ts + 1, start, end, key_fn, |_| {
        PessimisticAction::SkipCheck
    })
}

fn build_txn_files_with_first_chunk_id(
    dfs: &Arc<dyn Dfs>,
    chunk_id: u64,
    start: usize,
    end: usize,
) -> Vec<u64> {
    build_txn_files_with_pessimistic_action(dfs, chunk_id, start, end, |_| {
        PessimisticAction::SkipCheck
    })
}

fn build_txn_files_with_pessimistic_action(
    dfs: &Arc<dyn Dfs>,
    chunk_id: u64,
    start: usize,
    end: usize,
    pessimistic_actions_fn: impl FnMut(usize) -> PessimisticAction,
) -> Vec<u64> {
    build_txn_files_with_chunk_id_and_key_fn(
        dfs,
        chunk_id,
        start,
        end,
        i_to_tidb_key,
        pessimistic_actions_fn,
    )
}

fn build_txn_files_with_chunk_id_and_key_fn<F>(
    dfs: &Arc<dyn Dfs>,
    mut chunk_id: u64,
    start: usize,
    end: usize,
    key_fn: F,
    mut pessimistic_actions_fn: impl FnMut(usize) -> PessimisticAction,
) -> Vec<u64>
where
    F: Fn(usize) -> Vec<u8>,
{
    let mut chunk_ids = vec![];
    let mut txn_chunk_builder = TxnChunkBuilder::new(chunk_id, 64, None, TXN_FILE_FORMAT_V2);
    for i in start..end {
        let key = key_fn(i);
        let val = i_to_val(i);
        let pessimistic_action = pessimistic_actions_fn(i);
        txn_chunk_builder.add_entry(
            InnerKey::from_outer_key(&key),
            OP_PUT,
            pessimistic_action,
            &val,
        );
        if (i + 1) % 100 == 0 {
            let mut data_buf = vec![];
            txn_chunk_builder.finish(&mut data_buf);
            txn_chunk_builder = TxnChunkBuilder::new(chunk_id, 64, None, TXN_FILE_FORMAT_V2);
            let opts = dfs::Options::default().with_type(FileType::TxnChunk);
            dfs.get_runtime()
                .block_on(dfs.create(chunk_id, Bytes::from(data_buf), opts))
                .unwrap();
            chunk_ids.push(chunk_id);
            chunk_id += 1;
        }
    }
    chunk_ids
}

fn verify_range(client: &mut ClusterClient, start: usize, end: usize, version: u64) {
    let put_time = Instant::now();
    let gen_key = generate_keyspace_key(KEYSPACE_ID);
    for i in start..end {
        let key = gen_key(i);
        let val = i_to_val(i);
        let opt = RequestOptions::default();
        client
            .verify_key_value(&key, Some(&val), version, put_time, &opt)
            .unwrap();
    }
}

fn expect_err_msg(err: &client::Error, msg: &str) {
    let err_str = err.to_string();
    assert!(
        err_str.contains(msg),
        "err: {:?}({}), expect: {}",
        err,
        err_str,
        msg
    );
}

#[derive(Clone, Copy, Debug)]
enum ActiveResolveSplitMethod {
    TikvRpc,
    Pd,
}

fn assert_txn_file_lock_blocks_split(
    split_method: ActiveResolveSplitMethod,
    client: &mut ClusterClient,
    split_key: &[u8],
    blocked_reason: &str,
) -> usize {
    let initial_region_count = client.pd_client().get_regions_number();
    match split_method {
        ActiveResolveSplitMethod::TikvRpc => {
            let err = client.try_split(split_key, 1).unwrap_err();
            match err {
                client::Error::KeyErrors(key_errs) => {
                    assert_eq!(key_errs.len(), 1);
                    let lock = key_errs[0].get_locked();
                    assert_eq!(lock.key, split_key);
                    assert!(lock.is_txn_file);
                }
                _ => panic!("unexpected error: {:?}", err),
            }

            let split_succeeded_while_disabled = try_wait(
                || match client.try_split(split_key, 1) {
                    Ok(_) => true,
                    Err(client::Error::KeyErrors(_)) => false,
                    Err(err) => panic!("split failed: {:?}", err),
                },
                2,
            );
            assert!(
                !split_succeeded_while_disabled,
                "split succeeded while {}",
                blocked_reason
            );
        }
        ActiveResolveSplitMethod::Pd => {
            client.try_split_by_pd(split_key, 1).unwrap_err();
            assert_eq!(
                client.pd_client().get_regions_number(),
                initial_region_count,
                "pd region count changed while {}",
                blocked_reason
            );

            let split_succeeded_while_disabled =
                try_wait(|| client.try_split_by_pd(split_key, 1).is_ok(), 2);
            assert!(
                !split_succeeded_while_disabled,
                "pd-driven split succeeded while {}",
                blocked_reason
            );
            assert_eq!(
                client.pd_client().get_regions_number(),
                initial_region_count,
                "pd region count changed while {}",
                blocked_reason
            );
        }
    }
    initial_region_count
}

fn assert_txn_file_lock_allows_split(
    split_method: ActiveResolveSplitMethod,
    cluster: &ServerCluster,
    client: &mut ClusterClient,
    split_key: &[u8],
    initial_region_count: usize,
    unlocked_reason: &str,
) {
    match split_method {
        ActiveResolveSplitMethod::TikvRpc => {
            let ok = try_wait(
                || match client.try_split(split_key, 5) {
                    Ok(_) => true,
                    Err(client::Error::KeyErrors(_)) => false,
                    Err(err) => panic!("split failed: {:?}", err),
                },
                10,
            );
            assert!(ok, "split key failed after {}", unlocked_reason);
        }
        ActiveResolveSplitMethod::Pd => {
            let ok = try_wait(|| client.try_split_by_pd(split_key, 5).is_ok(), 10);
            assert!(ok, "pd-driven split failed after {}", unlocked_reason);
            cluster.wait_pd_region_count(initial_region_count + 1);
            assert_eq!(
                client.pd_client().get_regions_number(),
                initial_region_count + 1
            );
        }
    }
}