// Copyright 2025 PingCAP, Inc.

use std::{
    sync::mpsc::{self, TryRecvError},
    thread,
    time::Duration,
};

use kvproto::{
    deadlock::WaitForEntry,
    kvrpcpb::{ApiVersion, Assertion, Context, GetLockWaitInfoRequest, Op, WriteConflictReason},
};
use test_cloud_server::{
    ServerCluster,
    client::{MutateOptions, PrewriteExt, TxnMutations},
    util::Mutation,
};
use txn_types::{LockType, NULL_KEYSPACE_ID, WriteType};

use super::{helper::*, util::*};
use crate::{alloc_node_id_vec, i_to_key, i_to_key_with_keyspace, i_to_val};

fn get_lock_wait_info_entries(
    client: &test_cloud_server::client::ClusterClient,
    store_id: u64,
    api_version: ApiVersion,
    keyspace_id: u32,
) -> Vec<WaitForEntry> {
    let kv_client = client.get_kv_client(store_id);
    let mut ctx = Context::default();
    ctx.set_api_version(api_version);
    ctx.set_keyspace_id(keyspace_id);

    let mut req = GetLockWaitInfoRequest::default();
    req.set_context(ctx);

    kv_client
        .get_lock_wait_info(&req)
        .unwrap()
        .take_entries()
        .into_vec()
}

fn wait_for_lock_wait_info_entry(
    client: &test_cloud_server::client::ClusterClient,
    store_id: u64,
    api_version: ApiVersion,
    keyspace_id: u32,
    expected_key: &[u8],
) -> Vec<WaitForEntry> {
    let mut entries = Vec::new();
    for _ in 0..20 {
        entries = get_lock_wait_info_entries(client, store_id, api_version, keyspace_id);
        if entries.len() == 1 && entries[0].key == expected_key {
            return entries;
        }
        thread::sleep(Duration::from_millis(100));
    }
    panic!(
        "wait for lock wait info timeout, api_version {:?}, keyspace {}, entries {:?}",
        api_version, keyspace_id, entries
    );
}

fn acquire_shared_secondary(
    client: &mut test_cloud_server::client::ClusterClient,
    primary: &[u8],
    shared_key: &[u8],
    start_ts: txn_types::TimeStamp,
    for_update_ts: txn_types::TimeStamp,
) {
    let (_, primary_errors) = client
        .kv_pessimistic_lock(
            primary.to_vec().into(),
            vec![primary.to_vec().into()],
            start_ts,
            3000,
            for_update_ts,
            None,
        )
        .unwrap();
    assert!(primary_errors.is_empty(), "{:?}", primary_errors);

    let (_, shared_errors) = client
        .kv_shared_pessimistic_lock(
            primary.to_vec().into(),
            vec![shared_key.to_vec().into()],
            start_ts,
            3000,
            for_update_ts,
            None,
        )
        .unwrap();
    assert!(shared_errors.is_empty(), "{:?}", shared_errors);
}

#[test]
fn test_pes_acquire_lock_single_new_key_then_commit() {
    for pipelined in [true, false] {
        let mut cluster = ServerCluster::new(alloc_node_id_vec(1), |_, cfg| {
            cfg.pessimistic_txn.pipelined = pipelined;
        });
        cluster.wait_region_replicated(&[], 1);
        let mut client = cluster.new_client();
        let keyspace_id = api_version::ApiV2::get_u32_keyspace_id_by_key(&i_to_key(0)).unwrap();
        client.split_keyspace(keyspace_id);

        // Arrange
        let k1 = i_to_key(1);
        let v1 = i_to_val(1);

        // Act: T1 acquires pessimistic lock on K1
        let start_ts = client.get_ts();
        let for_update_ts = client.get_ts();
        let lock_ttl = 3000;

        let (_, key_errors) = client
            .kv_pessimistic_lock(
                k1.clone().into(),
                vec![k1.clone().into()],
                start_ts,
                lock_ttl,
                for_update_ts,
                None,
            )
            .expect("T1 pessimistic lock should succeed");

        assert!(key_errors.is_empty(), "No key errors expected");

        // Assert: K1 is locked by T1 with pessimistic lock type

        if pipelined {
            thread::sleep(Duration::from_millis(100));
        }
        let _ = must_locked_with_properties(&mut client, &k1, &k1, start_ts, LockType::Pessimistic);

        // prewrite the value and commit the transaction
        let commit_mutations = vec![new_put_mutation(k1.clone(), v1.clone())];
        let commit_ts = client
            .prewrite_and_commit(
                commit_mutations,
                MutateOptions::start_ts(start_ts),
                PrewriteExt::with_for_update_ts(start_ts),
            )
            .expect("T1 commit failed");

        // Assert: K1 is unlocked after commit
        must_unlocked(&mut client, &k1);

        must_written_with_properties(&mut client, &k1, start_ts, commit_ts, WriteType::Put);

        // Assert: T2 can read V1 from K1
        let t2_ts = client.get_ts();
        let val_read_by_t2 = client
            .get_key_version(&k1, t2_ts.into_inner())
            .expect("T2 Get failed")
            .expect("Value not found for K1 by T2");

        assert_eq!(val_read_by_t2, v1, "T2 Get(K1) should return V1");

        cluster.stop();
    }
}

