// Copyright 2025 PingCAP, Inc.

use std::{
    collections::HashMap,
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, Ordering},
        mpsc,
    },
    thread,
    time::Duration,
};

use bytes::Bytes;
use cloud_server::ASYNC_WRITE_CALLBACK_DROPPED_ERR_MSG;
use futures::executor::block_on;
use hyper::Method;
use kvproto::{
    deadlock::WaitForEntry,
    kvrpcpb::{
        self, ApiVersion, Context, GetLockWaitInfoRequest, Op, PessimisticRollbackRequest,
        WriteConflictReason,
    },
};
use matches::assert_matches;
use pd_client::PdClient;
use security::{RestfulClient, SecurityConfig, SecurityManager};
use test_cloud_server::{
    ServerCluster, alloc_node_id_vec,
    client::{ClusterClient, Error as ClientError, TxnMutations},
    try_wait,
    util::Mutation,
};
use test_pd_client::PdClientExt;
use tikv_util::config::ReadableDuration;
use txn_types::TimeStamp;

use crate::cases::{FailpointGuard, cfg_notify_once, i_to_key};

#[test]
fn test_undetermined_write_err() {
    let cluster = ServerCluster::new(alloc_node_id_vec(1), |_, _| {});
    cluster.wait_region_replicated(&[], 1);
    let mut client = cluster.new_client();

    let mut mutation = Mutation::default();
    let k1 = i_to_key(1);
    mutation.set_op(Op::Put);
    mutation.set_key(k1.clone());
    mutation.set_value(b"v".to_vec());
    let primary_key = k1;
    let start_ts = client.get_ts();
    let mutations = vec![mutation];

    fail::cfg("applied_cb_return_undetermined_err", "return()").unwrap();
    let err = client
        .kv_prewrite(
            primary_key.to_vec().into(),
            None,
            TxnMutations::from_normal(mutations.clone()),
            start_ts,
        )
        .unwrap_err();
    assert_matches!(err, ClientError::RegionError(ref region_err) if {
        assert!(region_err.has_undetermined_result());
        assert_eq!(
            region_err.get_undetermined_result().get_message(),
            ASYNC_WRITE_CALLBACK_DROPPED_ERR_MSG,
        );
        true
    });
    fail::remove("applied_cb_return_undetermined_err");
}

#[test]
fn test_raftkv_early_error_report() {
    test_util::init_log_for_test();
    let cluster = ServerCluster::new(alloc_node_id_vec(1), |_, _| {});
    cluster.wait_region_replicated(&[], 1);
    let mut client = cluster.new_client();

    let mut mutation = Mutation::default();
    let k1 = i_to_key(1);
    mutation.set_op(Op::Put);
    mutation.set_key(k1.clone());
    mutation.set_value(b"v".to_vec());
    let primary_key = k1;
    let start_ts = client.get_ts();
    let mutations = vec![mutation];

    fail::cfg("raftkv_early_error_report", "return").unwrap();
    let prewrite_resp = client
        .kv_prewrite_single_region_without_retry(
            primary_key.to_vec().into(),
            None,
            TxnMutations::from_normal(mutations.clone()),
            start_ts,
        )
        .unwrap();
    assert!(prewrite_resp.has_region_error(), "{:?}", prewrite_resp);
    assert!(
        prewrite_resp.get_region_error().has_region_not_found(),
        "{:?}",
        prewrite_resp
    );
    fail::remove("raftkv_early_error_report");
}

fn scan_locks_in_range(
    client: &mut ClusterClient,
    start_key: &[u8],
    end_key: &[u8],
    max_ts: u64,
    limit: u32,
) -> Vec<kvrpcpb::LockInfo> {
    let region = client.get_region_by_key(start_key);
    let (mut region_meta, leader) = block_on(client.pd_client().get_region_leader_by_id(region.id))
        .unwrap()
        .unwrap();
    let mut ctx = client.new_rpc_ctx(region.id, start_key).unwrap();
    ctx.set_region_epoch(region_meta.take_region_epoch());
    ctx.set_peer(leader);
    let kv_client = client.get_kv_client(ctx.get_peer().get_store_id());

    let mut req = kvrpcpb::ScanLockRequest::default();
    req.set_context(ctx);
    req.set_start_key(start_key.to_vec());
    req.set_end_key(end_key.to_vec());
    req.set_max_version(max_ts);
    req.set_limit(limit);

    let mut resp = kv_client.kv_scan_lock(&req).unwrap();
    assert!(
        !resp.has_region_error() && !resp.has_error(),
        "resp {:?}",
        resp
    );
    resp.take_locks().into_vec()
}

