// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

use std::{
    collections::{BTreeMap, Bound},
    fmt,
    sync::atomic::{AtomicU64, Ordering},
};

use kvproto::{
    errorpb::{self, EpochNotMatch, StaleCommand},
    metapb,
};
use lazy_static::lazy_static;
use parking_lot::RwLock;
use prometheus::{IntGauge, register_int_gauge};
use txn_types::{Key, Lock, PessimisticLock};

/// Transaction extensions related to a peer.
#[derive(Default)]
pub struct TxnExt {
    /// The max timestamp recorded in the concurrency manager is only updated at
    /// leader. So if a peer becomes leader from a follower, the max timestamp
    /// can be outdated. We need to update the max timestamp with a latest
    /// timestamp from PD before this peer can work. From the least significant
    /// to the most, 1 bit marks whether the timestamp is updated, 31 bits for
    /// the current epoch version, 32 bits for the current term. The version
    /// and term are stored to prevent stale UpdateMaxTimestamp task from
    /// marking the lowest bit.
    pub max_ts_sync_status: AtomicU64,

    /// The in-memory pessimistic lock table of the peer.
    pub pessimistic_locks: RwLock<PeerPessimisticLocks>,
}

impl TxnExt {
    pub fn is_max_ts_synced(&self) -> bool {
        self.max_ts_sync_status.load(Ordering::SeqCst) & 1 == 1
    }
}

impl fmt::Debug for TxnExt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut debug_struct = f.debug_struct("TxnExt");
        debug_struct.field("max_ts_sync_status", &self.max_ts_sync_status);
        if let Some(pessimistic_locks) = self.pessimistic_locks.try_read() {
            debug_struct.field("pessimistic_locks", &pessimistic_locks);
        } else {
            debug_struct.field("pessimistic_locks", &"(Locked)");
        }
        debug_struct.finish()
    }
}

lazy_static! {
    pub static ref INSTANCE_MEM_SIZE: IntGauge = register_int_gauge!(
        "tikv_pessimistic_lock_memory_size",
        "Total memory size of pessimistic locks in bytes."
    )
    .unwrap();
}

const DEFAULT_PEER_MEM_SIZE_LIMIT: usize = 512 << 10;
const DEFAULT_INSTANCE_MEM_SIZE_LIMIT: usize = 100 << 20;

/// Pessimistic locks of a region peer.
#[derive(PartialEq)]
pub struct PeerPessimisticLocks {
    /// The table that stores pessimistic locks.
    ///
    /// The bool marks an ongoing write request (which has been sent to the
    /// raftstore while not applied yet) will delete this lock. The lock will be
    /// really deleted after applying the write request. The flag will decide
    /// whether this lock should be migrated to other peers on leader or region
    /// changes:
    ///
    /// - **Transfer leader**: The lock with the deleted mark SHOULD NOT be
    ///   proposed before transferring leader. Considering the following cases:
    ///
    ///   1. Propose write -> propose locks -> apply write -> apply locks ->
    ///      transfer leader: The locks marked deleted will not be proposed. The
    ///      lock will be deleted when applying the write while not showing up
    ///      again after applying the locks.
    ///
    ///   2. Propose locks -> propose write -> transfer leader: No lock will be
    ///      lost in normal cases because the write request has been sent to the
    ///      raftstore, it is *likely* to be proposed successfully, while the
    ///      leader will need at least another round to receive the transfer
    ///      leader message from the transferee.
    ///
    /// - **Split region**: The lock with the deleted mark SHOULD be moved to
    ///   new regions on region split. Considering the following cases:
    ///
    ///   1. Propose write -> propose split -> apply write -> execute split: The
    ///      write will be applied earlier than split. So, the lock will be
    ///      deleted earlier than moving locks to new regions.
    ///
    ///   2. Propose split -> propose write -> ready split -> apply write: The
    ///      write will be skipped because its version is lower than the new
    ///      region. So, no lock should be deleted in this case.
    ///
    ///   3. Propose split -> ready split -> propose write: The write proposal
    ///      will be rejected because of version mismatch.
    ///
    /// - **Merge region**: The lock with the deleted mark SHOULD be included in
    ///   the catch up logs on region merge. Considering the following cases:
    ///
    ///   1. Propose write -> propose prepare merge -> apply write -> execute
    ///      merge: The locks marked deleted will be deleted when applying the
    ///      write request. So, the deleted locks will not be included again in
    ///      the commit merge request.
    ///
    ///   2. Propose prepare merge -> propose write -> execute merge -> apply
    ///      write: Applying the write will be skipped because of version
    ///      mismatch. So, no lock should be deleted. It's correct that we
    ///      include the locks that are marked deleted in the commit merge
    ///      request.
    map: BTreeMap<Key, (PessimisticLock, bool)>,
    /// Status of the pessimistic lock map.
    /// The map is writable only in the Normal state.
    pub status: LocksStatus,
    /// Refers to the Raft term in which the pessimistic lock table is valid.
    pub term: u64,
    /// Refers to the region version in which the pessimistic lock table is
    /// valid.
    pub version: u64,
    /// Estimated memory used by the pessimistic locks.
    pub memory_size: usize,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum LocksStatus {
    Normal,
    TransferringLeader,
    MergingRegion,
    NotLeader,
}

impl fmt::Debug for PeerPessimisticLocks {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("PeerPessimisticLocks")
            .field("count", &self.map.len())
            .field("memory_size", &self.memory_size)
            .field("status", &self.status)
            .field("term", &self.term)
            .field("version", &self.version)
            .finish()
    }
}

