// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.
// #[PerformanceCriticalPath]
//! Commands used in the transaction system
#[macro_use]
mod macros;
pub(crate) mod acquire_pessimistic_lock;
pub(crate) mod acquire_pessimistic_lock_resumed;
pub(crate) mod check_secondary_locks;
pub(crate) mod check_txn_status;
pub(crate) mod cleanup;
pub(crate) mod commit;
pub(crate) mod mvcc_by_key;
pub(crate) mod mvcc_by_start_ts;
pub(crate) mod pause;
pub(crate) mod pessimistic_rollback;
pub(crate) mod pessimistic_rollback_read_phase;
pub(crate) mod prewrite;
pub(crate) mod resolve_lock;
pub(crate) mod resolve_lock_async;
pub(crate) mod resolve_lock_lite;
pub(crate) mod resolve_lock_readphase;
pub(crate) mod rollback;
pub(crate) mod txn_file;
pub(crate) mod txn_heart_beat;

use std::{
    fmt::{self, Debug, Display, Formatter},
    iter,
    marker::PhantomData,
    ops::{Deref, DerefMut},
    sync::Arc,
};

pub use acquire_pessimistic_lock::AcquirePessimisticLock;
pub use acquire_pessimistic_lock_resumed::AcquirePessimisticLockResumed;
pub use check_secondary_locks::CheckSecondaryLocks;
pub use check_txn_status::CheckTxnStatus;
pub use cleanup::Cleanup;
use collections::HashMap;
pub use commit::Commit;
use concurrency_manager::{ConcurrencyManager, KeyHandleGuard};
use kvproto::kvrpcpb::*;
pub use mvcc_by_key::MvccByKey;
pub use mvcc_by_start_ts::MvccByStartTs;
pub use pause::Pause;
pub use pessimistic_rollback::PessimisticRollback;
pub use pessimistic_rollback_read_phase::PessimisticRollbackReadPhase;
pub use prewrite::{Prewrite, PrewritePessimistic, one_pc_commit};
pub use resolve_lock::{RESOLVE_LOCK_BATCH_SIZE, ResolveLock};
pub use resolve_lock_async::ResolveLockAsync;
pub use resolve_lock_lite::ResolveLockLite;
pub use resolve_lock_readphase::ResolveLockReadPhase;
pub use rollback::Rollback;
use tikv_util::{Either, deadline::Deadline};
use tracker::RequestType;
pub use txn_heart_beat::TxnHeartBeat;
use txn_types::{Key, LockInfoExt, TimeStamp, Value, Write};

use crate::storage::{
    Result as StorageResult, Snapshot, Statistics,
    kv::WriteData,
    lock_manager::{
        self, LockManager, LockWaitToken, WaitTimeout,
        lock_wait_context::LockWaitContextSharedState,
    },
    metrics,
    mvcc::{Lock as MvccLock, ReleasedLock, SharedLockUpgradeHint, SnapshotReader},
    txn::{
        ProcessResult, Result, commands::txn_file::TxnFileCommand, latch,
        txn_status_cache::TxnStatusCache,
    },
    types::{
        MvccInfo, PessimisticLockParameters, PessimisticLockResults, PrewriteResult,
        SecondaryLocksStatus, StorageCallbackType, TxnStatus,
    },
};

/// Store Transaction scheduler commands.
///
/// Learn more about our transaction system at
/// [Deep Dive TiKV: Distributed Transactions](https://tikv.org/docs/deep-dive/distributed-transaction/introduction/)
///
/// These are typically scheduled and used through the
/// [`Storage`](crate::storage::Storage) with functions like
/// [`prewrite`](prewrite::Prewrite) trait and are executed asynchronously.
pub enum Command {
    Prewrite(Prewrite),
    PrewritePessimistic(PrewritePessimistic),
    AcquirePessimisticLock(AcquirePessimisticLock),
    AcquirePessimisticLockResumed(AcquirePessimisticLockResumed),
    Commit(Commit),
    Cleanup(Cleanup),
    Rollback(Rollback),
    PessimisticRollback(PessimisticRollback),
    PessimisticRollbackReadPhase(PessimisticRollbackReadPhase),
    TxnHeartBeat(TxnHeartBeat),
    CheckTxnStatus(CheckTxnStatus),
    CheckSecondaryLocks(CheckSecondaryLocks),
    ResolveLockReadPhase(ResolveLockReadPhase),
    ResolveLockAsync(ResolveLockAsync),
    ResolveLock(ResolveLock),
    ResolveLockLite(ResolveLockLite),
    Pause(Pause),
    MvccByKey(MvccByKey),
    MvccByStartTs(MvccByStartTs),
    TxnFile(TxnFileCommand),
}

