assert_storage.rs

// Copyright 2017 TiKV Project Authors. Licensed under Apache-2.0.

use kvproto::kvrpcpb::{Context, LockInfo};

use test_raftstore::{Cluster, ServerCluster, SimulateEngine};
use tikv::storage::kv::{self, RocksEngine};
use tikv::storage::mvcc::{self, MAX_TXN_WRITE_SIZE};
use tikv::storage::txn;
use tikv::storage::{self, Engine, Key, KvPair, Mutation, Value};
use tikv_util::HandyRwLock;

use super::*;

#[derive(Clone)]
pub struct AssertionStorage<E: Engine> {
    pub store: SyncTestStorage<E>,
    pub ctx: Context,
}

impl Default for AssertionStorage<RocksEngine> {
    fn default() -> Self {
        AssertionStorage {
            ctx: Context::new(),
            store: SyncTestStorageBuilder::new().build().unwrap(),
        }
    }
}

impl AssertionStorage<SimulateEngine> {
    pub fn new_raft_storage_with_store_count(
        count: usize,
        key: &str,
    ) -> (Cluster<ServerCluster>, Self) {
        let (cluster, store, ctx) = new_raft_storage_with_store_count(count, key);
        let storage = Self { ctx, store };
        (cluster, storage)
    }

    pub fn update_with_key_byte(&mut self, cluster: &mut Cluster<ServerCluster>, key: &[u8]) {
        // ensure the leader of range which contains current key has been elected
        cluster.must_get(key);
        let region = cluster.get_region(key);
        let leader = cluster.leader_of_region(region.get_id()).unwrap();
        if leader.get_store_id() == self.ctx.get_peer().get_store_id() {
            return;
        }
        let engine = cluster.sim.rl().storages[&leader.get_id()].clone();
        self.ctx.set_region_id(region.get_id());
        self.ctx.set_region_epoch(region.get_region_epoch().clone());
        self.ctx.set_peer(leader.clone());
        self.store = SyncTestStorageBuilder::from_engine(engine).build().unwrap();
    }

    pub fn delete_ok_for_cluster(
        &mut self,
        cluster: &mut Cluster<ServerCluster>,
        key: &[u8],
        start_ts: u64,
        commit_ts: u64,
    ) {
        let mutations = vec![Mutation::Delete(Key::from_raw(key))];
        let commit_keys = vec![Key::from_raw(key)];
        self.two_pc_ok_for_cluster(cluster, mutations, key, commit_keys, start_ts, commit_ts);
    }

    fn get_from_custer(
        &mut self,
        cluster: &mut Cluster<ServerCluster>,
        key: &[u8],
        ts: u64,
    ) -> Option<Value> {
        for _ in 0..3 {
            let res = self.store.get(self.ctx.clone(), &Key::from_raw(key), ts);
            if let Ok(data) = res {
                return data;
            }
            self.expect_not_leader_or_stale_command(res.unwrap_err());
            self.update_with_key_byte(cluster, key);
        }
        panic!("failed with 3 try");
    }

    pub fn get_none_from_cluster(
        &mut self,
        cluster: &mut Cluster<ServerCluster>,
        key: &[u8],
        ts: u64,
    ) {
        assert_eq!(self.get_from_custer(cluster, key, ts), None);
    }

    pub fn put_ok_for_cluster(
        &mut self,
        cluster: &mut Cluster<ServerCluster>,
        key: &[u8],
        value: &[u8],
        start_ts: u64,
        commit_ts: u64,
    ) {
        let mutations = vec![Mutation::Put((Key::from_raw(key), value.to_vec()))];
        let commit_keys = vec![Key::from_raw(key)];
        self.two_pc_ok_for_cluster(cluster, mutations, key, commit_keys, start_ts, commit_ts);
    }

