txn: Error handling for pessimistic locks

src/request/plan.rs

@@ -16,6 +16,7 @@ use crate::{
     stats::tikv_stats,
     store::RegionStore,
     transaction::{resolve_locks, HasLocks},
+    util::iter::FlatMapOkIterExt,
     Error, Result,
 };
 
@@ -63,6 +64,7 @@ pub struct RetryableMultiRegion<P: Plan, PdC: PdClient> {
     pub(super) inner: P,
     pub pd_client: Arc<PdC>,
     pub backoff: Backoff,
+    pub preserve_results: bool,
 }
 
 impl<P: Plan + Shardable, PdC: PdClient> RetryableMultiRegion<P, PdC>
@@ -76,6 +78,7 @@ where
         current_plan: P,
         backoff: Backoff,
         permits: Arc<Semaphore>,
+        preserve_results: bool,
     ) -> Result<<Self as Plan>::Result> {
         let shards = current_plan.shards(&pd_client).collect::<Vec<_>>().await;
         let mut handles = Vec::new();
@@ -89,16 +92,29 @@ where
                 region_store,
                 backoff.clone(),
                 permits.clone(),
+                preserve_results,
             ));
             handles.push(handle);
         }
-        Ok(try_join_all(handles)
-            .await?
-            .into_iter()
-            .collect::<Result<Vec<_>>>()?
-            .into_iter()
-            .flatten()
-            .collect())
+
+        let results = try_join_all(handles).await?;
+        if preserve_results {
+            Ok(results
+                .into_iter()
+                .flat_map_ok(|x| x)
+                .map(|x| match x {
+                    Ok(r) => r,
+                    Err(e) => Err(e),
+                })
+                .collect())
+        } else {
+            Ok(results
+                .into_iter()
+                .collect::<Result<Vec<_>>>()?
+                .into_iter()
+                .flatten()
+                .collect())
+        }
     }
 
     #[async_recursion]
@@ -108,6 +124,7 @@ where
         region_store: RegionStore,
         mut backoff: Backoff,
         permits: Arc<Semaphore>,
+        preserve_results: bool,
     ) -> Result<<Self as Plan>::Result> {
         // limit concurrent requests
         let permit = permits.acquire().await.unwrap();
@@ -125,7 +142,8 @@ where
                     if !region_error_resolved {
                         futures_timer::Delay::new(duration).await;
                     }
-                    Self::single_plan_handler(pd_client, plan, backoff, permits).await
+                    Self::single_plan_handler(pd_client, plan, backoff, permits, preserve_results)
+                        .await
                 }
                 None => Err(Error::RegionError(e)),
             }
@@ -242,6 +260,7 @@ impl<P: Plan, PdC: PdClient> Clone for RetryableMultiRegion<P, PdC> {
             inner: self.inner.clone(),
             pd_client: self.pd_client.clone(),
             backoff: self.backoff.clone(),
+            preserve_results: self.preserve_results,
         }
     }
 }
@@ -263,6 +282,7 @@ where
             self.inner.clone(),
             self.backoff.clone(),
             concurrency_permits.clone(),
+            self.preserve_results,
         )
         .await
     }
@@ -556,6 +576,7 @@ mod test {
             },
             pd_client: Arc::new(MockPdClient::default()),
             backoff: Backoff::no_backoff(),
+            preserve_results: false,
         };
         assert!(plan.execute().await.is_err())
     }

src/request/plan_builder.rs

