// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.
// #[PerformanceCriticalPath]
use tikv_util::Either;
use txn_types::{Key, ReqType, TimeStamp};

use crate::storage::{
    ProcessResult, Snapshot, TxnStatus,
    kv::WriteData,
    lock_manager::LockManager,
    mvcc::{Error as MvccError, ErrorInner as MvccErrorInner, MvccTxn, SnapshotReader},
    txn::{
        Result,
        commands::{
            Command, CommandExt, ReaderWithStats, ReleasedLocks, ResponsePolicy, TypedCommand,
            WriteCommand, WriteContext, WriteResult,
        },
    },
};

command! {
    /// Heart beat of a transaction. It enlarges the primary lock's TTL.
    ///
    /// This is invoked on a transaction's primary lock. The lock may be generated by either
    /// [`AcquirePessimisticLock`](Command::AcquirePessimisticLock) or
    /// [`Prewrite`](Command::Prewrite).
    TxnHeartBeat:
        cmd_ty => TxnStatus,
        display => "kv::command::txn_heart_beat {} @ {} ttl {} | {:?}", (primary_key, start_ts, advise_ttl, ctx),
        content => {
            /// The primary key of the transaction.
            primary_key: Key,
            /// The transaction's start_ts.
            start_ts: TimeStamp,
            /// The new TTL that will be used to update the lock's TTL. If the lock's TTL is already
            /// greater than `advise_ttl`, nothing will happen.
            advise_ttl: u64,
            /// Used by file based transaction.
            is_txn_file: bool,
        }
}

impl CommandExt for TxnHeartBeat {
    ctx!();
    tag!(txn_heart_beat);
    request_type!(KvTxnHeartBeat);
    ts!(start_ts);
    write_bytes!(primary_key);
    gen_lock!(primary_key);
    can_build_txn_file!();
}

#[maybe_async::async_trait]
impl<S: Snapshot + 'static, L: LockManager> WriteCommand<S, L> for TxnHeartBeat {
    #[maybe_async]
    async fn process_write(self, snapshot: S, context: WriteContext<'_, L>) -> Result<WriteResult> {
        // TxnHeartBeat never remove locks. No need to wake up waiters.
        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,
        );
        fail_point!("txn_heart_beat", |err| Err(
            crate::storage::mvcc::Error::from(crate::storage::mvcc::txn::make_txn_error(
                err,
                &self.primary_key,
                self.start_ts,
            ))
            .into()
        ));

        let lock = match reader.load_lock(&self.primary_key)? {
            Some(Either::Left(mut lock)) if lock.ts == self.start_ts => {
                if !lock.is_txn_file {
                    if lock.ttl < self.advise_ttl {
                        lock.ttl = self.advise_ttl;
                        txn.put_lock(self.primary_key.clone(), &lock);
                    }
                } else {
                    // We do not try to update the txn file's ttl from this late received
                    // pessimistic lock txn heartbeat for simplicity. The txn
                    // file's ttl should be updated in the next round.
                    warn!("ignore late lock txn heartbeat after txn file is prewrited"; "lock_info" => ?lock);
                }

                lock
            }
            Some(Either::Right(shared_locks)) if shared_locks.contains_start_ts(self.start_ts) => {
                return Err(MvccError::from(MvccErrorInner::PrimaryMismatch(
                    shared_locks.into_lock_info(self.primary_key.into_raw()?),
                ))
                .into());
            }
            _ => {
                return Err(MvccError::from(MvccErrorInner::TxnNotFound {
                    start_ts: self.start_ts,
                    key: self.primary_key.into_raw()?,
                })
                .into());
            }
        };

        let pr = ProcessResult::TxnStatus {
            txn_status: TxnStatus::uncommitted(lock, false),
        };

        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::Heartbeat);
        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: 1,
            pr,
            lock_info: vec![],
            released_locks: ReleasedLocks::new(),
            new_acquired_locks: vec![],
            shared_lock_upgrade_hints,
            lock_guards: vec![],
            response_policy: ResponsePolicy::OnApplied,
            known_txn_status: vec![],
        })
    }
}

