Skip to content

Commit 605aed8

Browse files
iosmanthusiosmanthus
committed
fix TxnApiV1 test
Signed-off-by: iosmanthus <[email protected]>
1 parent 222908e commit 605aed8

14 files changed

+190
-139
lines changed

‎examples/transaction.rs

+13-7
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,15 @@
33
mod common;
44

55
use crate::common::parse_args;
6-
use tikv_client::{BoundRange, Config, Key, KvPair, TransactionClient as Client, Value};
6+
use tikv_client::{
7+
request::request_codec::{RequestCodec, TxnApiV1},
8+
BoundRange, Config, Key, KvPair, TransactionClient as Client, Value,
9+
};
710

8-
async fn puts(client: &Client, pairs: impl IntoIterator<Item = impl Into<KvPair>>) {
11+
async fn puts<C: RequestCodec>(
12+
client: &Client<C>,
13+
pairs: impl IntoIterator<Item = impl Into<KvPair>>,
14+
) {
915
let mut txn = client
1016
.begin_optimistic()
1117
.await
@@ -17,7 +23,7 @@ async fn puts(client: &Client, pairs: impl IntoIterator<Item = impl Into<KvPair>
1723
txn.commit().await.expect("Could not commit transaction");
1824
}
1925

20-
async fn get(client: &Client, key: Key) -> Option<Value> {
26+
async fn get<C: RequestCodec>(client: &Client<C>, key: Key) -> Option<Value> {
2127
let mut txn = client
2228
.begin_optimistic()
2329
.await
@@ -29,7 +35,7 @@ async fn get(client: &Client, key: Key) -> Option<Value> {
2935
res
3036
}
3137

32-
async fn key_exists(client: &Client, key: Key) -> bool {
38+
async fn key_exists<C: RequestCodec>(client: &Client<C>, key: Key) -> bool {
3339
let mut txn = client
3440
.begin_optimistic()
3541
.await
@@ -44,7 +50,7 @@ async fn key_exists(client: &Client, key: Key) -> bool {
4450
res
4551
}
4652

47-
async fn scan(client: &Client, range: impl Into<BoundRange>, limit: u32) {
53+
async fn scan<C: RequestCodec>(client: &Client<C>, range: impl Into<BoundRange>, limit: u32) {
4854
let mut txn = client
4955
.begin_optimistic()
5056
.await
@@ -56,7 +62,7 @@ async fn scan(client: &Client, range: impl Into<BoundRange>, limit: u32) {
5662
txn.commit().await.expect("Could not commit transaction");
5763
}
5864

59-
async fn dels(client: &Client, keys: impl IntoIterator<Item = Key>) {
65+
async fn dels<C: RequestCodec>(client: &Client<C>, keys: impl IntoIterator<Item = Key>) {
6066
let mut txn = client
6167
.begin_optimistic()
6268
.await
@@ -81,7 +87,7 @@ async fn main() {
8187
Config::default()
8288
};
8389

84-
let txn = Client::new_with_config(args.pd, config, None)
90+
let txn = Client::new_with_config(args.pd, config, TxnApiV1, None)
8591
.await
8692
.expect("Could not connect to tikv");
8793

‎src/lib.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,7 @@
9090
//! # })}
9191
//! ```
9292
93-
#![feature(specialization)]
93+
#![feature(min_specialization)]
9494
#![feature(explicit_generic_args_with_impl_trait)]
9595
#[macro_use]
9696
pub mod request;

‎src/pd/client.rs

+6-26
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
// Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.
22

3-
use std::{collections::HashMap, marker::PhantomData, sync::Arc, thread};
3+
use std::{collections::HashMap, sync::Arc, thread};
44

55
use async_trait::async_trait;
66
use futures::{prelude::*, stream::BoxStream};
@@ -14,7 +14,6 @@ use tikv_client_store::{KvClient, KvConnect, TikvConnect};
1414

1515
use crate::{
1616
compat::stream_fn,
17-
kv::codec,
1817
pd::{retry::RetryClientTrait, RetryClient},
1918
region::{RegionId, RegionVerId, RegionWithLeader},
2019
region_cache::RegionCache,
@@ -198,14 +197,6 @@ pub trait PdClient: Send + Sync + 'static {
198197
.boxed()
199198
}
200199

201-
fn decode_region(mut region: RegionWithLeader, enable_codec: bool) -> Result<RegionWithLeader> {
202-
if enable_codec {
203-
codec::decode_bytes_in_place(region.region.mut_start_key(), false)?;
204-
codec::decode_bytes_in_place(region.region.mut_end_key(), false)?;
205-
}
206-
Ok(region)
207-
}
208-
209200
async fn update_leader(&self, ver_id: RegionVerId, leader: metapb::Peer) -> Result<()>;
210201

211202
async fn invalidate_region_cache(&self, ver_id: RegionVerId);
@@ -219,9 +210,8 @@ pub struct PdRpcClient<C, KvC: KvConnect + Send + Sync + 'static = TikvConnect,
219210
pd: Arc<RetryClient<Cl>>,
220211
kv_connect: KvC,
221212
kv_client_cache: Arc<RwLock<HashMap<String, KvC::KvClient>>>,
222-
region_cache: RegionCache<RetryClient<Cl>>,
213+
region_cache: RegionCache<C, RetryClient<Cl>>,
223214
logger: Logger,
224-
codec: C,
225215
}
226216

227217
#[async_trait]
@@ -237,20 +227,11 @@ impl<C: RequestCodec, KvC: KvConnect + Send + Sync + 'static> PdClient for PdRpc
237227
}
238228

239229
async fn region_for_key(&self, key: &Key) -> Result<RegionWithLeader> {
240-
let enable_codec = self.enable_codec;
241-
let key = if enable_codec {
242-
key.to_encoded()
243-
} else {
244-
key.clone()
245-
};
246-
247-
let region = self.region_cache.get_region_by_key(&key).await?;
248-
Self::decode_region(region, enable_codec)
230+
self.region_cache.get_region_by_key(&key).await
249231
}
250232

251233
async fn region_for_id(&self, id: RegionId) -> Result<RegionWithLeader> {
252-
let region = self.region_cache.get_region_by_id(id).await?;
253-
Self::decode_region(region, self.enable_codec)
234+
self.region_cache.get_region_by_id(id).await
254235
}
255236

256237
async fn get_timestamp(self: Arc<Self>) -> Result<Timestamp> {
@@ -270,7 +251,7 @@ impl<C: RequestCodec, KvC: KvConnect + Send + Sync + 'static> PdClient for PdRpc
270251
}
271252

272253
fn get_request_codec(&self) -> Self::RequestCodec {
273-
todo!()
254+
self.region_cache.get_request_codec()
274255
}
275256
}
276257

@@ -338,9 +319,8 @@ impl<C, KvC: KvConnect + Send + Sync + 'static, Cl> PdRpcClient<C, KvC, Cl> {
338319
pd: pd.clone(),
339320
kv_client_cache,
340321
kv_connect: kv_connect(env, security_mgr),
341-
region_cache: RegionCache::new(pd),
322+
region_cache: RegionCache::new(codec, pd),
342323
logger,
343-
codec
344324
})
345325
}
346326

‎src/raw/client.rs

+11-11
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,12 @@ use tikv_client_common::Error;
99
use tikv_client_proto::metapb;
1010

1111
use crate::{
12-
Backoff,
1312
backoff::DEFAULT_REGION_BACKOFF,
14-
BoundRange,
15-
ColumnFamily,
1613
config::Config,
17-
Key, KvPair, pd::{PdClient, PdRpcClient}, raw::lowering::*, request::{Collect, CollectSingle, Plan, request_codec::RequestCodec}, Result, Value,
14+
pd::{PdClient, PdRpcClient},
15+
raw::lowering::*,
16+
request::{request_codec::RequestCodec, Collect, CollectSingle, Plan},
17+
Backoff, BoundRange, ColumnFamily, Key, KvPair, Result, Value,
1818
};
1919

2020
const MAX_RAW_KV_SCAN_LIMIT: u32 = 10240;
@@ -215,7 +215,7 @@ impl<C: RequestCodec, PdC: PdClient> Client<C, PdC> {
215215
/// ```
216216
pub async fn batch_get(
217217
&self,
218-
keys: impl IntoIterator<Item=impl Into<Key>>,
218+
keys: impl IntoIterator<Item = impl Into<Key>>,
219219
) -> Result<Vec<KvPair>> {
220220
debug!(self.logger, "invoking raw batch_get request");
221221
let request =
@@ -277,7 +277,7 @@ impl<C: RequestCodec, PdC: PdClient> Client<C, PdC> {
277277
/// ```
278278
pub async fn batch_put(
279279
&self,
280-
pairs: impl IntoIterator<Item=impl Into<KvPair>>,
280+
pairs: impl IntoIterator<Item = impl Into<KvPair>>,
281281
) -> Result<()> {
282282
debug!(self.logger, "invoking raw batch_put request");
283283
let request = new_raw_batch_put_request::<C>(
@@ -339,7 +339,7 @@ impl<C: RequestCodec, PdC: PdClient> Client<C, PdC> {
339339
/// let result: () = req.await.unwrap();
340340
/// # });
341341
/// ```
342-
pub async fn batch_delete(&self, keys: impl IntoIterator<Item=impl Into<Key>>) -> Result<()> {
342+
pub async fn batch_delete(&self, keys: impl IntoIterator<Item = impl Into<Key>>) -> Result<()> {
343343
debug!(self.logger, "invoking raw batch_delete request");
344344
self.assert_non_atomic()?;
345345
let request =
@@ -465,7 +465,7 @@ impl<C: RequestCodec, PdC: PdClient> Client<C, PdC> {
465465
/// ```
466466
pub async fn batch_scan(
467467
&self,
468-
ranges: impl IntoIterator<Item=impl Into<BoundRange>>,
468+
ranges: impl IntoIterator<Item = impl Into<BoundRange>>,
469469
each_limit: u32,
470470
) -> Result<Vec<KvPair>> {
471471
debug!(self.logger, "invoking raw batch_scan request");
@@ -497,7 +497,7 @@ impl<C: RequestCodec, PdC: PdClient> Client<C, PdC> {
497497
/// ```
498498
pub async fn batch_scan_keys(
499499
&self,
500-
ranges: impl IntoIterator<Item=impl Into<BoundRange>>,
500+
ranges: impl IntoIterator<Item = impl Into<BoundRange>>,
501501
each_limit: u32,
502502
) -> Result<Vec<Key>> {
503503
debug!(self.logger, "invoking raw batch_scan_keys request");
@@ -547,7 +547,7 @@ impl<C: RequestCodec, PdC: PdClient> Client<C, PdC> {
547547
&self,
548548
copr_name: impl Into<String>,
549549
copr_version_req: impl Into<String>,
550-
ranges: impl IntoIterator<Item=impl Into<BoundRange>>,
550+
ranges: impl IntoIterator<Item = impl Into<BoundRange>>,
551551
request_builder: impl Fn(metapb::Region, Vec<Range<Key>>) -> Vec<u8> + Send + Sync + 'static,
552552
) -> Result<Vec<(Vec<u8>, Vec<Range<Key>>)>> {
553553
let copr_version_req = copr_version_req.into();
@@ -593,7 +593,7 @@ impl<C: RequestCodec, PdC: PdClient> Client<C, PdC> {
593593

594594
async fn batch_scan_inner(
595595
&self,
596-
ranges: impl IntoIterator<Item=impl Into<BoundRange>>,
596+
ranges: impl IntoIterator<Item = impl Into<BoundRange>>,
597597
each_limit: u32,
598598
key_only: bool,
599599
) -> Result<Vec<KvPair>> {

‎src/raw/lowering.rs

+3-4
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,13 @@
44
//! types (i.e., the types from the client crate) and converts these to the types used in the
55
//! generated protobuf code, then calls the low-level ctor functions in the requests module.
66
7-
use std::{iter::Iterator, marker::PhantomData, ops::Range, sync::Arc};
7+
use std::{iter::Iterator, ops::Range, sync::Arc};
88

99
use tikv_client_proto::{kvrpcpb, metapb};
1010

1111
use crate::{
12-
raw::requests,
13-
request::{request_codec::RequestCodec, KvRequest},
14-
BoundRange, ColumnFamily, Key, KvPair, Value,
12+
raw::requests, request::request_codec::RequestCodec, BoundRange, ColumnFamily, Key, KvPair,
13+
Value,
1514
};
1615

1716
pub fn new_raw_get_request<C: RequestCodec>(

‎src/raw/requests.rs

+2-3
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,8 @@ use crate::{
1313
collect_first,
1414
pd::PdClient,
1515
request::{
16-
plan::ResponseWithShard,
17-
request_codec::{RawApiV1, RequestCodec},
18-
Collect, CollectSingle, DefaultProcessor, KvRequest, Merge, Process, Shardable, SingleKey,
16+
plan::ResponseWithShard, request_codec::RequestCodec, Collect, CollectSingle,
17+
DefaultProcessor, KvRequest, Merge, Process, Shardable, SingleKey,
1918
},
2019
store::{store_stream_for_keys, store_stream_for_ranges, RegionStore},
2120
transaction::HasLocks,

‎src/region_cache.rs

+42-22
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,22 @@
11
// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.
22

3-
use crate::{
4-
pd::{RetryClient, RetryClientTrait},
5-
region::{RegionId, RegionVerId, RegionWithLeader, StoreId},
6-
Key, Result,
7-
};
83
use std::{
94
collections::{BTreeMap, HashMap, HashSet},
105
sync::Arc,
116
};
7+
8+
use tokio::sync::{Notify, RwLock};
9+
1210
use tikv_client_common::Error;
1311
use tikv_client_pd::Cluster;
1412
use tikv_client_proto::metapb::{self, Store};
15-
use tokio::sync::{Notify, RwLock};
13+
14+
use crate::{
15+
pd::{RetryClient, RetryClientTrait},
16+
region::{RegionId, RegionVerId, RegionWithLeader, StoreId},
17+
request::request_codec::RequestCodec,
18+
Key, Result,
19+
};
1620

1721
const MAX_RETRY_WAITING_CONCURRENT_REQUEST: usize = 4;
1822

@@ -44,23 +48,25 @@ impl RegionCacheMap {
4448
}
4549
}
4650

47-
pub struct RegionCache<Client = RetryClient<Cluster>> {
51+
pub struct RegionCache<C, Client = RetryClient<Cluster>> {
4852
region_cache: RwLock<RegionCacheMap>,
4953
store_cache: RwLock<HashMap<StoreId, Store>>,
5054
inner_client: Arc<Client>,
55+
codec: C,
5156
}
5257

53-
impl<Client> RegionCache<Client> {
54-
pub fn new(inner_client: Arc<Client>) -> RegionCache<Client> {
58+
impl<C, Client> RegionCache<C, Client> {
59+
pub fn new(codec: C, inner_client: Arc<Client>) -> Self {
5560
RegionCache {
5661
region_cache: RwLock::new(RegionCacheMap::new()),
5762
store_cache: RwLock::new(HashMap::new()),
5863
inner_client,
64+
codec,
5965
}
6066
}
6167
}
6268

63-
impl<C: RetryClientTrait> RegionCache<C> {
69+
impl<C: RequestCodec, R: RetryClientTrait> RegionCache<C, R> {
6470
// Retrieve cache entry by key. If there's no entry, query PD and update cache.
6571
pub async fn get_region_by_key(&self, key: &Key) -> Result<RegionWithLeader> {
6672
let region_cache_guard = self.region_cache.read().await;
@@ -126,9 +132,14 @@ impl<C: RetryClientTrait> RegionCache<C> {
126132

127133
/// Force read through (query from PD) and update cache
128134
pub async fn read_through_region_by_key(&self, key: Key) -> Result<RegionWithLeader> {
129-
let region = self.inner_client.clone().get_region(key.into()).await?;
130-
self.add_region(region.clone()).await;
131-
Ok(region)
135+
let mut r = self
136+
.inner_client
137+
.clone()
138+
.get_region(self.codec.encode_pd_query(key).into())
139+
.await?;
140+
r.region = self.codec.decode_region(r.region)?;
141+
self.add_region(r.clone()).await;
142+
Ok(r)
132143
}
133144

134145
/// Force read through (query from PD) and update cache
@@ -140,7 +151,8 @@ impl<C: RetryClientTrait> RegionCache<C> {
140151
region_cache_guard.on_my_way_id.insert(id, notify.clone());
141152
}
142153

143-
let region = self.inner_client.clone().get_region_by_id(id).await?;
154+
let mut region = self.inner_client.clone().get_region_by_id(id).await?;
155+
region.region = self.codec.decode_region(region.region)?;
144156
self.add_region(region.clone()).await;
145157

146158
// notify others
@@ -226,27 +238,35 @@ impl<C: RetryClientTrait> RegionCache<C> {
226238
cache.key_to_ver_id.remove(&start_key);
227239
}
228240
}
241+
242+
pub fn get_request_codec(&self) -> C {
243+
self.codec.clone()
244+
}
229245
}
230246

231247
#[cfg(test)]
232248
mod test {
233-
use super::RegionCache;
234-
use crate::{
235-
pd::RetryClientTrait,
236-
region::{RegionId, RegionWithLeader},
237-
Key, Result,
238-
};
239-
use async_trait::async_trait;
240249
use std::{
241250
collections::{BTreeMap, HashMap, HashSet},
242251
sync::{
243252
atomic::{AtomicU64, Ordering::SeqCst},
244253
Arc,
245254
},
246255
};
256+
257+
use async_trait::async_trait;
258+
use tokio::sync::Mutex;
259+
247260
use tikv_client_common::Error;
248261
use tikv_client_proto::metapb;
249-
use tokio::sync::Mutex;
262+
263+
use crate::{
264+
pd::RetryClientTrait,
265+
region::{RegionId, RegionWithLeader},
266+
Key, Result,
267+
};
268+
269+
use super::RegionCache;
250270

251271
#[derive(Default)]
252272
struct MockRetryClient {

0 commit comments

Comments
 (0)