#[test]
fn test_pes_acquire_lock_key_locked_wait_and_write_conflict_after_release() {
    for pipelined in [true, false] {
        let mut cluster = ServerCluster::new(alloc_node_id_vec(1), |_, cfg| {
            cfg.pessimistic_txn.pipelined = pipelined;
        });
        cluster.wait_region_replicated(&[], 1);
        let mut client = cluster.new_client();
        let keyspace_id = api_version::ApiV2::get_u32_keyspace_id_by_key(&i_to_key(0)).unwrap();
        client.split_keyspace(keyspace_id);

        // Arrange: Key K1
        let k1 = i_to_key(2);

        // T_other acquires a pessimistic lock on K1
        let start_ts_other = client.get_ts();
        let for_update_ts_other = client.get_ts();
        let lock_ttl = 10000;

        client
            .kv_pessimistic_lock(
                k1.clone().into(),
                vec![k1.clone().into()],
                start_ts_other,
                lock_ttl,
                for_update_ts_other,
                None,
            )
            .expect("T_other pessimistic lock should succeed");

        // Verify K1 is locked by T_other
        if pipelined {
            thread::sleep(Duration::from_millis(100));
        }
        let _ = must_locked_with_properties(
            &mut client,
            &k1,
            &k1,
            start_ts_other,
            LockType::Pessimistic,
        );

        // Create a cloned client for T1 in a separate thread
        let client_clone = client.clone();
        let k1_clone = k1.clone();

        // Start T1 that will try to acquire the same lock and wait
        let t1_handle = thread::spawn(move || {
            let mut t1_client = client_clone;
            let start_ts_t1 = t1_client.get_ts();
            let for_update_ts_t1 = t1_client.get_ts();

            // T1 will wait up to 5 seconds for the lock
            let wait_timeout = Some(5000);

            let result = t1_client.kv_pessimistic_lock(
                k1_clone.clone().into(),
                vec![k1_clone.clone().into()],
                start_ts_t1,
                lock_ttl,
                for_update_ts_t1,
                wait_timeout,
            );

            (t1_client, start_ts_t1, result)
        });

        // Give T1 some time to start waiting for the lock
        thread::sleep(Duration::from_millis(500));

        // T_other releases its lock (e.g., by rollback)
        client
            .kv_rollback(
                TxnMutations::from_normal(vec![Mutation {
                    op: Op::PessimisticLock,
                    key: k1.clone().into(),
                    value: vec![].into(),
                    ..Default::default()
                }]),
                start_ts_other,
            )
            .expect("T_other rollback should succeed");

        // Wait for T1 to complete and get its result
        let (_, _, t1_result) = t1_handle.join().unwrap();

        // Assert: T1's lock acquisition returns write conflict error
        let (_, key_errors) =
            t1_result.expect("T1 pessimistic lock should return write conflict error");
        assert!(
            key_errors[0].has_conflict(),
            "T1 pessimistic lock should return write conflict error"
        );
        cluster.stop();
    }
}

#[test]
fn test_pes_acquire_lock_key_locked_wait_and_timeout() {
    for pipelined in [true, false] {
        let mut cluster = ServerCluster::new(alloc_node_id_vec(1), |_, cfg| {
            cfg.pessimistic_txn.pipelined = pipelined;
        });
        cluster.wait_region_replicated(&[], 1);
        let mut client = cluster.new_client();
        let keyspace_id = api_version::ApiV2::get_u32_keyspace_id_by_key(&i_to_key(0)).unwrap();
        client.split_keyspace(keyspace_id);

        // Arrange: Key K1
        let k1 = i_to_key(3);

        // T_other acquires a pessimistic lock on K1
        let start_ts_other = client.get_ts();
        let for_update_ts_other = client.get_ts();
        let lock_ttl = 10000;

        client
            .kv_pessimistic_lock(
                k1.clone().into(),
                vec![k1.clone().into()],
                start_ts_other,
                lock_ttl,
                for_update_ts_other,
                None,
            )
            .expect("T_other pessimistic lock should succeed");

        // Verify K1 is locked by T_other

        if pipelined {
            thread::sleep(Duration::from_millis(100));
        }
        let _ = must_locked_with_properties(
            &mut client,
            &k1,
            &k1,
            start_ts_other,
            LockType::Pessimistic,
        );

        // T1 tries to acquire the same lock with a short timeout
        let start_ts_t1 = client.get_ts();
        let for_update_ts_t1 = client.get_ts();

        // T1 will only wait 100ms for the lock
        let wait_timeout = Some(100);

        let (_lock_result, errors) = client
            .kv_pessimistic_lock(
                k1.clone().into(),
                vec![k1.clone().into()],
                start_ts_t1,
                lock_ttl,
                for_update_ts_t1,
                wait_timeout,
            )
            .unwrap();

        // Assert: T1's lock acquisition failed with timeout
        assert!(
            errors.len() == 1,
            "T1 pessimistic lock should return one error"
        );
        assert!(
            errors[0].has_locked(),
            "T1 pessimistic lock should return locked error, but got {:?}",
            errors[0]
        );

        // Assert: K1 is still locked by T_other
        let _ = must_locked_with_properties(
            &mut client,
            &k1,
            &k1,
            start_ts_other,
            LockType::Pessimistic,
        );

        cluster.stop();
    }
}

