// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.
// #[PerformanceCriticalPath]
use protobuf::Message;
use resource_metering::record_network_out_bytes;
use tikv_util::{Either, box_err};
use txn_types::{Key, Lock, ReqType, WriteType};

use crate::storage::{
    ProcessResult, Snapshot,
    kv::WriteData,
    lock_manager::LockManager,
    mvcc::{LockType, MvccTxn, SnapshotReader, TimeStamp, TxnCommitRecord},
    txn::{
        Result,
        actions::check_txn_status::{collapse_prev_rollback, make_rollback},
        commands::{
            Command, CommandExt, ReaderWithStats, ReleasedLocks, ResponsePolicy, TypedCommand,
            WriteCommand, WriteContext, WriteResult,
        },
    },
    types::SecondaryLocksStatus,
};

command! {
    /// Check secondary locks of an async commit transaction.
    ///
    /// If all prewritten locks exist, the lock information is returned.
    /// Otherwise, it returns the commit timestamp of the transaction.
    ///
    /// If the lock does not exist or is a pessimistic lock, to prevent the
    /// status being changed, a rollback may be written.
    CheckSecondaryLocks:
        cmd_ty => SecondaryLocksStatus,
        display => "kv::command::CheckSecondaryLocks {:?} keys@{} | {:?}", (keys, start_ts, ctx),
        content => {
            /// The keys of secondary locks.
            keys: Vec<Key>,
            /// The start timestamp of the transaction.
            start_ts: txn_types::TimeStamp,
        }
}

impl CommandExt for CheckSecondaryLocks {
    ctx!();
    tag!(check_secondary_locks);
    request_type!(KvCheckSecondaryLocks);
    ts!(start_ts);
    write_bytes!(keys: multiple);
    gen_lock!(keys: multiple);
}

#[derive(Debug, PartialEq)]
enum SecondaryLockStatus {
    Locked(Lock),
    Committed(TimeStamp),
    RolledBack,
}

