// Copyright 2024 PingCAP, Inc.

use std::{cmp::Ordering, convert::TryFrom, iter::Iterator as StdIterator, mem, ops::Deref};

use api_version::{ApiV2, api_v2::KEYSPACE_PREFIX_LEN};
use concurrency_manager::ConcurrencyManager;
use kvengine::{
    Iterator, LOCK_CF, SnapAccess, UserMeta, WRITE_CF,
    table::{
        BoundedDataSet, InnerKey, PessimisticAction,
        txn_file::{OP_CHECK_NOT_EXIST, OP_INSERT, TxnCtx, TxnFile, TxnFileId, TxnFileIterator},
    },
    txn_chunk_manager::TxnChunkManager,
};
use kvenginepb::TxnFileRef;
use kvproto::{
    errorpb,
    kvrpcpb::{CommandPri, WriteConflictReason},
};
use log_wrappers::Value as LogValue;
use protobuf::Message;
use raftstore::coprocessor::{RegionInfoAccessor, RegionInfoProvider};
use tikv_kv::{Snapshot, TxnFileWriteData, WriteData};
use tikv_util::{
    Either,
    metrics::{ServerIsBusyType, new_server_is_busy},
};
use txn_types::{Key, LockType, TimeStamp, Write, WriteType};

use crate::storage::{
    ProcessResult, TxnStatus, kv,
    lock_manager::LockManager,
    mvcc::{
        self, CloudReader, Error as MvccError, ErrorInner, MvccTxn, Result as MvccResult,
        SharedLockUpgradeHint, TxnCommitRecord,
        metrics::MVCC_COMMIT_REJECT_BY_BACKUP_TS_COUNTER_VEC,
    },
    txn::{
        Command, Error, Lock, Result,
        commands::{
            CommandExt, Prewrite, PrewritePessimistic, ReleasedLocks, ResolveLock,
            ResolveLockReadPhase, ResponsePolicy, WriteCommand, WriteContext, WriteResult,
            check_txn_status::update_max_ts_from_check_txn_status,
        },
        latch::TxnFileLatchTarget,
    },
    types,
};

const TXN_FILE_LOCK_LIMIT_REACHED_REASON: &str = "txn file lock limit reached";

pub(crate) fn reject_new_txn_file_prewrite_by_lock_limit(
    start_ts: u64,
    region_id: u64,
    visible_lock_txn_files: &[TxnFile],
    threshold: usize,
) -> Option<kv::Error> {
    let visible_lock_count = visible_lock_txn_files.len();
    if visible_lock_count < threshold
        || visible_lock_txn_files
            .iter()
            .any(|txn_file| txn_file.start_ts() == start_ts)
    {
        return None;
    }

    crate::storage::metrics::TXN_FILE_PREWRITE_REJECT_BY_LOCK_LIMIT_COUNTER.inc();
    debug!("reject txn-file prewrite by lock limit";
        "region_id" => region_id,
        "start_ts" => start_ts,
        "visible_lock_count" => visible_lock_count,
        "threshold" => threshold,
    );
    let mut err = errorpb::Error::default();
    err.set_server_is_busy(new_server_is_busy(
        ServerIsBusyType::TxnFile,
        TXN_FILE_LOCK_LIMIT_REACHED_REASON,
    ));
    Some(kv::ErrorInner::Request(err).into())
}

enum TxnFileLockMergeResult {
    NotFound,
    FoundEqual,
    FoundUpdated,
}

#[derive(Clone, Debug)]
struct TxnFileSnapshot {
    txn_file_ref: TxnFileRef,
    txn_file: TxnFile,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum TxnFileSnapshotCompare {
    Same,
    ShardVersionMismatch,
    Mismatch,
}

#[derive(Clone)]
pub(crate) struct TxnFileCommandContext {
    region_info_accessor: Option<RegionInfoAccessor>,
    max_txn_file_locks_per_region: usize,
}

impl TxnFileCommandContext {
    pub(crate) fn new(
        region_info_accessor: Option<RegionInfoAccessor>,
        max_txn_file_locks_per_region: usize,
    ) -> Self {
        Self {
            region_info_accessor,
            max_txn_file_locks_per_region,
        }
    }
}

pub struct TxnFileCommand {
    start_ts: u64,
    txn_file_snapshot: Option<TxnFileSnapshot>,
    modified: bool,
    pub lock_prefix: Option<txn_types::Lock>,
    pub inner_cmd: Option<Box<Command>>,

    context: TxnFileCommandContext,
}

impl TxnFileCommand {
    fn ensure_no_shared_lock_upgrade_hints(
        shared_lock_upgrade_hints: &[SharedLockUpgradeHint],
    ) -> Result<()> {
        if !shared_lock_upgrade_hints.is_empty() {
            return Err(box_err!(
                "txn-file command emitted shared-lock upgrade hints without per-key latches"
            ));
        }
        Ok(())
    }

    ///   Return a  `TypedCommand`  that encapsulates the result of executing
    /// this command.
    fn new(
        inner_cmd: Command,
        start_ts: u64,
        txn_file_snapshot: Option<TxnFileSnapshot>,
        context: TxnFileCommandContext,
    ) -> Result<Command> {
        let lock_prefix = txn_file_snapshot.as_ref().map_or(Ok(None), |snapshot| {
            Self::build_lock_prefix(&snapshot.txn_file_ref)
        })?;
        debug!(
            "new txn file command {:?}, start_ts {}, txn_file_snapshot {:?}, lock {:?}",
            inner_cmd, start_ts, txn_file_snapshot, lock_prefix
        );
        Ok(Command::TxnFile(TxnFileCommand {
            inner_cmd: Some(Box::new(inner_cmd)),
            modified: false,
            start_ts,
            lock_prefix,
            txn_file_snapshot,
            context,
        }))
    }

    pub fn priority(&self) -> CommandPri {
        // Force to low priority for prewrite as it's of high throughput but latency
        // insensitive.
        // Force to high priority for
        // rollback/check_txn_status/resolve_lock_lite/resolve_lock to improve the
        // efficiency of cleaning up txn file locks.
        // Ref: Command::need_flow_control().
        match self.inner_cmd.as_ref().unwrap().as_ref() {
            Command::Prewrite(_) | Command::PrewritePessimistic(_) => CommandPri::Low,
            Command::Rollback(_)
            | Command::CheckTxnStatus(_)
            | Command::ResolveLockLite(_)
            | Command::ResolveLock(_) => CommandPri::High,
            cmd => cmd.priority(),
        }
    }

    fn build_prewrite_txn_file_ref(req: &Prewrite, snap: &SnapAccess) -> TxnFileRef {
        let mut lock = txn_types::Lock::new(
            LockType::Put,
            req.primary.clone(),
            req.start_ts,
            req.lock_ttl,
            None,
            TimeStamp::zero(),
            req.txn_size,
            req.min_commit_ts,
        )
        .set_txn_source(req.get_ctx().get_txn_source());
        lock.is_txn_file = true;
        let lock_prefix = lock.to_bytes();
        let mut txn_file_ref = TxnFileRef::new();
        txn_file_ref.set_lock_val_prefix(lock_prefix);
        txn_file_ref.set_shard_ver(req.get_ctx().get_region_epoch().get_version());
        txn_file_ref.set_start_ts(req.start_ts.into_inner());
        txn_file_ref.set_chunk_ids(req.txn_file_chunks.clone());
        // Prepend keyspace to generate TxnFileRef the same as txn chunk files.
        txn_file_ref.set_inner_lower_bound(Self::prepend_keyspace(snap.get_inner_start(), snap));
        txn_file_ref.set_inner_upper_bound(Self::prepend_keyspace(snap.get_inner_end(), snap));
        txn_file_ref
    }

    fn build_prewrite_pessimistic_txn_file_ref(
        req: &PrewritePessimistic,
        snap: &SnapAccess,
    ) -> TxnFileRef {
        let mut lock = txn_types::Lock::new(
            LockType::Put,
            req.primary.clone(),
            req.start_ts,
            req.lock_ttl,
            None,
            req.for_update_ts,
            req.txn_size,
            req.min_commit_ts,
        )
        .set_txn_source(req.get_ctx().get_txn_source());
        lock.is_txn_file = true;
        let lock_prefix = lock.to_bytes();
        let mut txn_file_ref = TxnFileRef::new();
        txn_file_ref.set_lock_val_prefix(lock_prefix);
        txn_file_ref.set_shard_ver(req.get_ctx().get_region_epoch().get_version());
        txn_file_ref.set_start_ts(req.start_ts.into_inner());
        txn_file_ref.set_chunk_ids(req.txn_file_chunks.clone());
        // Prepend keyspace to generate TxnFileRef the same as txn chunk files.
        txn_file_ref.set_inner_lower_bound(Self::prepend_keyspace(snap.get_inner_start(), snap));
        txn_file_ref.set_inner_upper_bound(Self::prepend_keyspace(snap.get_inner_end(), snap));
        txn_file_ref
    }

    fn is_prewrite_cmd(cmd: &Command) -> bool {
        matches!(cmd, Command::Prewrite(_) | Command::PrewritePessimistic(_))
    }

    fn try_build_prewrite_txn_file_ref(cmd: &Command, snap: &SnapAccess) -> Result<TxnFileRef> {
        debug_assert!(cmd.can_build_txn_file());
        Ok(match cmd {
            Command::Prewrite(req) => {
                if !req.mutations.is_empty() {
                    return Err(box_err!(
                        "invalid txn file prewrite with mutations {:?}",
                        cmd
                    ));
                }
                Self::build_prewrite_txn_file_ref(req, snap)
            }
            Command::PrewritePessimistic(req) => {
                if !req.mutations.is_empty() {
                    return Err(box_err!(
                        "invalid txn file pessimistic prewrite with mutations {:?}",
                        cmd
                    ));
                }
                Self::build_prewrite_pessimistic_txn_file_ref(req, snap)
            }
            _ => unreachable!("unexpected command {:?}", cmd),
        })
    }

    pub(crate) fn try_from<S: Snapshot>(
        cmd: Command,
        snapshot: S,
        txn_chunk_manager: TxnChunkManager,
        context: TxnFileCommandContext,
    ) -> Result<Command> {
        let snap = snapshot.get_kvengine_snap().unwrap();
        let start_ts = Self::command_start_ts(&cmd);
        let txn_file_snapshot = if Self::is_prewrite_cmd(&cmd) {
            let txn_file_ref = Self::try_build_prewrite_txn_file_ref(&cmd, snap)?;
            let mut txn_chunks = Vec::with_capacity(txn_file_ref.chunk_ids.len());
            let txn_ctx = TxnCtx::from_txn_file_ref(&txn_file_ref);
            for &chunk_id in &txn_file_ref.chunk_ids {
                let txn_chunk = txn_chunk_manager
                    .get(chunk_id)
                    .ok_or_else(|| -> Error { box_err!("txn chunk not prepared: {}", chunk_id) })?;
                txn_chunks.push(txn_chunk);
            }
            let txn_file_id = TxnFileId::new(cmd.ctx().region_id, txn_file_ref.shard_ver, start_ts);
            let txn_file = TxnFile::new(txn_file_id, txn_chunks, txn_ctx).unwrap();
            Some(TxnFileSnapshot {
                txn_file_ref,
                txn_file,
            })
        } else {
            Self::build_txn_file_snapshot(snap, start_ts)
        };
        Self::new(cmd, start_ts, txn_file_snapshot, context)
    }

    fn build_txn_file_snapshot(snap: &SnapAccess, start_ts: u64) -> Option<TxnFileSnapshot> {
        let txn_file = snap.get_lock_txn_file(start_ts)?;
        let mut txn_file_ref = TxnFileRef::new();
        txn_file_ref.set_start_ts(start_ts);
        txn_file_ref.set_shard_ver(snap.get_version());
        txn_file_ref.set_chunk_ids(txn_file.chunk_ids());
        txn_file_ref.set_inner_lower_bound(Self::prepend_keyspace(txn_file.lower_bound(), snap));
        txn_file_ref.set_inner_upper_bound(Self::prepend_keyspace(txn_file.upper_bound(), snap));
        txn_file_ref.set_lock_val_prefix(txn_file.get_lock_val_prefix().to_vec());
        Some(TxnFileSnapshot {
            txn_file_ref,
            txn_file,
        })
    }