#[test]
fn test_shared_lock_lifecycle_with_multiple_holders() {
    test_util::init_log_for_test();
    for pipelined in [true, false] {
        let mut cluster = ServerCluster::new(alloc_node_id_vec(1), |_, cfg| {
            cfg.pessimistic_txn.pipelined = pipelined;
        });
        cluster.wait_region_replicated(&[], 1);
        let mut client = cluster.new_client();
        let keyspace_id = api_version::ApiV2::get_u32_keyspace_id_by_key(&i_to_key(0)).unwrap();
        client.split_keyspace(keyspace_id);

        let key = i_to_key(20);
        let pk1 = i_to_key(21);
        let pk2 = i_to_key(22);

        let start_ts_1 = client.get_ts();
        let for_update_ts_1 = client.get_ts();
        client
            .kv_pessimistic_lock(
                pk1.clone().into(),
                vec![pk1.clone().into()],
                start_ts_1,
                3000,
                for_update_ts_1,
                None,
            )
            .unwrap();
        client
            .kv_shared_pessimistic_lock(
                pk1.clone().into(),
                vec![key.clone().into()],
                start_ts_1,
                3000,
                for_update_ts_1,
                None,
            )
            .unwrap();

        let start_ts_2 = client.get_ts();
        let for_update_ts_2 = client.get_ts();
        client
            .kv_pessimistic_lock(
                pk2.clone().into(),
                vec![pk2.clone().into()],
                start_ts_2,
                3000,
                for_update_ts_2,
                None,
            )
            .unwrap();
        client
            .kv_shared_pessimistic_lock(
                pk2.clone().into(),
                vec![key.clone().into()],
                start_ts_2,
                3000,
                for_update_ts_2,
                None,
            )
            .unwrap();

        if pipelined {
            thread::sleep(Duration::from_millis(100));
        }

        let mut txn_1 = client.begin_transaction(Some(start_ts_1));
        client
            .kv_prewrite_ext(
                pk1.clone().into(),
                Some(&vec![key.clone().into()]),
                TxnMutations::from_normal(vec![
                    new_lock_mutation(pk1.clone()),
                    new_shared_lock_mutation(key.clone()),
                ]),
                &mut txn_1,
                PrewriteExt {
                    for_update_ts: for_update_ts_1,
                    ..Default::default()
                },
            )
            .unwrap();
        let mut txn_2 = client.begin_transaction(Some(start_ts_2));
        client
            .kv_prewrite_ext(
                pk2.clone().into(),
                Some(&vec![key.clone().into()]),
                TxnMutations::from_normal(vec![
                    new_lock_mutation(pk2.clone()),
                    new_shared_lock_mutation(key.clone()),
                ]),
                &mut txn_2,
                PrewriteExt {
                    for_update_ts: for_update_ts_2,
                    ..Default::default()
                },
            )
            .unwrap();

        let commit_ts_1 = client.get_ts();
        client
            .kv_commit(
                TxnMutations::from_normal(vec![
                    new_lock_mutation(pk1.clone()),
                    new_shared_lock_mutation(key.clone()),
                ]),
                start_ts_1,
                commit_ts_1,
                false,
            )
            .unwrap();

        let start_ts_3 = client.get_ts();
        let for_update_ts_3 = client.get_ts();
        let (_results, key_errors) = client
            .kv_pessimistic_lock(
                key.clone().into(),
                vec![key.clone().into()],
                start_ts_3,
                3000,
                for_update_ts_3,
                Some(1),
            )
            .unwrap();
        assert_eq!(key_errors.len(), 1);
        assert!(key_errors[0].has_locked(), "expected shared lock conflict");

        let commit_ts_2 = client.get_ts();
        client
            .kv_commit(
                TxnMutations::from_normal(vec![
                    new_lock_mutation(pk2.clone()),
                    new_shared_lock_mutation(key.clone()),
                ]),
                start_ts_2,
                commit_ts_2,
                false,
            )
            .unwrap();

        let start_ts_4 = client.get_ts();
        let for_update_ts_4 = client.get_ts();
        let (_results, key_errors) = client
            .kv_pessimistic_lock(
                key.clone().into(),
                vec![key.clone().into()],
                start_ts_4,
                3000,
                for_update_ts_4,
                None,
            )
            .unwrap();
        assert!(
            key_errors.is_empty(),
            "key should be lockable after both shared holders commit"
        );

        must_written_with_properties(&mut client, &key, start_ts_1, commit_ts_1, WriteType::Lock);
        must_written_with_properties(&mut client, &key, start_ts_2, commit_ts_2, WriteType::Lock);
        cluster.stop();
    }
}

#[test]
fn test_shared_lock_survives_restart() {
    test_util::init_log_for_test();
    for pipelined in [true, false] {
        let node_ids = alloc_node_id_vec(1);
        let node_id = node_ids[0];
        let mut cluster = ServerCluster::new(node_ids, |_, cfg| {
            cfg.pessimistic_txn.pipelined = pipelined;
        });
        cluster.wait_region_replicated(&[], 1);
        let mut client = cluster.new_client();
        let keyspace_id = api_version::ApiV2::get_u32_keyspace_id_by_key(&i_to_key(0)).unwrap();
        client.split_keyspace(keyspace_id);

        let key = i_to_key(24);
        let pk1 = i_to_key(25);
        let pk2 = i_to_key(26);

        let start_ts_1 = client.get_ts();
        let for_update_ts_1 = client.get_ts();
        client
            .kv_pessimistic_lock(
                pk1.clone().into(),
                vec![pk1.clone().into()],
                start_ts_1,
                3000,
                for_update_ts_1,
                None,
            )
            .unwrap();
        client
            .kv_shared_pessimistic_lock(
                pk1.clone().into(),
                vec![key.clone().into()],
                start_ts_1,
                3000,
                for_update_ts_1,
                None,
            )
            .unwrap();

        let start_ts_2 = client.get_ts();
        let for_update_ts_2 = client.get_ts();
        client
            .kv_pessimistic_lock(
                pk2.clone().into(),
                vec![pk2.clone().into()],
                start_ts_2,
                3000,
                for_update_ts_2,
                None,
            )
            .unwrap();
        client
            .kv_shared_pessimistic_lock(
                pk2.clone().into(),
                vec![key.clone().into()],
                start_ts_2,
                3000,
                for_update_ts_2,
                None,
            )
            .unwrap();

        if pipelined {
            thread::sleep(Duration::from_millis(100));
        }

        let mut txn_1 = client.begin_transaction(Some(start_ts_1));
        client
            .kv_prewrite_ext(
                pk1.clone().into(),
                Some(&vec![key.clone().into()]),
                TxnMutations::from_normal(vec![
                    new_lock_mutation(pk1.clone()),
                    new_shared_lock_mutation(key.clone()),
                ]),
                &mut txn_1,
                PrewriteExt {
                    for_update_ts: for_update_ts_1,
                    ..Default::default()
                },
            )
            .unwrap();
        let mut txn_2 = client.begin_transaction(Some(start_ts_2));
        client
            .kv_prewrite_ext(
                pk2.clone().into(),
                Some(&vec![key.clone().into()]),
                TxnMutations::from_normal(vec![
                    new_lock_mutation(pk2.clone()),
                    new_shared_lock_mutation(key.clone()),
                ]),
                &mut txn_2,
                PrewriteExt {
                    for_update_ts: for_update_ts_2,
                    ..Default::default()
                },
            )
            .unwrap();

        let commit_ts_1 = client.get_ts();
        client
            .kv_commit(
                TxnMutations::from_normal(vec![
                    new_lock_mutation(pk1.clone()),
                    new_shared_lock_mutation(key.clone()),
                ]),
                start_ts_1,
                commit_ts_1,
                false,
            )
            .unwrap();

        cluster.restart_node(node_id, Duration::from_secs(0), false, |_, _| {});
        thread::sleep(Duration::from_secs(1));
        cluster.wait_region_replicated(&key, 1);

        let mut restarted_client = cluster.new_client();
        let start_ts_3 = restarted_client.get_ts();
        let for_update_ts_3 = restarted_client.get_ts();
        let (_results, key_errors) = restarted_client
            .kv_pessimistic_lock(
                key.clone().into(),
                vec![key.clone().into()],
                start_ts_3,
                3000,
                for_update_ts_3,
                Some(1),
            )
            .unwrap();
        assert_eq!(key_errors.len(), 1);
        assert!(key_errors[0].has_locked(), "expected persisted shared lock");

        let commit_ts_2 = restarted_client.get_ts();
        restarted_client
            .kv_commit(
                TxnMutations::from_normal(vec![
                    new_lock_mutation(pk2.clone()),
                    new_shared_lock_mutation(key.clone()),
                ]),
                start_ts_2,
                commit_ts_2,
                false,
            )
            .unwrap();

        let start_ts_4 = restarted_client.get_ts();
        let for_update_ts_4 = restarted_client.get_ts();
        let (_results, key_errors) = restarted_client
            .kv_pessimistic_lock(
                key.clone().into(),
                vec![key.clone().into()],
                start_ts_4,
                3000,
                for_update_ts_4,
                None,
            )
            .unwrap();
        assert!(
            key_errors.is_empty(),
            "key should be lockable after restart and final shared-holder commit"
        );

        must_written_with_properties(
            &mut restarted_client,
            &key,
            start_ts_1,
            commit_ts_1,
            WriteType::Lock,
        );
        must_written_with_properties(
            &mut restarted_client,
            &key,
            start_ts_2,
            commit_ts_2,
            WriteType::Lock,
        );
        cluster.stop();
    }
}

