txn: Resolve locks for async commit (region level) (#378)

Signed-off-by: Ping Yu <yuping@pingcap.com>

Author: pingyu

src/request/mod.rs

@@ -15,7 +15,7 @@ pub use self::{
         ResponseWithShard, RetryableMultiRegion,
     },
     plan_builder::{PlanBuilder, SingleKey},
-    shard::Shardable,
+    shard::{HasNextBatch, NextBatch, Shardable},
 };
 
 pub mod plan;

src/request/plan.rs

@@ -12,10 +12,10 @@ use tokio::sync::Semaphore;
 use crate::{
     backoff::Backoff,
     pd::PdClient,
-    request::{KvRequest, Shardable},
+    request::{shard::HasNextBatch, KvRequest, NextBatch, Shardable},
     stats::tikv_stats,
     store::RegionStore,
-    transaction::{resolve_locks, HasLocks},
+    transaction::{resolve_locks, HasLocks, ResolveLocksContext, ResolveLocksOptions},
     util::iter::FlatMapOkIterExt,
     Error, Result,
 };
@@ -442,6 +442,169 @@ where
     }
 }
 
+#[derive(Default)]
+pub struct CleanupLocksResult {
+    pub region_error: Option<errorpb::Error>,
+    pub key_error: Option<Vec<Error>>,
+    pub meet_locks: usize,
+    // TODO: pub resolved_locks: usize,
+}
+
+impl Clone for CleanupLocksResult {
+    fn clone(&self) -> Self {
+        Self {
+            meet_locks: self.meet_locks,
+            ..Default::default() // Ignore errors, which should be extracted by `extract_error()`.
+        }
+    }
+}
+
+impl HasRegionError for CleanupLocksResult {
+    fn region_error(&mut self) -> Option<errorpb::Error> {
+        self.region_error.take()
+    }
+}
+
+impl HasKeyErrors for CleanupLocksResult {
+    fn key_errors(&mut self) -> Option<Vec<Error>> {
+        self.key_error.take()
+    }
+}
+
+impl Merge<CleanupLocksResult> for Collect {
+    type Out = CleanupLocksResult;
+
+    fn merge(&self, input: Vec<Result<CleanupLocksResult>>) -> Result<Self::Out> {
+        input
+            .into_iter()
+            .fold(Ok(CleanupLocksResult::default()), |acc, x| {
+                Ok(CleanupLocksResult {
+                    meet_locks: acc.unwrap().meet_locks + x?.meet_locks,
+                    ..Default::default()
+                })
+            })
+    }
+}
+
+pub struct CleanupLocks<P: Plan, PdC: PdClient> {
+    pub logger: slog::Logger,
+    pub inner: P,
+    pub ctx: ResolveLocksContext,
+    pub options: ResolveLocksOptions,
+    pub store: Option<RegionStore>,
+    pub pd_client: Arc<PdC>,
+    pub backoff: Backoff,
+}
+
+impl<P: Plan, PdC: PdClient> Clone for CleanupLocks<P, PdC> {
+    fn clone(&self) -> Self {
+        CleanupLocks {
+            logger: self.logger.clone(),
+            inner: self.inner.clone(),
+            ctx: self.ctx.clone(),
+            options: self.options,
+            store: None,
+            pd_client: self.pd_client.clone(),
+            backoff: self.backoff.clone(),
+        }
+    }
+}
+
+#[async_trait]
+impl<P: Plan + Shardable + NextBatch, PdC: PdClient> Plan for CleanupLocks<P, PdC>
+where
+    P::Result: HasLocks + HasNextBatch + HasKeyErrors + HasRegionError,
+{
+    type Result = CleanupLocksResult;
+
+    async fn execute(&self) -> Result<Self::Result> {
+        let mut result = CleanupLocksResult::default();
+        let mut inner = self.inner.clone();
+        let mut lock_resolver =
+            crate::transaction::LockResolver::new(self.logger.clone(), self.ctx.clone());
+        let region = &self.store.as_ref().unwrap().region_with_leader;
+        let mut has_more_batch = true;
+
+        while has_more_batch {
+            let mut scan_lock_resp = inner.execute().await?;
+
+            // Propagate errors to `retry_multi_region` for retry.
+            if let Some(e) = scan_lock_resp.key_errors() {
+                info!(
+                    self.logger,
+                    "CleanupLocks::execute, inner key errors:{:?}", e
+                );
+                result.key_error = Some(e);
+                return Ok(result);
+            } else if let Some(e) = scan_lock_resp.region_error() {
+                info!(
+                    self.logger,
+                    "CleanupLocks::execute, inner region error:{}",
+                    e.get_message()
+                );
+                result.region_error = Some(e);
+                return Ok(result);
+            }
+
+            // Iterate to next batch of inner.
+            match scan_lock_resp.has_next_batch() {
+                Some(range) if region.contains(range.0.as_ref()) => {
+                    debug!(self.logger, "CleanupLocks::execute, next range:{:?}", range);
+                    inner.next_batch(range);
+                }
+                _ => has_more_batch = false,
+            }
+
+            let mut locks = scan_lock_resp.take_locks();
+            if locks.is_empty() {
+                break;
+            }
+            if locks.len() < self.options.batch_size as usize {
+                has_more_batch = false;
+            }
+
+            if self.options.async_commit_only {
+                locks = locks
+                    .into_iter()
+                    .filter(|l| l.get_use_async_commit())
+                    .collect::<Vec<_>>();
+            }
+            debug!(
+                self.logger,
+                "CleanupLocks::execute, meet locks:{}",
+                locks.len()
+            );
+            result.meet_locks += locks.len();
+
+            match lock_resolver
+                .cleanup_locks(self.store.clone().unwrap(), locks, self.pd_client.clone())
+                .await
+            {
+                Ok(()) => {}
+                Err(Error::ExtractedErrors(mut errors)) => {
+                    // Propagate errors to `retry_multi_region` for retry.
+                    if let Error::RegionError(e) = errors.pop().unwrap() {
+                        result.region_error = Some(*e);
+                    } else {
+                        result.key_error = Some(errors);
+                    }
+                    return Ok(result);
+                }
+                Err(e) => {
+                    return Err(e);
+                }
+            }
+
+            // TODO: improve backoff
+            // if self.backoff.is_none() {
+            //     return Err(Error::ResolveLockError);
+            // }
+        }
+
+        Ok(result)
+    }
+}
+
 /// When executed, the plan extracts errors from its inner plan, and returns an
 /// `Err` wrapping the error.
 ///