    fn two_pc_ok_for_cluster(
        &mut self,
        cluster: &mut Cluster<ServerCluster>,
        prewrite_mutations: Vec<Mutation>,
        key: &[u8],
        commit_keys: Vec<Key>,
        start_ts: u64,
        commit_ts: u64,
    ) {
        let retry_time = 3;
        let mut success = false;
        for _ in 0..retry_time {
            let res = self.store.prewrite(
                self.ctx.clone(),
                prewrite_mutations.clone(),
                key.to_vec(),
                start_ts,
            );
            if res.is_ok() {
                success = true;
                break;
            }
            self.expect_not_leader_or_stale_command(res.unwrap_err());
            self.update_with_key_byte(cluster, key)
        }
        assert!(success);

        success = false;
        for _ in 0..retry_time {
            let res = self
                .store
                .commit(self.ctx.clone(), commit_keys.clone(), start_ts, commit_ts);
            if res.is_ok() {
                success = true;
                break;
            }
            self.expect_not_leader_or_stale_command(res.unwrap_err());
            self.update_with_key_byte(cluster, key)
        }
        assert!(success);
    }

    pub fn gc_ok_for_cluster(
        &mut self,
        cluster: &mut Cluster<ServerCluster>,
        region_key: &[u8],
        safe_point: u64,
    ) {
        for _ in 0..3 {
            let ret = self.store.gc(self.ctx.clone(), safe_point);
            if ret.is_ok() {
                return;
            }
            self.expect_not_leader_or_stale_command(ret.unwrap_err());
            self.update_with_key_byte(cluster, region_key);
        }
        panic!("failed with 3 retry!");
    }

    pub fn test_txn_store_gc3_for_cluster(
        &mut self,
        cluster: &mut Cluster<ServerCluster>,
        key_prefix: u8,
    ) {
        let key_len = 10_000;
        let key = vec![key_prefix; 1024];
        for k in 1u64..(MAX_TXN_WRITE_SIZE / key_len * 2) as u64 {
            self.put_ok_for_cluster(cluster, &key, b"", k * 10, k * 10 + 5);
        }

        self.delete_ok_for_cluster(cluster, &key, 1000, 1050);
        self.get_none_from_cluster(cluster, &key, 2000);
        self.gc_ok_for_cluster(cluster, &key, 2000);
        self.get_none_from_cluster(cluster, &key, 3000);
    }
}

impl<E: Engine> AssertionStorage<E> {
    pub fn get_none(&self, key: &[u8], ts: u64) {
        let key = Key::from_raw(key);
        assert_eq!(self.store.get(self.ctx.clone(), &key, ts).unwrap(), None);
    }

    pub fn get_err(&self, key: &[u8], ts: u64) {
        let key = Key::from_raw(key);
        assert!(self.store.get(self.ctx.clone(), &key, ts).is_err());
    }

    pub fn get_ok(&self, key: &[u8], ts: u64, expect: &[u8]) {
        let key = Key::from_raw(key);
        assert_eq!(
            self.store.get(self.ctx.clone(), &key, ts).unwrap().unwrap(),
            expect
        );
    }

    pub fn batch_get_ok(&self, keys: &[&[u8]], ts: u64, expect: Vec<&[u8]>) {
        let keys: Vec<Key> = keys.iter().map(|x| Key::from_raw(x)).collect();
        let result: Vec<Vec<u8>> = self
            .store
            .batch_get(self.ctx.clone(), &keys, ts)
            .unwrap()
            .into_iter()
            .map(|x| x.unwrap().1)
            .collect();
        let expect: Vec<Vec<u8>> = expect.into_iter().map(|x| x.to_vec()).collect();
        assert_eq!(result, expect);
    }

    fn expect_not_leader_or_stale_command(&self, err: storage::Error) {
        match err {
            storage::Error::Txn(txn::Error::Mvcc(mvcc::Error::Engine(kv::Error::Request(
                ref e,
            ))))
            | storage::Error::Txn(txn::Error::Engine(kv::Error::Request(ref e)))
            | storage::Error::Engine(kv::Error::Request(ref e)) => {
                assert!(
                    e.has_not_leader() | e.has_stale_command(),
                    "invalid error {:?}",
                    e
                );
            }
            _ => {
                panic!(
                    "expect not leader error or stale command, but got {:?}",
                    err
                );
            }
        }
    }