#[test]
fn test_shared_and_exclusive_pessimistic_locks_conflict() {
    test_util::init_log_for_test();
    let mut cluster = ServerCluster::new(alloc_node_id_vec(1), |_, _| {});
    cluster.wait_region_replicated(&[], 1);
    let mut client = cluster.new_client();
    let keyspace_id = api_version::ApiV2::get_u32_keyspace_id_by_key(&i_to_key(0)).unwrap();
    client.split_keyspace(keyspace_id);

    let key = i_to_key(30);
    let pk = i_to_key(31);

    let start_ts_1 = client.get_ts();
    let for_update_ts_1 = client.get_ts();
    client
        .kv_pessimistic_lock(
            pk.clone().into(),
            vec![key.clone().into()],
            start_ts_1,
            3000,
            for_update_ts_1,
            None,
        )
        .unwrap();

    let start_ts_2 = client.get_ts();
    let for_update_ts_2 = client.get_ts();
    let (_results, key_errors) = client
        .kv_shared_pessimistic_lock(
            pk.clone().into(),
            vec![key.clone().into()],
            start_ts_2,
            3000,
            for_update_ts_2,
            Some(1),
        )
        .unwrap();
    assert_eq!(key_errors.len(), 1);
    assert!(
        key_errors[0].has_locked(),
        "shared request should be blocked by exclusive lock"
    );

    client
        .kv_rollback(
            TxnMutations::from_normal(vec![new_pessimistic_lock_mutation(key.clone())]),
            start_ts_1,
        )
        .unwrap();

    let start_ts_2 = client.get_ts();
    let for_update_ts_2 = client.get_ts();
    client
        .kv_shared_pessimistic_lock(
            pk.clone().into(),
            vec![key.clone().into()],
            start_ts_2,
            3000,
            for_update_ts_2,
            None,
        )
        .unwrap();

    let start_ts_3 = client.get_ts();
    let for_update_ts_3 = client.get_ts();
    let (_results, key_errors) = client
        .kv_pessimistic_lock(
            key.clone().into(),
            vec![key.clone().into()],
            start_ts_3,
            3000,
            for_update_ts_3,
            Some(1),
        )
        .unwrap();
    assert_eq!(key_errors.len(), 1);
    assert!(
        key_errors[0].has_locked(),
        "exclusive request should be blocked by shared lock"
    );
    assert!(
        !key_errors[0].has_conflict(),
        "exclusive request should not retry instead of waiting: {:?}",
        key_errors[0]
    );

    cluster.stop();
}

#[test]
fn test_external_exclusive_nowait_on_shared_lock_returns_locked() {
    test_util::init_log_for_test();
    let mut cluster = ServerCluster::new(alloc_node_id_vec(1), |_, _| {});
    cluster.wait_region_replicated(&[], 1);
    let mut client = cluster.new_client();
    let keyspace_id = api_version::ApiV2::get_u32_keyspace_id_by_key(&i_to_key(0)).unwrap();
    client.split_keyspace(keyspace_id);

    let key = i_to_key(32);
    let primary = i_to_key(33);
    let shared_start_ts = client.get_ts();
    let shared_for_update_ts = client.get_ts();
    client
        .kv_shared_pessimistic_lock(
            primary.clone().into(),
            vec![key.clone().into()],
            shared_start_ts,
            3000,
            shared_for_update_ts,
            None,
        )
        .unwrap();

    let start_ts = client.get_ts();
    let for_update_ts = client.get_ts();
    let (_, key_errors) = client
        .kv_pessimistic_lock(
            key.clone().into(),
            vec![key.clone().into()],
            start_ts,
            3000,
            for_update_ts,
            Some(-1),
        )
        .unwrap();
    assert_eq!(key_errors.len(), 1);
    assert!(
        key_errors[0].has_locked(),
        "NOWAIT exclusive request should be blocked by shared lock: {:?}",
        key_errors[0]
    );
    assert!(
        !key_errors[0].has_conflict(),
        "NOWAIT exclusive request should not retry instead of waiting: {:?}",
        key_errors[0]
    );

    cluster.stop();
}