/// A `Command` with its return type, reified as the generic parameter `T`.
///
/// Incoming grpc requests (like `CommitRequest`, `PrewriteRequest`) are
/// converted to this type via a series of transformations. That process is
/// described below using `CommitRequest` as an example:
/// 1. A `CommitRequest` is handled by the `future_commit` method in kv.rs,
///    where it needs to be transformed to a `TypedCommand` before being passed
///    to the `storage.sched_txn_command` method.
/// 2. The `From<CommitRequest>` impl for `TypedCommand` gets chosen, and its
///    generic parameter indicates that the result type for this instance of
///    `TypedCommand` is going to be `TxnStatus` - one of the variants of the
///    `StorageCallback` enum.
/// 3. In the above `from` method, the details of the commit request are
///    captured by creating an instance of the struct
///    `storage::txn::commands::commit::Command` via its `new` method.
/// 4. This struct is wrapped in a variant of the enum
///    `storage::txn::commands::Command`. This enum exists to facilitate generic
///    operations over different commands.
/// 5. Finally, the `Command` enum variant for `Commit` is converted to the
///    `TypedCommand` using the `From<Command>` impl for `TypedCommand`.
///
/// For other requests, see the corresponding `future_` method, the `From` trait
/// implementation and so on.
pub struct TypedCommand<T> {
    pub cmd: Command,

    /// Track the type of the command's return value.
    _pd: PhantomData<T>,
}

impl<T: StorageCallbackType> From<Command> for TypedCommand<T> {
    fn from(cmd: Command) -> TypedCommand<T> {
        TypedCommand {
            cmd,
            _pd: PhantomData,
        }
    }
}

impl<T> From<TypedCommand<T>> for Command {
    fn from(t: TypedCommand<T>) -> Command {
        t.cmd
    }
}

impl From<PrewriteRequest> for TypedCommand<PrewriteResult> {
    fn from(mut req: PrewriteRequest) -> Self {
        let for_update_ts = req.get_for_update_ts();
        let secondary_keys = if req.get_use_async_commit() {
            Some(req.get_secondaries().into())
        } else {
            None
        };
        if for_update_ts == 0 {
            // Disallow skip_constraint_check for optimistic transactions
            let mut skip_constraint_check = req.get_skip_constraint_check();
            if skip_constraint_check {
                warn!(
                    "skip_constraint_check is not allowed for optimistic transactions; \
                     start_ts={}",
                    req.get_start_version()
                );
                // Override the flag to false for optimistic transactions
                skip_constraint_check = false;
            }

            Prewrite::new(
                req.take_mutations().into_iter().map(Into::into).collect(),
                req.take_primary_lock(),
                req.get_start_version().into(),
                req.get_lock_ttl(),
                skip_constraint_check,
                req.get_txn_size(),
                req.get_min_commit_ts().into(),
                req.get_max_commit_ts().into(),
                secondary_keys,
                req.get_try_one_pc(),
                req.get_assertion_level(),
                req.get_txn_file_chunks().into(),
                req.take_context(),
            )
        } else {
            let pessimistic_actions = req.take_pessimistic_actions();
            let mutations = req
                .take_mutations()
                .into_iter()
                .map(Into::into)
                .zip(pessimistic_actions)
                .collect();
            PrewritePessimistic::new(
                mutations,
                req.take_primary_lock(),
                req.get_start_version().into(),
                req.get_lock_ttl(),
                for_update_ts.into(),
                req.get_txn_size(),
                req.get_min_commit_ts().into(),
                req.get_max_commit_ts().into(),
                secondary_keys,
                req.get_try_one_pc(),
                req.get_assertion_level(),
                req.get_txn_file_chunks().into(),
                req.take_context(),
            )
        }
    }
}

impl From<PessimisticLockRequest> for TypedCommand<StorageResult<PessimisticLockResults>> {
    fn from(mut req: PessimisticLockRequest) -> Self {
        let keys = req
            .take_mutations()
            .into_iter()
            .map(|x| match x.get_op() {
                Op::PessimisticLock | Op::SharedPessimisticLock => (
                    Key::from_raw(x.get_key()),
                    x.get_assertion() == Assertion::NotExist,
                    x.get_op() == Op::SharedPessimisticLock,
                ),
                _ => panic!("mismatch Op in pessimistic lock mutations"),
            })
            .collect();

        let allow_lock_with_conflict = match req.get_wake_up_mode() {
            PessimisticLockWakeUpMode::WakeUpModeNormal => false,
            PessimisticLockWakeUpMode::WakeUpModeForceLock => true,
        };

        AcquirePessimisticLock::new(
            keys,
            req.take_primary_lock(),
            req.get_start_version().into(),
            req.get_lock_ttl(),
            req.get_is_first_lock(),
            req.get_for_update_ts().into(),
            WaitTimeout::from_encoded(req.get_wait_timeout()),
            req.get_return_values(),
            req.get_min_commit_ts().into(),
            req.get_check_existence(),
            req.get_lock_only_if_exists(),
            allow_lock_with_conflict,
            req.take_context(),
        )
    }
}