#[maybe_async::async_trait]
impl<S: Snapshot + 'static, L: LockManager> WriteCommand<S, L> for CheckSecondaryLocks {
    #[maybe_async]
    async fn process_write(self, snapshot: S, context: WriteContext<'_, L>) -> Result<WriteResult> {
        // It is not allowed for commit to overwrite a protected rollback. So we update
        // max_ts to prevent this case from happening.
        context
            .concurrency_manager
            .update_max_ts(self.start_ts, || {
                format!("check_secondary_locks-{}", self.start_ts)
            })?;

        let mut txn = MvccTxn::new(self.start_ts, context.concurrency_manager);
        let mut reader = ReaderWithStats::new(
            SnapshotReader::new_with_ctx(self.start_ts, snapshot, &self.ctx),
            context.statistics,
        );
        let mut released_locks = ReleasedLocks::new();
        let mut result = SecondaryLocksStatus::Locked(Vec::new());
        let mut result_size: u64 = 0;
        for key in self.keys {
            let mut released_lock = None;
            let mut mismatch_lock = None;
            // Checks whether the given secondary lock exists.
            let (status, need_rollback, rollback_overlapped_write) = match reader.load_lock(&key)? {
                // The lock exists, the lock information is returned.
                Some(Either::Left(lock)) if lock.ts == self.start_ts => {
                    if lock.lock_type == LockType::Pessimistic {
                        let for_update_ts = lock.for_update_ts;
                        let ttl = lock.ttl;
                        let use_async_commit = lock.use_async_commit;
                        let min_commit_ts = lock.min_commit_ts;
                        released_lock = txn.unlock_key(key.clone(), true, TimeStamp::zero());
                        tikv_util::txn_info!(trace_event::types::Category::WriteDetails,
                            "check_secondary_locks rolled back pessimistic lock";
                            "key" => %key,
                            "start_ts" => self.start_ts,
                            "ttl" => ttl,
                            "for_update_ts" => for_update_ts,
                            "use_async_commit" => use_async_commit,
                            "min_commit_ts" => min_commit_ts,
                            "request_source" => %self.ctx.get_request_source(),
                        );
                        let overlapped_write =
                            reader.get_txn_commit_record(&key).await?.unwrap_none();
                        (SecondaryLockStatus::RolledBack, true, overlapped_write)
                    } else {
                        (SecondaryLockStatus::Locked(lock), false, None)
                    }
                }
                // Async commit transactions don't write shared locks. If the
                // target transaction is found inside SharedLocks, reject it
                // explicitly instead of turning it into a protected rollback.
                // This intentionally diverges from upstream because cloud
                // engine rejects shared-lock async-commit at prewrite, so
                // reaching this path indicates an invariant violation.
                Some(Either::Right(mut shared_locks))
                    if shared_locks.contains_start_ts(self.start_ts) =>
                {
                    let lock = shared_locks
                        .get_lock(&self.start_ts)
                        .map_err(crate::storage::mvcc::Error::from)?
                        .ok_or_else(|| {
                            crate::storage::txn::ErrorInner::Other(box_err!(
                                "missing shared lock holder after start-ts check"
                            ))
                        })?;
                    debug_assert!(matches!(
                        lock.lock_type,
                        LockType::Lock | LockType::Pessimistic
                    ));
                    return Err(crate::storage::txn::ErrorInner::Other(box_err!(
                        "unexpected check-secondary-locks on shared lock holder {:?}: shared lock transactions cannot use async-commit",
                        lock.lock_type
                    ))
                    .into());
                }
                // SharedLocks owned by other transactions don't decide this
                // async-commit transaction's status, so check write CF directly.
                Some(Either::Right(_)) => match reader.get_txn_commit_record(&key).await? {
                    TxnCommitRecord::SingleRecord { commit_ts, write } => {
                        let status = if write.write_type != WriteType::Rollback {
                            SecondaryLockStatus::Committed(commit_ts)
                        } else {
                            SecondaryLockStatus::RolledBack
                        };
                        (status, false, None)
                    }
                    TxnCommitRecord::OverlappedRollback { .. } => {
                        (SecondaryLockStatus::RolledBack, false, None)
                    }
                    TxnCommitRecord::None { overlapped_write } => {
                        (SecondaryLockStatus::RolledBack, true, overlapped_write)
                    }
                },
                // Searches the write CF for the commit record of the lock and returns the commit
                // timestamp (0 if the lock is not committed).
                l => {
                    mismatch_lock = l.and_then(|lock_or_shared| lock_or_shared.left());
                    match reader.get_txn_commit_record(&key).await? {
                        TxnCommitRecord::SingleRecord { commit_ts, write } => {
                            let status = if write.write_type != WriteType::Rollback {
                                SecondaryLockStatus::Committed(commit_ts)
                            } else {
                                SecondaryLockStatus::RolledBack
                            };
                            // We needn't write a rollback once there is a write record for it:
                            // If it's a committed record, it cannot be changed.
                            // If it's a rollback record, it either comes from another
                            // check_secondary_lock (thus protected) or the client stops commit
                            // actively. So we don't need to make it protected again.
                            (status, false, None)
                        }
                        TxnCommitRecord::OverlappedRollback { .. } => {
                            (SecondaryLockStatus::RolledBack, false, None)
                        }
                        TxnCommitRecord::None { overlapped_write } => {
                            (SecondaryLockStatus::RolledBack, true, overlapped_write)
                        }
                    }
                }
            };
            // If the lock does not exist or is a pessimistic lock, to prevent the
            // status being changed, a rollback may be written and this rollback
            // needs to be protected.
            if need_rollback {
                let has_mismatch_lock = mismatch_lock.is_some();
                if let Some(ref l) = mismatch_lock {
                    txn.mark_rollback_on_mismatching_lock(&key, l.clone(), true);
                }
                // We must protect this rollback in case this rollback is collapsed and a stale
                // acquire_pessimistic_lock and prewrite succeed again.
                let has_overlapped_write = rollback_overlapped_write.is_some();
                if let Some(write) = make_rollback(self.start_ts, true, rollback_overlapped_write) {
                    tikv_util::txn_info!(trace_event::types::Category::WriteDetails,
                        "check_secondary_locks wrote protected rollback";
                        "key" => %key,
                        "start_ts" => self.start_ts,
                        "has_overlapped_write" => has_overlapped_write,
                        "has_mismatch_lock" => has_mismatch_lock,
                        "mismatch_lock" => ?mismatch_lock,
                        "request_source" => %self.ctx.get_request_source(),
                    );
                    txn.put_write(key.clone(), self.start_ts, write.as_ref().to_bytes());
                    collapse_prev_rollback(&mut txn, &mut reader, &key)?;
                }
            }
            released_locks.push(released_lock);
            match status {
                SecondaryLockStatus::Locked(lock) => {
                    let lock_info = lock.into_lock_info(key.to_raw()?);
                    result_size += lock_info.compute_size() as u64;
                    result.push(lock_info);
                }
                SecondaryLockStatus::Committed(commit_ts) => {
                    result = SecondaryLocksStatus::Committed(commit_ts);
                    break;
                }
                SecondaryLockStatus::RolledBack => {
                    result = SecondaryLocksStatus::RolledBack;
                    break;
                }
            }
        }

        record_network_out_bytes(result_size);
        let write_result_known_txn_status =
            if let SecondaryLocksStatus::Committed(commit_ts) = &result {
                vec![(self.start_ts, *commit_ts)]
            } else {
                vec![]
            };
        let mut rows = 0;
        if let SecondaryLocksStatus::RolledBack = &result {
            // One row is mutated only when a secondary lock is rolled back.
            rows = 1;
        }
        let pr = ProcessResult::SecondaryLocksStatus { status: result };
        let shared_lock_upgrade_hints = txn.take_shared_lock_upgrade_hints();
        let shared_lock_lifecycle = txn.has_shared_lock_lifecycle();
        let mut write_data = WriteData::from_modifies(txn.into_modifies());
        write_data.set_req_type(ReqType::CheckSecondaryLocks);
        write_data.set_shared_lock_lifecycle(shared_lock_lifecycle);
        write_data.set_allowed_on_disk_almost_full();
        Ok(WriteResult {
            ctx: self.ctx,
            to_be_write: write_data,
            rows,
            pr,
            lock_info: vec![],
            released_locks,
            new_acquired_locks: vec![],
            shared_lock_upgrade_hints,
            lock_guards: vec![],
            response_policy: ResponsePolicy::OnApplied,
            known_txn_status: write_result_known_txn_status,
        })
    }
}