#[test]
fn test_external_exclusive_waits_for_shared_holders_then_retries_to_acquire_lock() {
    test_util::init_log_for_test();
    let mut cluster = ServerCluster::new(alloc_node_id_vec(1), |_, cfg| {
        // Keep this protocol test synchronous. Pipelined outcome handling has
        // separate failpoint coverage, while this test focuses on the two RPC
        // phases required before an external exclusive waiter can really wait.
        cfg.pessimistic_txn.pipelined = false;
    });
    cluster.wait_region_replicated(&[], 1);
    let mut client = cluster.new_client();
    let keyspace_id = api_version::ApiV2::get_u32_keyspace_id_by_key(&i_to_key(0)).unwrap();
    client.split_keyspace(keyspace_id);

    let shared_key = i_to_key(34);
    let primary_1 = i_to_key(35);
    let primary_2 = i_to_key(36);
    let primary_x = i_to_key(37);
    let start_1 = client.get_ts();
    let for_update_1 = client.get_ts();
    let start_2 = client.get_ts();
    let for_update_2 = client.get_ts();

    // Create two independent shared holders with real primary locks so holder
    // rollback follows the same protocol as a client transaction.
    acquire_shared_secondary(&mut client, &primary_1, &shared_key, start_1, for_update_1);
    acquire_shared_secondary(&mut client, &primary_2, &shared_key, start_2, for_update_2);

    let start_x = client.get_ts();
    let for_update_x = client.get_ts();
    let (_, primary_x_errors) = client
        .kv_pessimistic_lock(
            primary_x.clone().into(),
            vec![primary_x.clone().into()],
            start_x,
            3000,
            for_update_x,
            None,
        )
        .unwrap();
    assert!(primary_x_errors.is_empty(), "{:?}", primary_x_errors);

    // The first external-X RPC does not enter lock wait. It returns the
    // original shared conflict after durably changing the wrapper to
    // shrink-only, which makes the following RPC eligible for ordinary wait.
    let (_, first_errors) = client
        .kv_pessimistic_lock(
            primary_x.clone().into(),
            vec![shared_key.clone().into()],
            start_x,
            3000,
            for_update_x,
            Some(30_000),
        )
        .unwrap();
    assert_eq!(first_errors.len(), 1, "{first_errors:?}");
    assert!(first_errors[0].has_locked(), "{:?}", first_errors);
    assert!(
        get_lock_wait_info_entries(&client, client.get_stores()[0], ApiVersion::V2, keyspace_id)
            .is_empty(),
        "the shrink RPC must not install the external waiter"
    );

    // The second RPC observes shrink-only and must reach a real lock-wait
    // entry. Its result is delivered through a bounded channel so holder
    // releases can be checked without relying on scheduling sleeps.
    let mut wait_client = client.clone();
    let wait_key = shared_key.clone();
    let wait_primary = primary_x.clone();
    let (wait_tx, wait_rx) = mpsc::sync_channel(1);
    let wait_task = thread::spawn(move || {
        let result = wait_client.kv_pessimistic_lock(
            wait_primary.into(),
            vec![wait_key.into()],
            start_x,
            3000,
            for_update_x,
            Some(30_000),
        );
        wait_tx.send(result).unwrap();
    });
    let store_id = client.get_stores()[0];
    let entries =
        wait_for_lock_wait_info_entry(&client, store_id, ApiVersion::V2, keyspace_id, &shared_key);
    assert_eq!(entries[0].txn, start_x.into_inner());

    // Removing one holder is only a partial shrink. The real wait entry and
    // its RPC must remain pending until the final shared holder is removed.
    let rollback_errors = client
        .kv_pessimistic_rollback(vec![shared_key.clone()], start_1, for_update_1)
        .unwrap();
    assert!(rollback_errors.is_empty(), "{:?}", rollback_errors);
    wait_for_lock_wait_info_entry(&client, store_id, ApiVersion::V2, keyspace_id, &shared_key);
    assert!(matches!(wait_rx.try_recv(), Err(TryRecvError::Empty)));

    // The final release wakes the legacy waiter with a directed retry. A fresh
    // third RPC then re-reads durable state and acquires the ordinary exclusive
    // pessimistic lock.
    let rollback_errors = client
        .kv_pessimistic_rollback(vec![shared_key.clone()], start_2, for_update_2)
        .unwrap();
    assert!(rollback_errors.is_empty(), "{:?}", rollback_errors);
    let (_, second_errors) = wait_rx
        .recv_timeout(Duration::from_secs(10))
        .expect("final shared-holder release should wake external X")
        .unwrap();
    assert_eq!(second_errors.len(), 1, "{second_errors:?}");
    assert!(second_errors[0].has_conflict(), "{:?}", second_errors);
    assert_eq!(
        second_errors[0].get_conflict().get_reason(),
        WriteConflictReason::PessimisticRetry
    );
    wait_task.join().unwrap();

    let (_, third_errors) = client
        .kv_pessimistic_lock(
            primary_x.clone().into(),
            vec![shared_key.clone().into()],
            start_x,
            3000,
            for_update_x,
            None,
        )
        .unwrap();
    assert!(third_errors.is_empty(), "{:?}", third_errors);
    must_locked_with_properties(
        &mut client,
        &shared_key,
        &primary_x,
        start_x,
        LockType::Pessimistic,
    );

    cluster.stop();
}

