From 19f829d46cbd7fe51c0dbe521b6163b2385ae410 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Mon, 28 Aug 2023 17:11:36 +0800 Subject: [PATCH] fix set_context Signed-off-by: Ping Yu --- src/raw/requests.rs | 4 ++-- src/request/shard.rs | 6 +++--- src/store/request.rs | 8 ++++++++ src/transaction/requests.rs | 13 +++++++------ 4 files changed, 20 insertions(+), 11 deletions(-) diff --git a/src/raw/requests.rs b/src/raw/requests.rs index ff1cc14f..23bfce73 100644 --- a/src/raw/requests.rs +++ b/src/raw/requests.rs @@ -162,7 +162,7 @@ impl Shardable for kvrpcpb::RawBatchPutRequest { } fn apply_shard(&mut self, shard: Self::Shard, store: &RegionStore) -> Result<()> { - self.context = Some(store.region_with_leader.context()?); + self.set_context(store.region_with_leader.context()?); self.pairs = shard; Ok(()) } @@ -293,7 +293,7 @@ impl Shardable for kvrpcpb::RawBatchScanRequest { } fn apply_shard(&mut self, shard: Self::Shard, store: &RegionStore) -> Result<()> { - self.context = Some(store.region_with_leader.context()?); + self.set_context(store.region_with_leader.context()?); self.ranges = shard; Ok(()) } diff --git a/src/request/shard.rs b/src/request/shard.rs index aaefab72..7c78743d 100644 --- a/src/request/shard.rs +++ b/src/request/shard.rs @@ -163,7 +163,7 @@ macro_rules! shardable_key { mut shard: Self::Shard, store: &$crate::store::RegionStore, ) -> $crate::Result<()> { - self.context = Some(store.region_with_leader.context()?); + self.set_context(store.region_with_leader.context()?); assert!(shard.len() == 1); self.key = shard.pop().unwrap(); Ok(()) @@ -196,7 +196,7 @@ macro_rules! shardable_keys { shard: Self::Shard, store: &$crate::store::RegionStore, ) -> $crate::Result<()> { - self.context = Some(store.region_with_leader.context()?); + self.set_context(store.region_with_leader.context()?); self.keys = shard.into_iter().map(Into::into).collect(); Ok(()) } @@ -225,7 +225,7 @@ macro_rules! shardable_range { shard: Self::Shard, store: &$crate::store::RegionStore, ) -> $crate::Result<()> { - self.context = Some(store.region_with_leader.context()?); + self.set_context(store.region_with_leader.context()?); self.start_key = shard.0.into(); self.end_key = shard.1.into(); diff --git a/src/store/request.rs b/src/store/request.rs index a26e57e0..2f1a31a0 100644 --- a/src/store/request.rs +++ b/src/store/request.rs @@ -21,6 +21,8 @@ pub trait Request: Any + Sync + Send + 'static { ) -> Result>; fn label(&self) -> &'static str; fn as_any(&self) -> &dyn Any; + /// Set the context for the request. + /// Should always use `set_context` other than modify the `self.context` directly. fn set_context(&mut self, context: kvrpcpb::Context); fn set_api_version(&mut self, api_version: kvrpcpb::ApiVersion); } @@ -53,7 +55,13 @@ macro_rules! impl_request { } fn set_context(&mut self, context: kvrpcpb::Context) { + let api_version = self + .context + .as_ref() + .map(|c| c.api_version) + .unwrap_or_default(); self.context = Some(context); + self.set_api_version(kvrpcpb::ApiVersion::from_i32(api_version).unwrap()); } fn set_api_version(&mut self, api_version: kvrpcpb::ApiVersion) { diff --git a/src/transaction/requests.rs b/src/transaction/requests.rs index 9cc4993f..bd3852e3 100644 --- a/src/transaction/requests.rs +++ b/src/transaction/requests.rs @@ -38,6 +38,7 @@ use crate::shardable_range; use crate::store::store_stream_for_keys; use crate::store::store_stream_for_range; use crate::store::RegionStore; +use crate::store::Request; use crate::timestamp::TimestampExt; use crate::transaction::HasLocks; use crate::util::iter::FlatMapOkIterExt; @@ -294,7 +295,7 @@ impl Shardable for kvrpcpb::PrewriteRequest { } fn apply_shard(&mut self, shard: Self::Shard, store: &RegionStore) -> Result<()> { - self.context = Some(store.region_with_leader.context()?); + self.set_context(store.region_with_leader.context()?); // Only need to set secondary keys if we're sending the primary key. if self.use_async_commit && !self.mutations.iter().any(|m| m.key == self.primary_lock) { @@ -361,7 +362,7 @@ impl Shardable for kvrpcpb::CommitRequest { } fn apply_shard(&mut self, shard: Self::Shard, store: &RegionStore) -> Result<()> { - self.context = Some(store.region_with_leader.context()?); + self.set_context(store.region_with_leader.context()?); self.keys = shard.into_iter().map(Into::into).collect(); Ok(()) } @@ -452,7 +453,7 @@ impl Shardable for kvrpcpb::PessimisticLockRequest { } fn apply_shard(&mut self, shard: Self::Shard, store: &RegionStore) -> Result<()> { - self.context = Some(store.region_with_leader.context()?); + self.set_context(store.region_with_leader.context()?); self.mutations = shard; Ok(()) } @@ -553,7 +554,7 @@ impl Shardable for kvrpcpb::ScanLockRequest { } fn apply_shard(&mut self, shard: Self::Shard, store: &RegionStore) -> Result<()> { - self.context = Some(store.region_with_leader.context()?); + self.set_context(store.region_with_leader.context()?); self.start_key = shard.0; Ok(()) } @@ -614,7 +615,7 @@ impl Shardable for kvrpcpb::TxnHeartBeatRequest { } fn apply_shard(&mut self, mut shard: Self::Shard, store: &RegionStore) -> Result<()> { - self.context = Some(store.region_with_leader.context()?); + self.set_context(store.region_with_leader.context()?); assert!(shard.len() == 1); self.primary_lock = shard.pop().unwrap(); Ok(()) @@ -672,7 +673,7 @@ impl Shardable for kvrpcpb::CheckTxnStatusRequest { } fn apply_shard(&mut self, mut shard: Self::Shard, store: &RegionStore) -> Result<()> { - self.context = Some(store.region_with_leader.context()?); + self.set_context(store.region_with_leader.context()?); assert!(shard.len() == 1); self.primary_key = shard.pop().unwrap(); Ok(())