#[cfg(test)]
pub mod tests {
    use cloud_server::mock_kv_engine::TestKvEngine;
    use concurrency_manager::ConcurrencyManager;
    use kvproto::kvrpcpb::Context;
    use tikv_util::deadline::Deadline;

    use super::*;
    use crate::storage::{
        Engine,
        kv::TestEngineBuilder,
        lock_manager::MockLockManager,
        mvcc::tests::*,
        txn::{
            commands::WriteCommand, scheduler::DEFAULT_EXECUTION_DURATION_LIMIT, tests::*,
            txn_status_cache::TxnStatusCache,
        },
    };

    #[cfg(test)]
    fn write_shared_locks<E: Engine>(engine: &E, key: &[u8], start_ts: impl Into<TimeStamp>) {
        write_shared_locks_with_type(engine, key, start_ts, LockType::Lock);
    }

    #[cfg(test)]
    fn write_shared_locks_with_type<E: Engine>(
        engine: &E,
        key: &[u8],
        start_ts: impl Into<TimeStamp>,
        lock_type: LockType,
    ) {
        let mut shared_locks = txn_types::SharedLocks::new();
        let start_ts = start_ts.into();
        shared_locks
            .insert_lock(Lock::new(
                lock_type,
                b"primary".to_vec(),
                start_ts,
                3000,
                None,
                TimeStamp::zero(),
                1,
                TimeStamp::zero(),
            ))
            .unwrap();
        let mut data = tikv_kv::WriteData::from_modifies(vec![tikv_kv::Modify::Put(
            engine_traits::CF_LOCK,
            Key::from_raw(key),
            shared_locks.to_bytes(),
        )]);
        data.extra.req_type = txn_types::ReqType::PessimisticLock;
        data.set_shared_lock_lifecycle(true);
        write(engine, &Context::default(), data);
    }