#[test]
fn test_removed_pending_holder_returns_shared_lock_lost() {
    test_util::init_log_for_test();

    // Start the test cluster and acquire U and H as shared holders of the same key.
    let mut cluster = ServerCluster::new(alloc_node_id_vec(1), |_, cfg| {
        cfg.pessimistic_txn.pipelined = false;
    });
    cluster.wait_region_replicated(&[], 1);
    let mut client = cluster.new_client();
    let keyspace_id = api_version::ApiV2::get_u32_keyspace_id_by_key(&i_to_key(0)).unwrap();
    client.split_keyspace(keyspace_id);

    let shared_key = i_to_key(70);
    let shared_key_end = i_to_key(71);
    let primary_u = i_to_key(72);
    let primary_h = i_to_key(73);
    let start_u = client.get_ts();
    let for_update_u = client.get_ts();
    let start_h = client.get_ts();
    let for_update_h = client.get_ts();

    acquire_shared_secondary(&mut client, &primary_u, &shared_key, start_u, for_update_u);
    acquire_shared_secondary(&mut client, &primary_h, &shared_key, start_h, for_update_h);

    // Start U's pending upgrade and observe its entry in the lock wait table.
    let (_, first_errors) = client
        .kv_pessimistic_lock(
            primary_u.clone().into(),
            vec![shared_key.clone().into()],
            start_u,
            3000,
            for_update_u,
            None,
        )
        .unwrap();
    assert_eq!(first_errors.len(), 1, "{:?}", first_errors);
    assert!(first_errors[0].has_conflict(), "{:?}", first_errors);
    assert_eq!(
        first_errors[0].get_conflict().get_reason(),
        WriteConflictReason::PessimisticRetry
    );

    let mut upgrade_client = client.clone();
    let upgrade_key = shared_key.clone();
    let upgrade_primary = primary_u.clone();
    let (upgrade_tx, upgrade_rx) = mpsc::sync_channel(1);
    let upgrade_task = thread::spawn(move || {
        let result = upgrade_client.kv_pessimistic_lock(
            upgrade_primary.into(),
            vec![upgrade_key.into()],
            start_u,
            3000,
            for_update_u,
            Some(30_000),
        );
        upgrade_tx.send(result).unwrap();
    });

    let store_id = client.get_stores()[0];
    wait_for_lock_wait_info_entry(&client, store_id, ApiVersion::V2, keyspace_id, &shared_key);

    // Remove U's shared holder through its pessimistic rollback.
    let rollback_errors = client
        .kv_pessimistic_rollback(vec![shared_key.clone()], start_u, for_update_u)
        .unwrap();
    assert!(rollback_errors.is_empty(), "{:?}", rollback_errors);

    // Verify that the pending upgrade receives the typed shared-lock-loss response.
    let (_, upgrade_errors) = upgrade_rx
        .recv_timeout(Duration::from_secs(10))
        .expect("holder removal should finish pending upgrader")
        .unwrap();
    assert_eq!(upgrade_errors.len(), 1, "{:?}", upgrade_errors);
    assert!(
        upgrade_errors[0].has_shared_lock_lost(),
        "{:?}",
        upgrade_errors
    );
    assert!(
        upgrade_errors[0].get_abort().is_empty(),
        "{:?}",
        upgrade_errors
    );
    let lost = upgrade_errors[0].get_shared_lock_lost();
    assert_eq!(lost.get_key(), shared_key);
    assert_eq!(lost.get_start_ts(), start_u.into_inner());
    upgrade_task.join().unwrap();

    // Verify H remains durable and the completed wait entry has been removed.
    let max_ts = client.get_ts();
    let locks = client
        .kv_scan_lock(&shared_key, &shared_key_end, max_ts, 10)
        .unwrap();
    assert_eq!(locks.len(), 1, "{:?}", locks);
    assert_eq!(locks[0].get_lock_type(), Op::SharedLock);
    let holders = locks[0]
        .get_shared_lock_infos()
        .iter()
        .map(|holder| holder.get_lock_version())
        .collect::<Vec<_>>();
    assert_eq!(holders, vec![start_h.into_inner()]);
    assert!(get_lock_wait_info_entries(&client, store_id, ApiVersion::V2, keyspace_id).is_empty());

    cluster.stop();
}

#[test]
fn test_pes_rollback_after_lock_and_prewrite() {
    for pipelined in [true, false] {
        let mut cluster = ServerCluster::new(alloc_node_id_vec(1), |_, cfg| {
            cfg.pessimistic_txn.pipelined = pipelined;
        });
        cluster.wait_region_replicated(&[], 1);
        let mut client = cluster.new_client();
        let keyspace_id = api_version::ApiV2::get_u32_keyspace_id_by_key(&i_to_key(0)).unwrap();
        client.split_keyspace(keyspace_id);

        // Arrange
        let k1 = i_to_key(20);
        let v1 = i_to_val(20);

        // Act: Transaction T1
        // 1. T1 acquires pessimistic lock on K1 and stages write of V1
        let t1_start_ts = client.get_ts();
        let t1_for_update_ts = client.get_ts();
        let lock_ttl = 3000; // ms

        let t1_mutations = vec![new_pessimistic_lock_mutation(k1.clone())];

        let (_lock_result, key_errors) = client
            .kv_pessimistic_lock(
                k1.clone().into(),
                vec![k1.clone().into()],
                t1_start_ts,
                lock_ttl,
                t1_for_update_ts,
                None,
            )
            .expect("T1: pessimistic_lock failed");
        assert!(
            key_errors.is_empty(),
            "T1: pessimistic_lock resulted in key_errors: {:?}",
            key_errors
        );

        // Verify K1 is locked by T1 (pessimistically)
        if pipelined {
            thread::sleep(Duration::from_millis(100));
        }
        let _ =
            must_locked_with_properties(&mut client, &k1, &k1, t1_start_ts, LockType::Pessimistic);

        // T1 prewrites the value
        let prewrite_mutations = vec![new_put_mutation(k1.clone(), v1.clone())];
        let mut prewrite_txn = client.begin_transaction(Some(t1_start_ts));
        client
            .kv_prewrite_ext(
                k1.clone().into(),
                None,
                TxnMutations::from_normal(prewrite_mutations),
                &mut prewrite_txn,
                PrewriteExt {
                    resolve_lock: false,
                    for_update_ts: t1_start_ts,
                    ..Default::default()
                },
            )
            .expect("T1 prewrite failed");

        // 2. T1 rolls back
        let t1_rollback_txns = TxnMutations::from_normal(t1_mutations);
        client
            .kv_rollback(t1_rollback_txns, t1_start_ts)
            .expect("T1: kv_rollback failed");

        // K1 should be unlocked after T1's rollback
        must_unlocked(&mut client, &k1);

        // Subsequent transaction T2 Get(K1) gets nothing
        let t2_start_ts = client.get_ts();
        let val_read_by_t2 = client
            .get_key_version(&k1, t2_start_ts.into_inner())
            .expect("T2: Get_key_version failed");

        assert!(
            val_read_by_t2.is_none(),
            "T2: Get(K1) should return None (KeyNotFound) as K1 was new and T1 rolled back, but got {:?}",
            val_read_by_t2
        );

        // Subsequent transaction T3 LockKeys(K1) succeeds (T1's lock is released)
        let t3_start_ts = client.get_ts();
        let t3_for_update_ts = client.get_ts();
        let (_t3_lock_result, t3_key_errors) = client
            .kv_pessimistic_lock(
                k1.clone().into(),
                vec![k1.clone().into()],
                t3_start_ts,
                lock_ttl,
                t3_for_update_ts,
                None,
            )
            .expect("T3: pessimistic_lock failed");
        assert!(
            t3_key_errors.is_empty(),
            "T3: pessimistic_lock should succeed without errors, but got: {:?}",
            t3_key_errors
        );

        // Verify T3's lock and then clean it up (good practice)
        if pipelined {
            thread::sleep(Duration::from_millis(100));
        }
        let _lock_info_t3 =
            must_locked_with_properties(&mut client, &k1, &k1, t3_start_ts, LockType::Pessimistic);

        cluster.stop();
    }
}

