From b4a4467521a517e6b16e33909818e628193fc01c Mon Sep 17 00:00:00 2001
From: yangjundong <1047934838@qq.com>
Date: Tue, 26 Nov 2024 05:38:24 +0000
Subject: [PATCH] [feat][store] Record region size when auto split

---
 src/split/split_checker.cc | 23 +++++++++++++----------
 src/split/split_checker.h  |  8 ++++----
 2 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/src/split/split_checker.cc b/src/split/split_checker.cc
index faf3e95d1..cc5152271 100644
--- a/src/split/split_checker.cc
+++ b/src/split/split_checker.cc
@@ -98,11 +98,10 @@ void MergedIterator::Next(IteratorPtr iter, int iter_pos) {
 
 // base physics key, contain key of multi version.
 std::string HalfSplitChecker::SplitKey(store::RegionPtr region, const pb::common::Range& range,
-                                       const std::vector<std::string>& cf_names, uint32_t& count) {
+                                       const std::vector<std::string>& cf_names, uint32_t& count, int64_t& size) {
   MergedIterator iter(raw_engine_, cf_names, range.end_key());
   iter.Seek(range.start_key());
 
-  int64_t size = 0;
   int64_t chunk_size = 0;
   std::string prev_key;
   std::vector<std::string> keys;
@@ -136,11 +135,10 @@ std::string HalfSplitChecker::SplitKey(store::RegionPtr region, const pb::common
 
 // base physics key, contain key of multi version.
 std::string SizeSplitChecker::SplitKey(store::RegionPtr region, const pb::common::Range& range,
-                                       const std::vector<std::string>& cf_names, uint32_t& count) {
+                                       const std::vector<std::string>& cf_names, uint32_t& count, int64_t& size) {
   MergedIterator iter(raw_engine_, cf_names, range.end_key());
   iter.Seek(range.start_key());
 
-  int64_t size = 0;
   std::string prev_key;
   std::string split_key;
   bool is_split = false;
@@ -168,11 +166,10 @@ std::string SizeSplitChecker::SplitKey(store::RegionPtr region, const pb::common
 
 // base logic key, ignore key of multi version.
 std::string KeysSplitChecker::SplitKey(store::RegionPtr region, const pb::common::Range& range,
-                                       const std::vector<std::string>& cf_names, uint32_t& count) {
+                                       const std::vector<std::string>& cf_names, uint32_t& count, int64_t& size) {
   MergedIterator iter(raw_engine_, cf_names, range.end_key());
   iter.Seek(range.start_key());
 
-  int64_t size = 0;
   int64_t split_key_count = 0;
   std::string prev_key;
   std::string split_key;
@@ -264,7 +261,8 @@ void SplitCheckTask::SplitCheck() {
   DINGO_LOG(INFO) << fmt::format("[split.check][region({})] Will check SplitKey for raw_range{} cf_names({})",
                                  region_->Id(), Helper::RangeToString(plain_range), Helper::VectorToString(cf_names));
   uint32_t key_count = 0;
-  std::string encode_split_key = split_checker_->SplitKey(region_, encode_range, cf_names, key_count);
+  int64_t size = 0;
+  std::string encode_split_key = split_checker_->SplitKey(region_, encode_range, cf_names, key_count, size);
 
   int64_t ts = 0;
   std::string plain_split_key;
@@ -273,9 +271,14 @@ void SplitCheckTask::SplitCheck() {
   }
 
   // Update region key count metrics.
-  if (region_metrics_ != nullptr && key_count > 0) {
-    region_metrics_->SetKeyCount(key_count);
-    region_metrics_->SetNeedUpdateKeyCount(false);
+  if (region_metrics_ != nullptr) {
+    if (key_count > 0) {
+      region_metrics_->SetKeyCount(key_count);
+      region_metrics_->SetNeedUpdateKeyCount(false);
+    }
+    if (size > 0) {
+      region_metrics_->SetRegionSize(size);
+    }
   }
 
   bool need_split = true;
diff --git a/src/split/split_checker.h b/src/split/split_checker.h
index c012812f6..3f5ed0832 100644
--- a/src/split/split_checker.h
+++ b/src/split/split_checker.h
@@ -85,7 +85,7 @@ class SplitChecker {
 
   // Calculate region split key.
   virtual std::string SplitKey(store::RegionPtr region, const pb::common::Range& range,
-                               const std::vector<std::string>& cf_names, uint32_t& count) = 0;
+                               const std::vector<std::string>& cf_names, uint32_t& count, int64_t& size) = 0;
 
  private:
   Policy policy_;
@@ -103,7 +103,7 @@ class HalfSplitChecker : public SplitChecker {
 
   // base physics key, contain key of multi version.
   std::string SplitKey(store::RegionPtr region, const pb::common::Range& range,
-                       const std::vector<std::string>& cf_names, uint32_t& count) override;
+                       const std::vector<std::string>& cf_names, uint32_t& count, int64_t& size) override;
 
  private:
   // Split region when exceed the split_threshold_size.
@@ -125,7 +125,7 @@ class SizeSplitChecker : public SplitChecker {
 
   // base physics key, contain key of multi version.
   std::string SplitKey(store::RegionPtr region, const pb::common::Range& range,
-                       const std::vector<std::string>& cf_names, uint32_t& count) override;
+                       const std::vector<std::string>& cf_names, uint32_t& count, int64_t& size) override;
 
  private:
   // Split when region exceed the split_size.
@@ -147,7 +147,7 @@ class KeysSplitChecker : public SplitChecker {
 
   // base logic key, ignore key of multi version.
   std::string SplitKey(store::RegionPtr region, const pb::common::Range& range,
-                       const std::vector<std::string>& cf_names, uint32_t& count) override;
+                       const std::vector<std::string>& cf_names, uint32_t& count, int64_t& size) override;
 
  private:
   // Split when region key number exceed split_key_number.