tikv
/
client-rust
mirror of https://github.com/tikv/client-rust.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

mock tikv: remove unnecessary clones; delete requests won't return errors 

Signed-off-by: ekexium <ekexium@gmail.com>
This commit is contained in:
ekexium 2020-09-14 14:33:06 +08:00
parent 1bf9374939
commit 194acb633c
5 changed files with 29 additions and 61 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
Cargo.lock generated
View File

@ -713,15 +713,6 @@ dependencies = [
"either",
]
[[package]]
name = "itertools"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "284f18f85651fe11e8a991b2adb42cb078325c996ed026d994719efcfca1d54b"
dependencies = [
"either",
]
[[package]]
name = "itoa"
version = "0.4.6"
@ -907,7 +898,6 @@ dependencies = [
"derive-new",
"futures 0.3.5",
"grpcio",
"itertools 0.9.0",
"kvproto",
"log",
"tikv-client-common",
@ -1232,7 +1222,7 @@ checksum = "02b10678c913ecbd69350e8535c3aef91a8676c0773fc1d7b95cdd196d7f2f26"
dependencies = [
"bytes",
"heck",
"itertools 0.8.2",
"itertools",
"log",
"multimap",
"petgraph",
@ -1249,7 +1239,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "537aa19b95acde10a12fec4301466386f757403de4cd4e5b4fa78fb5ecb18f72"
dependencies = [
"anyhow",
"itertools 0.8.2",
"itertools",
"proc-macro2 1.0.19",
"quote 1.0.7",
"syn 1.0.35",

1
mock-tikv/Cargo.toml
View File

@ -10,4 +10,3 @@ kvproto = { git = "https://github.com/pingcap/kvproto.git", rev = "1e28226154c37
derive-new = "0.5.8"
tikv-client-common = { path = "../tikv-client-common"}
log = "0.4"
itertools = "0.9.0"

34
mock-tikv/src/server.rs
View File

@ -4,7 +4,6 @@ use crate::{spawn_unary_success, KvStore};
use derive_new::new;
use futures::{FutureExt, TryFutureExt};
use grpcio::{Environment, Server, ServerBuilder};
use itertools::Itertools;
use kvproto::{kvrpcpb::*, tikvpb::*};
use std::sync::Arc;
@ -189,11 +188,11 @@ impl Tikv for MockTikv {
fn raw_batch_get(
&mut self,
ctx: grpcio::RpcContext,
req: kvproto::kvrpcpb::RawBatchGetRequest,
mut req: kvproto::kvrpcpb::RawBatchGetRequest,
sink: grpcio::UnarySink<kvproto::kvrpcpb::RawBatchGetResponse>,
) {
let mut resp = kvproto::kvrpcpb::RawBatchGetResponse::default();
resp.set_pairs(self.inner.raw_batch_get(req.get_keys()));
resp.set_pairs(self.inner.raw_batch_get(req.take_keys()));
spawn_unary_success!(ctx, req, resp, sink);
}
@ -212,10 +211,10 @@ impl Tikv for MockTikv {
fn raw_batch_put(
&mut self,
ctx: grpcio::RpcContext,
req: kvproto::kvrpcpb::RawBatchPutRequest,
mut req: kvproto::kvrpcpb::RawBatchPutRequest,
sink: grpcio::UnarySink<kvproto::kvrpcpb::RawBatchPutResponse>,
) {
let pairs = req.get_pairs();
let pairs = req.take_pairs();
self.inner.raw_batch_put(pairs);
let resp = RawBatchPutResponse::default();
spawn_unary_success!(ctx, req, resp, sink);
@ -228,33 +227,20 @@ impl Tikv for MockTikv {
sink: grpcio::UnarySink<kvproto::kvrpcpb::RawDeleteResponse>,
) {
let key = req.get_key();
let res = self.inner.raw_delete(key);
let mut resp = RawDeleteResponse::default();
if res.is_none() {
resp.set_error("Key not exist".to_owned());
}
self.inner.raw_delete(key);
let resp = RawDeleteResponse::default();
spawn_unary_success!(ctx, req, resp, sink);
}
fn raw_batch_delete(
&mut self,
ctx: grpcio::RpcContext,
req: kvproto::kvrpcpb::RawBatchDeleteRequest,
mut req: kvproto::kvrpcpb::RawBatchDeleteRequest,
sink: grpcio::UnarySink<kvproto::kvrpcpb::RawBatchDeleteResponse>,
) {
let keys: &[Vec<u8>] = req.get_keys();
let res = self.inner.raw_batch_delete(keys);
let mut resp = RawBatchDeleteResponse::default();
if res.is_err() {
resp.set_error(format!(
"Non-existent keys:[{}]",
res.err()
.unwrap()
.iter()
.map(|k| std::str::from_utf8(k).unwrap())
.join(", ")
));
}
let keys = req.take_keys();
self.inner.raw_batch_delete(keys);
let resp = RawBatchDeleteResponse::default();
spawn_unary_success!(ctx, req, resp, sink);
}

38
mock-tikv/src/store.rs
View File

@ -29,14 +29,14 @@ impl KvStore {
data.get(key).map(|v| v.to_vec())
}
pub fn raw_batch_get(&self, keys: &[Vec<u8>]) -> Vec<KvPair> {
pub fn raw_batch_get(&self, keys: Vec<Vec<u8>>) -> Vec<KvPair> {
let data = self.data.read().unwrap();
keys.iter()
keys.into_iter()
.filter_map(|key| {
if data.contains_key(key) {
if data.contains_key(&key) {
let mut pair = KvPair::default();
pair.set_value(data.get(key).map(|v| v.to_vec()).unwrap_or_else(Vec::new));
pair.set_key(key.to_vec());
pair.set_value(data.get(&key).unwrap().to_vec());
pair.set_key(key);
Some(pair)
} else {
None
@ -50,30 +50,26 @@ impl KvStore {
data.insert(key, value);
}
pub fn raw_batch_put(&self, pairs: &[KvPair]) {
pub fn raw_batch_put(&self, pairs: Vec<KvPair>) {
let mut data = self.data.write().unwrap();
for pair in pairs {
data.insert(pair.get_key().to_vec(), pair.get_value().to_vec());
}
data.extend(
pairs
.into_iter()
.map(|mut pair| (pair.take_key(), pair.take_value())),
);
}
// if success, return the key deleted
pub fn raw_delete(&self, key: &[u8]) -> Option<Vec<u8>> {
pub fn raw_delete(&self, key: &[u8]) {
let mut data = self.data.write().unwrap();
data.remove(key)
data.remove(key);
}
// if any of the key does not exist, return non-existent keys; delete other keys
pub fn raw_batch_delete(&self, keys: &[Vec<u8>]) -> Result<(), Vec<Vec<u8>>> {
pub fn raw_batch_delete(&self, keys: Vec<Vec<u8>>) {
let mut data = self.data.write().unwrap();
let non_exist: Vec<_> = keys
.iter()
.filter_map(|k| data.remove(k).xor(Some(k.to_vec())))
.collect();
if non_exist.is_empty() {
Ok(())
} else {
Err(non_exist)
}
keys.iter().for_each(|k| {
data.remove(k);
});
}
}

3
tests/mock_tikv_tests.rs
View File

@ -54,9 +54,6 @@ mod test {
let res = client.delete("k3".to_owned()).await;
assert!(res.is_ok());
let res = client.delete("key-not-exist".to_owned()).await;
assert!(res.is_err());
let res = client.get("k3".to_owned()).await;
assert_eq!(res.unwrap(), None);