    fn expect_invalid_tso_err(&self, resp: Result<(), storage::Error>, sts: u64, cmt_ts: u64) {
        assert!(resp.is_err());
        let err = resp.unwrap_err();
        match err {
            storage::Error::Txn(txn::Error::InvalidTxnTso {
                start_ts,
                commit_ts,
            }) => {
                assert_eq!(sts, start_ts);
                assert_eq!(cmt_ts, commit_ts);
            }
            _ => {
                panic!("expect invalid tso error, but got {:?}", err);
            }
        }
    }

    pub fn put_ok(&self, key: &[u8], value: &[u8], start_ts: u64, commit_ts: u64) {
        self.store
            .prewrite(
                self.ctx.clone(),
                vec![Mutation::Put((Key::from_raw(key), value.to_vec()))],
                key.to_vec(),
                start_ts,
            )
            .unwrap();
        self.store
            .commit(
                self.ctx.clone(),
                vec![Key::from_raw(key)],
                start_ts,
                commit_ts,
            )
            .unwrap();
    }

    pub fn delete_ok(&self, key: &[u8], start_ts: u64, commit_ts: u64) {
        self.store
            .prewrite(
                self.ctx.clone(),
                vec![Mutation::Delete(Key::from_raw(key))],
                key.to_vec(),
                start_ts,
            )
            .unwrap();
        self.store
            .commit(
                self.ctx.clone(),
                vec![Key::from_raw(key)],
                start_ts,
                commit_ts,
            )
            .unwrap();
    }

    pub fn scan_ok(
        &self,
        start_key: &[u8],
        limit: usize,
        ts: u64,
        expect: Vec<Option<(&[u8], &[u8])>>,
    ) {
        let key_address = Key::from_raw(start_key);
        let result = self
            .store
            .scan(self.ctx.clone(), key_address, None, limit, false, ts)
            .unwrap();
        let result: Vec<Option<KvPair>> = result.into_iter().map(Result::ok).collect();
        let expect: Vec<Option<KvPair>> = expect
            .into_iter()
            .map(|x| x.map(|(k, v)| (k.to_vec(), v.to_vec())))
            .collect();
        assert_eq!(result, expect);
    }

    pub fn reverse_scan_ok(
        &self,
        start_key: &[u8],
        limit: usize,
        ts: u64,
        expect: Vec<Option<(&[u8], &[u8])>>,
    ) {
        let key_address = Key::from_raw(start_key);
        let result = self
            .store
            .reverse_scan(self.ctx.clone(), key_address, None, limit, false, ts)
            .unwrap();
        let result: Vec<Option<KvPair>> = result.into_iter().map(Result::ok).collect();
        let expect: Vec<Option<KvPair>> = expect
            .into_iter()
            .map(|x| x.map(|(k, v)| (k.to_vec(), v.to_vec())))
            .collect();
        assert_eq!(result, expect);
    }

    pub fn scan_key_only_ok(
        &self,
        start_key: &[u8],
        limit: usize,
        ts: u64,
        expect: Vec<Option<&[u8]>>,
    ) {
        let key_address = Key::from_raw(start_key);
        let result = self
            .store
            .scan(self.ctx.clone(), key_address, None, limit, true, ts)
            .unwrap();
        let result: Vec<Option<KvPair>> = result.into_iter().map(Result::ok).collect();
        let expect: Vec<Option<KvPair>> = expect
            .into_iter()
            .map(|x| x.map(|k| (k.to_vec(), vec![])))
            .collect();
        assert_eq!(result, expect);
    }

    pub fn prewrite_ok(&self, mutations: Vec<Mutation>, primary: &[u8], start_ts: u64) {
        self.store
            .prewrite(self.ctx.clone(), mutations, primary.to_vec(), start_ts)
            .unwrap();
    }