    #[cfg(test)]
    fn check_secondary_err<E: Engine>(
        engine: &mut E,
        key: &[u8],
        lock_ts: impl Into<TimeStamp>,
    ) -> crate::storage::txn::Error {
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let lock_ts = lock_ts.into();
        let cm = ConcurrencyManager::new(lock_ts);
        let command = crate::storage::txn::commands::CheckSecondaryLocks {
            ctx: Context::default(),
            keys: vec![Key::from_raw(key)],
            start_ts: lock_ts,
            deadline: Deadline::from_now(DEFAULT_EXECUTION_DURATION_LIMIT),
        };
        match command.process_write(
            snapshot,
            WriteContext {
                lock_mgr: &MockLockManager::new(),
                concurrency_manager: cm,
                extra_op: Default::default(),
                extra_cf_scan_mode: Default::default(),
                statistics: &mut Default::default(),
                async_apply_prewrite: false,
                txn_status_cache: &TxnStatusCache::new_for_test(),
            },
        ) {
            Ok(_) => panic!("expected check secondary locks to fail"),
            Err(err) => err,
        }
    }

    pub fn must_success<E: Engine>(
        engine: &mut E,
        key: &[u8],
        lock_ts: impl Into<TimeStamp>,
        expect_status: SecondaryLocksStatus,
    ) {
        let ctx = Context::default();
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let lock_ts = lock_ts.into();
        let cm = ConcurrencyManager::new(lock_ts);
        let command = crate::storage::txn::commands::CheckSecondaryLocks {
            ctx: ctx.clone(),
            keys: vec![Key::from_raw(key)],
            start_ts: lock_ts,
            deadline: Deadline::from_now(DEFAULT_EXECUTION_DURATION_LIMIT),
        };
        let result = command
            .process_write(
                snapshot,
                WriteContext {
                    lock_mgr: &MockLockManager::new(),
                    concurrency_manager: cm,
                    extra_op: Default::default(),
                    extra_cf_scan_mode: Default::default(),
                    statistics: &mut Default::default(),
                    async_apply_prewrite: false,
                    txn_status_cache: &TxnStatusCache::new_for_test(),
                },
            )
            .unwrap();
        if let ProcessResult::SecondaryLocksStatus { status } = result.pr {
            assert_eq!(status, expect_status);
            let shared_lock_lifecycle = result.to_be_write.extra.shared_lock_lifecycle;
            let mut data = tikv_kv::WriteData::from_modifies(result.to_be_write.modifies);
            data.extra.req_type = txn_types::ReqType::CheckSecondaryLocks;
            data.set_shared_lock_lifecycle(shared_lock_lifecycle);
            write(engine, &ctx, data);
        } else {
            unreachable!();
        }
    }

    #[test]
    fn test_check_secondary_locks_rejects_target_shared_lock() {
        let mut engine = TestKvEngine::new().unwrap();
        let key = b"k-shared-secondary";
        write_shared_locks(&engine, key, 10);

        let err = check_secondary_err(&mut engine, key, 10);
        match err {
            crate::storage::txn::Error(box crate::storage::txn::ErrorInner::Other(err)) => {
                let message = err.to_string();
                assert!(
                    message.contains("unexpected check-secondary-locks on shared lock"),
                    "unexpected error: {:?}",
                    message
                );
            }
            err => panic!("unexpected error: {:?}", err),
        }
        must_not_have_write(&mut engine, key, 10);
    }

    #[test]
    fn test_check_secondary_locks_rejects_target_shared_pessimistic_lock() {
        let mut engine = TestKvEngine::new().unwrap();
        let key = b"k-shared-pessimistic-secondary";
        write_shared_locks_with_type(&engine, key, 10, LockType::Pessimistic);

        let err = check_secondary_err(&mut engine, key, 10);
        match err {
            crate::storage::txn::Error(box crate::storage::txn::ErrorInner::Other(err)) => {
                let message = err.to_string();
                assert!(
                    message.contains("unexpected check-secondary-locks on shared lock holder"),
                    "unexpected error: {:?}",
                    message
                );
                assert!(
                    message.contains("Pessimistic"),
                    "unexpected error: {:?}",
                    message
                );
            }
            err => panic!("unexpected error: {:?}", err),
        }
        must_not_have_write(&mut engine, key, 10);
    }