@@ -113,13 +113,31 @@ where
     pub fn retry_multi_region(
         self,
         backoff: Backoff,
+    ) -> PlanBuilder<PdC, RetryableMultiRegion<P, PdC>, Targetted> {
+        self.make_retry_multi_region(backoff, false)
+    }
+
+    /// Preserve all results, even some of them are Err.
+    /// To pass all responses to merge, and handle partial successful results correctly.
+    pub fn retry_multi_region_preserve_results(
+        self,
+        backoff: Backoff,
+    ) -> PlanBuilder<PdC, RetryableMultiRegion<P, PdC>, Targetted> {
+        self.make_retry_multi_region(backoff, true)
+    }
+
+    fn make_retry_multi_region(
+        self,
+        backoff: Backoff,
+        preserve_results: bool,
     ) -> PlanBuilder<PdC, RetryableMultiRegion<P, PdC>, Targetted> {
         PlanBuilder {
             pd_client: self.pd_client.clone(),
             plan: RetryableMultiRegion {
                 inner: self.plan,
                 pd_client: self.pd_client,
                 backoff,
+                preserve_results,
             },
             phantom: PhantomData,
         }

src/transaction/requests.rs

@@ -16,6 +16,7 @@ use crate::{
 use either::Either;
 use futures::stream::BoxStream;
 use std::{collections::HashMap, iter, sync::Arc};
+use tikv_client_common::Error::PessimisticLockError;
 use tikv_client_proto::{
     kvrpcpb::{self, TxnHeartBeatResponse},
     pdpb::Timestamp,
@@ -379,35 +380,51 @@ impl Merge<ResponseWithShard<kvrpcpb::PessimisticLockResponse, Vec<kvrpcpb::Muta
             Result<ResponseWithShard<kvrpcpb::PessimisticLockResponse, Vec<kvrpcpb::Mutation>>>,
         >,
     ) -> Result<Self::Out> {
-        input
-            .into_iter()
-            .flat_map_ok(|ResponseWithShard(mut resp, mutations)| {
-                let values = resp.take_values();
-                let values_len = values.len();
-                let not_founds = resp.take_not_founds();
-                let kvpairs = mutations
-                    .into_iter()
-                    .map(|m| m.key)
-                    .zip(values)
-                    .map(KvPair::from);
-                assert_eq!(kvpairs.len(), values_len);
-                if not_founds.is_empty() {
-                    // Legacy TiKV does not distiguish not existing key and existing key
-                    // that with empty value. We assume that key does not exist if value
-                    // is empty.
-                    Either::Left(kvpairs.filter(|kvpair| !kvpair.value().is_empty()))
-                } else {
-                    assert_eq!(kvpairs.len(), not_founds.len());
-                    Either::Right(kvpairs.zip(not_founds).filter_map(|(kvpair, not_found)| {
-                        if not_found {
-                            None
-                        } else {
-                            Some(kvpair)
-                        }
-                    }))
-                }
+        let (success, mut errors): (Vec<_>, Vec<_>) = input.into_iter().partition(Result::is_ok);
+        if let Some(first_err) = errors.pop() {
+            let success_keys = success
+                .into_iter()
+                .map(Result::unwrap)
+                .flat_map(|ResponseWithShard(_resp, mutations)| {
+                    mutations.into_iter().map(|m| m.key)
+                })
+                .collect();
+            Err(PessimisticLockError {
+                inner: Box::new(first_err.unwrap_err()),
+                success_keys,
             })
-            .collect()
+        } else {
+            Ok(success
+                .into_iter()
+                .map(Result::unwrap)
+                .flat_map(|ResponseWithShard(mut resp, mutations)| {
+                    let values = resp.take_values();
+                    let values_len = values.len();
+                    let not_founds = resp.take_not_founds();
+                    let kvpairs = mutations
+                        .into_iter()
+                        .map(|m| m.key)
+                        .zip(values)
+                        .map(KvPair::from);
+                    assert_eq!(kvpairs.len(), values_len);
+                    if not_founds.is_empty() {
+                        // Legacy TiKV does not distinguish not existing key and existing key
+                        // that with empty value. We assume that key does not exist if value
+                        // is empty.
+                        Either::Left(kvpairs.filter(|kvpair| !kvpair.value().is_empty()))
+                    } else {
+                        assert_eq!(kvpairs.len(), not_founds.len());
+                        Either::Right(kvpairs.zip(not_founds).filter_map(|(kvpair, not_found)| {
+                            if not_found {
+                                None
+                            } else {
+                                Some(kvpair)
+                            }
+                        }))
+                    }
+                })
+                .collect())
+        }
     }
 }
 
@@ -653,3 +670,108 @@ impl HasLocks for kvrpcpb::PrewriteResponse {
             .collect()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        request::{plan::Merge, CollectWithShard, ResponseWithShard},
+        KvPair,
+    };
+    use tikv_client_common::Error::{PessimisticLockError, ResolveLockError};
+    use tikv_client_proto::kvrpcpb;
+
+    #[tokio::test]
+    async fn test_merge_pessimistic_lock_response() {
+        let (key1, key2, key3, key4) = (b"key1", b"key2", b"key3", b"key4");
+        let (value1, value4) = (b"value1", b"value4");
+        let value_empty = b"";
+
+        let resp1 = ResponseWithShard(
+            kvrpcpb::PessimisticLockResponse {
+                values: vec![value1.to_vec()],
+                ..Default::default()
+            },
+            vec![kvrpcpb::Mutation {
+                op: kvrpcpb::Op::PessimisticLock as i32,
+                key: key1.to_vec(),
+                ..Default::default()
+            }],
+        );
+
+        let resp_empty_value = ResponseWithShard(
+            kvrpcpb::PessimisticLockResponse {
+                values: vec![value_empty.to_vec()],
+                ..Default::default()
+            },
+            vec![kvrpcpb::Mutation {
+                op: kvrpcpb::Op::PessimisticLock as i32,
+                key: key2.to_vec(),
+                ..Default::default()
+            }],
+        );
+
+        let resp_not_found = ResponseWithShard(
+            kvrpcpb::PessimisticLockResponse {
+                values: vec![value_empty.to_vec(), value4.to_vec()],
+                not_founds: vec![true, false],
+                ..Default::default()
+            },
+            vec![
+                kvrpcpb::Mutation {
+                    op: kvrpcpb::Op::PessimisticLock as i32,
+                    key: key3.to_vec(),
+                    ..Default::default()
+                },
+                kvrpcpb::Mutation {
+                    op: kvrpcpb::Op::PessimisticLock as i32,
+                    key: key4.to_vec(),
+                    ..Default::default()
+                },
+            ],
+        );
+
+        let merger = CollectWithShard {};
+        {
+            // empty values & not founds are filtered.
+            let input = vec![
+                Ok(resp1.clone()),
+                Ok(resp_empty_value.clone()),
+                Ok(resp_not_found.clone()),
+            ];
+            let result = merger.merge(input);
+            println!("result: {:?}", result);
+
+            assert_eq!(
+                result.unwrap(),
+                vec![
+                    KvPair::new(key1.to_vec(), value1.to_vec()),
+                    KvPair::new(key4.to_vec(), value4.to_vec()),
+                ]
+            );
+        }
+        {
+            let input = vec![
+                Ok(resp1),
+                Ok(resp_empty_value),
+                Err(ResolveLockError),
+                Ok(resp_not_found),
+            ];
+            let result = merger.merge(input);
+            println!("result: {:?}", result);
+
+            if let PessimisticLockError {
+                inner,
+                success_keys,
+            } = result.unwrap_err()
+            {
+                assert!(matches!(*inner, ResolveLockError));
+                assert_eq!(
+                    success_keys,
+                    vec![key1.to_vec(), key2.to_vec(), key3.to_vec(), key4.to_vec()]
+                );
+            } else {
+                panic!();
+            }
+        }
+    }
+}