impl Default for PeerPessimisticLocks {
    fn default() -> Self {
        PeerPessimisticLocks {
            map: BTreeMap::default(),
            status: LocksStatus::Normal,
            term: 0,
            version: 0,
            memory_size: 0,
        }
    }
}

impl PeerPessimisticLocks {
    /// Rejects reads whose snapshot was taken under a different leader term or
    /// region epoch than the current lock table.
    pub fn check_term_version_status(
        &self,
        expected_term: u64,
        expected_version: u64,
    ) -> std::result::Result<(), errorpb::Error> {
        if expected_term != 0 && self.term != expected_term {
            let mut err = errorpb::Error::default();
            err.set_stale_command(StaleCommand::default());
            return Err(err);
        }
        if expected_version != 0 && self.version != expected_version {
            let mut err = errorpb::Error::default();
            err.set_epoch_not_match(EpochNotMatch::default());
            return Err(err);
        }
        Ok(())
    }

    /// Inserts pessimistic locks using default memory limits.
    ///
    /// This is a convenience wrapper for tests only. Production code should use
    /// `insert_with_limits` with actual memory limits from configuration.
    pub fn insert<P: PessimisticLockPair>(&mut self, pairs: Vec<P>) -> Result<(), Vec<P>> {
        self.insert_with_limits(
            pairs,
            DEFAULT_PEER_MEM_SIZE_LIMIT,
            DEFAULT_INSTANCE_MEM_SIZE_LIMIT,
        )
    }

    /// Inserts pessimistic locks into the map with custom memory limits.
    ///
    /// Returns whether the operation succeeds.
    pub fn insert_with_limits<P: PessimisticLockPair>(
        &mut self,
        pairs: Vec<P>,
        peer_mem_size_limit: usize,
        instance_mem_size_limit: usize,
    ) -> Result<(), Vec<P>> {
        let mut incr = 0;
        // Pre-check the memory limit of pessimistic locks.
        for pair in &pairs {
            let (key, lock) = pair.as_pair();
            // If the key already exists in the map, it's an overwrite.
            // The primary lock does not change during an overwrite, so we don't need to
            // update the memory size.
            if !self.map.contains_key(key) {
                incr += key.len() + lock.memory_size();
            }
        }
        if self.memory_size + incr > peer_mem_size_limit
            || INSTANCE_MEM_SIZE.get() as usize + incr > instance_mem_size_limit
        {
            return Err(pairs);
        }
        // Insert after check has passed.
        for pair in pairs {
            let (key, lock) = pair.into_pair();
            self.map.insert(key, (lock, false));
        }
        self.memory_size += incr;
        INSTANCE_MEM_SIZE.add(incr as i64);
        Ok(())
    }

    pub fn remove(&mut self, key: &Key) {
        if let Some((lock, _)) = self.map.remove(key) {
            let desc = key.len() + lock.memory_size();
            self.memory_size -= desc;
            INSTANCE_MEM_SIZE.sub(desc as i64);
        }
    }