fn new_shared_lock_test_cluster_with_pipelined(pipelined: bool) -> (ServerCluster, ClusterClient) {
    let cluster = ServerCluster::new(alloc_node_id_vec(1), move |_, 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);
    (cluster, client)
}

fn new_shared_lock_test_cluster() -> (ServerCluster, ClusterClient) {
    new_shared_lock_test_cluster_with_pipelined(false)
}

fn cfg_block_once_with_timeout(
    name: &'static str,
) -> (mpsc::Receiver<()>, mpsc::SyncSender<()>, FailpointGuard) {
    let (entered_tx, entered_rx) = mpsc::sync_channel(1);
    let (release_tx, release_rx) = mpsc::sync_channel(1);
    let release_rx = Mutex::new(Some(release_rx));
    let guard = FailpointGuard::new().cfg_callback(name, move || {
        let Some(release_rx) = release_rx.lock().unwrap().take() else {
            return;
        };
        if entered_tx.send(()).is_err() {
            return;
        }
        // Dropping the test-side sender during unwinding also releases the
        // blocked worker. Avoid panicking in that worker so an earlier, more
        // useful test failure is not obscured by barrier cleanup.
        let _ = release_rx.recv_timeout(Duration::from_secs(10));
    });
    (entered_rx, release_tx, guard)
}

fn acquire_shared_secondary(
    client: &mut ClusterClient,
    primary: &[u8],
    shared_key: &[u8],
    start_ts: TimeStamp,
    for_update_ts: 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);
}

fn mark_shared_wrapper_shrink_only(
    client: &mut ClusterClient,
    primary: &[u8],
    shared_key: &[u8],
    start_ts: TimeStamp,
    for_update_ts: TimeStamp,
) {
    let (_, errors) = client
        .kv_pessimistic_lock(
            primary.to_vec().into(),
            vec![shared_key.to_vec().into()],
            start_ts,
            3000,
            for_update_ts,
            None,
        )
        .unwrap();
    assert_eq!(errors.len(), 1, "{errors:?}");
    assert!(errors[0].has_conflict(), "{:?}", errors);
    assert_eq!(
        errors[0].get_conflict().get_reason(),
        WriteConflictReason::PessimisticRetry
    );
}

fn raw_pessimistic_rollback(
    client: &mut ClusterClient,
    key: &[u8],
    start_ts: TimeStamp,
    for_update_ts: TimeStamp,
) -> kvrpcpb::PessimisticRollbackResponse {
    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());
    let mut req = PessimisticRollbackRequest::default();
    req.set_context(ctx);
    req.set_keys(vec![key.to_vec()].into());
    req.set_start_version(start_ts.into_inner());
    req.set_for_update_ts(for_update_ts.into_inner());
    kv_client.kv_pessimistic_rollback(&req).unwrap()
}

fn must_shared_lock_holders(
    client: &mut ClusterClient,
    key: &[u8],
    end_key: &[u8],
    expected: &[TimeStamp],
) {
    let max_ts = client.get_ts().into_inner();
    let locks = scan_locks_in_range(client, key, end_key, max_ts, 10);
    assert_eq!(locks.len(), 1, "locks {locks:?}");
    assert_eq!(locks[0].get_lock_type(), Op::SharedLock);
    let mut actual = locks[0]
        .get_shared_lock_infos()
        .iter()
        .map(|holder| holder.get_lock_version())
        .collect::<Vec<_>>();
    actual.sort_unstable();
    let mut expected = expected
        .iter()
        .map(|ts| ts.into_inner())
        .collect::<Vec<_>>();
    expected.sort_unstable();
    assert_eq!(actual, expected);
}