impl From<CommitRequest> for TypedCommand<TxnStatus> {
    fn from(mut req: CommitRequest) -> Self {
        let keys = req.get_keys().iter().map(|x| Key::from_raw(x)).collect();

        Commit::new(
            keys,
            req.get_start_version().into(),
            req.get_commit_version().into(),
            req.get_use_async_commit(),
            req.get_is_txn_file(),
            req.take_context(),
        )
    }
}

impl From<CleanupRequest> for TypedCommand<()> {
    fn from(mut req: CleanupRequest) -> Self {
        Cleanup::new(
            Key::from_raw(req.get_key()),
            req.get_start_version().into(),
            req.get_current_ts().into(),
            req.take_context(),
        )
    }
}

impl From<BatchRollbackRequest> for TypedCommand<()> {
    fn from(mut req: BatchRollbackRequest) -> Self {
        let keys = req.get_keys().iter().map(|x| Key::from_raw(x)).collect();
        Rollback::new(
            keys,
            req.get_start_version().into(),
            req.get_is_txn_file(),
            req.take_context(),
        )
    }
}

impl From<PessimisticRollbackRequest> for TypedCommand<Vec<StorageResult<()>>> {
    fn from(mut req: PessimisticRollbackRequest) -> Self {
        // If the keys are empty, try to scan locks with specified `start_ts` and
        // `for_update_ts`, and then pass them to a new pessimitic rollback
        // command to clean up, just like resolve lock with read phase.
        if req.get_keys().is_empty() {
            PessimisticRollbackReadPhase::new(
                req.get_start_version().into(),
                req.get_for_update_ts().into(),
                None,
                req.take_context(),
            )
        } else {
            let keys = req.get_keys().iter().map(|x| Key::from_raw(x)).collect();
            PessimisticRollback::new(
                keys,
                req.get_start_version().into(),
                req.get_for_update_ts().into(),
                None,
                req.take_context(),
            )
        }
    }
}

impl From<TxnHeartBeatRequest> for TypedCommand<TxnStatus> {
    fn from(mut req: TxnHeartBeatRequest) -> Self {
        TxnHeartBeat::new(
            Key::from_raw(req.get_primary_lock()),
            req.get_start_version().into(),
            req.get_advise_lock_ttl(),
            req.get_is_txn_file(),
            req.take_context(),
        )
    }
}

impl From<CheckTxnStatusRequest> for TypedCommand<TxnStatus> {
    fn from(mut req: CheckTxnStatusRequest) -> Self {
        CheckTxnStatus::new(
            Key::from_raw(req.get_primary_key()),
            req.get_lock_ts().into(),
            req.get_caller_start_ts().into(),
            req.get_current_ts().into(),
            req.get_rollback_if_not_exist(),
            req.get_force_sync_commit(),
            req.get_resolving_pessimistic_lock(),
            req.get_is_txn_file(),
            req.get_verify_is_primary(),
            req.take_context(),
        )
    }
}

impl From<CheckSecondaryLocksRequest> for TypedCommand<SecondaryLocksStatus> {
    fn from(mut req: CheckSecondaryLocksRequest) -> Self {
        CheckSecondaryLocks::new(
            req.take_keys()
                .into_iter()
                .map(|k| Key::from_raw(&k))
                .collect(),
            req.get_start_version().into(),
            req.take_context(),
        )
    }
}

impl From<ResolveLockRequest> for TypedCommand<()> {
    fn from(mut req: ResolveLockRequest) -> Self {
        let resolve_keys: Vec<Key> = req
            .get_keys()
            .iter()
            .map(|key| Key::from_raw(key))
            .collect();
        let txn_file_status: HashMap<TimeStamp, TimeStamp> = req
            .get_txn_infos()
            .iter()
            .filter(|info| info.is_txn_file)
            .map(|info| (info.txn.into(), info.status.into()))
            .collect();
        let txn_status = if req.get_start_version() > 0 {
            iter::once((
                req.get_start_version().into(),
                req.get_commit_version().into(),
            ))
            .collect()
        } else {
            req.take_txn_infos()
                .into_iter()
                .filter(|info| !info.is_txn_file)
                .map(|info| (info.txn.into(), info.status.into()))
                .collect()
        };
        let is_async_region_level_resolve = req.get_is_async()
            && resolve_keys.is_empty()
            && !req.get_is_txn_file()
            && txn_file_status.is_empty();
        if is_async_region_level_resolve {
            ResolveLockAsync::new(txn_status, req.take_context())
        } else if !txn_file_status.is_empty() {
            ResolveLock::new(
                txn_status,
                None,
                vec![],
                txn_file_status,
                false,
                req.take_context(),
            )
        } else if resolve_keys.is_empty() && !req.get_is_txn_file() {
            ResolveLockReadPhase::new(txn_status, None, false, req.take_context())
        } else {
            let start_ts: TimeStamp = req.get_start_version().into();
            assert!(!start_ts.is_zero());
            let commit_ts = req.get_commit_version().into();
            ResolveLockLite::new(
                start_ts,
                commit_ts,
                resolve_keys,
                req.get_is_txn_file(),
                req.take_context(),
            )
        }
    }
}