    pub fn clear(&mut self) {
        self.map = BTreeMap::default();
        INSTANCE_MEM_SIZE.sub(self.memory_size as i64);
        self.memory_size = 0;
    }

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

    pub fn len(&self) -> usize {
        self.map.len()
    }

    pub fn is_writable(&self) -> bool {
        self.status == LocksStatus::Normal
    }

    pub fn get(&self, key: &Key) -> Option<&(PessimisticLock, bool)> {
        self.map.get(key)
    }

    pub fn get_mut(&mut self, key: &Key) -> Option<&mut (PessimisticLock, bool)> {
        self.map.get_mut(key)
    }

    /// Scans in-memory pessimistic locks in the given key range.
    ///
    /// Returns locks that satisfy the filter function. The `limit` parameter
    /// specifies the maximum number of locks to return (0 means unlimited).
    ///
    /// Returns `(locks, has_remain)` where `has_remain` indicates whether there
    /// MAY be remaining locks that can be scanned.
    pub fn scan_locks<F>(
        &self,
        start: Option<&Key>,
        end: Option<&Key>,
        filter: F,
        limit: usize,
    ) -> (Vec<(Key, Lock)>, bool)
    where
        F: Fn(&Key, &PessimisticLock) -> bool,
    {
        let mut locks = Vec::new();
        let mut iter = self.map.range((
            start.map_or(Bound::Unbounded, Bound::Included),
            end.map_or(Bound::Unbounded, Bound::Excluded),
        ));
        while let Some((key, (lock, _))) = iter.next() {
            if filter(key, lock) {
                locks.push((key.clone(), lock.to_lock()));
            }
            if limit > 0 && locks.len() >= limit {
                return (locks, iter.next().is_some());
            }
        }
        (locks, false)
    }

    /// Group pessimistic locks in the original region to the split regions.
    ///
    /// The given regions MUST be sorted by key in the ascending order. The
    /// returned `HashMap`s are in the same order of the given regions.
    ///
    /// The locks belonging to the derived region will be kept in the given
    /// `locks` map, and the corresponding position in the returned `Vec`
    /// will be an empty map.
    pub fn group_by_regions(
        &mut self,
        regions: &[metapb::Region],
        derived: &metapb::Region,
    ) -> Vec<PeerPessimisticLocks> {
        // Assert regions are sorted by key in ascending order.
        if cfg!(debug_assertions) {
            for (r1, r2) in regions.iter().zip(regions.iter().skip(1)) {
                assert!(r1.get_start_key() < r2.get_start_key());
            }
        }

        let mut res: Vec<PeerPessimisticLocks> = regions
            .iter()
            .map(|_| PeerPessimisticLocks::default())
            .collect();

        // Locks that are marked deleted still need to be moved to the new regions,
        // and the deleted mark should also be cleared.
        // Refer to the comment in `PeerPessimisticLocks` for details.
        // There is no drain_filter for BTreeMap, so extra clone is needed.
        let mut removed_locks = Vec::new();
        let (start_key, end_key) = (derived.get_start_key(), derived.get_end_key());
        self.map.retain(|key, value| {
            let key_ref = key.as_encoded().as_slice();
            if key_ref < start_key || (!end_key.is_empty() && key_ref >= end_key) {
                removed_locks.push((key.clone(), value.clone()));
                false
            } else {
                true
            }
        });

        for (key, (lock, _)) in removed_locks {
            let idx = regions
                .binary_search_by_key(&&**key.as_encoded(), |region| region.get_start_key())
                .unwrap_or_else(|idx| idx.saturating_sub(1));
            let size = key.len() + lock.memory_size();
            self.memory_size -= size;
            res[idx].map.insert(key, (lock, false));
            res[idx].memory_size += size;
        }
        res
    }

    #[cfg(test)]
    fn from_locks(locks: impl IntoIterator<Item = (Key, (PessimisticLock, bool))>) -> Self {
        let mut res = PeerPessimisticLocks::default();
        for (key, (locks, is_deleted)) in locks {
            res.memory_size += key.len() + locks.memory_size();
            res.map.insert(key, (locks, is_deleted));
        }
        res
    }
}