fn get_lock_wait_info_entries(
    client: &ClusterClient,
    store_id: u64,
    keyspace_id: u32,
) -> Vec<WaitForEntry> {
    let kv_client = client.get_kv_client(store_id);
    let mut ctx = Context::default();
    ctx.set_api_version(ApiVersion::V2);
    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_entry(
    client: &ClusterClient,
    store_id: u64,
    keyspace_id: u32,
    key: &[u8],
    should_exist: bool,
) {
    for attempt in 0..50 {
        let entries = get_lock_wait_info_entries(client, store_id, keyspace_id);
        let exists = entries.iter().any(|entry| entry.key == key);
        if exists == should_exist {
            return;
        }
        if attempt == 49 {
            panic!(
                "wait-entry state did not become {} for {:?}: {:?}",
                should_exist, key, entries
            );
        }
        thread::sleep(Duration::from_millis(100));
    }
}

fn assert_pessimistic_retry(errors: &[kvrpcpb::KeyError]) {
    assert_eq!(errors.len(), 1, "{errors:?}");
    assert!(errors[0].has_conflict(), "{:?}", errors);
    assert!(!errors[0].has_shared_lock_lost(), "{:?}", errors);
    assert!(errors[0].get_abort().is_empty(), "{:?}", errors);
    assert_eq!(
        errors[0].get_conflict().get_reason(),
        WriteConflictReason::PessimisticRetry
    );
}

fn assert_shared_lock_lost(errors: &[kvrpcpb::KeyError], start_ts: TimeStamp, key: &[u8]) {
    assert_eq!(errors.len(), 1, "{:?}", errors);
    assert!(errors[0].has_shared_lock_lost(), "{:?}", errors);
    assert!(errors[0].get_abort().is_empty(), "{:?}", errors);
    let lost = errors[0].get_shared_lock_lost();
    assert_eq!(lost.get_start_ts(), start_ts.into_inner());
    assert_eq!(lost.get_key(), key);
}

fn assert_key_is_locked(errors: &[kvrpcpb::KeyError]) {
    assert_eq!(errors.len(), 1, "{errors:?}");
    assert!(errors[0].has_locked(), "{:?}", errors);
}

#[test]
fn test_matching_upgrade_hint_retries_before_determined_write_failure() {
    test_util::init_log_for_test();
    let (mut cluster, mut client) = new_shared_lock_test_cluster();
    let shared_key = i_to_key(600);
    let shared_key_end = i_to_key(601);
    let primary_u = i_to_key(610);
    let primary_h = i_to_key(611);
    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);
    mark_shared_wrapper_shrink_only(&mut client, &primary_u, &shared_key, start_u, for_update_u);

    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];
    let keyspace_id = api_version::ApiV2::get_u32_keyspace_id_by_key(&shared_key).unwrap();
    wait_for_lock_wait_entry(&client, store_id, keyspace_id, &shared_key, true);

    // The hint is consumed before the determined engine error is reported, while
    // Finished remains paused with the same-key latch held.
    let failure = FailpointGuard::new().cfg("raftkv_early_error_report", "return");
    let (finish_entered, release_finish, _finish_guard) =
        cfg_block_once_with_timeout("scheduler_async_write_finish");
    let mut rollback_client = client.clone();
    let rollback_key = shared_key.clone();
    let (rollback_tx, rollback_rx) = mpsc::sync_channel(1);
    let rollback_task = thread::spawn(move || {
        let response =
            raw_pessimistic_rollback(&mut rollback_client, &rollback_key, start_h, for_update_h);
        rollback_tx.send(response).unwrap();
    });
    finish_entered
        .recv_timeout(Duration::from_secs(10))
        .expect("failed holder removal should reach Finished");

    let (_, upgrade_errors) = upgrade_rx
        .recv_timeout(Duration::from_secs(10))
        .expect("matching hint should retry U before write completion")
        .unwrap();
    assert_pessimistic_retry(&upgrade_errors);
    upgrade_task.join().unwrap();
    wait_for_lock_wait_entry(&client, store_id, keyspace_id, &shared_key, false);

    // A fresh retry cannot re-read the old wrapper until the failed writer
    // releases the latch.
    let mut retry_client = client.clone();
    let retry_key = shared_key.clone();
    let retry_primary = primary_u.clone();
    let (retry_tx, retry_rx) = mpsc::sync_channel(1);
    let retry_task = thread::spawn(move || {
        let result = retry_client.kv_pessimistic_lock(
            retry_primary.into(),
            vec![retry_key.into()],
            start_u,
            3000,
            for_update_u,
            Some(30_000),
        );
        retry_tx.send(result).unwrap();
    });
    assert_matches!(
        retry_rx.recv_timeout(Duration::from_millis(500)),
        Err(mpsc::RecvTimeoutError::Timeout)
    );
    wait_for_lock_wait_entry(&client, store_id, keyspace_id, &shared_key, false);

    release_finish
        .send(())
        .expect("Finished pause should still be installed");
    let failed = rollback_rx
        .recv_timeout(Duration::from_secs(10))
        .expect("failed rollback should return after Finished");
    assert!(failed.has_region_error(), "{:?}", failed);
    assert!(
        failed.get_region_error().has_region_not_found(),
        "{:?}",
        failed
    );
    rollback_task.join().unwrap();
    must_shared_lock_holders(
        &mut client,
        &shared_key,
        &shared_key_end,
        &[start_u, start_h],
    );

    // The retry re-reads the unchanged wrapper and waits again. A successful
    // removal then emits a new matching hint and completes that exact waiter.
    wait_for_lock_wait_entry(&client, store_id, keyspace_id, &shared_key, true);
    drop(failure);
    let succeeded = raw_pessimistic_rollback(&mut client, &shared_key, start_h, for_update_h);
    assert!(!succeeded.has_region_error(), "{:?}", succeeded);
    assert!(succeeded.get_errors().is_empty(), "{:?}", succeeded);
    must_shared_lock_holders(&mut client, &shared_key, &shared_key_end, &[start_u]);
    let (_, retry_errors) = retry_rx
        .recv_timeout(Duration::from_secs(10))
        .expect("successful removal should retry the fresh waiter")
        .unwrap();
    assert_pessimistic_retry(&retry_errors);
    retry_task.join().unwrap();
    wait_for_lock_wait_entry(&client, store_id, keyspace_id, &shared_key, false);
    cluster.stop();
}
#[test]
fn test_determined_full_shared_wrapper_delete_preserves_durable_lock() {
    test_util::init_log_for_test();
    let (mut cluster, mut client) = new_shared_lock_test_cluster();
    let shared_key = i_to_key(640);
    let shared_key_end = i_to_key(641);
    let primary_h = i_to_key(650);
    let primary_x = i_to_key(651);
    let start_h = client.get_ts();
    let for_update_h = client.get_ts();

    // Create the sole durable shared holder. A pending shared upgrader cannot
    // naturally coexist with this topology: the sole holder's exclusive
    // request takes the fast upgrade path instead of entering lock wait.
    acquire_shared_secondary(&mut client, &primary_h, &shared_key, start_h, for_update_h);
    must_shared_lock_holders(&mut client, &shared_key, &shared_key_end, &[start_h]);

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

    // The first external-X request writes shrink-only and returns KeyIsLocked;
    // the second request then installs a real ordinary waiter. Keeping that
    // waiter makes this integration exercise the existing released_locks
    // recovery in addition to the durable full-delete failure.
    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_key_is_locked(&first_errors);

    let mut wait_client = client.clone();
    let wait_primary = primary_x.clone();
    let wait_key = shared_key.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 keyspace_id = api_version::ApiV2::get_u32_keyspace_id_by_key(&shared_key).unwrap();
    wait_for_lock_wait_entry(&client, store_id, keyspace_id, &shared_key, true);

    // This rollback would delete the whole wrapper, but the determined
    // RaftKV failure happens before proposal. Durable storage must retain H.
    // The legacy waiter may already receive the early released_locks retry;
    // the same-key latch makes that safe by delaying its re-read until this
    // failed writer has completed.
    let failure = FailpointGuard::new().cfg("raftkv_early_error_report", "return");
    let failed = raw_pessimistic_rollback(&mut client, &shared_key, start_h, for_update_h);
    assert!(failed.has_region_error(), "{:?}", failed);
    assert!(
        failed.get_region_error().has_region_not_found(),
        "{:?}",
        failed
    );
    must_shared_lock_holders(&mut client, &shared_key, &shared_key_end, &[start_h]);
    let (_, errors) = wait_rx
        .recv_timeout(Duration::from_secs(10))
        .expect("early full-delete notification should retry external X")
        .unwrap();
    assert_pessimistic_retry(&errors);
    wait_task.join().unwrap();

    // Reissuing the waiter after the determined failure must observe the old
    // durable shrink-only wrapper and enter real wait again. This is the
    // released_locks recovery contract: retry re-reads storage rather than
    // trusting the optimistic release notification.
    let mut wait_client = client.clone();
    let wait_primary = primary_x.clone();
    let wait_key = shared_key.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();
    });
    wait_for_lock_wait_entry(&client, store_id, keyspace_id, &shared_key, true);

    // The successful retry may still issue released_locks before its engine
    // write is known durable; the same-key latch keeps X's next re-read behind
    // this rollback. Once the synchronous rollback returns, first confirm the
    // wrapper is durably deleted, then consume the legacy PessimisticRetry,
    // which may already have arrived.
    drop(failure);
    let succeeded = raw_pessimistic_rollback(&mut client, &shared_key, start_h, for_update_h);
    assert!(!succeeded.has_region_error(), "{:?}", succeeded);
    assert!(succeeded.get_errors().is_empty(), "{:?}", succeeded);
    let max_ts = client.get_ts().into_inner();
    let locks = scan_locks_in_range(&mut client, &shared_key, &shared_key_end, max_ts, 10);
    assert!(locks.is_empty(), "wrapper should be deleted: {:?}", locks);

    let (_, errors) = wait_rx
        .recv_timeout(Duration::from_secs(10))
        .expect("successful full-delete retry should notify external X")
        .unwrap();
    assert_pessimistic_retry(&errors);
    wait_task.join().unwrap();
    wait_for_lock_wait_entry(&client, store_id, keyspace_id, &shared_key, false);
    cluster.stop();
}

