Merge pull request #245 from ekexium/fix-result-haserror

nrc · web-flow · commit 47478cb59faa · 2021-03-05T13:49:11.000+13:00
Fix HasError and HasLocks
diff --git a/src/raw/client.rs b/src/raw/client.rs
@@ -179,6 +179,7 @@ impl Client {
             .await?
             .resolve_lock(OPTIMISTIC_BACKOFF)
             .retry_region(DEFAULT_REGION_BACKOFF)
+            .extract_error()
             .plan();
         plan.execute().await?;
         Ok(())
@@ -210,6 +211,7 @@ impl Client {
             .resolve_lock(OPTIMISTIC_BACKOFF)
             .multi_region()
             .retry_region(DEFAULT_REGION_BACKOFF)
+            .extract_error()
             .plan();
         plan.execute().await?;
         Ok(())
@@ -239,6 +241,7 @@ impl Client {
             .await?
             .resolve_lock(OPTIMISTIC_BACKOFF)
             .retry_region(DEFAULT_REGION_BACKOFF)
+            .extract_error()
             .plan();
         plan.execute().await?;
         Ok(())
@@ -268,6 +271,7 @@ impl Client {
             .resolve_lock(OPTIMISTIC_BACKOFF)
             .multi_region()
             .retry_region(DEFAULT_REGION_BACKOFF)
+            .extract_error()
             .plan();
         plan.execute().await?;
         Ok(())
@@ -294,6 +298,7 @@ impl Client {
             .resolve_lock(OPTIMISTIC_BACKOFF)
             .multi_region()
             .retry_region(DEFAULT_REGION_BACKOFF)
+            .extract_error()
             .plan();
         plan.execute().await?;
         Ok(())
diff --git a/src/request/mod.rs b/src/request/mod.rs
@@ -10,8 +10,8 @@ use tikv_client_store::{HasError, Request};
 
 pub use self::{
     plan::{
-        Collect, CollectError, DefaultProcessor, Dispatch, Merge, MergeResponse, MultiRegion, Plan,
-        Process, ProcessResponse, ResolveLock, RetryRegion,
+        Collect, CollectError, DefaultProcessor, Dispatch, ExtractError, Merge, MergeResponse,
+        MultiRegion, Plan, Process, ProcessResponse, ResolveLock, RetryRegion,
     },
     plan_builder::{PlanBuilder, SingleKey},
     shard::Shardable,
@@ -66,15 +66,17 @@ mod test {
     use crate::{
         mock::{MockKvClient, MockPdClient},
         store::store_stream_for_keys,
+        transaction::lowering::new_commit_request,
         Error, Key, Result,
     };
     use futures::executor;
     use grpcio::CallOption;
     use std::{
         any::Any,
+        iter,
         sync::{Arc, Mutex},
     };
-    use tikv_client_proto::{kvrpcpb, tikvpb::TikvClient};
+    use tikv_client_proto::{kvrpcpb, pdpb::Timestamp, tikvpb::TikvClient};
     use tikv_client_store::HasRegionError;
 
     #[test]
@@ -167,10 +169,54 @@ mod test {
             .resolve_lock(Backoff::no_jitter_backoff(1, 1, 3))
             .multi_region()
             .retry_region(Backoff::no_jitter_backoff(1, 1, 3))
+            .extract_error()
             .plan();
         let _ = executor::block_on(async { plan.execute().await });
 
         // Original call plus the 3 retries
         assert_eq!(*invoking_count.lock().unwrap(), 4);
     }
+
+    #[test]
+    fn test_extract_error() {
+        let pd_client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook(
+            |_: &dyn Any| {
+                Ok(Box::new(kvrpcpb::CommitResponse {
+                    region_error: None,
+                    error: Some(kvrpcpb::KeyError {
+                        locked: None,
+                        retryable: String::new(),
+                        abort: String::new(),
+                        conflict: None,
+                        already_exist: None,
+                        deadlock: None,
+                        commit_ts_expired: None,
+                        txn_not_found: None,
+                        commit_ts_too_large: None,
+                    }),
+                    commit_version: 0,
+                }) as Box<dyn Any>)
+            },
+        )));
+
+        let key: Key = "key".to_owned().into();
+        let req = new_commit_request(iter::once(key), Timestamp::default(), Timestamp::default());
+
+        // does not extract error
+        let plan = crate::request::PlanBuilder::new(pd_client.clone(), req.clone())
+            .resolve_lock(OPTIMISTIC_BACKOFF)
+            .multi_region()
+            .retry_region(OPTIMISTIC_BACKOFF)
+            .plan();
+        assert!(executor::block_on(async { plan.execute().await }).is_ok());
+
+        // extract error
+        let plan = crate::request::PlanBuilder::new(pd_client.clone(), req)
+            .resolve_lock(OPTIMISTIC_BACKOFF)
+            .multi_region()
+            .retry_region(OPTIMISTIC_BACKOFF)
+            .extract_error()
+            .plan();
+        assert!(executor::block_on(async { plan.execute().await }).is_err());
+    }
 }
diff --git a/src/request/plan.rs b/src/request/plan.rs
@@ -249,6 +249,42 @@ where
     }
 }
 
+pub struct ExtractError<P: Plan> {
+    pub inner: P,
+}
+
+impl<P: Plan> Clone for ExtractError<P> {
+    fn clone(&self) -> Self {
+        ExtractError {
+            inner: self.inner.clone(),
+        }
+    }
+}
+
+/// When executed, the plan extracts errors from its inner plan, and
+/// returns an `Err` wrapping the error.
+///
+/// The errors come from two places: `Err` from inner plans, and `Ok(response)`
+/// where `response` contains unresolved errors (`error` and `region_error`).
+#[async_trait]
+impl<P: Plan> Plan for ExtractError<P>
+where
+    P::Result: HasError,
+{
+    type Result = P::Result;
+
+    async fn execute(&self) -> Result<Self::Result> {
+        let mut result = self.inner.execute().await?;
+        if let Some(error) = result.error() {
+            Err(error)
+        } else if let Some(error) = result.region_error() {
+            Err(error)
+        } else {
+            Ok(result)
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
diff --git a/src/request/plan_builder.rs b/src/request/plan_builder.rs
@@ -4,8 +4,8 @@ use crate::{
     backoff::Backoff,
     pd::PdClient,
     request::{
-        DefaultProcessor, Dispatch, KvRequest, Merge, MergeResponse, MultiRegion, Plan, Process,
-        ProcessResponse, ResolveLock, RetryRegion, Shardable,
+        DefaultProcessor, Dispatch, ExtractError, KvRequest, Merge, MergeResponse, MultiRegion,
+        Plan, Process, ProcessResponse, ResolveLock, RetryRegion, Shardable,
     },
     store::Store,
     transaction::HasLocks,
@@ -161,6 +161,19 @@ impl<PdC: PdClient, R: KvRequest> PlanBuilder<PdC, Dispatch<R>, NoTarget> {
     }
 }
 
+impl<PdC: PdClient, P: Plan> PlanBuilder<PdC, P, Targetted>
+where
+    P::Result: HasError,
+{
+    pub fn extract_error(self) -> PlanBuilder<PdC, ExtractError<P>, Targetted> {
+        PlanBuilder {
+            pd_client: self.pd_client,
+            plan: ExtractError { inner: self.plan },
+            phantom: self.phantom,
+        }
+    }
+}
+
 fn set_single_region_store<PdC: PdClient, R: KvRequest>(
     mut plan: Dispatch<R>,
     store: Store,
diff --git a/src/transaction/lock.rs b/src/transaction/lock.rs
@@ -104,6 +104,7 @@ async fn resolve_lock_with_retry(
             .await?
             .resolve_lock(Backoff::no_backoff())
             .retry_region(Backoff::no_backoff())
+            .extract_error()
             .plan();
         match plan.execute().await {
             Ok(_) => {
diff --git a/src/transaction/requests.rs b/src/transaction/requests.rs
@@ -13,6 +13,47 @@ use futures::stream::BoxStream;
 use std::{collections::HashMap, iter, sync::Arc};
 use tikv_client_proto::{kvrpcpb, pdpb::Timestamp};
 
+// implement HasLocks for a response type that has a `pairs` field,
+// where locks can be extracted from both the `pairs` and `error` fields
+macro_rules! pair_locks {
+    ($response_type:ty) => {
+        impl HasLocks for $response_type {
+            fn take_locks(&mut self) -> Vec<kvrpcpb::LockInfo> {
+                if self.pairs.is_empty() {
+                    self.error
+                        .as_mut()
+                        .and_then(|error| error.locked.take())
+                        .into_iter()
+                        .collect()
+                } else {
+                    self.pairs
+                        .iter_mut()
+                        .filter_map(|pair| {
+                            pair.error.as_mut().and_then(|error| error.locked.take())
+                        })
+                        .collect()
+                }
+            }
+        }
+    };
+}
+
+// implement HasLocks for a response type that does not have a `pairs` field,
+// where locks are only extracted from the `error` field
+macro_rules! error_locks {
+    ($response_type:ty) => {
+        impl HasLocks for $response_type {
+            fn take_locks(&mut self) -> Vec<kvrpcpb::LockInfo> {
+                self.error
+                    .as_mut()
+                    .and_then(|error| error.locked.take())
+                    .into_iter()
+                    .collect()
+            }
+        }
+    };
+}
+
 pub fn new_get_request(key: Vec<u8>, timestamp: u64) -> kvrpcpb::GetRequest {
     let mut req = kvrpcpb::GetRequest::default();
     req.set_key(key);
@@ -43,16 +84,6 @@ impl Process<kvrpcpb::GetResponse> for DefaultProcessor {
     }
 }
 
-impl HasLocks for kvrpcpb::GetResponse {
-    fn take_locks(&mut self) -> Vec<kvrpcpb::LockInfo> {
-        self.error
-            .as_mut()
-            .and_then(|error| error.locked.take())
-            .into_iter()
-            .collect()
-    }
-}
-
 pub fn new_batch_get_request(keys: Vec<Vec<u8>>, timestamp: u64) -> kvrpcpb::BatchGetRequest {
     let mut req = kvrpcpb::BatchGetRequest::default();
     req.set_keys(keys);
@@ -77,15 +108,6 @@ impl Merge<kvrpcpb::BatchGetResponse> for Collect {
     }
 }
 
-impl HasLocks for kvrpcpb::BatchGetResponse {
-    fn take_locks(&mut self) -> Vec<kvrpcpb::LockInfo> {
-        self.pairs
-            .iter_mut()
-            .filter_map(|pair| pair.error.as_mut().and_then(|error| error.locked.take()))
-            .collect()
-    }
-}
-
 pub fn new_scan_request(
     start_key: Vec<u8>,
     end_key: Vec<u8>,
@@ -119,15 +141,6 @@ impl Merge<kvrpcpb::ScanResponse> for Collect {
     }
 }
 
-impl HasLocks for kvrpcpb::ScanResponse {
-    fn take_locks(&mut self) -> Vec<kvrpcpb::LockInfo> {
-        self.pairs
-            .iter_mut()
-            .filter_map(|pair| pair.error.as_mut().and_then(|error| error.locked.take()))
-            .collect()
-    }
-}
-
 pub fn new_resolve_lock_request(
     start_version: u64,
     commit_version: u64,
@@ -547,13 +560,32 @@ pub struct SecondaryLocksStatus {
     pub commit_ts: Option<Timestamp>,
 }
 
-impl HasLocks for kvrpcpb::CommitResponse {}
+impl HasLocks for kvrpcpb::PessimisticRollbackResponse {
+    fn take_locks(&mut self) -> Vec<kvrpcpb::LockInfo> {
+        self.errors
+            .iter_mut()
+            .filter_map(|error| error.locked.take())
+            .collect()
+    }
+}
+
+impl HasLocks for kvrpcpb::PessimisticLockResponse {
+    fn take_locks(&mut self) -> Vec<kvrpcpb::LockInfo> {
+        self.errors
+            .iter_mut()
+            .filter_map(|error| error.locked.take())
+            .collect()
+    }
+}
+
+pair_locks!(kvrpcpb::BatchGetResponse);
+pair_locks!(kvrpcpb::ScanResponse);
+error_locks!(kvrpcpb::GetResponse);
+error_locks!(kvrpcpb::ResolveLockResponse);
+error_locks!(kvrpcpb::CommitResponse);
+error_locks!(kvrpcpb::BatchRollbackResponse);
+error_locks!(kvrpcpb::TxnHeartBeatResponse);
+error_locks!(kvrpcpb::CheckTxnStatusResponse);
+error_locks!(kvrpcpb::CheckSecondaryLocksResponse);
 impl HasLocks for kvrpcpb::CleanupResponse {}
-impl HasLocks for kvrpcpb::BatchRollbackResponse {}
-impl HasLocks for kvrpcpb::PessimisticRollbackResponse {}
-impl HasLocks for kvrpcpb::ResolveLockResponse {}
 impl HasLocks for kvrpcpb::ScanLockResponse {}
-impl HasLocks for kvrpcpb::PessimisticLockResponse {}
-impl HasLocks for kvrpcpb::TxnHeartBeatResponse {}
-impl HasLocks for kvrpcpb::CheckTxnStatusResponse {}
-impl HasLocks for kvrpcpb::CheckSecondaryLocksResponse {}
diff --git a/src/transaction/transaction.rs b/src/transaction/transaction.rs
@@ -901,6 +901,7 @@ impl Committer {
             .multi_region()
             .retry_region(self.options.retry_options.region_backoff.clone())
             .merge(CollectError)
+            .extract_error()
             .plan();
         let response = plan.execute().await?;
 
@@ -939,6 +940,7 @@ impl Committer {
             .resolve_lock(self.options.retry_options.lock_backoff.clone())
             .multi_region()
             .retry_region(self.options.retry_options.region_backoff.clone())
+            .extract_error()
             .plan();
         plan.execute()
             .inspect_err(|e| {
@@ -975,6 +977,7 @@ impl Committer {
             .resolve_lock(self.options.retry_options.lock_backoff)
             .multi_region()
             .retry_region(self.options.retry_options.region_backoff)
+            .extract_error()
             .plan();
         plan.execute().await?;
         Ok(())
@@ -995,6 +998,7 @@ impl Committer {
                     .resolve_lock(self.options.retry_options.lock_backoff)
                     .multi_region()
                     .retry_region(self.options.retry_options.region_backoff)
+                    .extract_error()
                     .plan();
                 plan.execute().await?;
             }
@@ -1004,6 +1008,7 @@ impl Committer {
                     .resolve_lock(self.options.retry_options.lock_backoff)
                     .multi_region()
                     .retry_region(self.options.retry_options.region_backoff)
+                    .extract_error()
                     .plan();
                 plan.execute().await?;
             }
diff --git a/tikv-client-common/src/errors.rs b/tikv-client-common/src/errors.rs
@@ -69,6 +69,8 @@ pub enum Error {
     KvError { message: String },
     #[error("{}", message)]
     InternalError { message: String },
+    #[error("{0}")]
+    StringError(String),
 }
 
 impl From<tikv_client_proto::errorpb::Error> for Error {
diff --git a/tikv-client-store/src/errors.rs b/tikv-client-store/src/errors.rs

Original file line number	Diff line number	Diff line change
`@@ -69,6 +69,8 @@ pub enum Error {`
`69`	`69`	`KvError { message: String },`
`70`	`70`	`#[error("{}", message)]`
`71`	`71`	`InternalError { message: String },`
	`72`	`+ #[error("{0}")]`
	`73`	`+ StringError(String),`
`72`	`74`	`}`
`73`	`75`
`74`	`76`	`impl From<tikv_client_proto::errorpb::Error> for Error {`