impl<'a> IntoIterator for &'a PeerPessimisticLocks {
    type Item = (&'a Key, &'a (PessimisticLock, bool));
    type IntoIter = std::collections::btree_map::Iter<'a, Key, (PessimisticLock, bool)>;

    fn into_iter(self) -> Self::IntoIter {
        self.map.iter()
    }
}

impl<'a> IntoIterator for &'a mut PeerPessimisticLocks {
    type Item = (&'a Key, &'a mut (PessimisticLock, bool));
    type IntoIter = std::collections::btree_map::IterMut<'a, Key, (PessimisticLock, bool)>;

    fn into_iter(self) -> Self::IntoIter {
        self.map.iter_mut()
    }
}

impl Drop for PeerPessimisticLocks {
    fn drop(&mut self) {
        INSTANCE_MEM_SIZE.sub(self.memory_size as i64);
    }
}

pub trait PessimisticLockPair {
    fn as_pair(&self) -> (&Key, &PessimisticLock);

    fn into_pair(self) -> (Key, PessimisticLock);
}

impl PessimisticLockPair for (Key, PessimisticLock) {
    fn as_pair(&self) -> (&Key, &PessimisticLock) {
        (&self.0, &self.1)
    }

    fn into_pair(self) -> (Key, PessimisticLock) {
        self
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Mutex;

    use tikv_util::defer;

    use super::*;

    lazy_static! {
        static ref TEST_MUTEX: Mutex<()> = Mutex::new(());
    }

    fn lock(primary: &[u8]) -> PessimisticLock {
        PessimisticLock {
            primary: primary.to_vec().into_boxed_slice(),
            start_ts: 100.into(),
            ttl: 3000,
            for_update_ts: 110.into(),
            min_commit_ts: 110.into(),
            last_change_ts: 105.into(),
            versions_to_last_change: 2,
        }
    }

    #[test]
    fn test_memory_size() {
        let _guard = TEST_MUTEX.lock().unwrap();

        let mut locks1 = PeerPessimisticLocks::default();
        let mut locks2 = PeerPessimisticLocks::default();
        let k1 = Key::from_raw(b"k1");
        let k2 = Key::from_raw(b"k22");
        let k3 = Key::from_raw(b"k333");
        let peer_mem_size_limit = 512 << 10;
        let instance_mem_size_limit = 100 << 20;

        // Test the memory size of peer pessimistic locks after inserting.
        locks1
            .insert_with_limits(
                vec![(k1.clone(), lock(b"k1"))],
                peer_mem_size_limit,
                instance_mem_size_limit,
            )
            .unwrap();
        assert_eq!(locks1.get(&k1), Some(&(lock(b"k1"), false)));
        assert_eq!(locks1.memory_size, k1.len() + lock(b"k1").memory_size());
        locks1
            .insert_with_limits(
                vec![(k2.clone(), lock(b"k1"))],
                peer_mem_size_limit,
                instance_mem_size_limit,
            )
            .unwrap();
        assert_eq!(locks1.get(&k2), Some(&(lock(b"k1"), false)));
        assert_eq!(
            locks1.memory_size,
            k1.len() + k2.len() + 2 * lock(b"k1").memory_size()
        );

        // Test the global memory size after inserting.
        locks2
            .insert_with_limits(
                vec![(k3.clone(), lock(b"k1"))],
                peer_mem_size_limit,
                instance_mem_size_limit,
            )
            .unwrap();
        assert_eq!(locks2.get(&k3), Some(&(lock(b"k1"), false)));
        assert_eq!(
            INSTANCE_MEM_SIZE.get() as usize,
            locks1.memory_size + locks2.memory_size
        );

        // Test the memory size after replacing, it should not change.
        locks1
            .insert_with_limits(
                vec![(k2.clone(), lock(b"k2"))],
                peer_mem_size_limit,
                instance_mem_size_limit,
            )
            .unwrap();
        assert_eq!(locks1.get(&k2), Some(&(lock(b"k2"), false)));
        assert_eq!(
            locks1.memory_size,
            k1.len() + k2.len() + 2 * lock(b"k1").memory_size()
        );
        assert_eq!(
            INSTANCE_MEM_SIZE.get() as usize,
            locks1.memory_size + locks2.memory_size
        );

        // Test the memory size after removing.
        locks1.remove(&k1);
        assert!(locks1.get(&k1).is_none());
        assert_eq!(locks1.memory_size, k2.len() + lock(b"k2").memory_size());
        assert_eq!(
            INSTANCE_MEM_SIZE.get() as usize,
            locks1.memory_size + locks2.memory_size
        );

        // Test the memory size after clearing.
        locks2.clear();
        assert!(locks2.is_empty());
        assert_eq!(locks2.memory_size, 0);
        assert_eq!(INSTANCE_MEM_SIZE.get() as usize, locks1.memory_size);

        // Test the global memory size after dropping.
        drop(locks1);
        drop(locks2);
        assert_eq!(INSTANCE_MEM_SIZE.get(), 0);
    }

    #[test]
    fn test_insert_checking_memory_limit() {
        let _guard = TEST_MUTEX.lock().unwrap();
        defer!(INSTANCE_MEM_SIZE.set(0));
        let peer_mem_size_limit = 512 << 10;
        let instance_mem_size_limit = 100 << 20;

        let mut locks = PeerPessimisticLocks::default();
        locks
            .insert_with_limits(
                vec![(Key::from_raw(b"k1"), lock(&[0; 512000]))],
                peer_mem_size_limit,
                instance_mem_size_limit,
            )
            .unwrap();

        // Exceeding the region limit
        locks
            .insert_with_limits(
                vec![(Key::from_raw(b"k2"), lock(&[0; 32000]))],
                peer_mem_size_limit,
                instance_mem_size_limit,
            )
            .unwrap_err();
        assert!(locks.get(&Key::from_raw(b"k2")).is_none());

        // Not exceeding the region limit, but exceeding the global limit
        INSTANCE_MEM_SIZE.set((instance_mem_size_limit - 10) as i64);
        let res = locks.insert_with_limits(
            vec![(Key::from_raw(b"k2"), lock(b"abc"))],
            peer_mem_size_limit,
            instance_mem_size_limit,
        );
        res.unwrap_err();
        assert!(locks.get(&Key::from_raw(b"k2")).is_none());
    }

    #[test]
    fn test_group_locks_by_regions() {
        fn lock(key: &[u8], deleted: bool) -> (Key, (PessimisticLock, bool)) {
            (
                Key::from_raw(key),
                (
                    PessimisticLock {
                        primary: key.to_vec().into_boxed_slice(),
                        start_ts: 10.into(),
                        ttl: 1000,
                        for_update_ts: 10.into(),
                        min_commit_ts: 20.into(),
                        last_change_ts: 5.into(),
                        versions_to_last_change: 2,
                    },
                    deleted,
                ),
            )
        }
        fn region(start_key: &[u8], end_key: &[u8]) -> metapb::Region {
            let mut region = metapb::Region::default();
            region.set_start_key(start_key.to_vec());
            region.set_end_key(end_key.to_vec());
            region
        }
        let _guard = TEST_MUTEX.lock().unwrap();
        defer!(INSTANCE_MEM_SIZE.set(0));

        let mut original = PeerPessimisticLocks::from_locks(vec![
            lock(b"a", true),
            lock(b"c", false),
            lock(b"e", true),
            lock(b"g", false),
            lock(b"i", false),
        ]);
        let regions = vec![
            region(b"", b"b"),  // test leftmost region
            region(b"b", b"c"), // no lock inside
            region(b"c", b"d"), // test key equals to start_key
            region(b"d", b"h"), // test multiple locks inside
            region(b"h", b""),  // test rightmost region
        ];
        let output = original.group_by_regions(&regions, &regions[4]);
        let expected: Vec<_> = vec![
            vec![lock(b"a", false)],
            vec![],
            vec![lock(b"c", false)],
            vec![lock(b"e", false), lock(b"g", false)],
            vec![], // the position of the derived region is empty
        ]
        .into_iter()
        .map(PeerPessimisticLocks::from_locks)
        .collect();
        assert_eq!(output, expected);
        // The lock that belongs to the derived region is kept in the original map.
        assert_eq!(
            original,
            PeerPessimisticLocks::from_locks(vec![lock(b"i", false)])
        );
    }
}