fix set_context

Signed-off-by: Ping Yu <[email protected]>

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(())
     }

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();

src/store/request.rs

@@ -21,6 +21,8 @@ pub trait Request: Any + Sync + Send + 'static {
     ) -> Result<Box<dyn Any>>;
     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) {

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(())