impl From<MvccGetByKeyRequest> for TypedCommand<MvccInfo> {
    fn from(mut req: MvccGetByKeyRequest) -> Self {
        MvccByKey::new(Key::from_raw(req.get_key()), req.take_context())
    }
}

impl From<MvccGetByStartTsRequest> for TypedCommand<Option<(Key, MvccInfo)>> {
    fn from(mut req: MvccGetByStartTsRequest) -> Self {
        MvccByStartTs::new(req.get_start_ts().into(), req.take_context())
    }
}

/// Represents for a scheduler command, when should the response sent to the
/// client. For most cases, the response should be sent after the result being
/// successfully applied to the storage (if needed). But in some special cases,
/// some optimizations allows the response to be returned at an earlier phase.
///
/// Note that this doesn't affect latch releasing. The latch and the memory lock
/// (if any) are always released after applying, regardless of when the response
/// is sent.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum ResponsePolicy {
    /// Return the response to the client when the command has finished
    /// applying.
    OnApplied,
    /// Return the response after finishing Raft committing.
    OnCommitted,
    /// Return the response after finishing raft proposing.
    OnProposed,
}

pub struct WriteResult {
    pub ctx: Context,
    pub to_be_write: WriteData,
    pub rows: usize,
    pub pr: ProcessResult,
    pub lock_info: Vec<WriteResultLockInfo>,
    pub released_locks: ReleasedLocks,
    pub new_acquired_locks: Vec<LockInfo>,
    /// Hints to consume while affected per-key latches remain held, before
    /// submitting the corresponding writes to the engine.
    pub shared_lock_upgrade_hints: Vec<SharedLockUpgradeHint>,
    pub lock_guards: Vec<KeyHandleGuard>,
    pub response_policy: ResponsePolicy,
    /// The txn status that can be inferred by the successful writing. This will
    /// be used to update the cache.
    ///
    /// Currently only commit_ts of committed transactions will be collected.
    /// Rolled-back transactions may also be collected in the future.
    pub known_txn_status: Vec<(TimeStamp, TimeStamp)>,
}

pub struct WriteResultLockInfo {
    pub lock_digest: lock_manager::LockDigest,
    pub key: Key,
    pub should_not_exist: bool,
    pub lock_info_pb: LockInfo,
    pub parameters: PessimisticLockParameters,
    pub hash_for_latch: u64,
    /// Whether the pending request is trying to acquire a shared lock.
    pub is_shared_lock_request: bool,
    /// Whether this wait is a shared-holder upgrade on the same key.
    pub is_lock_upgrade: bool,
    /// If a request is woken up after waiting for some lock, and it encounters
    /// another lock again after resuming, this field will carry the token
    /// that was already allocated before.
    pub lock_wait_token: LockWaitToken,
    /// For resumed pessimistic lock requests, this is needed to check if it's
    /// canceled outside.
    pub req_states: Option<Arc<LockWaitContextSharedState>>,
}

impl WriteResultLockInfo {
    pub fn new(
        lock_info_pb: LockInfo,
        parameters: PessimisticLockParameters,
        key: Key,
        should_not_exist: bool,
        is_shared_lock_request: bool,
    ) -> Self {
        let lock = lock_manager::LockDigest {
            ts: lock_info_pb.get_lock_version().into(),
            hash: key.gen_hash(),
        };
        let hash_for_latch = latch::Lock::hash(&key);
        let is_lock_upgrade = !is_shared_lock_request
            && lock_info_pb.is_shared_lock()
            && lock_info_pb
                .iter_locks()
                .any(|lock| TimeStamp::from(lock.get_lock_version()) == parameters.start_ts);
        Self {
            lock_digest: lock,
            key,
            should_not_exist,
            lock_info_pb,
            parameters,
            hash_for_latch,
            is_shared_lock_request,
            is_lock_upgrade,
            lock_wait_token: LockWaitToken(None),
            req_states: None,
        }
    }
}