#[test]
fn test_lock_waiting() {
    let cases = vec![
        case!(lock!(10, 10); ok!()),
        case!(lock!(10, 10, async, 100); ok!()),
        // should wait 1s, test 900ms
        case!(lock!(10, 10), lock!(20, 20, async, 900); timeout!()),
        // 1s timeout, write conflict
        case!(lock!(10, 10), lock!(20, 20, async, 2000); fail!()),
        case!(concurrent!(
            sequence!(lock!(10, 10), sleep!(500), rollback!(10); ok!()),
            // fail => it is woken up
            sequence!(sleep!(300), lock!(20, 20, async, 500); fail!())
        ); ok!()),
        // 2 txns waiting
        case!(concurrent!(
            sequence!(lock!(10, 10), sleep!(500), rollback!(10); ok!()),
            sequence!(sleep!(300), lock!(20, 20, async, 500); fail!()),
            sequence!(sleep!(300), lock!(30, 30, async, 150); timeout!())
        ); ok!()),
    ];
    run_test_cases(cases);
}

#[test]
fn test_get_lock_wait_info_filters_by_keyspace() {
    test_util::init_log_for_test();
    const KEYSPACE_1: u32 = 42;
    const KEYSPACE_2: u32 = 43;

    for pipelined in [true, false] {
        let mut cluster = ServerCluster::new(alloc_node_id_vec(1), |_, cfg| {
            cfg.pessimistic_txn.pipelined = pipelined;
        });
        cluster.wait_region_replicated(&[], 1);

        let mut client = cluster.new_client();
        client.split_keyspace(KEYSPACE_1);
        client.split_keyspace(KEYSPACE_2);

        let key_1 = i_to_key_with_keyspace(KEYSPACE_1)(1);
        let key_2 = i_to_key_with_keyspace(KEYSPACE_2)(2);
        let lock_ttl = 10_000;

        let txn_1_start_ts = client.get_ts();
        let txn_1_for_update_ts = client.get_ts();
        client
            .kv_pessimistic_lock(
                key_1.clone().into(),
                vec![key_1.clone().into()],
                txn_1_start_ts,
                lock_ttl,
                txn_1_for_update_ts,
                None,
            )
            .unwrap();

        let txn_2_start_ts = client.get_ts();
        let txn_2_for_update_ts = client.get_ts();
        client
            .kv_pessimistic_lock(
                key_2.clone().into(),
                vec![key_2.clone().into()],
                txn_2_start_ts,
                lock_ttl,
                txn_2_for_update_ts,
                None,
            )
            .unwrap();

        if pipelined {
            thread::sleep(Duration::from_millis(100));
        }
        must_locked_with_properties(
            &mut client,
            &key_1,
            &key_1,
            txn_1_start_ts,
            LockType::Pessimistic,
        );
        must_locked_with_properties(
            &mut client,
            &key_2,
            &key_2,
            txn_2_start_ts,
            LockType::Pessimistic,
        );

        let txn_3 = {
            let mut txn_3_client = client.clone();
            let key = key_1.clone();
            thread::spawn(move || {
                let start_ts = txn_3_client.get_ts();
                let for_update_ts = txn_3_client.get_ts();
                let res = txn_3_client.kv_pessimistic_lock(
                    key.clone().into(),
                    vec![key.clone().into()],
                    start_ts,
                    lock_ttl,
                    for_update_ts,
                    Some(5000),
                );
                (start_ts, res)
            })
        };

        let txn_4 = {
            let mut txn_4_client = client.clone();
            let key = key_2.clone();
            thread::spawn(move || {
                let start_ts = txn_4_client.get_ts();
                let for_update_ts = txn_4_client.get_ts();
                let res = txn_4_client.kv_pessimistic_lock(
                    key.clone().into(),
                    vec![key.clone().into()],
                    start_ts,
                    lock_ttl,
                    for_update_ts,
                    Some(5000),
                );
                (start_ts, res)
            })
        };

        let store_id = client.get_stores()[0];
        let entries_1 =
            wait_for_lock_wait_info_entry(&client, store_id, ApiVersion::V2, KEYSPACE_1, &key_1);
        let entries_2 =
            wait_for_lock_wait_info_entry(&client, store_id, ApiVersion::V2, KEYSPACE_2, &key_2);
        let entries_v1 = get_lock_wait_info_entries(&client, store_id, ApiVersion::V1, 0);

        assert_eq!(entries_1[0].key, key_1);
        assert_eq!(entries_1[0].wait_for_txn, txn_1_start_ts.into_inner());
        assert_eq!(entries_2[0].key, key_2);
        assert_eq!(entries_2[0].wait_for_txn, txn_2_start_ts.into_inner());
        assert!(
            entries_v1.is_empty(),
            "api v1 wait entries: {:?}",
            entries_v1
        );
        assert!(
            get_lock_wait_info_entries(&client, store_id, ApiVersion::V2, NULL_KEYSPACE_ID)
                .is_empty()
        );

        client
            .kv_rollback(
                TxnMutations::from_normal(vec![Mutation {
                    op: Op::PessimisticLock,
                    key: key_1.clone().into(),
                    value: vec![].into(),
                    assertion: Assertion::None,
                }]),
                txn_1_start_ts,
            )
            .unwrap();
        client
            .kv_rollback(
                TxnMutations::from_normal(vec![Mutation {
                    op: Op::PessimisticLock,
                    key: key_2.clone().into(),
                    value: vec![].into(),
                    assertion: Assertion::None,
                }]),
                txn_2_start_ts,
            )
            .unwrap();

        let (_, txn_3_res) = txn_3.join().unwrap();
        let (_, txn_4_res) = txn_4.join().unwrap();
        let (_, txn_3_errors) = txn_3_res.unwrap();
        let (_, txn_4_errors) = txn_4_res.unwrap();
        assert_eq!(txn_3_errors.len(), 1);
        assert_eq!(txn_4_errors.len(), 1);
        assert!(txn_3_errors[0].has_conflict(), "{:?}", txn_3_errors);
        assert!(txn_4_errors[0].has_conflict(), "{:?}", txn_4_errors);

        cluster.stop();
    }
}

