Skip to content

Commit 5c5809a

Browse files
iosmanthusiosmanthus
committed
introduce impl_kv_request to simplify code
Signed-off-by: iosmanthus <[email protected]>
1 parent 317116b commit 5c5809a

File tree

4 files changed

+97
-100
lines changed

4 files changed

+97
-100
lines changed

‎src/raw/requests.rs

Lines changed: 20 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ pub fn new_raw_get_request(key: Vec<u8>, cf: Option<ColumnFamily>) -> kvrpcpb::R
3232
req
3333
}
3434

35-
impl_kv_request_for_single_key_op!(kvrpcpb::RawGetRequest, kvrpcpb::RawGetResponse);
35+
impl_kv_request!(kvrpcpb::RawGetRequest, key; kvrpcpb::RawGetResponse;);
3636
shardable_key!(kvrpcpb::RawGetRequest);
3737
collect_first!(kvrpcpb::RawGetResponse);
3838

@@ -66,8 +66,7 @@ pub fn new_raw_batch_get_request(
6666
req
6767
}
6868

69-
impl_kv_request_for_batch_get!(kvrpcpb::RawBatchGetRequest, kvrpcpb::RawBatchGetResponse);
70-
69+
impl_kv_request!(kvrpcpb::RawBatchGetRequest, keys; kvrpcpb::RawBatchGetResponse;);
7170
shardable_keys!(kvrpcpb::RawBatchGetRequest);
7271

