From babbc22e4cb3a8260dfd59685c2081c601402331 Mon Sep 17 00:00:00 2001
From: "qi.xu" <tonxuqi@outlook.com>
Date: Mon, 13 Sep 2021 03:53:39 +0800
Subject: [PATCH] fix the bucket keys generation bug (should use origin key)

---
 .../src/coprocessor/split_check/half.rs          |  4 ++--
 .../src/coprocessor/split_check/size.rs          |  2 +-
 components/raftstore/src/store/peer.rs           |  2 +-
 .../raftstore/src/store/worker/split_check.rs    | 16 ++++++++--------
 4 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/components/raftstore/src/coprocessor/split_check/half.rs b/components/raftstore/src/coprocessor/split_check/half.rs
index 7e5df3315cd..f1d7a2da95e 100644
--- a/components/raftstore/src/coprocessor/split_check/half.rs
+++ b/components/raftstore/src/coprocessor/split_check/half.rs
@@ -236,9 +236,9 @@ mod tests {
         // so bucket key will be all these keys
         let mut exp_bucket_keys = vec![];
         for i in 0..11 {
-            let k = format!("{:04}", i).into_bytes();
+            let k = format!("{:04}", i).into_bytes(); 
+            exp_bucket_keys.push(Key::from_raw(&k).as_encoded().clone());
             let k = keys::data_key(Key::from_raw(&k).as_encoded());
-            exp_bucket_keys.push(k.clone());
             engine.put_cf(CF_DEFAULT, &k, &k).unwrap();
             // Flush for every key so that we can know the exact middle key.
             engine.flush_cf(CF_DEFAULT, true).unwrap();
diff --git a/components/raftstore/src/coprocessor/split_check/size.rs b/components/raftstore/src/coprocessor/split_check/size.rs
index 1116c080abe..2e0d5566115 100644
--- a/components/raftstore/src/coprocessor/split_check/size.rs
+++ b/components/raftstore/src/coprocessor/split_check/size.rs
@@ -488,7 +488,7 @@ pub mod tests {
             engine.put_cf(data_cf, &s, &s).unwrap();
             if i % 10 == 0 && i > 0 {
                 if i < 40 {
-                    exp_bucket_keys.push(s.clone());
+                    exp_bucket_keys.push(keys::origin_key(&s).to_vec());
                 }
                 engine.flush_cf(data_cf, true).unwrap();
             }
diff --git a/components/raftstore/src/store/peer.rs b/components/raftstore/src/store/peer.rs
index 6383ad68415..e2373348c64 100644
--- a/components/raftstore/src/store/peer.rs
+++ b/components/raftstore/src/store/peer.rs
@@ -3524,7 +3524,7 @@ where
 
     pub fn heartbeat_pd<T>(&mut self, ctx: &PollContext<EK, ER, T>) {
         let mut region = self.region().clone();
-        let mut region_buckets = Vec::with_capacity(self.region_buckets.len());
+        let mut region_buckets = vec![metapb::RegionBucket::default();self.region_buckets.len()];
         region_buckets.clone_from_slice(&self.region_buckets); // TODO: remove the clone if possible
         region.set_region_bucket(protobuf::RepeatedField::from(region_buckets));
         if region.region_bucket.len() != 0 {
diff --git a/components/raftstore/src/store/worker/split_check.rs b/components/raftstore/src/store/worker/split_check.rs
index 934a57c7858..0e4c464a434 100644
--- a/components/raftstore/src/store/worker/split_check.rs
+++ b/components/raftstore/src/store/worker/split_check.rs
@@ -224,7 +224,7 @@ where
             CheckPolicy::Approximate => {
                 if host.enable_region_bucket() {
                     let mut bucket_keys = match host.approximate_bucket_keys(region, &tablet) {
-                        Ok(keys) => keys.into_iter().map(|k| k.to_vec()).collect(),
+                        Ok(keys) => keys.into_iter().map(|k| keys::origin_key(&k).to_vec()).collect(),
                         Err(e) => {
                             error!(%e;
                                 "failed to get approximate bucket key";
@@ -235,8 +235,8 @@ where
                     };
                     info!("starting approximate_bucket_keys {}", bucket_keys.len());
                     if bucket_keys.len() > 0 {
-                        bucket_keys.insert(0, start_key.clone()); //
-                        bucket_keys.push(end_key.clone());
+                        bucket_keys.insert(0, keys::origin_key(&start_key).to_vec()); //
+                        bucket_keys.push(keys::origin_end_key(&end_key).to_vec());
                         let mut i = 0;
                         let mut region_buckets = vec![];
                         while i < bucket_keys.len() - 1 {
@@ -313,7 +313,7 @@ where
             let mut size = 0;
             let mut keys = 0;
             let mut bucket_size: u64 = 0;
-            let mut bucket_start_key = start_key.to_vec();
+            let mut bucket_start_key = keys::origin_key(&start_key).to_vec();
             while let Some(e) = iter.next() {
                 if host.on_kv(region, &e) {
                     return;
@@ -325,16 +325,16 @@ where
                     if bucket_size >= host.region_bucket_size() {
                         let mut region_bucket = RegionBucket::default();
                         region_bucket.start_key = bucket_start_key;
-                        region_bucket.end_key = e.key().to_vec();
+                        region_bucket.end_key = keys::origin_key(&e.key()).to_vec();
                         region_buckets.push(region_bucket);
                         bucket_size = 0;
-                        bucket_start_key = e.key().to_vec();
+                        bucket_start_key = keys::origin_key(&e.key()).to_vec();
                     }
                 }
             }
             let buckets_len = region_buckets.len();
-            if buckets_len > 1 {
-                region_buckets[buckets_len - 1].end_key = end_key.to_vec();
+            if buckets_len >= 1 {
+                region_buckets[buckets_len - 1].end_key = keys::origin_end_key(&end_key).to_vec();
             }
 
             // if we scan the whole range, we can update approximate size and keys with accurate value.