    #[test]
    #[ignore = "CSE incompatibility; see #6158"]
    fn test_check_async_commit_secondary_locks() {
        let mut engine = TestEngineBuilder::new().build().unwrap();
        let mut engine_clone = engine.clone();
        let ctx = Context::default();
        let cm = ConcurrencyManager::new(1.into());

        let mut check_secondary = |key, ts| {
            let snapshot = engine_clone.snapshot(Default::default()).unwrap();
            let key = Key::from_raw(key);
            let ts = TimeStamp::new(ts);
            let command = crate::storage::txn::commands::CheckSecondaryLocks {
                ctx: Default::default(),
                keys: vec![key],
                start_ts: ts,
                deadline: Deadline::from_now(DEFAULT_EXECUTION_DURATION_LIMIT),
            };
            let result = command
                .process_write(
                    snapshot,
                    WriteContext {
                        lock_mgr: &MockLockManager::new(),
                        concurrency_manager: cm.clone(),
                        extra_op: Default::default(),
                        extra_cf_scan_mode: Default::default(),
                        statistics: &mut Default::default(),
                        async_apply_prewrite: false,
                        txn_status_cache: &TxnStatusCache::new_for_test(),
                    },
                )
                .unwrap();
            if !result.to_be_write.modifies.is_empty() {
                engine_clone.write(&ctx, result.to_be_write).unwrap();
            }
            if let ProcessResult::SecondaryLocksStatus { status } = result.pr {
                status
            } else {
                unreachable!();
            }
        };

        must_prewrite_lock(&mut engine, b"k1", b"key", 1);
        must_commit(&mut engine, b"k1", 1, 3);
        must_rollback(&mut engine, b"k1", 5, false);
        must_prewrite_lock(&mut engine, b"k1", b"key", 7);
        must_commit(&mut engine, b"k1", 7, 9);

        // Lock CF has no lock
        //
        // LOCK CF       | WRITE CF
        // --------------+---------------------
        //               | 9: start_ts = 7
        //               | 5: rollback
        //               | 3: start_ts = 1

        assert_eq!(
            check_secondary(b"k1", 7),
            SecondaryLocksStatus::Committed(9.into())
        );
        must_get_commit_ts(&mut engine, b"k1", 7, 9);
        assert_eq!(check_secondary(b"k1", 5), SecondaryLocksStatus::RolledBack);
        must_get_rollback_ts(&mut engine, b"k1", 5);
        assert_eq!(
            check_secondary(b"k1", 1),
            SecondaryLocksStatus::Committed(3.into())
        );
        must_get_commit_ts(&mut engine, b"k1", 1, 3);
        assert_eq!(check_secondary(b"k1", 6), SecondaryLocksStatus::RolledBack);
        must_get_rollback_protected(&mut engine, b"k1", 6, true);

        // ----------------------------

        must_acquire_pessimistic_lock(&mut engine, b"k1", b"key", 11, 11);

        // Lock CF has a pessimistic lock
        //
        // LOCK CF       | WRITE CF
        // ------------------------------------
        // ts = 11 (pes) | 9: start_ts = 7
        //               | 5: rollback
        //               | 3: start_ts = 1

        let status = check_secondary(b"k1", 11);
        assert_eq!(status, SecondaryLocksStatus::RolledBack);
        must_get_rollback_protected(&mut engine, b"k1", 11, true);

        // ----------------------------

        must_prewrite_lock(&mut engine, b"k1", b"key", 13);

        // Lock CF has an optimistic lock
        //
        // LOCK CF       | WRITE CF
        // ------------------------------------
        // ts = 13 (opt) | 11: rollback
        //               |  9: start_ts = 7
        //               |  5: rollback
        //               |  3: start_ts = 1

        match check_secondary(b"k1", 13) {
            SecondaryLocksStatus::Locked(_) => {}
            res => panic!("unexpected lock status: {:?}", res),
        }
        must_locked(&mut engine, b"k1", 13);

        // ----------------------------

        must_commit(&mut engine, b"k1", 13, 15);

        // Lock CF has an optimistic lock
        //
        // LOCK CF       | WRITE CF
        // ------------------------------------
        //               | 15: start_ts = 13
        //               | 11: rollback
        //               |  9: start_ts = 7
        //               |  5: rollback
        //               |  3: start_ts = 1

        match check_secondary(b"k1", 14) {
            SecondaryLocksStatus::RolledBack => {}
            res => panic!("unexpected lock status: {:?}", res),
        }
        must_get_rollback_protected(&mut engine, b"k1", 14, true);

        match check_secondary(b"k1", 15) {
            SecondaryLocksStatus::RolledBack => {}
            res => panic!("unexpected lock status: {:?}", res),
        }
        must_get_overlapped_rollback(&mut engine, b"k1", 15, 13, WriteType::Lock, Some(0));
    }
}