#[test]
fn test_undetermined_mismatching_upgrade_hint_returns_shared_lock_lost() {
    test_util::init_log_for_test();
    let (mut cluster, mut client) = new_shared_lock_test_cluster();
    let shared_key = i_to_key(620);
    let shared_key_end = i_to_key(621);
    let primary_u = i_to_key(630);
    let primary_h = i_to_key(631);
    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);
    mark_shared_wrapper_shrink_only(&mut client, &primary_u, &shared_key, start_u, for_update_u);

    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];
    let keyspace_id = api_version::ApiV2::get_u32_keyspace_id_by_key(&shared_key).unwrap();
    wait_for_lock_wait_entry(&client, store_id, keyspace_id, &shared_key, true);

    // The write applies but reports an undetermined result. The speculative
    // mismatching hint fences U before the outcome is known.
    let failure = FailpointGuard::new().cfg("applied_cb_return_undetermined_err", "return()");
    let response = raw_pessimistic_rollback(&mut client, &shared_key, start_u, for_update_u);
    assert!(response.has_region_error(), "{:?}", response);
    assert!(
        response.get_region_error().has_undetermined_result(),
        "{:?}",
        response
    );
    assert_eq!(
        response
            .get_region_error()
            .get_undetermined_result()
            .get_message(),
        ASYNC_WRITE_CALLBACK_DROPPED_ERR_MSG
    );

    let (_, upgrade_errors) = upgrade_rx
        .recv_timeout(Duration::from_secs(10))
        .expect("mismatching hint should fence U")
        .unwrap();
    assert_shared_lock_lost(&upgrade_errors, start_u, &shared_key);
    upgrade_task.join().unwrap();
    wait_for_lock_wait_entry(&client, store_id, keyspace_id, &shared_key, false);

    drop(failure);
    must_shared_lock_holders(&mut client, &shared_key, &shared_key_end, &[start_h]);
    let cleanup_errors = client
        .kv_pessimistic_rollback(vec![primary_u], start_u, for_update_u)
        .unwrap();
    assert!(cleanup_errors.is_empty(), "{:?}", cleanup_errors);
    cluster.stop();
}