    pub fn prewrite_err(&self, mutations: Vec<Mutation>, primary: &[u8], start_ts: u64) {
        self.store
            .prewrite(self.ctx.clone(), mutations, primary.to_vec(), start_ts)
            .unwrap_err();
    }

    pub fn prewrite_locked(
        &self,
        mutations: Vec<Mutation>,
        primary: &[u8],
        start_ts: u64,
        expect_locks: Vec<(&[u8], &[u8], u64)>,
    ) {
        let res = self
            .store
            .prewrite(self.ctx.clone(), mutations, primary.to_vec(), start_ts)
            .unwrap();
        let locks: Vec<(&[u8], &[u8], u64)> = res
            .iter()
            .filter_map(|x| {
                if let Err(storage::Error::Txn(txn::Error::Mvcc(mvcc::Error::KeyIsLocked {
                    ref key,
                    ref primary,
                    ts,
                    ..
                }))) = *x
                {
                    Some((key.as_ref(), primary.as_ref(), ts))
                } else {
                    None
                }
            })
            .collect();
        assert_eq!(expect_locks, locks);
    }

    pub fn prewrite_conflict(
        &self,
        mutations: Vec<Mutation>,
        cur_primary: &[u8],
        cur_start_ts: u64,
        confl_key: &[u8],
        confl_ts: u64,
    ) {
        let err = self
            .store
            .prewrite(
                self.ctx.clone(),
                mutations,
                cur_primary.to_vec(),
                cur_start_ts,
            )
            .unwrap_err();

        match err {
            storage::Error::Txn(txn::Error::Mvcc(mvcc::Error::WriteConflict {
                start_ts,
                conflict_start_ts,
                ref key,
                ref primary,
                ..
            })) => {
                assert_eq!(cur_start_ts, start_ts);
                assert_eq!(confl_ts, conflict_start_ts);
                assert_eq!(key.to_owned(), confl_key.to_owned());
                assert_eq!(primary.to_owned(), cur_primary.to_owned());
            }
            _ => {
                panic!("expect conflict error, but got {:?}", err);
            }
        }
    }

    pub fn commit_ok(&self, keys: Vec<&[u8]>, start_ts: u64, commit_ts: u64) {
        let keys: Vec<Key> = keys.iter().map(|x| Key::from_raw(x)).collect();
        self.store
            .commit(self.ctx.clone(), keys, start_ts, commit_ts)
            .unwrap();
    }

    pub fn commit_with_illegal_tso(&self, keys: Vec<&[u8]>, start_ts: u64, commit_ts: u64) {
        let keys: Vec<Key> = keys.iter().map(|x| Key::from_raw(x)).collect();
        let resp = self
            .store
            .commit(self.ctx.clone(), keys, start_ts, commit_ts);
        self.expect_invalid_tso_err(resp, start_ts, commit_ts);
    }

    pub fn cleanup_ok(&self, key: &[u8], start_ts: u64) {
        self.store
            .cleanup(self.ctx.clone(), Key::from_raw(key), start_ts)
            .unwrap();
    }

    pub fn cleanup_err(&self, key: &[u8], start_ts: u64) {
        assert!(self
            .store
            .cleanup(self.ctx.clone(), Key::from_raw(key), start_ts)
            .is_err());
    }

    pub fn rollback_ok(&self, keys: Vec<&[u8]>, start_ts: u64) {
        let keys: Vec<Key> = keys.iter().map(|x| Key::from_raw(x)).collect();
        self.store
            .rollback(self.ctx.clone(), keys, start_ts)
            .unwrap();
    }

    pub fn rollback_err(&self, keys: Vec<&[u8]>, start_ts: u64) {
        let keys: Vec<Key> = keys.iter().map(|x| Key::from_raw(x)).collect();
        assert!(self
            .store
            .rollback(self.ctx.clone(), keys, start_ts)
            .is_err());
    }