#[derive(Default)]
pub struct ReleasedLocks(Vec<ReleasedLock>);

impl ReleasedLocks {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn push(&mut self, lock: Option<ReleasedLock>) {
        if let Some(lock) = lock {
            self.0.push(lock);
        }
    }

    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    pub fn clear(&mut self) {
        self.0.clear()
    }

    pub fn into_iter(self) -> impl Iterator<Item = ReleasedLock> {
        self.0.into_iter()
    }
}

type LockWritesVals = (
    Option<MvccLock>,
    Vec<(TimeStamp, Write)>,
    Vec<(TimeStamp, Value)>,
);

#[maybe_async::both]
pub(crate) async fn find_mvcc_infos_by_key<S: Snapshot>(
    reader: &mut SnapshotReader<S>,
    key: &Key,
    ts: TimeStamp,
) -> Result<LockWritesVals> {
    let mut values = vec![];
    let lock = reader
        .load_lock(key)?
        .and_then(|lock_or_shared| match lock_or_shared {
            Either::Left(lock) => Some(lock),
            // This debug surface returns one representative holder for shared locks.
            Either::Right(mut shared_locks) => {
                let first_ts = shared_locks.iter_ts().next().cloned();
                first_ts.and_then(|ts| shared_locks.get_lock(&ts).ok().flatten().cloned())
            }
        });
    let mut writes = reader.cloud_reader.scan_write_for_key(key, ts).await?;
    writes.extend(reader.cloud_reader.get_extras(key)?);
    for (ts, v) in reader.scan_values_in_default(key)? {
        values.push((ts, v));
    }
    Ok((lock, writes, values))
}

pub trait CommandExt: Display {
    fn tag(&self) -> metrics::CommandKind;

    fn request_type(&self) -> RequestType {
        RequestType::Unknown
    }

    fn get_ctx(&self) -> &Context;

    fn get_ctx_mut(&mut self) -> &mut Context;

    fn deadline(&self) -> Deadline;

    fn incr_cmd_metric(&self) {
        metrics::KV_COMMAND_COUNTER_VEC_STATIC.get(self.tag()).inc();
    }

    fn ts(&self) -> TimeStamp {
        TimeStamp::zero()
    }

    fn readonly(&self) -> bool {
        false
    }

    fn is_sys_cmd(&self) -> bool {
        false
    }

    fn can_be_pipelined(&self) -> bool {
        false
    }

    fn write_bytes(&self) -> usize;

    fn gen_lock(&self) -> latch::Lock;

    fn can_build_txn_file(&self) -> bool {
        false
    }
}

pub struct RawExt {
    pub ts: TimeStamp,
    pub key_guard: KeyHandleGuard,
}

pub struct WriteContext<'a, L: LockManager> {
    pub lock_mgr: &'a L,
    pub concurrency_manager: ConcurrencyManager,
    pub extra_op: ExtraOp,
    /// EXTRA_CF scan mode captured when the scheduler starts this command.
    pub extra_cf_scan_mode: crate::storage::config::ExtraCfScanMode,
    pub statistics: &'a mut Statistics,
    pub async_apply_prewrite: bool,
    pub txn_status_cache: &'a TxnStatusCache,
}

pub struct ReaderWithStats<'a, S: Snapshot> {
    reader: SnapshotReader<S>,
    statistics: &'a mut Statistics,
}

impl<'a, S: Snapshot> ReaderWithStats<'a, S> {
    fn new(reader: SnapshotReader<S>, statistics: &'a mut Statistics) -> Self {
        Self { reader, statistics }
    }
}

impl<S: Snapshot> Deref for ReaderWithStats<'_, S> {
    type Target = SnapshotReader<S>;

    fn deref(&self) -> &Self::Target {
        &self.reader
    }
}

impl<S: Snapshot> DerefMut for ReaderWithStats<'_, S> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.reader
    }
}

impl<S: Snapshot> Drop for ReaderWithStats<'_, S> {
    fn drop(&mut self) {
        self.statistics.add(&self.reader.take_statistics())
    }
}

impl Command {
    // These two are for backward compatibility, after some other refactors are done
    // we can remove Command totally and use `&dyn CommandExt` instead
    fn command_ext(&self) -> &dyn CommandExt {
        match &self {
            Command::Prewrite(t) => t,
            Command::PrewritePessimistic(t) => t,
            Command::AcquirePessimisticLock(t) => t,
            Command::AcquirePessimisticLockResumed(t) => t,
            Command::Commit(t) => t,
            Command::Cleanup(t) => t,
            Command::Rollback(t) => t,
            Command::PessimisticRollback(t) => t,
            Command::PessimisticRollbackReadPhase(t) => t,
            Command::TxnHeartBeat(t) => t,
            Command::CheckTxnStatus(t) => t,
            Command::CheckSecondaryLocks(t) => t,
            Command::ResolveLockReadPhase(t) => t,
            Command::ResolveLockAsync(t) => t,
            Command::ResolveLock(t) => t,
            Command::ResolveLockLite(t) => t,
            Command::Pause(t) => t,
            Command::MvccByKey(t) => t,
            Command::MvccByStartTs(t) => t,
            Command::TxnFile(t) => t,
        }
    }