#[test]
fn test_mismatching_upgrade_hint_can_abort_on_determined_write_failure() {
    test_util::init_log_for_test();
    let (mut cluster, mut client) = new_shared_lock_test_cluster();
    let shared_key = i_to_key(660);
    let shared_key_end = i_to_key(661);
    let primary_u = i_to_key(670);
    let primary_h = i_to_key(671);
    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);
    mark_shared_wrapper_shrink_only(&mut client, &primary_u, &shared_key, start_u, for_update_u);

    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];
    let keyspace_id = api_version::ApiV2::get_u32_keyspace_id_by_key(&shared_key).unwrap();
    wait_for_lock_wait_entry(&client, store_id, keyspace_id, &shared_key, true);

    // A determined failure leaves {U, H} durable, but the speculative mismatch
    // may conservatively abort U before the engine rejects the write.
    let failure = FailpointGuard::new().cfg("raftkv_early_error_report", "return");
    let failed = raw_pessimistic_rollback(&mut client, &shared_key, start_u, for_update_u);
    assert!(failed.has_region_error(), "{:?}", failed);
    assert!(
        failed.get_region_error().has_region_not_found(),
        "{:?}",
        failed
    );
    let (_, upgrade_errors) = upgrade_rx
        .recv_timeout(Duration::from_secs(10))
        .expect("mismatching hint should conservatively fence U")
        .unwrap();
    assert_shared_lock_lost(&upgrade_errors, start_u, &shared_key);
    upgrade_task.join().unwrap();
    wait_for_lock_wait_entry(&client, store_id, keyspace_id, &shared_key, false);
    must_shared_lock_holders(
        &mut client,
        &shared_key,
        &shared_key_end,
        &[start_u, start_h],
    );

    // Roll back the already-fenced transaction after removing the failure.
    drop(failure);
    let cleanup_errors = client
        .kv_pessimistic_rollback(
            vec![primary_u.clone(), shared_key.clone()],
            start_u,
            for_update_u,
        )
        .unwrap();
    assert!(cleanup_errors.is_empty(), "{:?}", cleanup_errors);
    must_shared_lock_holders(&mut client, &shared_key, &shared_key_end, &[start_h]);
    cluster.stop();
}