    fn prepend_keyspace(key: InnerKey<'_>, snap: &SnapAccess) -> Vec<u8> {
        if let Some(keyspace_id) = snap.prepend_keyspace_id() {
            let mut buf = ApiV2::get_txn_keyspace_prefix(keyspace_id);
            buf.extend_from_slice(key.deref());
            buf
        } else {
            key.to_vec()
        }
    }

    fn validate_and_refresh_non_prewrite_lock(
        &mut self,
        cmd: &Command,
        snap: &SnapAccess,
    ) -> Result<()> {
        if Self::is_prewrite_cmd(cmd) {
            return Ok(());
        }
        let current_snapshot = Self::build_txn_file_snapshot(snap, self.start_ts);
        match Self::compare_optional_txn_file_snapshot(
            self.txn_file_snapshot.as_ref(),
            current_snapshot.as_ref(),
        ) {
            TxnFileSnapshotCompare::Same => {
                self.txn_file_snapshot = current_snapshot;
                self.lock_prefix = self
                    .txn_file_snapshot
                    .as_ref()
                    .map_or(Ok(None), |snapshot| {
                        Self::build_lock_prefix(&snapshot.txn_file_ref)
                    })?;
                Ok(())
            }
            TxnFileSnapshotCompare::ShardVersionMismatch => {
                Err(self.epoch_not_match_region_error(cmd.ctx().get_region_id()))
            }
            TxnFileSnapshotCompare::Mismatch => {
                Err(Self::stale_command_region_error(self.start_ts))
            }
        }
    }

    fn get_conflict_lock(
        &self,
        snap_access: &SnapAccess,
    ) -> Option<(Vec<u8>, txn_types::LockOrSharedLocks)> {
        if let Some((key, conflict_lock)) = snap_access.get_txn_file_conflict_lock(self.txn_file())
        {
            return Some((key, conflict_lock));
        }
        None
    }

    #[maybe_async::both]
    async fn get_conflict_write(&self, snap_access: &SnapAccess) -> Option<(Vec<u8>, UserMeta)> {
        if let Some((key, um)) = snap_access
            .get_txn_file_conflict_write(self.txn_file())
            .await
        {
            return Some((key, um));
        }
        None
    }

    // TODO: choose best match algorithm based on stats.
    #[maybe_async::both]
    async fn check_constraint(&self, snap_access: &SnapAccess) -> MvccResult<()> {
        if let Some(already_exist_key) = snap_access
            .check_txn_file_constraint(self.txn_file(), self.ts().into_inner())
            .await
        {
            return Err(ErrorInner::AlreadyExist {
                key: already_exist_key,
            }
            .into());
        }
        Ok(())
    }

    #[maybe_async::both]
    #[allow(dead_code)]
    async fn check_constraint_by_linear_match(&self, snap_access: &SnapAccess) -> MvccResult<()> {
        let mut txn_file_iter = TxnFileIterator::new(self.txn_file().clone(), false);
        let prefix = if snap_access.get_keyspace_id() > 0 {
            ApiV2::get_keyspace_prefix(snap_access.get_start_key()).unwrap()
        } else {
            &[]
        };
        let prefix_len = prefix.len();
        let mut start_key = prefix.to_vec();
        start_key.extend_from_slice(self.txn_file().smallest().deref());
        let mut end_key = prefix.to_vec();
        end_key.extend_from_slice(self.txn_file().biggest().deref());
        end_key.push(0);
        let mut snap_iter = snap_access
            .new_iterator(WRITE_CF, false, false, Some(self.ts().into_inner()), true)
            .await;
        snap_iter.set_range(start_key.into(), end_key.into()).await;
        txn_file_iter.rewind();
        while txn_file_iter.valid() && snap_iter.valid() {
            let txn_key = txn_file_iter.key();
            let snap_key = snap_iter.key();
            let txn_op = txn_file_iter.get_op();
            match &snap_key[prefix_len..].cmp(txn_key.deref()) {
                Ordering::Less => {
                    snap_iter.next().await;
                }
                Ordering::Equal => {
                    if !snap_iter.val().is_empty()
                        && (txn_op == OP_INSERT || txn_op == OP_CHECK_NOT_EXIST)
                    {
                        return Err(ErrorInner::AlreadyExist {
                            key: snap_iter.key().to_vec(),
                        }
                        .into());
                    }
                    snap_iter.next().await;
                    txn_file_iter.next();
                }
                Ordering::Greater => {
                    txn_file_iter.next();
                }
            }
        }
        Ok(())
    }

    fn try_merge_txn_file_locks(
        &mut self,
        snap_access: &SnapAccess,
    ) -> MvccResult<TxnFileLockMergeResult> {
        if let Some(existed) = snap_access.get_lock_txn_file(self.txn_file().start_ts()) {
            // The only scene to merge txn file locks is region merge after parts of txn
            // chunks have been prewrite.
            if existed.lower_bound() < self.txn_file().lower_bound()
                || self.txn_file().upper_bound() < existed.upper_bound()
                || self.txn_file().id().shard_id != existed.id().shard_id
                || self.txn_file().shard_ver() < existed.shard_ver()
            {
                error!("unexpected txn file locks of same transaction"; "start_ts" => self.ts(), "txn_file" => ?self.txn_file(), "existed" => ?existed);
                return Err(box_err!(
                    "unexpected txn file locks of same transaction, start_ts: {}",
                    self.ts()
                ));
            }

            if self.txn_file().chunk_ids() == existed.chunk_ids()
                && self.txn_file().lower_bound() == existed.lower_bound()
                && self.txn_file().upper_bound() == existed.upper_bound()
            {
                info!("txn file locks are equal"; "start_ts" => self.ts(), "txn_file" => ?self.txn_file(), "existed" => ?existed);
                return Ok(TxnFileLockMergeResult::FoundEqual);
            }

            let merged_chunks = TxnFile::merge_chunks(self.txn_file(), &existed);
            let mut txn_ctx = self.txn_file().txn_ctx().clone();
            // Use existed `lock_val_prefix` as it would be updated, e.g, by
            // check_txn_status.
            txn_ctx.set_lock_val_prefix(existed.lock_val_prefix());
            let new_txn_file = TxnFile::new(self.txn_file().id(), merged_chunks, txn_ctx)?;
            new_txn_file.validate()?;

            debug!("merge txn file locks"; "start_ts" => self.ts(), "req" => ?self.txn_file(), "existed" => ?existed, "merged" => ?new_txn_file);
            let new_chunk_ids = new_txn_file.chunk_ids();
            let new_lock_val_prefix = new_txn_file.get_lock_val_prefix().to_vec();
            let snapshot = self.txn_file_snapshot_mut()?;
            snapshot.txn_file = new_txn_file;
            snapshot.txn_file_ref.set_chunk_ids(new_chunk_ids);
            snapshot
                .txn_file_ref
                .set_lock_val_prefix(new_lock_val_prefix);
            return Ok(TxnFileLockMergeResult::FoundUpdated);
        }
        Ok(TxnFileLockMergeResult::NotFound)
    }

    #[maybe_async::both]
    async fn process_prewrite(
        &mut self,
        snap_access: &SnapAccess,
        concurrency_manager: &ConcurrencyManager,
    ) -> MvccResult<ProcessResult> {
        self.txn_file().validate()?;

        if !self.txn_file().has_data_in_bound(snap_access.data_bound()) {
            error!("process_prewrite: txn file has no data in region"; "start_ts" => self.ts());
            return Err(box_err!("txn file has no data in region"));
        }

        let merge_result = self.try_merge_txn_file_locks(snap_access)?;
        if matches!(merge_result, TxnFileLockMergeResult::FoundEqual) {
            info!("process_prewrite: ignore duplicated prewrite"; "start_ts" => self.ts());
        } else {
            if matches!(merge_result, TxnFileLockMergeResult::NotFound) {
                if let Some(err) = reject_new_txn_file_prewrite_by_lock_limit(
                    self.txn_file().start_ts(),
                    self.txn_file().id().shard_id,
                    snap_access.get_lock_txn_files(),
                    self.context.max_txn_file_locks_per_region,
                ) {
                    return Err(err.into());
                }
            }
            if let Some((key, conflict_lock)) = self.get_conflict_lock(snap_access) {
                match conflict_lock {
                    Either::Left(conflict_lock) => {
                        if conflict_lock.ts != self.ts() {
                            return Err(
                                ErrorInner::KeyIsLocked(conflict_lock.into_lock_info(key)).into()
                            );
                        }
                        error!("unexpected lock of same transaction"; "start_ts" => self.ts(), "conflict_lock" => ?conflict_lock, "key" => &LogValue::key(&key));
                        return Err(box_err!(
                            "unexpected lock of same transaction, start_ts: {}",
                            self.ts()
                        ));
                    }
                    Either::Right(shared_locks) => {
                        return Err(
                            ErrorInner::KeyIsLocked(shared_locks.into_lock_info(key)).into()
                        );
                    }
                }
            }
            if let Some(commit_ts) = self.check_txn_commit_record(snap_access).await? {
                // already committed
                let lock = self.lock_prefix.as_ref().unwrap();
                let start_ts = lock.ts;
                let key = lock.primary.clone();
                return Err(ErrorInner::Committed {
                    start_ts,
                    commit_ts,
                    key,
                }
                .into());
            }
            if let Some((key, conflict_um)) = self.get_conflict_write(snap_access).await {
                let lock_prefix = self.lock_prefix.as_ref().unwrap();
                return Err(ErrorInner::WriteConflict {
                    start_ts: self.ts(),
                    conflict_start_ts: conflict_um.start_ts.into(),
                    conflict_commit_ts: conflict_um.commit_ts.into(),
                    key,
                    primary: lock_prefix.primary.clone(),
                    reason: Default::default(),
                }
                .into());
            }
            if self.txn_file().get_inserts() + self.txn_file().get_check_not_exists() > 0 {
                self.check_constraint(snap_access).await?;
                let start_ts = self.lock_prefix.as_ref().unwrap().ts;
                concurrency_manager.update_max_ts(start_ts, || format!("prewrite-{}", start_ts))?;
            }
            self.modified = true;
        }
        Ok(ProcessResult::PrewriteResult {
            result: types::PrewriteResult {
                locks: vec![],
                min_commit_ts: Default::default(),
                one_pc_commit_ts: Default::default(),
            },
        })
    }

    #[maybe_async::both]
    async fn process_pessimistic_prewrite(
        &mut self,
        snap_access: &SnapAccess,
        concurrency_manager: &ConcurrencyManager,
        extra_cf_scan_mode: crate::storage::config::ExtraCfScanMode,
    ) -> MvccResult<ProcessResult> {
        self.txn_file().validate()?;

        info!(
            "txn file process_pessimistic_prewrite, for update ts: {}",
            self.lock_prefix.as_ref().unwrap().for_update_ts
        );

        if !self.txn_file().has_data_in_bound(snap_access.data_bound()) {
            error!("process_pessimistic_prewrite: txn file has no data in region"; "start_ts" => self.ts());
            return Err(box_err!("txn file has no data in region"));
        }

        if !self.txn_file().support_pessimistic_txn() {
            error!(
                "process_pessimistic_prewrite: the txn file format version does not support pessimistic txn"
            );
            return Err(box_err!(
                "txn file format outdated, does not support pessimistic txn"
            ));
        }

        let merge_result = self.try_merge_txn_file_locks(snap_access)?;
        if matches!(merge_result, TxnFileLockMergeResult::FoundEqual) {
            info!("process_prewrite: ignore duplicated prewrite"; "start_ts" => self.ts());
        } else {
            if matches!(merge_result, TxnFileLockMergeResult::NotFound) {
                if let Some(err) = reject_new_txn_file_prewrite_by_lock_limit(
                    self.txn_file().start_ts(),
                    self.txn_file().id().shard_id,
                    snap_access.get_lock_txn_files(),
                    self.context.max_txn_file_locks_per_region,
                ) {
                    return Err(err.into());
                }
            }
            if let Some(commit_ts) = self.check_txn_commit_record(snap_access).await? {
                // already committed
                let lock = self.lock_prefix.as_ref().unwrap();
                let start_ts = lock.ts;
                let key = lock.primary.clone();
                return Err(ErrorInner::Committed {
                    start_ts,
                    commit_ts,
                    key,
                }
                .into());
            }

            let start_ts = self.ts();
            let for_update_ts = self.lock_prefix.as_ref().unwrap().for_update_ts;
            let reader =
                CloudReader::new_with_scan_mode(snap_access.clone(), true, extra_cf_scan_mode);
            let current_raw_key = ApiV2::get_keyspace_prefix_by_id(snap_access.get_keyspace_id());
            let mut checker = TxnFilePessimisticActionChecker {
                err: None,
                start_ts,
                reader,
                for_update_ts,
                primary_key: self.lock_prefix.as_ref().unwrap().primary.clone(),
                current_raw_key,
            };
            self.txn_file()
                .iter_check_entries(snap_access.data_bound(), &mut checker)
                .await;
            if let Some(err) = checker.err.take() {
                return Err(err);
            }
            concurrency_manager.update_max_ts(start_ts, || format!("prewrite-{}", start_ts))?;
            self.modified = true;
        }

        Ok(ProcessResult::PrewriteResult {
            result: types::PrewriteResult {
                locks: vec![],
                min_commit_ts: Default::default(),
                one_pc_commit_ts: Default::default(),
            },
        })
    }