    fn command_ext_mut(&mut self) -> &mut dyn CommandExt {
        match self {
            Command::Prewrite(t) => t,
            Command::PrewritePessimistic(t) => t,
            Command::AcquirePessimisticLock(t) => t,
            Command::AcquirePessimisticLockResumed(t) => t,
            Command::Commit(t) => t,
            Command::Cleanup(t) => t,
            Command::Rollback(t) => t,
            Command::PessimisticRollback(t) => t,
            Command::PessimisticRollbackReadPhase(t) => t,
            Command::TxnHeartBeat(t) => t,
            Command::CheckTxnStatus(t) => t,
            Command::CheckSecondaryLocks(t) => t,
            Command::ResolveLockReadPhase(t) => t,
            Command::ResolveLockAsync(t) => t,
            Command::ResolveLock(t) => t,
            Command::ResolveLockLite(t) => t,
            Command::Pause(t) => t,
            Command::MvccByKey(t) => t,
            Command::MvccByStartTs(t) => t,
            Command::TxnFile(t) => t,
        }
    }

    #[maybe_async::both]
    pub(super) async fn process_read<S: Snapshot + 'static>(
        self,
        snapshot: S,
        statistics: &mut Statistics,
    ) -> Result<ProcessResult> {
        match self {
            Command::ResolveLockReadPhase(t) => t.process_read(snapshot, statistics).await,
            Command::ResolveLockAsync(t) => t.process_read(snapshot, statistics).await,
            Command::MvccByKey(t) => t.process_read(snapshot, statistics).await,
            Command::MvccByStartTs(t) => t.process_read(snapshot, statistics).await,
            Command::PessimisticRollbackReadPhase(t) => t.process_read(snapshot, statistics).await,
            _ => panic!("unsupported read command"),
        }
    }

    #[maybe_async::both]
    pub(crate) async fn process_write<S: Snapshot + 'static, L: LockManager>(
        self,
        snapshot: S,
        context: WriteContext<'_, L>,
    ) -> Result<WriteResult> {
        match self {
            Command::Prewrite(t) => t.process_write(snapshot, context).await,
            Command::PrewritePessimistic(t) => t.process_write(snapshot, context).await,
            Command::AcquirePessimisticLock(t) => t.process_write(snapshot, context).await,
            Command::AcquirePessimisticLockResumed(t) => t.process_write(snapshot, context).await,
            Command::Commit(t) => t.process_write(snapshot, context).await,
            Command::Cleanup(t) => t.process_write(snapshot, context).await,
            Command::Rollback(t) => t.process_write(snapshot, context).await,
            Command::PessimisticRollback(t) => t.process_write(snapshot, context).await,
            Command::ResolveLock(t) => t.process_write(snapshot, context).await,
            Command::ResolveLockLite(t) => t.process_write(snapshot, context).await,
            Command::TxnHeartBeat(t) => t.process_write(snapshot, context).await,
            Command::CheckTxnStatus(t) => t.process_write(snapshot, context).await,
            Command::CheckSecondaryLocks(t) => t.process_write(snapshot, context).await,
            Command::Pause(t) => t.process_write(snapshot, context).await,
            Command::TxnFile(t) => t.process_write(snapshot, context).await,
            _ => panic!("unsupported write command"),
        }
    }

    pub fn readonly(&self) -> bool {
        self.command_ext().readonly()
    }

    pub fn incr_cmd_metric(&self) {
        self.command_ext().incr_cmd_metric()
    }

    pub fn priority(&self) -> CommandPri {
        if self.command_ext().is_sys_cmd() {
            return CommandPri::High;
        }
        if let Command::TxnFile(txn_file) = self {
            return txn_file.priority();
        }
        self.command_ext().get_ctx().get_priority()
    }

    pub fn need_flow_control(&self) -> bool {
        !self.readonly() && self.priority() != CommandPri::High
    }

    pub fn tag(&self) -> metrics::CommandKind {
        self.command_ext().tag()
    }

    pub fn request_type(&self) -> RequestType {
        self.command_ext().request_type()
    }

    pub fn ts(&self) -> TimeStamp {
        self.command_ext().ts()
    }

    pub fn write_bytes(&self) -> usize {
        self.command_ext().write_bytes()
    }

    pub fn gen_lock(&self) -> latch::Lock {
        self.command_ext().gen_lock()
    }

    pub fn can_be_pipelined(&self) -> bool {
        self.command_ext().can_be_pipelined()
    }

    pub fn ctx(&self) -> &Context {
        self.command_ext().get_ctx()
    }

    pub fn ctx_mut(&mut self) -> &mut Context {
        self.command_ext_mut().get_ctx_mut()
    }

    pub fn deadline(&self) -> Deadline {
        self.command_ext().deadline()
    }

    pub fn can_build_txn_file(&self) -> bool {
        self.command_ext().can_build_txn_file()
    }
}