fn send_async_region_resolve_lock(
    client: &mut ClusterClient,
    representative_key: &[u8],
    start_ts: u64,
    commit_ts: u64,
) {
    let region = client.get_region_by_key(representative_key);
    let (mut region_meta, leader) = block_on(client.pd_client().get_region_leader_by_id(region.id))
        .unwrap()
        .unwrap();
    let mut ctx = client.new_rpc_ctx(region.id, representative_key).unwrap();
    ctx.set_region_epoch(region_meta.take_region_epoch());
    ctx.set_peer(leader);
    let kv_client = client.get_kv_client(ctx.get_peer().get_store_id());

    let txn_info = kvrpcpb::TxnInfo {
        txn: start_ts,
        status: commit_ts,
        ..Default::default()
    };
    let mut req = kvrpcpb::ResolveLockRequest::default();
    req.set_context(ctx);
    req.set_is_async(true);
    req.set_txn_infos(vec![txn_info].into());

    let resp = kv_client.kv_resolve_lock(&req).unwrap();
    assert!(
        !resp.has_region_error() && !resp.has_error(),
        "resp {:?}",
        resp
    );
}

fn update_async_resolve_lock_batch_interval(status_addr: String, interval: &str) {
    let security_mgr = Arc::new(SecurityManager::new(&SecurityConfig::default()).unwrap());
    let rt = tokio::runtime::Runtime::new().unwrap();
    let change = HashMap::from([(
        "storage.scheduler-async-resolve-lock-batch-interval".to_owned(),
        interval.to_owned(),
    )]);
    rt.block_on(async move {
        RestfulClient::new("async-resolve-lock-config", vec![status_addr], security_mgr)
            .unwrap()
            .request(
                "config",
                Method::POST,
                Some(Bytes::from(serde_json::to_vec(&change).unwrap())),
            )
            .await
            .unwrap();
    });
}