    #[maybe_async::both]
    async fn check_txn_commit_record(
        &mut self,
        snap_access: &SnapAccess,
    ) -> MvccResult<Option<TimeStamp>> {
        let lock_prefix = self.lock_prefix.as_ref().unwrap();
        let primary = lock_prefix.primary.as_slice();
        if snap_access.get_start_key() <= primary && primary < snap_access.get_end_key() {
            let mut cloud_reader = CloudReader::new(snap_access.clone(), true);
            let primary = Key::from_raw(&lock_prefix.primary);
            if let TxnCommitRecord::SingleRecord { commit_ts, write } = cloud_reader
                .get_txn_commit_record(&primary, lock_prefix.ts)
                .await?
            {
                if write.write_type == WriteType::Rollback {
                    return Err(ErrorInner::WriteConflict {
                        start_ts: write.start_ts,
                        conflict_start_ts: write.start_ts,
                        conflict_commit_ts: write.start_ts,
                        key: lock_prefix.primary.clone(),
                        primary: lock_prefix.primary.clone(),
                        reason: WriteConflictReason::SelfRolledBack,
                    }
                    .into());
                } else {
                    // already committed
                    return Ok(Some(commit_ts));
                }
            }
        }
        Ok(None)
    }

    // Return Ok(None) when all keys are out of region, and the status of txn is
    // unknown.
    #[maybe_async::both]
    async fn check_txn_commit_record_by_keys(
        &self,
        keys: &[Key],
        start_ts: TimeStamp,
        commit_ts: TimeStamp,
        snap_access: &SnapAccess,
    ) -> MvccResult<Option<TimeStamp>> {
        let mut reader = CloudReader::new(snap_access.clone(), true);
        for key in keys {
            let raw_key = key.to_raw()?;
            if snap_access.get_start_key() <= raw_key.as_slice()
                && raw_key.as_slice() < snap_access.get_end_key()
            {
                return match reader.get_txn_commit_record(key, start_ts).await?.info() {
                    Some((_, WriteType::Rollback)) | None => {
                        // None: related Rollback has been collapsed.
                        // Rollback: rollback by concurrent transaction.
                        info!(
                            "txn conflict (lock not found)";
                            "key" => %key,
                            "start_ts" => start_ts,
                            "commit_ts" => commit_ts,
                        );
                        Err(ErrorInner::TxnLockNotFound {
                            start_ts,
                            commit_ts,
                            key: raw_key,
                        }
                        .into())
                    }
                    // Committed by concurrent transaction.
                    Some((commit_ts, WriteType::Put))
                    | Some((commit_ts, WriteType::Delete))
                    | Some((commit_ts, WriteType::Lock)) => Ok(Some(commit_ts)),
                };
            }
        }
        Ok(None)
    }

    /// Check the transaction status when a "not primary txn file lock" commit
    /// to primary region (i.e. the region contains primary key).
    ///
    /// This condition will happen when the primary txn file lock is committed
    /// or rollback, and then the primary region is merged.
    ///
    /// See https://github.com/tidbcloud/cloud-storage-engine/issues/1800.
    #[maybe_async::both]
    async fn check_commit_primary_region(
        &self,
        lock: &txn_types::Lock,
        lock_txn_file: &TxnFile,
        start_ts: TimeStamp,
        commit_ts: TimeStamp,
        snap_access: &SnapAccess,
    ) -> MvccResult<bool /* is_primary */> {
        if !snap_access.key_is_in_range(&lock.primary) {
            return Ok(false);
        }

        let is_primary_txn_file =
            lock_txn_file.contains_key(InnerKey::from_outer_key(&lock.primary));
        if !is_primary_txn_file {
            let txn_lock_not_found = || -> mvcc::Error {
                ErrorInner::TxnLockNotFound {
                    start_ts,
                    commit_ts,
                    key: vec![],
                }
                .into()
            };

            let mut cloud_reader = CloudReader::new(snap_access.clone(), true);
            let record = cloud_reader
                .get_txn_commit_record(&Key::from_raw(&lock.primary), start_ts)
                .await?
                .info();
            return match record {
                Some((_, WriteType::Rollback)) | None => Err(txn_lock_not_found()),
                Some((committed_ts, WriteType::Put | WriteType::Delete | WriteType::Lock)) => {
                    debug_assert_eq!(
                        committed_ts, commit_ts,
                        "commit_ts mismatch: start_ts {}, commit_ts {}, committed_ts {}",
                        start_ts, commit_ts, committed_ts
                    );
                    Ok(is_primary_txn_file)
                }
            };
        }
        Ok(is_primary_txn_file)
    }

    #[maybe_async::both]
    async fn process_commit(
        &mut self,
        snap_access: &SnapAccess,
        commit_ts: TimeStamp,
        keys: Option<Vec<Key>>,
        txn: &mut MvccTxn,
    ) -> MvccResult<ProcessResult> {
        let start_ts = self.ts();
        if commit_ts <= start_ts {
            return Err(box_err!(
                "Invalid transaction tso with start_ts:{start_ts}, commit_ts:{commit_ts}"
            ));
        }
        if let Some(lock_txn_file) = snap_access.get_lock_txn_file(self.start_ts) {
            // Txn-file metadata stores only ordinary lock prefixes. Shared-lock
            // wrappers from other transactions are handled in the live lock-CF
            // conflict path instead of here.
            let lock = txn_types::parse_non_shared_lock(lock_txn_file.get_lock_val_prefix())?;
            let is_primary = self
                .check_commit_primary_region(
                    &lock,
                    &lock_txn_file,
                    start_ts,
                    commit_ts,
                    snap_access,
                )
                .await?;

            debug_assert_eq!(lock.ts, self.ts());
            let mut min_commit_ts = lock.min_commit_ts;
            let mut reject_by_backup_ts = false;
            if is_primary {
                if let Some(backup_ts) = txn.backup_ts.as_ref() {
                    let backup_ts = backup_ts.get();
                    if min_commit_ts < backup_ts {
                        min_commit_ts = backup_ts;
                        reject_by_backup_ts = true;
                    }
                    txn.backup_ts_checked = true;
                }
            }
            if commit_ts < min_commit_ts {
                info!(
                    "trying to commit with smaller commit_ts than min_commit_ts";
                    "lock" => ?lock,
                    "start_ts" => start_ts,
                    "commit_ts" => commit_ts,
                    "min_commit_ts" => min_commit_ts,
                    "reject_by_backup_ts" => reject_by_backup_ts,
                );
                if reject_by_backup_ts {
                    MVCC_COMMIT_REJECT_BY_BACKUP_TS_COUNTER_VEC.txn_file.inc();
                }
                return Err(ErrorInner::CommitTsExpired {
                    start_ts,
                    commit_ts,
                    key: lock.primary.clone(),
                    min_commit_ts,
                }
                .into());
            }
            self.set_committed_or_rolled_back(commit_ts.into_inner())?;
            return Ok(ProcessResult::TxnStatus {
                txn_status: TxnStatus::committed(commit_ts),
            });
        } else if let Some(keys) = keys.as_ref() {
            if keys.is_empty() {
                // There is no way to check committed when there is no key in commit request.
                // So we just return committed.
                info!("process_commit: no keys in txn file commit request"; "start_ts" => self.ts(), "commit_ts" => commit_ts);
                return Ok(ProcessResult::TxnStatus {
                    txn_status: TxnStatus::committed(commit_ts),
                });
            }

            if let Some(_new_commit_ts) = self
                .check_txn_commit_record_by_keys(keys.as_slice(), start_ts, commit_ts, snap_access)
                .await?
            {
                // Normal txn return the commit_ts in request other than the committed record.
                // We keep the same here.
                return Ok(ProcessResult::TxnStatus {
                    txn_status: TxnStatus::committed(commit_ts),
                });
            }
        }
        Err(ErrorInner::TxnLockNotFound {
            start_ts: self.ts(),
            commit_ts,
            key: vec![],
        }
        .into())
    }

    #[maybe_async::both]
    async fn process_rollback(
        &mut self,
        snap_access: &SnapAccess,
        keys: Option<Vec<Key>>,
    ) -> MvccResult<ProcessResult> {
        let start_ts = self.ts();
        if let Some(lock_txn_file) = snap_access.get_lock_txn_file(self.start_ts) {
            // TODO: remove this check when it's stable enough.
            let lock = txn_types::parse_non_shared_lock(lock_txn_file.get_lock_val_prefix())?;
            debug_assert_eq!(lock.ts, start_ts);

            self.set_committed_or_rolled_back(0)?;
            Ok(ProcessResult::Res)
        } else if let Some(keys) = keys {
            let mut cloud_reader = CloudReader::new(snap_access.clone(), true);
            for key in keys {
                match self
                    .process_check_txn_status_missing_lock(&mut cloud_reader, key.clone())
                    .await?
                {
                    TxnStatus::RolledBack | TxnStatus::LockNotExist => {}
                    TxnStatus::Committed { commit_ts } => {
                        return Err(ErrorInner::Committed {
                            start_ts: self.ts(),
                            commit_ts,
                            key: key.into_raw()?,
                        }
                        .into());
                    }
                    _ => unreachable!(),
                }
            }
            // Note: normal txn will proceed to write the rollback record when lock not
            // exist. But we can not do so here because when lock file not
            // exist, we don't know all the keys.
            Ok(ProcessResult::Res)
        } else {
            Err(ErrorInner::TxnLockNotFound {
                start_ts,
                commit_ts: Default::default(),
                key: vec![],
            }
            .into())
        }
    }

    #[maybe_async::both]
    async fn process_check_txn_status(
        &mut self,
        snap_access: &SnapAccess,
        primary_key: Key,
        lock_ts: TimeStamp,
        current_ts: TimeStamp,
        caller_start_ts: TimeStamp,
        concurrency_manager: &ConcurrencyManager,
    ) -> MvccResult<ProcessResult> {
        update_max_ts_from_check_txn_status(
            concurrency_manager,
            lock_ts,
            current_ts,
            caller_start_ts,
        )?;
        let mut reader = CloudReader::new(snap_access.clone(), true);
        if let Some(lock_or_shared) = reader.load_lock(&primary_key).unwrap() {
            match lock_or_shared {
                Either::Left(lock) if lock.ts == self.start_ts.into() => {
                    return self.process_check_txn_status_lock_exists(
                        primary_key,
                        lock,
                        current_ts,
                        caller_start_ts,
                    );
                }
                Either::Right(shared_locks) => {
                    let raw_key = primary_key.clone().into_raw()?;
                    return Err(
                        ErrorInner::KeyIsLocked(shared_locks.into_lock_info(raw_key)).into(),
                    );
                }
                Either::Left(_) => {}
            }
        }
        let txn_status = self
            .process_check_txn_status_missing_lock(&mut reader, primary_key)
            .await?;

        // Resolve txn file locks merged from other regions.
        // See https://github.com/tidbcloud/cloud-storage-engine/issues/2182.
        if let Some(lock_txn_file) = snap_access.get_lock_txn_file(self.start_ts) {
            match &txn_status {
                TxnStatus::Committed { commit_ts } => {
                    info!("process_check_txn_status: commit txn file lock";
                        "start_ts" => self.ts(), "commit_ts" => commit_ts, "lock" => ?lock_txn_file);
                    self.set_committed_or_rolled_back(commit_ts.into_inner())?;
                }
                TxnStatus::RolledBack => {
                    info!("process_check_txn_status: rollback txn file lock"; "start_ts" => self.ts(), "lock" => ?lock_txn_file);
                    self.set_committed_or_rolled_back(0)?;
                }
                TxnStatus::LockNotExist => {
                    warn!("process_check_txn_status: lock may be lost"; "start_ts" => self.ts(), "lock" => ?lock_txn_file);
                    debug_assert!(false, "lock may be lost");
                    self.set_committed_or_rolled_back(0)?;
                }
                _ => unreachable!("unexpected txn status {:?}", txn_status),
            }
        }

        Ok(ProcessResult::TxnStatus { txn_status })
    }