impl Display for Command {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.command_ext().fmt(f)
    }
}

impl Debug for Command {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.command_ext().fmt(f)
    }
}

/// Commands that do not need to modify the database during execution will
/// implement this trait.
#[maybe_async::async_trait]
pub trait ReadCommand<S: Snapshot>: CommandExt {
    #[maybe_async]
    async fn process_read(self, snapshot: S, statistics: &mut Statistics) -> Result<ProcessResult>;
}

#[macro_export]
macro_rules! command_process_read {
    ($cmd:expr, $snapshot:expr, $statistics:expr) => {{
        if $snapshot.is_sync() {
            $cmd.process_read($snapshot, $statistics)
        } else {
            $cmd.process_read_async($snapshot, $statistics).await
        }
    }};
}

/// Commands that need to modify the database during execution will implement
/// this trait.
#[maybe_async::async_trait]
pub trait WriteCommand<S: Snapshot, L: LockManager>: CommandExt {
    #[maybe_async]
    async fn process_write(self, snapshot: S, context: WriteContext<'_, L>) -> Result<WriteResult>;
}

#[macro_export]
macro_rules! command_process_write {
    ($cmd:expr, $snapshot:expr, $context:expr, $sample:expr) => {{
        if $snapshot.is_sync() {
            if let Some(sample) = $sample.as_mut() {
                let _guard = sample.observe_cpu();
                $cmd.process_write($snapshot, $context)
            } else {
                $cmd.process_write($snapshot, $context)
            }
        } else {
            let task = async { $cmd.process_write_async($snapshot, $context).await };
            if let Some(sample) = $sample.as_mut() {
                let (cpu_time, res) = sample.observe_cpu_async(task).await;
                sample.add_cpu_time(cpu_time);
                res
            } else {
                task.await
            }
        }
    }};
    ($cmd:expr, $snapshot:expr, $context:expr) => {{
        if $snapshot.is_sync() {
            $cmd.process_write($snapshot, $context)
        } else {
            let task = async { $cmd.process_write_async($snapshot, $context).await };
            task.await
        }
    }};
}

#[cfg(test)]
pub mod test_util {
    use txn_types::Mutation;

    use super::*;
    use crate::storage::{
        Engine, MockLockManager,
        mvcc::{Error as MvccError, ErrorInner as MvccErrorInner},
        txn::{Error, ErrorInner, Result},
    };

    // Some utils for tests that may be used in multiple source code files.

    pub fn prewrite_command<E: Engine>(
        engine: &mut E,
        cm: ConcurrencyManager,
        statistics: &mut Statistics,
        cmd: TypedCommand<PrewriteResult>,
    ) -> Result<PrewriteResult> {
        let snap = engine.snapshot(Default::default())?;
        let context = WriteContext {
            lock_mgr: &MockLockManager::new(),
            concurrency_manager: cm,
            extra_op: ExtraOp::Noop,
            extra_cf_scan_mode: Default::default(),
            statistics,
            async_apply_prewrite: false,
            txn_status_cache: &TxnStatusCache::new_for_test(),
        };
        let ret = cmd.cmd.process_write(snap, context)?;
        let res = match ret.pr {
            ProcessResult::PrewriteResult {
                result: PrewriteResult { locks, .. },
            } if !locks.is_empty() => {
                let info = LockInfo::default();
                return Err(Error::from(ErrorInner::Mvcc(MvccError::from(
                    MvccErrorInner::KeyIsLocked(info),
                ))));
            }
            ProcessResult::PrewriteResult { result } => result,
            _ => unreachable!(),
        };
        let ctx = Context::default();
        if !ret.to_be_write.modifies.is_empty() {
            engine.write(&ctx, ret.to_be_write).unwrap();
        }
        Ok(res)
    }