#[test]
fn test_shared_lock_upgrade_fast_path_and_first_retry() {
    test_util::init_log_for_test();
    for pipelined in [true, false] {
        let mut cluster = ServerCluster::new(alloc_node_id_vec(1), |_, cfg| {
            cfg.pessimistic_txn.pipelined = pipelined;
        });
        cluster.wait_region_replicated(&[], 1);
        let mut client = cluster.new_client();
        let keyspace_id = api_version::ApiV2::get_u32_keyspace_id_by_key(&i_to_key(0)).unwrap();
        client.split_keyspace(keyspace_id);

        // Verify that a sole shared-lock holder upgrades without waiting.
        let key_fast = i_to_key(40);
        let pk_fast = i_to_key(41);
        let start_ts_fast = client.get_ts();
        let for_update_ts_fast = client.get_ts();

        client
            .kv_pessimistic_lock(
                pk_fast.clone().into(),
                vec![pk_fast.clone().into()],
                start_ts_fast,
                3000,
                for_update_ts_fast,
                None,
            )
            .unwrap();
        client
            .kv_shared_pessimistic_lock(
                pk_fast.clone().into(),
                vec![key_fast.clone().into()],
                start_ts_fast,
                3000,
                for_update_ts_fast,
                None,
            )
            .unwrap();

        if pipelined {
            thread::sleep(Duration::from_millis(100));
        }

        let (_, fast_errors) = client
            .kv_pessimistic_lock(
                pk_fast.clone().into(),
                vec![key_fast.clone().into()],
                start_ts_fast,
                3000,
                for_update_ts_fast,
                None,
            )
            .unwrap();
        assert!(fast_errors.is_empty(), "{:?}", fast_errors);
        if pipelined {
            thread::sleep(Duration::from_millis(100));
        }
        must_locked_with_properties(
            &mut client,
            &key_fast,
            &pk_fast,
            start_ts_fast,
            LockType::Pessimistic,
        );
        let shared_after_upgrade_start_ts = client.get_ts();
        let shared_after_upgrade_for_update_ts = client.get_ts();
        let (_, shared_after_upgrade_errors) = client
            .kv_shared_pessimistic_lock(
                pk_fast.clone().into(),
                vec![key_fast.clone().into()],
                shared_after_upgrade_start_ts,
                3000,
                shared_after_upgrade_for_update_ts,
                Some(1),
            )
            .unwrap();
        assert_eq!(
            shared_after_upgrade_errors.len(),
            1,
            "{:?}",
            shared_after_upgrade_errors
        );
        assert!(
            shared_after_upgrade_errors[0].has_locked(),
            "{:?}",
            shared_after_upgrade_errors
        );
        let upgraded_lock = shared_after_upgrade_errors[0].get_locked();
        assert_eq!(upgraded_lock.get_lock_type(), Op::PessimisticLock);
        assert_eq!(upgraded_lock.get_shared_lock_infos().len(), 0);

        // Verify that a multi-holder upgrade returns PessimisticRetry and leaves no
        // lock-wait entry behind.
        let key_retry = i_to_key(42);
        let pk_retry_1 = i_to_key(43);
        let pk_retry_2 = i_to_key(44);
        let start_ts_retry_1 = client.get_ts();
        let for_update_ts_retry_1 = client.get_ts();
        client
            .kv_pessimistic_lock(
                pk_retry_1.clone().into(),
                vec![pk_retry_1.clone().into()],
                start_ts_retry_1,
                3000,
                for_update_ts_retry_1,
                None,
            )
            .unwrap();
        client
            .kv_shared_pessimistic_lock(
                pk_retry_1.clone().into(),
                vec![key_retry.clone().into()],
                start_ts_retry_1,
                3000,
                for_update_ts_retry_1,
                None,
            )
            .unwrap();

        let start_ts_retry_2 = client.get_ts();
        let for_update_ts_retry_2 = client.get_ts();
        client
            .kv_pessimistic_lock(
                pk_retry_2.clone().into(),
                vec![pk_retry_2.clone().into()],
                start_ts_retry_2,
                3000,
                for_update_ts_retry_2,
                None,
            )
            .unwrap();
        client
            .kv_shared_pessimistic_lock(
                pk_retry_2.clone().into(),
                vec![key_retry.clone().into()],
                start_ts_retry_2,
                3000,
                for_update_ts_retry_2,
                None,
            )
            .unwrap();

        if pipelined {
            thread::sleep(Duration::from_millis(100));
        }

        let store_id = client.get_stores()[0];
        let (_, retry_errors) = client
            .kv_pessimistic_lock(
                pk_retry_1.clone().into(),
                vec![key_retry.clone().into()],
                start_ts_retry_1,
                3000,
                for_update_ts_retry_1,
                None,
            )
            .unwrap();
        assert_eq!(retry_errors.len(), 1, "{:?}", retry_errors);
        assert!(retry_errors[0].has_conflict(), "{:?}", retry_errors);
        assert_eq!(
            retry_errors[0].get_conflict().get_reason(),
            WriteConflictReason::PessimisticRetry
        );
        assert!(
            get_lock_wait_info_entries(&client, store_id, ApiVersion::V2, keyspace_id).is_empty()
        );

        cluster.stop();
    }
}

#[test]
fn test_deadlock_detection() {
    let k1 = i_to_key(100);
    let k2 = i_to_key(101);

    let k3 = i_to_key(102);
    let k4 = i_to_key(103);
    let k5 = i_to_key(104);

    let cases = vec![
        case!(concurrent!(
            sequence!(lock!(10, 10, k1), sleep!(200), lock!(10, 10, k2); fail!()),
            sequence!(lock!(20, 20, k2), sleep!(400), lock!(20, 20, k1); deadlock!())
        ); ok!()),
        case!(concurrent!(
            sequence!(lock!(10, 10, k3), sleep!(200), lock!(10, 10, k4); fail!()),
            sequence!(lock!(20, 20, k4), sleep!(300), lock!(20, 20, k5); fail!()),
            sequence!(lock!(30, 30, k5), sleep!(400), lock!(30, 30, k3); deadlock!())
        ); ok!()),
    ];

    run_test_cases(cases);
}