    fn process_check_txn_status_lock_exists(
        &mut self,
        primary_key: Key,
        mut lock: txn_types::Lock,
        current_ts: TimeStamp,
        caller_start_ts: TimeStamp,
    ) -> MvccResult<ProcessResult> {
        if !primary_key.is_encoded_from(&lock.primary) {
            return Err(
                ErrorInner::PrimaryMismatch(lock.into_lock_info(primary_key.into_raw()?)).into(),
            );
        }

        let txn_status = if lock.ts.physical() + lock.ttl < current_ts.physical() {
            // rollback.
            self.set_committed_or_rolled_back(0)?;
            TxnStatus::TtlExpire
        } else {
            // If lock.min_commit_ts is 0, it's not a large transaction and we can't push
            // forward its min_commit_ts otherwise the transaction can't be committed by
            // old version TiDB during rolling update.
            if !lock.min_commit_ts.is_zero()
                && !caller_start_ts.is_max()
                // Push forward the min_commit_ts so that reading won't be blocked by locks.
                && caller_start_ts >= lock.min_commit_ts
            {
                lock.min_commit_ts = caller_start_ts.next();

                if lock.min_commit_ts < current_ts {
                    lock.min_commit_ts = current_ts;
                }

                lock.short_value = None;
                self.set_lock_val_prefix(lock.to_bytes())?;
                self.modified = true;
            }

            // As long as the primary lock's min_commit_ts > caller_start_ts, locks belong
            // to the same transaction can't block reading. Return MinCommitTsPushed
            // result to the client to let it bypass locks.
            let min_commit_ts_pushed = (!caller_start_ts.is_zero() && lock.min_commit_ts > caller_start_ts)
                // If the caller_start_ts is max, it's a point get in the autocommit transaction.
                // We don't push forward lock's min_commit_ts and the point get can ignore the lock
                // next time because it's not committed yet.
                || caller_start_ts.is_max();

            TxnStatus::uncommitted(lock, min_commit_ts_pushed)
        };
        Ok(ProcessResult::TxnStatus { txn_status })
    }

    #[maybe_async::both]
    async fn process_check_txn_status_missing_lock(
        &mut self,
        reader: &mut CloudReader,
        primary_key: Key,
    ) -> MvccResult<TxnStatus> {
        let txn_status = match reader
            .get_txn_commit_record(&primary_key, self.start_ts.into())
            .await?
        {
            TxnCommitRecord::SingleRecord { commit_ts, write } => {
                if write.write_type == WriteType::Rollback {
                    TxnStatus::RolledBack
                } else {
                    TxnStatus::committed(commit_ts)
                }
            }
            TxnCommitRecord::None { .. } => TxnStatus::LockNotExist,
            _ => unreachable!(),
        };
        Ok(txn_status)
    }

    fn process_txn_heartbeat(
        &mut self,
        snap_access: &SnapAccess,
        primary_key: Vec<u8>,
        advise_ttl: u64,
    ) -> MvccResult<ProcessResult> {
        let item = snap_access.get(LOCK_CF, &primary_key, 0);
        let start_ts = TimeStamp::new(self.start_ts);
        if item.is_value_empty() {
            return Err(ErrorInner::TxnNotFound {
                start_ts,
                key: primary_key,
            }
            .into());
        }
        let mut lock = txn_types::parse_non_shared_lock(item.get_value())
            .expect("invariant: txn-file primary lock is non-shared");
        if lock.ts != start_ts {
            return Err(ErrorInner::TxnNotFound {
                start_ts,
                key: primary_key,
            }
            .into());
        }
        if lock.ttl < advise_ttl {
            lock.ttl = advise_ttl;
            lock.short_value = None;
            self.set_lock_val_prefix(lock.to_bytes())?;
            self.modified = true;
        }
        Ok(ProcessResult::TxnStatus {
            txn_status: TxnStatus::uncommitted(lock, false),
        })
    }

    #[maybe_async::both]
    async fn process_resolve_txn_lock(
        &mut self,
        snap_access: &SnapAccess,
        commit_ts: TimeStamp,
        resolve_keys: Option<Vec<Key>>,
        txn: &mut MvccTxn,
    ) -> MvccResult<ProcessResult> {
        let res = if commit_ts.is_zero() {
            self.process_rollback(snap_access, resolve_keys).await
        } else {
            self.process_commit(snap_access, commit_ts, resolve_keys, txn)
                .await
        };
        match res {
            Ok(_) => Ok(ProcessResult::Res),
            Err(err @ mvcc::Error(box ErrorInner::TxnLockNotFound { .. })) => {
                // If the lock is not found, it means the transaction is already committed or
                // rollbacked by others. We can safely ignore it.
                info!("process_resolve_txn_lock: txn lock not found"; "err" => ?err, "start_ts" => self.ts());
                Ok(ProcessResult::Res)
            }
            Err(err) => Err(err),
        }
    }

    #[maybe_async::both]
    async fn process_resolve_lock(
        &mut self,
        snap_access: &SnapAccess,
        mut resolve_lock: ResolveLock,
        txn: &mut MvccTxn,
    ) -> MvccResult<ProcessResult> {
        let resolved_commit_ts = resolve_lock
            .txn_file_status
            .remove(&self.ts())
            .expect("txn-file resolve lock must contain the current txn status");
        self.process_resolve_txn_lock(snap_access, resolved_commit_ts, None, txn)
            .await?;
        if resolve_lock.txn_file_status.is_empty() && resolve_lock.txn_status.is_empty() {
            // All locks are resolved.
            return Ok(ProcessResult::Res);
        }
        if resolve_lock.txn_file_status.is_empty() {
            // Only non-txn-file locks.
            let next_cmd = ResolveLockReadPhase {
                ctx: resolve_lock.ctx.clone(),
                deadline: resolve_lock.deadline,
                txn_status: resolve_lock.txn_status,
                scan_key: None,
                is_background: resolve_lock.is_background,
            };
            return Ok(ProcessResult::NextCommand {
                cmd: Command::ResolveLockReadPhase(next_cmd),
            });
        }
        // More locks to resolve,
        Ok(ProcessResult::NextCommand {
            cmd: Command::ResolveLock(resolve_lock),
        })
    }

    // `commit_ts > 0`: committed, `commit_ts == 0`: rollbacked.
    fn set_committed_or_rolled_back(&mut self, commit_ts: u64) -> MvccResult<()> {
        let user_meta = UserMeta::new(self.start_ts, commit_ts);
        self.set_user_meta_and_lock_val_prefix(user_meta, vec![])?;
        self.modified = true;
        Ok(())
    }
}

impl CommandExt for TxnFileCommand {
    fn get_ctx(&self) -> &crate::storage::Context {
        self.inner_cmd.as_ref().unwrap().ctx()
    }
    fn get_ctx_mut(&mut self) -> &mut crate::storage::Context {
        self.inner_cmd.as_mut().unwrap().ctx_mut()
    }
    fn deadline(&self) -> ::tikv_util::deadline::Deadline {
        self.inner_cmd.as_ref().unwrap().deadline()
    }
    tag!(txn_file);

    fn ts(&self) -> TimeStamp {
        TimeStamp::new(self.start_ts)
    }

    fn write_bytes(&self) -> usize {
        let mut bytes = self
            .txn_file_snapshot
            .as_ref()
            .map_or(0, |snapshot| snapshot.txn_file_ref.compute_size() as usize);
        match self.inner_cmd.as_ref().unwrap().as_ref() {
            Command::Prewrite(_) | Command::PrewritePessimistic(_) => {
                bytes += self.txn_file().size();
            }
            _ => {}
        }
        bytes
    }

    fn gen_lock(&self) -> Lock {
        Lock {
            required_hashes: vec![],
            keyspace_id: self.get_ctx().get_keyspace_id(),
            owned_count: 0,
            region_id: self.inner_cmd.as_ref().unwrap().ctx().region_id,
            start_ts: self.start_ts,
            checked_txn_cid: 0,
            txn_file: Some(match self.txn_file_snapshot.as_ref() {
                Some(snapshot) => TxnFileLatchTarget::Present(snapshot.txn_file.clone()),
                None => TxnFileLatchTarget::Missing {
                    start_ts: self.start_ts,
                },
            }),
            count_added: false,
            deadline: Some(self.deadline()),
        }
    }
}

#[maybe_async::async_trait]
impl<S: Snapshot + 'static, L: LockManager> WriteCommand<S, L> for TxnFileCommand {
    #[maybe_async]
    async fn process_write(
        mut self,
        snapshot: S,
        context: WriteContext<'_, L>,
    ) -> crate::storage::txn::Result<WriteResult> {
        let snap = snapshot.get_kvengine_snap().unwrap();
        let write_bytes = self.write_bytes();
        let cmd = mem::take(&mut self.inner_cmd).unwrap();
        let ctx = cmd.ctx().clone();
        debug!("txn file process write"; "cmd" => ?cmd, "start_ts" => self.start_ts, "txn_file_snapshot" => ?self.txn_file_snapshot, "ctx" => ?ctx, "snap" => ?snap);

        self.validate_and_refresh_non_prewrite_lock(&cmd, snap)?;

        let mut txn = if matches!(cmd.as_ref(), Command::Commit(_)) {
            MvccTxn::new_with_backup_ts(self.ts(), context.concurrency_manager, false)
        } else {
            MvccTxn::new(self.ts(), context.concurrency_manager)
        };
        let pr = match *cmd {
            Command::Prewrite(_) => {
                self.process_prewrite(snap, &txn.concurrency_manager)
                    .await?
            }
            Command::PrewritePessimistic(_) => {
                self.process_pessimistic_prewrite(
                    snap,
                    &txn.concurrency_manager,
                    context.extra_cf_scan_mode,
                )
                .await?
            }
            Command::Commit(commit) => {
                let pr = self
                    .process_commit(snap, commit.commit_ts, Some(commit.keys), &mut txn)
                    .await?;
                fail::fail_point!("txn::txn_file_process_commit_finish");
                pr
            }
            Command::Rollback(rollback) => self.process_rollback(snap, Some(rollback.keys)).await?,
            Command::TxnHeartBeat(txn_heartbeat) => {
                let primary_key = txn_heartbeat.primary_key.to_raw().unwrap();
                self.process_txn_heartbeat(snap, primary_key, txn_heartbeat.advise_ttl)?
            }
            Command::CheckTxnStatus(check_txn_status) => {
                self.process_check_txn_status(
                    snap,
                    check_txn_status.primary_key,
                    check_txn_status.lock_ts,
                    check_txn_status.current_ts,
                    check_txn_status.caller_start_ts,
                    &txn.concurrency_manager,
                )
                .await?
            }
            Command::ResolveLock(resolve) => {
                self.process_resolve_lock(snap, resolve, &mut txn).await?
            }
            Command::ResolveLockLite(resolve) => {
                self.process_resolve_txn_lock(
                    snap,
                    resolve.commit_ts,
                    Some(resolve.resolve_keys),
                    &mut txn,
                )
                .await?
            }
            _ => {
                return Err(box_err!("unsupported txn file command"));
            }
        };
        let mut write_data = WriteData::default();
        if self.modified {
            let Some(snapshot) = self.txn_file_snapshot.as_ref() else {
                debug_assert!(
                    false,
                    "missing txn-file lock must not emit txn-file metadata"
                );
                return Err(Self::stale_command_region_error(self.start_ts));
            };
            write_data.txn_file = Some(TxnFileWriteData {
                txn_file_ref: snapshot.txn_file_ref.clone(),
                write_bytes,
            });
            write_data.backup_ts_checked = txn.take_checked_backup_ts();
        }
        let shared_lock_upgrade_hints = txn.take_shared_lock_upgrade_hints();
        let hint_check = Self::ensure_no_shared_lock_upgrade_hints(&shared_lock_upgrade_hints);
        debug_assert!(
            hint_check.is_ok(),
            "txn-file shared-lock upgrade hint invariant violated"
        );
        hint_check?;
        let result = WriteResult {
            ctx,
            to_be_write: write_data,
            rows: 1,
            pr,
            lock_info: vec![],
            released_locks: ReleasedLocks::new(),
            new_acquired_locks: vec![],
            shared_lock_upgrade_hints: vec![],
            lock_guards: vec![],
            response_policy: ResponsePolicy::OnApplied,
            // TODO: fix this
            known_txn_status: vec![],
        };
        Ok(result)
    }
}