7372
impl Merge<kvrpcpb::RawBatchGetResponse> for Collect {
@@ -96,7 +95,7 @@ pub fn new_raw_put_request(
9695
req
9796
}
9897

99-
impl_kv_request!(kvrpcpb::RawPutRequest, kvrpcpb::RawPutResponse; encode=key);
98+
impl_kv_request!(kvrpcpb::RawPutRequest, key; kvrpcpb::RawPutResponse;);
10099
shardable_key!(kvrpcpb::RawPutRequest);
101100
collect_first!(kvrpcpb::RawPutResponse);
102101

@@ -119,15 +118,7 @@ pub fn new_raw_batch_put_request(
119118
req
120119
}
121120

122-
impl<C: RequestCodec> KvRequest<C> for kvrpcpb::RawBatchPutRequest {
123-
type Response = kvrpcpb::RawBatchPutResponse;
124-
125-
fn encode_request(mut self, codec: &C) -> Self {
126-
*self.mut_pairs() = codec.encode_pairs(self.take_pairs());
127-
128-
self
129-
}
130-
}
121+
impl_kv_request!(kvrpcpb::RawBatchPutRequest, pairs; kvrpcpb::RawBatchPutResponse;);
131122

132123
impl Shardable for kvrpcpb::RawBatchPutRequest {
133124
type Shard = Vec<kvrpcpb::KvPair>;
@@ -164,7 +155,7 @@ pub fn new_raw_delete_request(
164155
req
165156
}
166157

167-
impl_kv_request_for_single_key_op!(kvrpcpb::RawDeleteRequest, kvrpcpb::RawDeleteResponse);
158+
impl_kv_request!(kvrpcpb::RawDeleteRequest, key; kvrpcpb::RawDeleteResponse;);
168159
shardable_key!(kvrpcpb::RawDeleteRequest);
169160
collect_first!(kvrpcpb::RawDeleteResponse);
170161
impl SingleKey for kvrpcpb::RawDeleteRequest {
@@ -184,16 +175,7 @@ pub fn new_raw_batch_delete_request(
184175
req
185176
}
186177

187-
impl<C: RequestCodec> KvRequest<C> for kvrpcpb::RawBatchDeleteRequest {
188-
type Response = kvrpcpb::RawBatchDeleteResponse;
189-
190-
fn encode_request(mut self, codec: &C) -> Self {
191-
*self.mut_keys() = codec.encode_keys(self.take_keys());
192-
193-
self
194-
}
195-
}
196-
178+
impl_kv_request!(kvrpcpb::RawBatchDeleteRequest, keys; kvrpcpb::RawBatchDeleteResponse;);
197179
shardable_keys!(kvrpcpb::RawBatchDeleteRequest);
198180

199181
pub fn new_raw_delete_range_request(
@@ -209,18 +191,10 @@ pub fn new_raw_delete_range_request(
209191
req
210192
}
211193

212-
impl<C: RequestCodec> KvRequest<C> for kvrpcpb::RawDeleteRangeRequest {
213-
type Response = kvrpcpb::RawDeleteRangeResponse;
214-
215-
fn encode_request(mut self, codec: &C) -> Self {
216-
let (start, end) = (self.take_start_key(), self.take_end_key());
217-
218-
self.set_start_key(codec.encode_key(start));
219-
self.set_end_key(codec.encode_key(end));
220-
221-
self
222-
}
223-
}
194+
impl_kv_request!(
195+
kvrpcpb::RawDeleteRangeRequest;
196+
kvrpcpb::RawDeleteRangeResponse;
197+
);
224198

225199
shardable_range!(kvrpcpb::RawDeleteRangeRequest);
226200

@@ -241,7 +215,7 @@ pub fn new_raw_scan_request(
241215
req
242216
}
243217

244-
impl_kv_request_for_scan_op!(kvrpcpb::RawScanRequest, kvrpcpb::RawScanResponse, kvs);
218+
impl_kv_request!(kvrpcpb::RawScanRequest; kvrpcpb::RawScanResponse, kvs;);
245219
shardable_range!(kvrpcpb::RawScanRequest);
246220

247221
impl Merge<kvrpcpb::RawScanResponse> for Collect {
@@ -270,15 +244,10 @@ pub fn new_raw_batch_scan_request(
270244
req
271245
}
272246

273-
impl<C: RequestCodec> KvRequest<C> for kvrpcpb::RawBatchScanRequest {
274-
type Response = kvrpcpb::RawBatchScanResponse;
275-
276-
fn encode_request(mut self, codec: &C) -> Self {
277-
*self.mut_ranges() = codec.encode_ranges(self.take_ranges());
278-
279-
self
280-
}
281-
}
247+
impl_kv_request!(
248+
kvrpcpb::RawBatchScanRequest, ranges;
249+
kvrpcpb::RawBatchScanResponse, kvs;
250+
);
282251

283252
impl Shardable for kvrpcpb::RawBatchScanRequest {
284253
type Shard = Vec<kvrpcpb::KeyRange>;
@@ -325,7 +294,11 @@ pub fn new_cas_request(
325294
req
326295
}
327296

328-
impl_kv_request_for_single_key_op!(kvrpcpb::RawCasRequest, kvrpcpb::RawCasResponse);
297+
impl_kv_request!(
298+
kvrpcpb::RawCasRequest, key;
299+
kvrpcpb::RawCasResponse;
300+
);
301+
329302
shardable_key!(kvrpcpb::RawCasRequest);
330303
collect_first!(kvrpcpb::RawCasResponse);
331304
impl SingleKey for kvrpcpb::RawCasRequest {

‎src/request/codec.rs

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ pub trait RequestCodec: Sized + Clone + Sync + Send + 'static {
5858
Ok(())
5959
}
6060

61-
fn decode_key_error(&self, err: &mut kvrpcpb::KeyError) -> Result<()> {
61+
fn decode_error(&self, err: &mut kvrpcpb::KeyError) -> Result<()> {
6262
if err.has_locked() {
6363
let locked = err.mut_locked();
6464
self.decode_lock(locked)?;
@@ -97,9 +97,9 @@ pub trait RequestCodec: Sized + Clone + Sync + Send + 'static {
9797
Ok(())
9898
}
9999

100-
fn decode_key_errors(&self, errors: &mut [kvrpcpb::KeyError]) -> Result<()> {
100+
fn decode_errors(&self, errors: &mut [kvrpcpb::KeyError]) -> Result<()> {
101101
for err in errors.iter_mut() {
102-
self.decode_key_error(err)?;
102+
self.decode_error(err)?;
103103
}
104104

105105
Ok(())
@@ -127,6 +127,10 @@ pub trait RequestCodec: Sized + Clone + Sync + Send + 'static {
127127
Ok(())
128128
}
129129

130+
fn decode_kvs(&self, kvs: &mut [kvrpcpb::KvPair]) -> Result<()> {
131+
self.decode_pairs(kvs)
132+
}
133+
130134
fn encode_range(&self, start: Vec<u8>, end: Vec<u8>) -> (Vec<u8>, Vec<u8>) {
131135
(start, end)
132136
}

‎src/request/mod.rs

Lines changed: 47 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,36 @@ pub trait KvRequest<C>: Request + Sized + Clone + Sync + Send + 'static {
4141
}
4242
}
4343

44+
macro_rules! impl_decode_response {
45+
($($o:ident)*; $($e:ident)*) => {
46+
fn decode_response(&self, codec: &C, mut resp: Self::Response) -> Result<Self::Response> {
47+
$(
48+
paste::paste! {
49+
codec.[<decode_ $o>](resp.[<mut_ $o>]())?;
50+
}
51+
)*
52+
53+
// decode errors
54+
if resp.has_region_error() {
55+
codec.decode_region_error(resp.mut_region_error())?;
56+
}
57+
58+
$(
59+
paste::paste! {
60+
if resp.[<has_ $e>]() {
61+
codec.[<decode_ $e>](resp.[<mut_ $e>]())?;
62+
}
63+
}
64+
)*
65+
66+
Ok(resp)
67+
}
68+
};
69+
}
70+
4471
#[macro_export]
4572
macro_rules! impl_kv_request {
46-
($req:ty, $resp:ty; encode=$($i:ident),*) => {
73+
($req:ty $(,$i:ident)+; $resp:ty $(,$o:ident)*; $($e:ident),*) => {
4774
impl<C> KvRequest<C> for $req
4875
where C: RequestCodec
4976
{
@@ -54,18 +81,30 @@ macro_rules! impl_kv_request {
5481
paste::paste! {
5582
*self.[<mut_ $i>]() = codec.[<encode_ $i>](self.[<take_ $i>]());
5683
}
57-
58-
self
5984
)*
85+
86+
self
6087
}
6188

62-
fn decode_response(&self, codec: &C, mut resp: Self::Response) -> Result<Self::Response> {
63-
if resp.has_region_error() {
64-
codec.decode_region_error(resp.mut_region_error())?;
65-
}
89+
impl_decode_response!{$($o)*; $($e)*}
90+
}
91+
};
6692

67-
Ok(resp)
93+
($req:ty; $resp:ty $(,$o:ident)*;$($e:ident)*) => {
94+
impl<C> KvRequest<C> for $req
95+
where C: RequestCodec
96+
{
97+
type Response = $resp;
98+
99+
fn encode_request(mut self, codec: &C) -> Self {
100+
let (start, end) = codec.encode_range(self.take_start_key(), self.take_end_key());
101+
*self.mut_start_key() = start;
102+
*self.mut_end_key() = end;
103+
104+
self
68105
}
106+
107+
impl_decode_response!{$($o)*; $($e)*}
69108
}
70109
};
71110
}

0 commit comments

Comments
 (0)