src/request/plan_builder.rs

@@ -5,11 +5,12 @@ use crate::{
     backoff::Backoff,
     pd::PdClient,
     request::{
-        DefaultProcessor, Dispatch, ExtractError, KvRequest, Merge, MergeResponse, Plan, Process,
-        ProcessResponse, ResolveLock, RetryableMultiRegion, Shardable,
+        plan::CleanupLocks, shard::HasNextBatch, DefaultProcessor, Dispatch, ExtractError,
+        KvRequest, Merge, MergeResponse, NextBatch, Plan, Process, ProcessResponse, ResolveLock,
+        RetryableMultiRegion, Shardable,
     },
     store::RegionStore,
-    transaction::HasLocks,
+    transaction::{HasLocks, ResolveLocksContext, ResolveLocksOptions},
     Result,
 };
 use std::{marker::PhantomData, sync::Arc};
@@ -68,6 +69,32 @@ impl<PdC: PdClient, P: Plan, Ph: PlanBuilderPhase> PlanBuilder<PdC, P, Ph> {
         }
     }
 
+    pub fn cleanup_locks(
+        self,
+        logger: slog::Logger, // TODO: add logger to PlanBuilder.
+        ctx: ResolveLocksContext,
+        options: ResolveLocksOptions,
+        backoff: Backoff,
+    ) -> PlanBuilder<PdC, CleanupLocks<P, PdC>, Ph>
+    where
+        P: Shardable + NextBatch,
+        P::Result: HasLocks + HasNextBatch + HasRegionError + HasKeyErrors,
+    {
+        PlanBuilder {
+            pd_client: self.pd_client.clone(),
+            plan: CleanupLocks {
+                logger,
+                inner: self.plan,
+                ctx,
+                options,
+                store: None,
+                backoff,
+                pd_client: self.pd_client,
+            },
+            phantom: PhantomData,
+        }
+    }
+
     /// Merge the results of a request. Usually used where a request is sent to multiple regions
     /// to combine the responses from each region.
     pub fn merge<In, M: Merge<In>>(self, merge: M) -> PlanBuilder<PdC, MergeResponse<P, In, M>, Ph>

src/request/shard.rs

@@ -3,7 +3,7 @@
 use super::plan::PreserveShard;
 use crate::{
     pd::PdClient,
-    request::{Dispatch, KvRequest, Plan, ResolveLock},
+    request::{plan::CleanupLocks, Dispatch, KvRequest, Plan, ResolveLock},
     store::RegionStore,
     Result,
 };
@@ -38,6 +38,17 @@ pub trait Shardable {
     fn apply_shard(&mut self, shard: Self::Shard, store: &RegionStore) -> Result<()>;
 }
 
+// Use to iterate in a region for scan requests that have batch size limit.
+// HasNextBatch use to get the next batch according to previous response.
+pub trait HasNextBatch {
+    fn has_next_batch(&self) -> Option<(Vec<u8>, Vec<u8>)>;
+}
+
+// NextBatch use to change start key of request by result of `has_next_batch`.
+pub trait NextBatch {
+    fn next_batch(&mut self, _range: (Vec<u8>, Vec<u8>));
+}
+
 impl<Req: KvRequest + Shardable> Shardable for Dispatch<Req> {
     type Shard = Req::Shard;
 
@@ -54,6 +65,12 @@ impl<Req: KvRequest + Shardable> Shardable for Dispatch<Req> {
     }
 }
 
+impl<Req: KvRequest + NextBatch> NextBatch for Dispatch<Req> {
+    fn next_batch(&mut self, range: (Vec<u8>, Vec<u8>)) {
+        self.request.next_batch(range);
+    }
+}
+
 impl<P: Plan + Shardable> Shardable for PreserveShard<P> {
     type Shard = P::Shard;
 
@@ -74,6 +91,22 @@ impl<P: Plan + Shardable, PdC: PdClient> Shardable for ResolveLock<P, PdC> {
     impl_inner_shardable!();
 }
 
+impl<P: Plan + Shardable, PdC: PdClient> Shardable for CleanupLocks<P, PdC> {
+    type Shard = P::Shard;
+
+    fn shards(
+        &self,
+        pd_client: &Arc<impl PdClient>,
+    ) -> BoxStream<'static, Result<(Self::Shard, RegionStore)>> {
+        self.inner.shards(pd_client)
+    }
+
+    fn apply_shard(&mut self, shard: Self::Shard, store: &RegionStore) -> Result<()> {
+        self.store = Some(store.clone());
+        self.inner.apply_shard(shard, store)
+    }
+}
+
 #[macro_export]
 macro_rules! shardable_key {
     ($type_: ty) => {