impl std::fmt::Display for TxnFileCommand {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "kv::command::txn_file start_ts={} {:?} {:?}",
            self.start_ts,
            self.txn_file_snapshot,
            self.inner_cmd
                .as_ref()
                .map(|cmd| cmd.to_string())
                .unwrap_or_default(),
        )
    }
}

impl std::fmt::Debug for TxnFileCommand {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self)
    }
}

// Helpers:
impl TxnFileCommand {
    fn command_start_ts(cmd: &Command) -> u64 {
        match cmd {
            Command::Prewrite(req) => req.start_ts.into_inner(),
            Command::PrewritePessimistic(req) => req.start_ts.into_inner(),
            Command::Commit(req) => req.lock_ts.into_inner(),
            Command::Rollback(req) => req.start_ts.into_inner(),
            Command::TxnHeartBeat(req) => req.start_ts.into_inner(),
            Command::CheckTxnStatus(req) => req.ts().into_inner(),
            Command::ResolveLockLite(req) => req.ts().into_inner(),
            Command::ResolveLock(req) => req
                .txn_file_status
                .iter()
                .next()
                .map(|(txn_id, _)| txn_id.into_inner())
                .expect("txn-file resolve lock must contain txn_file_status"),
            _ => unreachable!("unexpected command {:?}", cmd),
        }
    }

    fn build_lock_prefix(txn_file_ref: &TxnFileRef) -> Result<Option<txn_types::Lock>> {
        if txn_file_ref.lock_val_prefix.is_empty() {
            return Ok(None);
        }

        // Txn-file transactions themselves stay non-shared in the current
        // rollout, so the persisted lock prefix must remain an ordinary
        // `Lock` even though conflict reads can now observe shared-lock
        // wrappers from other transactions in LOCK_CF.
        txn_types::parse_non_shared_lock(&txn_file_ref.lock_val_prefix)
            .map(Some)
            .map_err(|err| -> Error {
                debug_assert!(false, "invalid txn-file lock prefix: {:?}", err);
                box_err!("invalid txn-file lock prefix: {}", err)
            })
    }

    #[inline]
    fn txn_file_snapshot(&self) -> &TxnFileSnapshot {
        self.txn_file_snapshot
            .as_ref()
            .expect("txn-file snapshot is required for this command path")
    }

    fn txn_file_snapshot_mut(&mut self) -> MvccResult<&mut TxnFileSnapshot> {
        let start_ts = self.start_ts;
        self.txn_file_snapshot.as_mut().ok_or_else(|| {
            ErrorInner::TxnLockNotFound {
                start_ts: start_ts.into(),
                commit_ts: TimeStamp::zero(),
                key: vec![],
            }
            .into()
        })
    }

    #[inline]
    fn txn_file(&self) -> &TxnFile {
        &self.txn_file_snapshot().txn_file
    }

    fn set_lock_val_prefix(&mut self, lock_val_prefix: Vec<u8>) -> MvccResult<()> {
        let snapshot = self.txn_file_snapshot_mut()?;
        snapshot.txn_file_ref.set_lock_val_prefix(lock_val_prefix);
        snapshot.txn_file = snapshot.txn_file.clone_with_ref(&snapshot.txn_file_ref)?;
        Ok(())
    }

    fn set_user_meta_and_lock_val_prefix(
        &mut self,
        user_meta: UserMeta,
        lock_val_prefix: Vec<u8>,
    ) -> MvccResult<()> {
        let snapshot = self.txn_file_snapshot_mut()?;
        snapshot
            .txn_file_ref
            .set_user_meta(user_meta.to_array().to_vec());
        snapshot.txn_file_ref.set_lock_val_prefix(lock_val_prefix);
        snapshot.txn_file = snapshot.txn_file.clone_with_ref(&snapshot.txn_file_ref)?;
        Ok(())
    }

    fn compare_optional_txn_file_snapshot(
        lhs: Option<&TxnFileSnapshot>,
        rhs: Option<&TxnFileSnapshot>,
    ) -> TxnFileSnapshotCompare {
        match (lhs, rhs) {
            (None, None) => TxnFileSnapshotCompare::Same,
            (Some(lhs), Some(rhs)) => Self::compare_txn_file_snapshot(lhs, rhs),
            _ => TxnFileSnapshotCompare::Mismatch,
        }
    }

    fn compare_txn_file_snapshot(
        lhs: &TxnFileSnapshot,
        rhs: &TxnFileSnapshot,
    ) -> TxnFileSnapshotCompare {
        let lhs_file = &lhs.txn_file;
        let rhs_file = &rhs.txn_file;
        if lhs_file.start_ts() != rhs_file.start_ts()
            || lhs_file.id().shard_id != rhs_file.id().shard_id
        {
            return TxnFileSnapshotCompare::Mismatch;
        }
        if lhs.txn_file_ref.shard_ver != rhs.txn_file_ref.shard_ver {
            return TxnFileSnapshotCompare::ShardVersionMismatch;
        }
        if lhs_file.chunk_ids() == rhs_file.chunk_ids()
            && lhs_file.lower_bound() == rhs_file.lower_bound()
            && lhs_file.upper_bound() == rhs_file.upper_bound()
            && lhs_file.get_lock_val_prefix() == rhs_file.get_lock_val_prefix()
        {
            TxnFileSnapshotCompare::Same
        } else {
            TxnFileSnapshotCompare::Mismatch
        }
    }

    fn stale_command_region_error(start_ts: u64) -> Error {
        let mut err = errorpb::Error::default();
        err.set_message(format!(
            "txn-file command observed stale txn-file metadata for start_ts {}",
            start_ts
        ));
        err.set_stale_command(errorpb::StaleCommand::default());
        kv::Error::from(kv::ErrorInner::Request(err)).into()
    }

    fn epoch_not_match_region_error(&self, region_id: u64) -> Error {
        let mut err = errorpb::Error::default();
        let mut epoch_not_match = errorpb::EpochNotMatch::default();
        let current_region = self
            .context
            .region_info_accessor
            .as_ref()
            .and_then(|accessor| {
                accessor
                    .find_region_by_id(region_id)
                    .map(|info| info.region)
            });
        if let Some(region) = current_region {
            err.set_message(format!(
                "txn-file metadata region version changed for start_ts {} in region {}; current is {}",
                self.start_ts,
                region.get_id(),
                region.get_region_epoch().get_version(),
            ));
            epoch_not_match.mut_current_regions().push(region);
        } else {
            err.set_message(format!(
                "txn-file metadata region version changed for start_ts {} in region {}",
                self.start_ts, region_id
            ));
        }
        err.set_epoch_not_match(epoch_not_match);
        kv::Error::from(kv::ErrorInner::Request(err)).into()
    }
}

pub(crate) fn command_chunks_to_load(cmd: &Command) -> &[u64] {
    match cmd {
        Command::Prewrite(req) => &req.txn_file_chunks,
        Command::PrewritePessimistic(req) => &req.txn_file_chunks,
        _ => &[],
    }
}

struct TxnFilePessimisticActionChecker {
    err: Option<crate::storage::mvcc::Error>,
    start_ts: TimeStamp,
    for_update_ts: TimeStamp,
    reader: crate::storage::mvcc::reader::CloudReader,
    primary_key: Vec<u8>,
    current_raw_key: Vec<u8>,
}

impl TxnFilePessimisticActionChecker {
    fn set_current_inner_key(&mut self, key: InnerKey<'_>) {
        self.current_raw_key.truncate(KEYSPACE_PREFIX_LEN);
        self.current_raw_key.extend_from_slice(&key);
    }

    fn check_lock(&mut self, lock: txn_types::Lock, action: PessimisticAction) -> MvccResult<()> {
        if lock.ts != self.start_ts {
            // Abort on lock belonging to other transaction if
            // prewrites a pessimistic lock.
            if action == PessimisticAction::DoCheck {
                warn!(
                    "prewrite failed (pessimistic lock not found)";
                    "start_ts" => self.start_ts,
                    "key" => ?self.current_raw_key,
                    "lock_ts" => lock.ts
                );
                return Err(ErrorInner::PessimisticLockNotFound {
                    start_ts: self.start_ts,
                    key: self.current_raw_key.clone(),
                }
                .into());
            }

            let mut info = lock.into_lock_info(self.current_raw_key.clone());
            info.set_lock_ttl(0);
            return Err(ErrorInner::KeyIsLocked(info).into());
        }
        // if the lock type is not pessimistic, it must means the entry is already
        // prewritten, we can just skip it.
        Ok(())
    }

    fn write_conflict_error(
        &self,
        raw_key: Vec<u8>,
        write: &Write,
        commit_ts: TimeStamp,
        reason: WriteConflictReason,
    ) -> MvccResult<()> {
        Err(ErrorInner::WriteConflict {
            start_ts: self.start_ts,
            conflict_start_ts: write.start_ts,
            conflict_commit_ts: commit_ts,
            key: raw_key,
            primary: self.primary_key.clone(),
            reason,
        }
        .into())
    }

    #[maybe_async::both]
    async fn amend_pessimistic_lock(&mut self) -> MvccResult<()> {
        // try amend the lock
        if let Some(user_meta) = self.reader.find_extra_cf_conflict_record_with_raw_key(
            &self.current_raw_key,
            self.start_ts,
            self.start_ts,
        )? {
            if user_meta.is_rollback() {
                warn!(
                    "prewrite failed (pessimistic lock not found): transaction has been rolled back";
                    "start_ts" => self.start_ts,
                    "key" => ?log_wrappers::Value::key(&self.current_raw_key),
                );
            } else {
                warn!(
                    "prewrite failed (pessimistic lock not found): another transaction acquired lock";
                    "start_ts" => self.start_ts,
                    "conflict_commit_ts" => user_meta.commit_ts,
                    "key" => ?log_wrappers::Value::key(&self.current_raw_key),
                );
            }
            return Err(ErrorInner::PessimisticLockNotFound {
                start_ts: self.start_ts,
                key: self.current_raw_key.clone(),
            }
            .into());
        }

        let write = self
            .reader
            .seek_write_with_raw_key(&self.current_raw_key, TimeStamp::max())
            .await?;
        if let Some((commit_ts, _write)) = write.as_ref() {
            if *commit_ts >= self.start_ts {
                // The invariants of pessimistic locks are:
                //   1. lock's for_update_ts >= key's latest commit_ts
                //   2. lock's for_update_ts >= txn's start_ts
                //   3. If the data is changed after acquiring the pessimistic lock, key's new
                // commit_ts > lock's for_update_ts
                //
                // So, if the key's latest commit_ts is still less than or equal to lock's
                // for_update_ts, the data is not changed. However, we can't get lock's
                // for_update_ts in current implementation (txn's for_update_ts is updated for
                // each DML), we can only use txn's start_ts to check -- If the key's
                // commit_ts is less than txn's start_ts, it's less than
                // lock's for_update_ts too.
                warn!(
                    "prewrite failed (pessimistic lock not found)";
                    "start_ts" => self.start_ts,
                    "commit_ts" => *commit_ts,
                    "key" => ?self.current_raw_key
                );
                return Err(ErrorInner::PessimisticLockNotFound {
                    start_ts: self.start_ts,
                    key: self.current_raw_key.clone(),
                }
                .into());
            }
        }
        Ok(())
    }