#[test]
fn test_async_region_resolve_lock_online_batch_interval_update() {
    test_util::init_log_for_test();
    let node_id = 1;
    let mut cluster = ServerCluster::new(vec![node_id], |_, _| {});
    assert_eq!(
        cluster
            .get_node_config(node_id)
            .storage
            .scheduler_async_resolve_lock_batch_interval,
        ReadableDuration::millis(500)
    );
    cluster.wait_region_replicated(&[], 1);
    let mut client = cluster.new_client();

    update_async_resolve_lock_batch_interval(cluster.status_addr(node_id), "0s");
    let key = i_to_key(20);
    let value = b"async-resolve-fallback".to_vec();
    let start_ts = client.get_ts();
    let commit_ts = client.get_ts();
    let txn = TxnMutations::from_normal(vec![Mutation {
        key: key.clone().into(),
        value: value.clone().into(),
        op: Op::Put,
        ..Default::default()
    }]);
    client
        .kv_prewrite_with_retry(txn.primary(), None, txn, start_ts)
        .expect("kv_prewrite");

    let max_ts = client.get_ts().into_inner();
    assert_eq!(
        scan_locks_in_range(
            &mut client,
            key.as_slice(),
            i_to_key(21).as_slice(),
            max_ts,
            10,
        )
        .len(),
        1
    );

    let schedule_guard =
        FailpointGuard::new().cfg("before_async_batch_resolve_lock_schedule", "pause");
    send_async_region_resolve_lock(
        &mut client,
        &key,
        start_ts.into_inner(),
        commit_ts.into_inner(),
    );

    let max_ts = client.get_ts().into_inner();
    assert!(
        scan_locks_in_range(
            &mut client,
            key.as_slice(),
            i_to_key(21).as_slice(),
            max_ts,
            10,
        )
        .is_empty(),
        "async ResolveLock should synchronously clear the lock when async batching is disabled"
    );
    assert_eq!(
        client
            .get_key_version(&key, u64::MAX)
            .expect("get committed key should succeed"),
        Some(value)
    );
    drop(schedule_guard);

    update_async_resolve_lock_batch_interval(cluster.status_addr(node_id), "50ms");
    let key = i_to_key(21);
    let value = b"async-resolve-after-online-enable".to_vec();
    let start_ts = client.get_ts();
    let commit_ts = client.get_ts();
    let txn = TxnMutations::from_normal(vec![Mutation {
        key: key.clone().into(),
        value: value.clone().into(),
        op: Op::Put,
        ..Default::default()
    }]);
    client
        .kv_prewrite_with_retry(txn.primary(), None, txn, start_ts)
        .expect("kv_prewrite");

    let max_ts = client.get_ts().into_inner();
    assert_eq!(
        scan_locks_in_range(
            &mut client,
            key.as_slice(),
            i_to_key(22).as_slice(),
            max_ts,
            10,
        )
        .len(),
        1
    );

    let schedule_guard =
        FailpointGuard::new().cfg("before_async_batch_resolve_lock_schedule", "pause");
    send_async_region_resolve_lock(
        &mut client,
        &key,
        start_ts.into_inner(),
        commit_ts.into_inner(),
    );
    let max_ts = client.get_ts().into_inner();
    assert_eq!(
        scan_locks_in_range(
            &mut client,
            key.as_slice(),
            i_to_key(22).as_slice(),
            max_ts,
            10,
        )
        .len(),
        1,
        "async ResolveLock should stay pending while async scheduling is paused"
    );
    drop(schedule_guard);
    assert!(
        try_wait(
            || {
                let mut check_client = cluster.new_client();
                let max_ts = check_client.get_ts().into_inner();
                scan_locks_in_range(
                    &mut check_client,
                    key.as_slice(),
                    i_to_key(22).as_slice(),
                    max_ts,
                    10,
                )
                .is_empty()
            },
            10
        ),
        "async ResolveLock should clear the lock after async scheduling resumes"
    );
    assert_eq!(
        client
            .get_key_version(&key, u64::MAX)
            .expect("get committed key should succeed"),
        Some(value)
    );

    cluster.stop();
}