    pub fn scan_locks_ok(
        &self,
        max_ts: u64,
        start_key: Vec<u8>,
        limit: usize,
        expect: Vec<LockInfo>,
    ) {
        assert_eq!(
            self.store
                .scan_locks(self.ctx.clone(), max_ts, start_key, limit)
                .unwrap(),
            expect
        );
    }

    pub fn resolve_lock_ok(&self, start_ts: u64, commit_ts: Option<u64>) {
        self.store
            .resolve_lock(self.ctx.clone(), start_ts, commit_ts)
            .unwrap();
    }

    pub fn resolve_lock_batch_ok(
        &self,
        start_ts_1: u64,
        commit_ts_1: u64,
        start_ts_2: u64,
        commit_ts_2: u64,
    ) {
        self.store
            .resolve_lock_batch(
                self.ctx.clone(),
                vec![(start_ts_1, commit_ts_1), (start_ts_2, commit_ts_2)],
            )
            .unwrap();
    }

    pub fn resolve_lock_with_illegal_tso(&self, start_ts: u64, commit_ts: Option<u64>) {
        let resp = self
            .store
            .resolve_lock(self.ctx.clone(), start_ts, commit_ts);
        self.expect_invalid_tso_err(resp, start_ts, commit_ts.unwrap())
    }

    pub fn gc_ok(&self, safe_point: u64) {
        self.store.gc(self.ctx.clone(), safe_point).unwrap();
    }

    pub fn raw_get_ok(&self, cf: String, key: Vec<u8>, value: Option<Vec<u8>>) {
        assert_eq!(
            self.store.raw_get(self.ctx.clone(), cf, key).unwrap(),
            value
        );
    }

    pub fn raw_put_ok(&self, cf: String, key: Vec<u8>, value: Vec<u8>) {
        self.store
            .raw_put(self.ctx.clone(), cf, key, value)
            .unwrap();
    }

    pub fn raw_put_err(&self, cf: String, key: Vec<u8>, value: Vec<u8>) {
        self.store
            .raw_put(self.ctx.clone(), cf, key, value)
            .unwrap_err();
    }

    pub fn raw_delete_ok(&self, cf: String, key: Vec<u8>) {
        self.store.raw_delete(self.ctx.clone(), cf, key).unwrap()
    }

    pub fn raw_delete_err(&self, cf: String, key: Vec<u8>) {
        self.store
            .raw_delete(self.ctx.clone(), cf, key)
            .unwrap_err();
    }

    pub fn raw_scan_ok(
        &self,
        cf: String,
        start_key: Vec<u8>,
        limit: usize,
        expect: Vec<(&[u8], &[u8])>,
    ) {
        let result: Vec<KvPair> = self
            .store
            .raw_scan(self.ctx.clone(), cf, start_key, None, limit)
            .unwrap()
            .into_iter()
            .map(|x| x.unwrap())
            .collect();
        let expect: Vec<KvPair> = expect
            .into_iter()
            .map(|(k, v)| (k.to_vec(), v.to_vec()))
            .collect();
        assert_eq!(result, expect);
    }

    pub fn test_txn_store_gc(&self, key: &str) {
        let key_bytes = key.as_bytes();
        self.put_ok(key_bytes, b"v1", 5, 10);
        self.put_ok(key_bytes, b"v2", 15, 20);
        self.gc_ok(30);
        self.get_none(key_bytes, 15);
        self.get_ok(key_bytes, 25, b"v2");
    }

    pub fn test_txn_store_gc3(&self, key_prefix: u8) {
        let key_len = 10_000;
        let key = vec![key_prefix; 1024];
        for k in 1u64..(MAX_TXN_WRITE_SIZE / key_len * 2) as u64 {
            self.put_ok(&key, b"", k * 10, k * 10 + 5);
        }
        self.delete_ok(&key, 1000, 1050);
        self.get_none(&key, 2000);
        self.gc_ok(2000);
        self.get_none(&key, 3000);
    }
}