    pub fn prewrite<E: Engine>(
        engine: &mut E,
        statistics: &mut Statistics,
        mutations: Vec<Mutation>,
        primary: Vec<u8>,
        start_ts: u64,
        one_pc_max_commit_ts: Option<u64>,
    ) -> Result<PrewriteResult> {
        let cm = ConcurrencyManager::new(start_ts.into());
        prewrite_with_cm(
            engine,
            cm,
            statistics,
            mutations,
            primary,
            start_ts,
            one_pc_max_commit_ts,
        )
    }

    pub fn prewrite_with_cm<E: Engine>(
        engine: &mut E,
        cm: ConcurrencyManager,
        statistics: &mut Statistics,
        mutations: Vec<Mutation>,
        primary: Vec<u8>,
        start_ts: u64,
        one_pc_max_commit_ts: Option<u64>,
    ) -> Result<PrewriteResult> {
        let cmd = if let Some(max_commit_ts) = one_pc_max_commit_ts {
            Prewrite::with_1pc(
                mutations,
                primary,
                TimeStamp::from(start_ts),
                max_commit_ts.into(),
            )
        } else {
            Prewrite::with_defaults(mutations, primary, TimeStamp::from(start_ts))
        };
        prewrite_command(engine, cm, statistics, cmd)
    }

    pub fn pessimistic_prewrite<E: Engine>(
        engine: &mut E,
        statistics: &mut Statistics,
        mutations: Vec<(Mutation, PrewriteRequestPessimisticAction)>,
        primary: Vec<u8>,
        start_ts: u64,
        for_update_ts: u64,
        one_pc_max_commit_ts: Option<u64>,
    ) -> Result<PrewriteResult> {
        let cm = ConcurrencyManager::new(start_ts.into());
        pessimistic_prewrite_with_cm(
            engine,
            cm,
            statistics,
            mutations,
            primary,
            start_ts,
            for_update_ts,
            one_pc_max_commit_ts,
        )
    }

    pub fn pessimistic_prewrite_with_cm<E: Engine>(
        engine: &mut E,
        cm: ConcurrencyManager,
        statistics: &mut Statistics,
        mutations: Vec<(Mutation, PrewriteRequestPessimisticAction)>,
        primary: Vec<u8>,
        start_ts: u64,
        for_update_ts: u64,
        one_pc_max_commit_ts: Option<u64>,
    ) -> Result<PrewriteResult> {
        let cmd = if let Some(max_commit_ts) = one_pc_max_commit_ts {
            PrewritePessimistic::with_1pc(
                mutations,
                primary,
                start_ts.into(),
                for_update_ts.into(),
                max_commit_ts.into(),
            )
        } else {
            PrewritePessimistic::with_defaults(
                mutations,
                primary,
                start_ts.into(),
                for_update_ts.into(),
            )
        };
        prewrite_command(engine, cm, statistics, cmd)
    }

    pub fn commit<E: Engine>(
        engine: &mut E,
        statistics: &mut Statistics,
        keys: Vec<Key>,
        lock_ts: u64,
        commit_ts: u64,
    ) -> Result<()> {
        let ctx = Context::default();
        let snap = engine.snapshot(Default::default())?;
        let concurrency_manager = ConcurrencyManager::new(lock_ts.into());
        let cmd = Commit::new(
            keys,
            TimeStamp::from(lock_ts),
            TimeStamp::from(commit_ts),
            false,
            false,
            ctx,
        );

        let context = WriteContext {
            lock_mgr: &MockLockManager::new(),
            concurrency_manager,
            extra_op: ExtraOp::Noop,
            extra_cf_scan_mode: Default::default(),
            statistics,
            async_apply_prewrite: false,
            txn_status_cache: &TxnStatusCache::new_for_test(),
        };

        let ret = cmd.cmd.process_write(snap, context)?;
        let ctx = Context::default();
        engine.write(&ctx, ret.to_be_write).unwrap();
        Ok(())
    }

    pub fn rollback<E: Engine>(
        engine: &mut E,
        statistics: &mut Statistics,
        keys: Vec<Key>,
        start_ts: u64,
    ) -> Result<()> {
        let ctx = Context::default();
        let snap = engine.snapshot(Default::default())?;
        let concurrency_manager = ConcurrencyManager::new(start_ts.into());
        let cmd = Rollback::new(keys, TimeStamp::from(start_ts), false, ctx);
        let context = WriteContext {
            lock_mgr: &MockLockManager::new(),
            concurrency_manager,
            extra_op: ExtraOp::Noop,
            extra_cf_scan_mode: Default::default(),
            statistics,
            async_apply_prewrite: false,
            txn_status_cache: &TxnStatusCache::new_for_test(),
        };

        let ret = cmd.cmd.process_write(snap, context)?;
        let ctx = Context::default();
        engine.write(&ctx, ret.to_be_write).unwrap();
        Ok(())
    }
}