    #[maybe_async::both]
    async fn check_for_newer_version(
        &mut self,
        op: u8,
        action: PessimisticAction,
    ) -> MvccResult<()> {
        if let Some((ts, write)) = self
            .reader
            .get_extra_with_raw_key(&self.current_raw_key, self.start_ts)
        {
            return self.write_conflict_error(
                self.current_raw_key.clone(),
                &write,
                ts,
                WriteConflictReason::SelfRolledBack,
            );
        }
        let should_not_exist = matches!(op, OP_INSERT | OP_CHECK_NOT_EXIST);
        // new version should not exist for insert.
        if !should_not_exist {
            // DoConstraintCheck: For keys that do not acquire pessimistic locks
            let (check_ts, conflict_reason) = if action == PessimisticAction::DoConstraintCheck {
                (
                    self.start_ts,
                    Some(WriteConflictReason::NotLockedKeyConflict),
                )
            } else {
                (self.for_update_ts, None)
            };
            if let Some((commit_ts, write)) = self
                .reader
                .get_newer_with_raw_key(&self.current_raw_key, check_ts)
                .await?
            {
                if let Some(reason) = conflict_reason {
                    return self.write_conflict_error(
                        self.current_raw_key.clone(),
                        &write,
                        commit_ts,
                        reason,
                    );
                } else {
                    warn!("conflicting write was found, pessimistic lock must be lost for the corresponding row key";
                                "key" => ?self.current_raw_key,
                                "start_ts" => self.start_ts,
                                "for_update_ts" => check_ts,
                                "conflicting start_ts" => write.start_ts,
                                "conflicting commit_ts" => commit_ts);
                    return Err(ErrorInner::PessimisticLockNotFound {
                        start_ts: self.start_ts,
                        key: self.current_raw_key.clone(),
                    }
                    .into());
                }
            }
            return Ok(());
        }

        let mut seek_ts = TimeStamp::max();
        while let Some((commit_ts, write)) = self
            .reader
            .seek_write_with_raw_key(&self.current_raw_key, seek_ts)
            .await?
        {
            // If there's a write record whose commit_ts equals to our start ts, the current
            // transaction is ok to continue, unless the record means that the current
            // transaction has been rolled back.
            if commit_ts == self.start_ts
                && (write.write_type == WriteType::Rollback || write.has_overlapped_rollback)
            {
                // TODO: Maybe we need to add a new error for the rolled back case.
                self.write_conflict_error(
                    self.current_raw_key.clone(),
                    &write,
                    commit_ts,
                    WriteConflictReason::SelfRolledBack,
                )?;
            }
            if PessimisticAction::DoConstraintCheck == action {
                // Do the same as optimistic transactions if constraint checks are needed.
                if commit_ts > self.start_ts {
                    return self.write_conflict_error(
                        self.current_raw_key.clone(),
                        &write,
                        commit_ts,
                        WriteConflictReason::NotLockedKeyConflict,
                    );
                }
            }
            if commit_ts > self.for_update_ts {
                // Don't treat newer Rollback records as write conflicts. They can cause
                // false positive errors because they can be written even if the pessimistic
                // lock of the corresponding row key exists.
                // Rollback records are only used to prevent retried prewrite from
                // succeeding. Even if the Rollback record of the current transaction is
                // collapsed by a newer record, it is safe to prewrite this non-pessimistic
                // key because either the primary key is rolled back or it's protected
                // because it's written by CheckSecondaryLocks.
                if write.write_type == WriteType::Rollback {
                    seek_ts = commit_ts.prev();
                    continue;
                }

                warn!("conflicting write was found, pessimistic lock must be lost for the corresponding row key";
                    "key" => ?self.current_raw_key,
                    "start_ts" => self.start_ts,
                    "for_update_ts" => self.for_update_ts,
                    "conflicting start_ts" => write.start_ts,
                    "conflicting commit_ts" => commit_ts);
                return Err(ErrorInner::PessimisticLockNotFound {
                    start_ts: self.start_ts,
                    key: self.current_raw_key.clone(),
                }
                .into());
            }
            // Should check it when no lock exists, otherwise it can report error when there
            // is a lock belonging to a committed transaction which deletes the key.
            self.check_data_constraint(should_not_exist, &write, commit_ts)
                .await?;
            return Ok(());
        }

        Ok(())
    }

    // this is the same as `actions::check_data_constraint`, only the reader type is
    // different.
    #[maybe_async::both]
    async fn check_data_constraint(
        &mut self,
        should_not_exist: bool,
        write: &Write,
        write_commit_ts: TimeStamp,
    ) -> MvccResult<()> {
        // Here we assume `write` is the latest version of the key. So it should not
        // contain a GC fence ts. Otherwise, it must be an already-deleted version.
        let write_is_invalid =
            matches!(write.gc_fence, Some(gc_fence_ts) if !gc_fence_ts.is_zero());

        if !should_not_exist || write.write_type == WriteType::Delete || write_is_invalid {
            return Ok(());
        }

        // The current key exists under any of the following conditions:
        // 1.The current write type is `PUT`
        // 2.The current write type is `Rollback` or `Lock`, and the key have an older
        // version.
        if write.write_type == WriteType::Put
            || self
                .reader
                .seek_write_with_raw_key(&self.current_raw_key, write_commit_ts.prev())
                .await?
                .is_some()
        {
            return Err(ErrorInner::AlreadyExist {
                key: self.current_raw_key.clone(),
            }
            .into());
        }
        Ok(())
    }

    #[maybe_async::both]
    async fn check_pessimistic_action(&mut self, op: u8, action: u8) -> MvccResult<()> {
        // We need to verify the action byte again in case the file is corrupted.
        let action = PessimisticAction::try_from(action)
            .map_err(|e| MvccError::from(ErrorInner::Other(e.into())))?;
        if let Some(lock_or_shared) = self.reader.load_lock_with_raw_key(&self.current_raw_key)? {
            match lock_or_shared {
                Either::Left(lock) => self.check_lock(lock, action)?,
                Either::Right(shared_locks) => {
                    return Err(MvccError::from(ErrorInner::KeyIsLocked(
                        shared_locks.into_lock_info(self.current_raw_key.clone()),
                    )));
                }
            }
        } else if action == PessimisticAction::DoCheck {
            self.amend_pessimistic_lock().await?;
        }

        // whether to skip constrain check.
        // This is the same as `PrewriteMutation::skip_constraint_check`
        // NOTE:
        // 1. retry request is already handled in `process_pessimistic_prewrite`, we do
        //    not need to cosider the `retry request` logic here.
        // 2. We do not face the same issue on retry request as https://github.com/tikv/tikv/issues/11187
        //    because the txn_file is only write once, its retry do not actually write
        //    the key again to kvengine as normal txn, so there is no need to do
        //    constraint check as normal txn for all keys.

        // NOTE: while its safe to skip constraint check for txn-file request, but here
        // we keep the same behavior with normal prewrite to keep to logic
        // simpler.
        if matches!(
            action,
            PessimisticAction::DoConstraintCheck | PessimisticAction::SkipCheck
        ) {
            self.check_for_newer_version(op, action).await?;
        }
        Ok(())
    }
}