#[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>) {
        let mut shared_locks = txn_types::SharedLocks::new();
        let start_ts = start_ts.into();
        shared_locks
            .insert_lock(txn_types::Lock::new(
                txn_types::LockType::Lock,
                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 txn_heart_beat_err<E: Engine>(
        engine: &mut E,
        primary_key: &[u8],
        start_ts: impl Into<TimeStamp>,
        advise_ttl: u64,
    ) -> crate::storage::txn::Error {
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let start_ts = start_ts.into();
        let cm = ConcurrencyManager::new(start_ts);
        let command = crate::storage::txn::commands::TxnHeartBeat {
            ctx: Context::default(),
            primary_key: Key::from_raw(primary_key),
            start_ts,
            advise_ttl,
            is_txn_file: false,
            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 txn heartbeat to fail"),
            Err(err) => err,
        }
    }

    pub fn must_success<E: Engine>(
        engine: &mut E,
        primary_key: &[u8],
        start_ts: impl Into<TimeStamp>,
        advise_ttl: u64,
        expect_ttl: u64,
    ) {
        let ctx = Context::default();
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let start_ts = start_ts.into();
        let cm = ConcurrencyManager::new(start_ts);
        let command = crate::storage::txn::commands::TxnHeartBeat {
            ctx: Context::default(),
            primary_key: Key::from_raw(primary_key),
            start_ts,
            advise_ttl,
            is_txn_file: false,
            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::TxnStatus {
            txn_status: TxnStatus::Uncommitted { lock, .. },
        } = result.pr
        {
            let shared_lock_lifecycle = result.to_be_write.extra.shared_lock_lifecycle;
            let mut data = WriteData::from_modifies(result.to_be_write.modifies);
            data.extra.req_type = ReqType::Heartbeat;
            data.set_shared_lock_lifecycle(shared_lock_lifecycle);
            write(engine, &ctx, data);
            assert_eq!(lock.ttl, expect_ttl);
        } else {
            unreachable!();
        }
    }

    pub fn must_err<E: Engine>(
        engine: &mut E,
        primary_key: &[u8],
        start_ts: impl Into<TimeStamp>,
        advise_ttl: u64,
    ) {
        let ctx = Context::default();
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let start_ts = start_ts.into();
        let cm = ConcurrencyManager::new(start_ts);
        let command = crate::storage::txn::commands::TxnHeartBeat {
            ctx,
            primary_key: Key::from_raw(primary_key),
            start_ts,
            advise_ttl,
            is_txn_file: false,
            deadline: Deadline::from_now(DEFAULT_EXECUTION_DURATION_LIMIT),
        };
        assert!(
            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(),
                    },
                )
                .is_err()
        );
    }

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

        let err = txn_heart_beat_err(&mut engine, key, 5, 5000);
        match err {
            crate::storage::txn::Error(box crate::storage::txn::ErrorInner::Mvcc(
                crate::storage::mvcc::Error(box crate::storage::mvcc::ErrorInner::PrimaryMismatch(
                    lock_info,
                )),
            )) => {
                assert_eq!(lock_info.get_lock_type(), kvproto::kvrpcpb::Op::SharedLock);
                assert_eq!(lock_info.get_shared_lock_infos().len(), 1);
            }
            err => panic!("unexpected error: {:?}", err),
        }
    }

    #[test]
    fn test_txn_heart_beat() {
        let mut engine = TestEngineBuilder::new().build().unwrap();

        let (k, v) = (b"k1", b"v1");

        fn test(ts: u64, k: &[u8], engine: &mut impl Engine) {
            // Do nothing if advise_ttl is less smaller than current TTL.
            must_success(engine, k, ts, 90, 100);
            // Return the new TTL if the TTL when the TTL is updated.
            must_success(engine, k, ts, 110, 110);
            // The lock's TTL is updated and persisted into the db.
            must_success(engine, k, ts, 90, 110);
            // Heart beat another transaction's lock will lead to an error.
            must_err(engine, k, ts - 1, 150);
            must_err(engine, k, ts + 1, 150);
            // The existing lock is not changed.
            must_success(engine, k, ts, 90, 110);
        }

        // No lock.
        must_err(&mut engine, k, 5, 100);

        // Create a lock with TTL=100.
        // The initial TTL will be set to 0 after calling must_prewrite_put. Update it
        // first.
        must_prewrite_put(&mut engine, k, v, k, 5);
        must_locked(&mut engine, k, 5);
        must_success(&mut engine, k, 5, 100, 100);

        test(5, k, &mut engine);

        must_locked(&mut engine, k, 5);
        must_commit(&mut engine, k, 5, 10);
        must_unlocked(&mut engine, k);

        // No lock.
        must_err(&mut engine, k, 5, 100);
        must_err(&mut engine, k, 10, 100);

        must_acquire_pessimistic_lock(&mut engine, k, k, 8, 15);
        must_pessimistic_locked(&mut engine, k, 8, 15);
        must_success(&mut engine, k, 8, 100, 100);

        test(8, k, &mut engine);

        must_pessimistic_locked(&mut engine, k, 8, 15);
    }
}