#[test]
fn test_async_region_resolve_lock_retries_after_leader_change() {
    test_util::init_log_for_test();
    let nodes = alloc_node_id_vec(3);
    let mut cluster = ServerCluster::new(nodes, |_, _| {});
    for config in cluster.get_node_configs().values() {
        assert_eq!(
            config.storage.scheduler_async_resolve_lock_batch_interval,
            ReadableDuration::millis(500)
        );
    }
    cluster.wait_region_replicated(&[], 3);
    let pd_client = cluster.get_pd_client();
    pd_client.disable_default_operator();
    let mut client = cluster.new_client();

    // Create a real prewrite lock so the background worker resolves the lock
    // through the same scan-and-resolve path used in production.
    let key = i_to_key(20);
    let value = b"async-resolve-leader-change".to_vec();
    let start_ts = client.get_ts();
    let commit_ts = client.get_ts();
    let txn = TxnMutations::from_normal(vec![Mutation {
        key: key.clone().into(),
        value: value.clone().into(),
        op: Op::Put,
        ..Default::default()
    }]);
    client
        .kv_prewrite_with_retry(txn.primary(), None, txn, start_ts)
        .expect("kv_prewrite");
    let max_ts = client.get_ts().into_inner();
    assert_eq!(
        scan_locks_in_range(
            &mut client,
            key.as_slice(),
            i_to_key(21).as_slice(),
            max_ts,
            10,
        )
        .len(),
        1
    );

    let region = client.get_region_by_key(&key);
    let (_, old_leader) = block_on(pd_client.get_region_leader_by_id(region.id))
        .unwrap()
        .unwrap();
    let new_leader = region
        .peers
        .iter()
        .find(|peer| peer.get_id() != old_leader.get_id())
        .unwrap()
        .clone();

    // Pause after the request is queued in the original leader's background worker,
    // then transfer the leader before the worker schedules the actual resolve
    // command.
    let stale_task_failed = Arc::new(AtomicBool::new(false));
    let stale_task_failed_clone = stale_task_failed.clone();
    let _stale_task_failed_guard = FailpointGuard::new().cfg_callback(
        "async_batch_resolve_lock_background_task_failed",
        move || stale_task_failed_clone.store(true, Ordering::SeqCst),
    );
    let (schedule_hit, schedule_guard) =
        cfg_notify_once("before_async_batch_resolve_lock_schedule");
    send_async_region_resolve_lock(
        &mut client,
        &key,
        start_ts.into_inner(),
        commit_ts.into_inner(),
    );
    let resume_schedule =
        block_on(schedule_hit).expect("async resolve should reach the scheduling failpoint");

    pd_client.must_transfer_leader(region.id, new_leader);
    let _ = resume_schedule.send(());
    drop(schedule_guard);

    // Wait until the stale leader task has actually observed the region error,
    // then verify it did not resolve the lock under an outdated leader context.
    assert!(
        try_wait(|| stale_task_failed.load(Ordering::SeqCst), 10),
        "the stale-leader background resolve should fail before retry"
    );
    client.clear_region_cache();
    client.put_kv(22..23, i_to_key, |_| b"leader-ready".to_vec());
    let max_ts = client.get_ts().into_inner();
    assert_eq!(
        scan_locks_in_range(
            &mut client,
            key.as_slice(),
            i_to_key(21).as_slice(),
            max_ts,
            10,
        )
        .len(),
        1,
        "the stale-leader background resolve should not clear the lock"
    );

    // A later async resolve request against the new leader should still be able to
    // schedule a fresh background task and finish the transaction.
    let mut retry_client = cluster.new_client();
    send_async_region_resolve_lock(
        &mut retry_client,
        &key,
        start_ts.into_inner(),
        commit_ts.into_inner(),
    );
    assert!(
        try_wait(
            || {
                let mut check_client = cluster.new_client();
                let max_ts = check_client.get_ts().into_inner();
                scan_locks_in_range(
                    &mut check_client,
                    key.as_slice(),
                    i_to_key(21).as_slice(),
                    max_ts,
                    10,
                )
                .is_empty()
            },
            10
        ),
        "retry on the new leader should clear the lock"
    );
    assert_eq!(
        retry_client
            .get_key_version(&key, u64::MAX)
            .expect("get committed key should succeed"),
        Some(value)
    );

    cluster.stop();
}