#[maybe_async::async_trait]
impl kvengine::table::txn_file::EntryChecker for TxnFilePessimisticActionChecker {
    #[maybe_async]
    async fn check(&mut self, op: u8, action: u8, key: InnerKey<'_>) -> bool {
        self.set_current_inner_key(key);
        if let Err(e) = self.check_pessimistic_action(op, action).await {
            self.err = Some(e);
            return false;
        }
        true
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::{AtomicU64, Ordering};

    use bytes::Bytes;
    use collections::HashMap;
    use kvengine::{
        EncryptionState, GLOBAL_SHARD_END_KEY, IoType, TXN_FILE_REF, WriteBatch,
        dfs::{self, FileType},
        table::{
            InnerKey, TXN_FILE_FORMAT_V2, TxnChunkBuilder,
            txn_file::{OP_PUT, TxnCtx, TxnFile, TxnFileId},
        },
        test_engine::{ApplyTask, TestEngine, new_test_engine},
    };
    use kvenginepb::TxnFileRefs;
    use kvproto::metapb;
    use raft::StateRole;
    use raftstore::coprocessor::RegionInfo;
    use tikv_util::{deadline::Deadline, mpsc};

    use super::*;
    use crate::storage::txn::commands::Rollback;

    static NEXT_TEST_TXN_CHUNK_ID: AtomicU64 = AtomicU64::new(1_000_000);

    fn new_test_lock_txn_file(start_ts: u64) -> TxnFile {
        let txn_ctx = TxnCtx::new(
            Bytes::new(),
            Bytes::new(),
            0,
            InnerKey::from_inner_buf(b"a"),
            InnerKey::from_inner_buf(b"z"),
        );
        TxnFile::new(TxnFileId::new(1, 1, start_ts), vec![], txn_ctx).unwrap()
    }

    fn make_test_lock_prefix(primary: Vec<u8>, start_ts: u64, min_commit_ts: u64) -> Vec<u8> {
        let mut lock = txn_types::Lock::new(
            LockType::Put,
            primary,
            start_ts.into(),
            3000,
            None,
            TimeStamp::zero(),
            100,
            min_commit_ts.into(),
        );
        lock.is_txn_file = true;
        lock.to_bytes()
    }

    fn write_test_data(wb: WriteBatch, applier_tx: &mpsc::Sender<ApplyTask>) {
        let (result_tx, result_rx) = mpsc::bounded(1);
        applier_tx.send(ApplyTask::new_wb(wb, result_tx)).unwrap();
        result_rx.recv().unwrap().unwrap();
    }

    fn build_test_txn_chunk(engine: &TestEngine, start: usize, end: usize) -> u64 {
        let chunk_id = NEXT_TEST_TXN_CHUNK_ID.fetch_add(1, Ordering::Relaxed);
        let mut chunk_builder = TxnChunkBuilder::new(chunk_id, 64, None, TXN_FILE_FORMAT_V2);
        for i in start..end {
            let key = engine.key_builder().i_to_key(i);
            chunk_builder.add_entry(
                InnerKey::from_outer_key(&key),
                OP_PUT,
                PessimisticAction::SkipCheck,
                &key,
            );
        }

        let mut buf = vec![];
        chunk_builder.finish(&mut buf);
        let runtime = engine.fs.get_runtime();
        let fs = engine.fs.clone();
        let opts = dfs::Options::default().with_type(FileType::TxnChunk);
        runtime
            .block_on(fs.create(chunk_id, buf.into(), opts))
            .unwrap();
        engine
            .get_txn_chunk_manager()
            .prepare(chunk_id, EncryptionState::Disabled, IoType::Other)
            .unwrap();
        chunk_id
    }

    fn build_test_txn_file(
        engine: &TestEngine,
        start_ts: u64,
        chunk_ranges: &[(usize, usize)],
        lock_val_prefix: Vec<u8>,
        lower_bound: InnerKey<'_>,
        upper_bound: InnerKey<'_>,
    ) -> TxnFile {
        build_test_txn_file_with_id(
            engine,
            TxnFileId::new(1, 1, start_ts),
            chunk_ranges,
            lock_val_prefix,
            lower_bound,
            upper_bound,
        )
    }

    fn build_test_txn_file_with_id(
        engine: &TestEngine,
        txn_file_id: TxnFileId,
        chunk_ranges: &[(usize, usize)],
        lock_val_prefix: Vec<u8>,
        lower_bound: InnerKey<'_>,
        upper_bound: InnerKey<'_>,
    ) -> TxnFile {
        let txn_chunk_manager = engine.get_txn_chunk_manager();
        let chunks = chunk_ranges
            .iter()
            .map(|&(start, end)| {
                let chunk_id = build_test_txn_chunk(engine, start, end);
                txn_chunk_manager.get(chunk_id).unwrap()
            })
            .collect();
        let txn_ctx = TxnCtx::new(
            Bytes::new(),
            lock_val_prefix.into(),
            0,
            lower_bound,
            upper_bound,
        );
        TxnFile::new(txn_file_id, chunks, txn_ctx).unwrap()
    }

    fn txn_file_snapshot(txn_file: TxnFile) -> TxnFileSnapshot {
        let mut txn_file_ref = TxnFileRef::new();
        txn_file_ref.set_start_ts(txn_file.start_ts());
        txn_file_ref.set_shard_ver(txn_file.shard_ver());
        txn_file_ref.set_chunk_ids(txn_file.chunk_ids());
        txn_file_ref.set_inner_lower_bound(txn_file.lower_bound().to_vec());
        txn_file_ref.set_inner_upper_bound(txn_file.upper_bound().to_vec());
        txn_file_ref.set_lock_val_prefix(txn_file.get_lock_val_prefix().to_vec());
        TxnFileSnapshot {
            txn_file_ref,
            txn_file,
        }
    }

    fn apply_lock_txn_file(txn_file: &TxnFile, applier_tx: &mpsc::Sender<ApplyTask>) {
        let mut txn_file_refs = TxnFileRefs::new();
        txn_file_refs
            .mut_txn_file_refs()
            .push(txn_file.to_txn_file_ref(None).unwrap());
        let mut wb = WriteBatch::default();
        wb.reset(1);
        wb.set_property(TXN_FILE_REF, &txn_file_refs.write_to_bytes().unwrap());
        write_test_data(wb, applier_tx);
    }

    fn new_test_txn_file_command(txn_file: TxnFile) -> TxnFileCommand {
        let txn_file_ref = txn_file.to_txn_file_ref(None).unwrap();
        let start_ts = txn_file.start_ts();
        TxnFileCommand {
            start_ts,
            txn_file_snapshot: Some(TxnFileSnapshot {
                txn_file_ref,
                txn_file,
            }),
            context: TxnFileCommandContext::new(None, usize::MAX),
            modified: false,
            lock_prefix: None,
            inner_cmd: None,
        }
    }

    fn new_missing_test_txn_file_command(start_ts: u64) -> TxnFileCommand {
        TxnFileCommand {
            start_ts,
            txn_file_snapshot: None,
            context: TxnFileCommandContext::new(None, usize::MAX),
            modified: false,
            lock_prefix: None,
            inner_cmd: None,
        }
    }

    fn test_rollback_command(start_ts: u64) -> Command {
        let mut ctx = kvproto::kvrpcpb::Context::default();
        ctx.set_region_id(1);
        Command::Rollback(Rollback {
            ctx,
            deadline: Deadline::from_now(
                crate::storage::txn::scheduler::DEFAULT_EXECUTION_DURATION_LIMIT,
            ),
            keys: vec![],
            start_ts: start_ts.into(),
            is_txn_file: true,
        })
    }

    #[test]
    fn test_txn_file_rejects_shared_lock_upgrade_hints_without_per_key_latches() {
        let shared_lock_upgrade_hints = vec![SharedLockUpgradeHint {
            key: Key::from_raw(b"txn-file-shared-lock-upgrade-hint"),
            sole_holder_start_ts: 42.into(),
        }];

        let err = TxnFileCommand::ensure_no_shared_lock_upgrade_hints(&shared_lock_upgrade_hints)
            .unwrap_err();
        assert!(
            err.to_string()
                .contains("shared-lock upgrade hints without per-key latches")
        );

        let mut cmd = new_missing_test_txn_file_command(42);
        cmd.inner_cmd = Some(Box::new(test_rollback_command(42)));
        assert!(cmd.gen_lock().required_hashes.is_empty());
    }

    #[derive(Debug)]
    enum ExpectedRegionError {
        StaleCommand,
        EpochNotMatch { current_region_id: u64 },
        EpochNotMatchWithoutCurrentRegion,
    }

    fn assert_region_error(err: Error, expected: ExpectedRegionError) {
        let Error(inner) = err;
        let crate::storage::txn::ErrorInner::Engine(kv_err) = *inner else {
            panic!("expected engine region error, got {:?}", inner);
        };
        let crate::storage::kv::ErrorInner::Request(region_err) = *kv_err.0 else {
            panic!("expected request region error, got {:?}", kv_err);
        };
        match expected {
            ExpectedRegionError::StaleCommand => {
                assert!(
                    region_err.has_stale_command(),
                    "region_err: {:?}",
                    region_err
                );
                assert!(
                    region_err
                        .get_message()
                        .contains("txn-file command observed stale txn-file metadata"),
                    "region_err: {:?}",
                    region_err
                );
            }
            ExpectedRegionError::EpochNotMatch { current_region_id } => {
                assert!(
                    region_err.has_epoch_not_match(),
                    "region_err: {:?}",
                    region_err
                );
                let current_regions = region_err.get_epoch_not_match().get_current_regions();
                assert_eq!(current_regions.len(), 1, "region_err: {:?}", region_err);
                assert_eq!(current_regions[0].get_id(), current_region_id);
            }
            ExpectedRegionError::EpochNotMatchWithoutCurrentRegion => {
                assert!(
                    region_err.has_epoch_not_match(),
                    "region_err: {:?}",
                    region_err
                );
                assert!(
                    region_err
                        .get_epoch_not_match()
                        .get_current_regions()
                        .is_empty(),
                    "region_err: {:?}",
                    region_err
                );
                assert!(
                    region_err
                        .get_message()
                        .contains("txn-file metadata region version changed"),
                    "region_err: {:?}",
                    region_err
                );
            }
        }
    }

    #[test]
    fn test_txn_file_stale_metadata_detection_cases() {
        fn make_snapshot(
            start_ts: u64,
            shard_id: u64,
            shard_ver: u64,
            lock_prefix: Vec<u8>,
        ) -> TxnFileSnapshot {
            let txn_ctx = TxnCtx::new(
                Bytes::new(),
                lock_prefix.into(),
                0,
                InnerKey::from_inner_buf(b""),
                InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
            );
            let txn_file = TxnFile::new(
                TxnFileId::new(shard_id, shard_ver, start_ts),
                vec![],
                txn_ctx,
            )
            .unwrap();
            let mut txn_file_ref = TxnFileRef::new();
            txn_file_ref.set_start_ts(txn_file.start_ts());
            txn_file_ref.set_shard_ver(txn_file.shard_ver());
            txn_file_ref.set_inner_lower_bound(txn_file.lower_bound().to_vec());
            txn_file_ref.set_inner_upper_bound(txn_file.upper_bound().to_vec());
            txn_file_ref.set_lock_val_prefix(txn_file.get_lock_val_prefix().to_vec());
            TxnFileSnapshot {
                txn_file_ref,
                txn_file,
            }
        }

        let base = make_snapshot(42, 1, 1, vec![1]);
        let different_prefix = make_snapshot(42, 1, 1, vec![2]);
        let different_shard_id = make_snapshot(42, 2, 1, vec![1]);
        let different_shard_ver = make_snapshot(42, 1, 2, vec![1]);

        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(None, None),
            TxnFileSnapshotCompare::Same
        );
        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(None, Some(&base)),
            TxnFileSnapshotCompare::Mismatch
        );
        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(Some(&base), None),
            TxnFileSnapshotCompare::Mismatch
        );
        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(
                Some(&base),
                Some(&different_prefix)
            ),
            TxnFileSnapshotCompare::Mismatch
        );
        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(
                Some(&base),
                Some(&different_shard_id)
            ),
            TxnFileSnapshotCompare::Mismatch
        );
        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(
                Some(&base),
                Some(&different_shard_ver)
            ),
            TxnFileSnapshotCompare::ShardVersionMismatch
        );
    }

    #[test]
    fn test_reject_new_txn_file_prewrite_by_lock_limit() {
        let existing_locks = vec![new_test_lock_txn_file(10)];

        assert!(reject_new_txn_file_prewrite_by_lock_limit(20, 1, &existing_locks, 1).is_some());
        assert!(reject_new_txn_file_prewrite_by_lock_limit(10, 1, &existing_locks, 1).is_none());
        assert!(reject_new_txn_file_prewrite_by_lock_limit(20, 1, &existing_locks, 2).is_none());
        assert!(reject_new_txn_file_prewrite_by_lock_limit(20, 1, &[], 0).is_some());
    }

    #[test]
    fn test_optional_txn_file_snapshot_equality() {
        let (engine, _tx) = new_test_engine();
        let primary = engine.key_builder().i_to_outer_key(10);
        let base_prefix = make_test_lock_prefix(primary.clone(), 42, 43);
        let base = txn_file_snapshot(build_test_txn_file(
            &engine,
            42,
            &[(10, 20)],
            base_prefix.clone(),
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        ));

        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(None, None),
            TxnFileSnapshotCompare::Same
        );
        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(None, Some(&base)),
            TxnFileSnapshotCompare::Mismatch
        );
        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(Some(&base), None),
            TxnFileSnapshotCompare::Mismatch
        );
        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(Some(&base), Some(&base)),
            TxnFileSnapshotCompare::Same
        );

        let different_chunks = txn_file_snapshot(build_test_txn_file(
            &engine,
            42,
            &[(20, 30)],
            base_prefix.clone(),
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        ));
        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(
                Some(&base),
                Some(&different_chunks)
            ),
            TxnFileSnapshotCompare::Mismatch
        );

        let lower_bound_key = engine.key_builder().i_to_key(11);
        let different_lower = txn_file_snapshot(build_test_txn_file(
            &engine,
            42,
            &[(10, 20)],
            base_prefix.clone(),
            InnerKey::from_outer_key(&lower_bound_key),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        ));
        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(Some(&base), Some(&different_lower)),
            TxnFileSnapshotCompare::Mismatch
        );

        let upper_bound_key = engine.key_builder().i_to_key(19);
        let different_upper = txn_file_snapshot(build_test_txn_file(
            &engine,
            42,
            &[(10, 20)],
            base_prefix.clone(),
            InnerKey::from_inner_buf(b""),
            InnerKey::from_outer_key(&upper_bound_key),
        ));
        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(Some(&base), Some(&different_upper)),
            TxnFileSnapshotCompare::Mismatch
        );

        let different_prefix = txn_file_snapshot(build_test_txn_file(
            &engine,
            42,
            &[(10, 20)],
            make_test_lock_prefix(primary.clone(), 42, 88),
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        ));
        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(
                Some(&base),
                Some(&different_prefix)
            ),
            TxnFileSnapshotCompare::Mismatch
        );

        let different_shard_id = txn_file_snapshot(build_test_txn_file_with_id(
            &engine,
            TxnFileId::new(2, 1, 42),
            &[(10, 20)],
            base_prefix.clone(),
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        ));
        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(
                Some(&base),
                Some(&different_shard_id)
            ),
            TxnFileSnapshotCompare::Mismatch
        );

        let different_shard_ver = txn_file_snapshot(build_test_txn_file_with_id(
            &engine,
            TxnFileId::new(1, 2, 42),
            &[(10, 20)],
            base_prefix,
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        ));
        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(
                Some(&base),
                Some(&different_shard_ver)
            ),
            TxnFileSnapshotCompare::ShardVersionMismatch
        );

        let mut different_ref_shard_ver = base.clone();
        different_ref_shard_ver.txn_file_ref.set_shard_ver(2);
        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(
                Some(&base),
                Some(&different_ref_shard_ver)
            ),
            TxnFileSnapshotCompare::ShardVersionMismatch
        );
    }

    #[test]
    fn test_missing_txn_file_snapshot_build_state() {
        let (engine, _tx) = new_test_engine();
        let snap = engine.get_snap_access(1).unwrap();

        let missing = TxnFileCommand::build_txn_file_snapshot(&snap, 42);

        assert!(missing.is_none());
    }

    #[test]
    fn test_build_txn_file_snapshot_keeps_lock_state() {
        let (engine, applier_tx) = new_test_engine();
        let start_ts = 42;
        let primary = engine.key_builder().i_to_outer_key(10);
        let lock_prefix = make_test_lock_prefix(primary, start_ts, 43);
        let txn_file = build_test_txn_file(
            &engine,
            start_ts,
            &[(10, 20)],
            lock_prefix.clone(),
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        );
        apply_lock_txn_file(&txn_file, &applier_tx);
        let snap = engine.get_snap_access(1).unwrap();

        let snapshot = TxnFileCommand::build_txn_file_snapshot(&snap, start_ts).unwrap();

        assert!(snapshot.txn_file_ref.get_user_meta().is_empty());
        assert_eq!(
            snapshot.txn_file_ref.get_lock_val_prefix(),
            lock_prefix.as_slice()
        );
        assert_eq!(
            snapshot.txn_file.get_lock_val_prefix(),
            lock_prefix.as_slice()
        );
    }

    #[test]
    fn test_lock_prefix_update_keeps_snapshot_locked() {
        let (engine, applier_tx) = new_test_engine();
        let start_ts = 42;
        let primary = engine.key_builder().i_to_outer_key(10);
        let original_lock_prefix = make_test_lock_prefix(primary.clone(), start_ts, 43);
        let updated_lock_prefix = make_test_lock_prefix(primary, start_ts, 88);
        let txn_file = build_test_txn_file(
            &engine,
            start_ts,
            &[(10, 20)],
            original_lock_prefix,
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        );
        apply_lock_txn_file(&txn_file, &applier_tx);
        let snap = engine.get_snap_access(1).unwrap();
        let snapshot = TxnFileCommand::build_txn_file_snapshot(&snap, start_ts).unwrap();
        let mut cmd = TxnFileCommand {
            start_ts,
            txn_file_snapshot: Some(snapshot),
            context: TxnFileCommandContext::new(None, usize::MAX),
            modified: false,
            lock_prefix: None,
            inner_cmd: None,
        };

        cmd.set_lock_val_prefix(updated_lock_prefix.clone())
            .unwrap();

        assert!(
            cmd.txn_file_snapshot()
                .txn_file_ref
                .get_user_meta()
                .is_empty()
        );
        assert_eq!(
            cmd.txn_file_snapshot().txn_file_ref.get_lock_val_prefix(),
            updated_lock_prefix.as_slice()
        );
        assert_eq!(
            cmd.txn_file().get_lock_val_prefix(),
            updated_lock_prefix.as_slice()
        );
    }

    #[test]
    fn test_final_txn_file_state_clears_lock_prefix() {
        let (engine, _tx) = new_test_engine();
        let lock_prefix = make_test_lock_prefix(engine.key_builder().i_to_outer_key(10), 42, 43);
        let txn_file = build_test_txn_file(
            &engine,
            42,
            &[(10, 20)],
            lock_prefix,
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        );
        let mut cmd = new_test_txn_file_command(txn_file);

        cmd.set_committed_or_rolled_back(100).unwrap();
        let user_meta = UserMeta::from_slice(cmd.txn_file_snapshot().txn_file_ref.get_user_meta());
        assert_eq!(user_meta.start_ts, 42);
        assert_eq!(user_meta.commit_ts, 100);
        assert!(
            cmd.txn_file_snapshot()
                .txn_file_ref
                .get_lock_val_prefix()
                .is_empty()
        );
        assert!(cmd.txn_file().get_lock_val_prefix().is_empty());
        assert_eq!(cmd.txn_file().version(), 100);
    }

    #[test]
    fn test_missing_txn_file_snapshot_mutation_returns_error() {
        let mut cmd = new_missing_test_txn_file_command(42);

        assert!(cmd.set_lock_val_prefix(vec![1]).is_err());
        assert!(cmd.set_committed_or_rolled_back(100).is_err());
        assert!(!cmd.modified);
    }

    #[test]
    fn test_missing_txn_file_resolve_lock_uses_request_status() {
        let (engine, _tx) = new_test_engine();
        let snap = engine.get_snap_access(1).unwrap();
        let concurrency_manager = ConcurrencyManager::new(1.into());
        let mut txn = MvccTxn::new(TimeStamp::new(42), concurrency_manager);
        let mut txn_file_status = HashMap::default();
        txn_file_status.insert(TimeStamp::new(42), TimeStamp::new(100));
        let resolve_lock = ResolveLock {
            ctx: Default::default(),
            deadline: Deadline::from_now(
                crate::storage::txn::scheduler::DEFAULT_EXECUTION_DURATION_LIMIT,
            ),
            txn_status: Default::default(),
            scan_key: None,
            key_locks: vec![],
            txn_file_status,
            is_background: false,
        };
        let mut cmd = new_missing_test_txn_file_command(42);

        let result = cmd
            .process_resolve_lock(&snap, resolve_lock, &mut txn)
            .unwrap();

        assert!(matches!(result, ProcessResult::Res));
        assert!(!cmd.modified);
    }

    #[test]
    fn test_non_prewrite_rejects_stale_txn_file_snapshot() {
        let (engine, applier_tx) = new_test_engine();
        let primary = engine.key_builder().i_to_outer_key(10);
        let start_ts = 42;
        let cmd = test_rollback_command(start_ts);

        let current_snapshot = build_test_txn_file(
            &engine,
            start_ts,
            &[(10, 20)],
            make_test_lock_prefix(primary.clone(), start_ts, 43),
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        );
        apply_lock_txn_file(&current_snapshot, &applier_tx);
        let snap_with_lock = engine.get_snap_access(1).unwrap();

        let mut missing_build = new_missing_test_txn_file_command(start_ts);
        assert_region_error(
            missing_build
                .validate_and_refresh_non_prewrite_lock(&cmd, &snap_with_lock)
                .unwrap_err(),
            ExpectedRegionError::StaleCommand,
        );

        let stale_build_lock = build_test_txn_file(
            &engine,
            start_ts,
            &[(10, 20)],
            make_test_lock_prefix(primary.clone(), start_ts, 44),
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        );
        let mut changed_prefix = new_test_txn_file_command(stale_build_lock);
        assert_region_error(
            changed_prefix
                .validate_and_refresh_non_prewrite_lock(&cmd, &snap_with_lock)
                .unwrap_err(),
            ExpectedRegionError::StaleCommand,
        );

        let different_shard_ver = build_test_txn_file_with_id(
            &engine,
            TxnFileId::new(1, 2, start_ts),
            &[(10, 20)],
            make_test_lock_prefix(primary, start_ts, 43),
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        );
        assert_eq!(
            TxnFileCommand::compare_optional_txn_file_snapshot(
                Some(&txn_file_snapshot(current_snapshot)),
                Some(&txn_file_snapshot(different_shard_ver)),
            ),
            TxnFileSnapshotCompare::ShardVersionMismatch,
        );

        let mut current_region = metapb::Region::default();
        current_region.set_id(1);
        current_region.mut_region_epoch().set_version(1);
        let region_info_accessor =
            RegionInfoAccessor::new_with_regions_for_test(vec![RegionInfo::new(
                current_region,
                StateRole::Leader,
                0,
            )]);
        let stale_shard_version_lock = build_test_txn_file_with_id(
            &engine,
            TxnFileId::new(1, 2, start_ts),
            &[(10, 20)],
            make_test_lock_prefix(engine.key_builder().i_to_outer_key(10), start_ts, 43),
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        );
        let mut stale_shard_version = new_test_txn_file_command(stale_shard_version_lock);
        stale_shard_version.context =
            TxnFileCommandContext::new(Some(region_info_accessor), usize::MAX);
        assert_region_error(
            stale_shard_version
                .validate_and_refresh_non_prewrite_lock(&cmd, &snap_with_lock)
                .unwrap_err(),
            ExpectedRegionError::EpochNotMatch {
                current_region_id: 1,
            },
        );

        let stale_shard_version_lock_without_region = build_test_txn_file_with_id(
            &engine,
            TxnFileId::new(1, 2, start_ts),
            &[(10, 20)],
            make_test_lock_prefix(engine.key_builder().i_to_outer_key(10), start_ts, 43),
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        );
        let mut stale_shard_version_without_region =
            new_test_txn_file_command(stale_shard_version_lock_without_region);
        assert_region_error(
            stale_shard_version_without_region
                .validate_and_refresh_non_prewrite_lock(&cmd, &snap_with_lock)
                .unwrap_err(),
            ExpectedRegionError::EpochNotMatchWithoutCurrentRegion,
        );
    }

    #[test]
    fn test_non_prewrite_missing_txn_file_snapshot_validation_states() {
        let (engine, _tx) = new_test_engine();
        let snap_without_lock = engine.get_snap_access(1).unwrap();
        let start_ts = 42;
        let cmd = test_rollback_command(start_ts);

        let mut both_missing = new_missing_test_txn_file_command(start_ts);
        both_missing
            .validate_and_refresh_non_prewrite_lock(&cmd, &snap_without_lock)
            .unwrap();
        assert!(both_missing.txn_file_snapshot.is_none());
        assert!(both_missing.lock_prefix.is_none());

        let primary = engine.key_builder().i_to_outer_key(10);
        let build_lock = build_test_txn_file(
            &engine,
            start_ts,
            &[(10, 20)],
            make_test_lock_prefix(primary, start_ts, 43),
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        );
        let mut missing_process = new_test_txn_file_command(build_lock);
        assert_region_error(
            missing_process
                .validate_and_refresh_non_prewrite_lock(&cmd, &snap_without_lock)
                .unwrap_err(),
            ExpectedRegionError::StaleCommand,
        );
    }

    #[test]
    fn test_try_merge_txn_file_locks_not_found() {
        let (engine, _tx) = new_test_engine();
        let primary = engine.key_builder().i_to_outer_key(10);
        let txn_file = build_test_txn_file(
            &engine,
            42,
            &[(10, 20)],
            make_test_lock_prefix(primary, 42, 43),
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        );
        let mut cmd = new_test_txn_file_command(txn_file);
        let snap = engine.get_snap_access(1).unwrap();

        assert!(matches!(
            cmd.try_merge_txn_file_locks(&snap).unwrap(),
            TxnFileLockMergeResult::NotFound
        ));
    }

    #[test]
    fn test_try_merge_txn_file_locks_found_equal() {
        let (engine, tx) = new_test_engine();
        let primary = engine.key_builder().i_to_outer_key(10);
        let txn_file = build_test_txn_file(
            &engine,
            42,
            &[(10, 20)],
            make_test_lock_prefix(primary, 42, 43),
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        );
        apply_lock_txn_file(&txn_file, &tx);

        let mut cmd = new_test_txn_file_command(txn_file.clone());
        let snap = engine.get_snap_access(1).unwrap();

        assert!(matches!(
            cmd.try_merge_txn_file_locks(&snap).unwrap(),
            TxnFileLockMergeResult::FoundEqual
        ));
        assert_eq!(cmd.txn_file().chunk_ids(), txn_file.chunk_ids());
    }

    #[test]
    fn test_try_merge_txn_file_locks_found_updated() {
        let (engine, tx) = new_test_engine();
        let primary = engine.key_builder().i_to_outer_key(10);
        let request_lock_prefix = make_test_lock_prefix(primary.clone(), 42, 43);
        let existing_lock_prefix = make_test_lock_prefix(primary, 42, 88);
        let existing_txn_file = build_test_txn_file(
            &engine,
            42,
            &[(10, 20)],
            existing_lock_prefix.clone(),
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        );
        apply_lock_txn_file(&existing_txn_file, &tx);

        let request_txn_file = build_test_txn_file(
            &engine,
            42,
            &[(10, 20), (20, 30)],
            request_lock_prefix,
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        );
        let expected_chunk_ids = request_txn_file.chunk_ids();
        let mut cmd = new_test_txn_file_command(request_txn_file);
        let snap = engine.get_snap_access(1).unwrap();

        assert!(matches!(
            cmd.try_merge_txn_file_locks(&snap).unwrap(),
            TxnFileLockMergeResult::FoundUpdated
        ));
        assert_eq!(cmd.txn_file().chunk_ids(), expected_chunk_ids);
        assert_eq!(
            cmd.txn_file_snapshot().txn_file_ref.get_chunk_ids(),
            expected_chunk_ids.as_slice()
        );
        assert_eq!(
            cmd.txn_file().get_lock_val_prefix(),
            existing_lock_prefix.as_slice()
        );
        assert_eq!(
            cmd.txn_file_snapshot().txn_file_ref.get_lock_val_prefix(),
            existing_lock_prefix.as_slice()
        );
    }

    #[test]
    fn test_try_merge_txn_file_locks_rejects_unexpected_mismatch() {
        let (engine, tx) = new_test_engine();
        let primary = engine.key_builder().i_to_outer_key(10);
        let existing_txn_file = build_test_txn_file(
            &engine,
            42,
            &[(10, 20)],
            make_test_lock_prefix(primary.clone(), 42, 43),
            InnerKey::from_inner_buf(b""),
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        );
        apply_lock_txn_file(&existing_txn_file, &tx);

        let mismatched_lower_bound_key = engine.key_builder().i_to_key(15);
        let mismatched_lower_bound = InnerKey::from_outer_key(&mismatched_lower_bound_key);
        let request_txn_file = build_test_txn_file(
            &engine,
            42,
            &[(10, 20)],
            make_test_lock_prefix(primary, 42, 43),
            mismatched_lower_bound,
            InnerKey::from_inner_buf(GLOBAL_SHARD_END_KEY),
        );
        let mut cmd = new_test_txn_file_command(request_txn_file);
        let snap = engine.get_snap_access(1).unwrap();
        let err = match cmd.try_merge_txn_file_locks(&snap) {
            Ok(_) => panic!("expected mismatch error"),
            Err(err) => err,
        };

        assert!(
            err.to_string()
                .contains("unexpected txn file locks of same transaction")
        );
    }
}