From 4581618b09ec5e0e6633f919f2ad20a4ec83495e Mon Sep 17 00:00:00 2001
From: Yongqiang YANG <98214048+dataroaring@users.noreply.github.com>
Date: Tue, 2 Jan 2024 14:08:09 +0800
Subject: [PATCH] [improvement](disk) pick disk randomly when usage is less
 than 0.7 (#29368)

---
 be/src/olap/data_dir.cpp       |  3 +--
 be/src/olap/data_dir.h         |  7 +++++++
 be/src/olap/storage_engine.cpp | 38 +++++++++++++++++++++-------------
 3 files changed, 32 insertions(+), 16 deletions(-)

diff --git a/be/src/olap/data_dir.cpp b/be/src/olap/data_dir.cpp
index 154673dfb4ff95..d19e77f6e8e096 100644
--- a/be/src/olap/data_dir.cpp
+++ b/be/src/olap/data_dir.cpp
@@ -870,8 +870,7 @@ size_t DataDir::tablet_size() const {
 }
 
 bool DataDir::reach_capacity_limit(int64_t incoming_data_size) {
-    double used_pct = (_disk_capacity_bytes - _available_bytes + incoming_data_size) /
-                      (double)_disk_capacity_bytes;
+    double used_pct = get_usage(incoming_data_size);
     int64_t left_bytes = _available_bytes - incoming_data_size;
     if (used_pct >= config::storage_flood_stage_usage_percent / 100.0 &&
         left_bytes <= config::storage_flood_stage_left_capacity_bytes) {
diff --git a/be/src/olap/data_dir.h b/be/src/olap/data_dir.h
index 84b79a8b2876b6..f8a405d4ffd4fd 100644
--- a/be/src/olap/data_dir.h
+++ b/be/src/olap/data_dir.h
@@ -133,6 +133,13 @@ class DataDir {
 
     void disks_compaction_num_increment(int64_t delta);
 
+    double get_usage(int64_t incoming_data_size) const {
+        return _disk_capacity_bytes == 0
+                       ? 0
+                       : (_disk_capacity_bytes - _available_bytes + incoming_data_size) /
+                                 (double)_disk_capacity_bytes;
+    }
+
     // Move tablet to trash.
     Status move_to_trash(const std::string& tablet_path);
 
diff --git a/be/src/olap/storage_engine.cpp b/be/src/olap/storage_engine.cpp
index 4b3827daf29394..33844e899290d4 100644
--- a/be/src/olap/storage_engine.cpp
+++ b/be/src/olap/storage_engine.cpp
@@ -455,28 +455,38 @@ std::vector<DataDir*> StorageEngine::get_stores_for_create_tablet(
         std::lock_guard<std::mutex> l(_store_lock);
         for (auto& it : _store_map) {
             if (it.second->is_used()) {
-                if (_available_storage_medium_type_count == 1 ||
-                    it.second->storage_medium() == storage_medium) {
+                if ((_available_storage_medium_type_count == 1 ||
+                     it.second->storage_medium() == storage_medium) &&
+                    !it.second->reach_capacity_limit(0)) {
                     stores.push_back(it.second);
                 }
             }
         }
     }
-    std::random_device rd;
-    std::mt19937 g(rd());
-    std::shuffle(stores.begin(), stores.end(), g);
-    // Two random choices
-    for (int i = 0; i < stores.size(); i++) {
-        int j = i + 1;
-        if (j < stores.size()) {
-            if (stores[i]->tablet_size() > stores[j]->tablet_size()) {
-                std::swap(stores[i], stores[j]);
-            }
-            std::shuffle(stores.begin() + j, stores.end(), g);
-        } else {
+
+    std::sort(stores.begin(), stores.end(),
+              [](DataDir* a, DataDir* b) { return a->get_usage(0) < b->get_usage(0); });
+
+    size_t seventy_percent_index = stores.size();
+    size_t eighty_five_percent_index = stores.size();
+    for (size_t index = 0; index < stores.size(); index++) {
+        // If the usage of the store is less than 70%, we choose disk randomly.
+        if (stores[index]->get_usage(0) > 0.7 && seventy_percent_index == stores.size()) {
+            seventy_percent_index = index;
+        }
+        if (stores[index]->get_usage(0) > 0.85 && eighty_five_percent_index == stores.size()) {
+            eighty_five_percent_index = index;
             break;
         }
     }
+
+    std::random_device rd;
+    std::mt19937 g(rd());
+    std::shuffle(stores.begin(), stores.begin() + seventy_percent_index, g);
+    std::shuffle(stores.begin() + seventy_percent_index, stores.begin() + eighty_five_percent_index,
+                 g);
+    std::shuffle(stores.begin() + eighty_five_percent_index, stores.end(), g);
+
     return stores;
 }