From 1d98913f90379345e80e75c746b8ffa34e2cbb0f Mon Sep 17 00:00:00 2001
From: Azat Khuzhin <a3at.mail@gmail.com>
Date: Wed, 18 May 2022 16:54:51 +0300
Subject: [PATCH] Disks: Reduce number of statfs() calls for least_used disk
 load balancing policy

Signed-off-by: Azat Khuzhin <a.khuzhin@semrush.com>
---
 src/Disks/VolumeJBOD.cpp |  8 ++++----
 src/Disks/VolumeJBOD.h   | 25 ++++++++++++++++++++-----
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/src/Disks/VolumeJBOD.cpp b/src/Disks/VolumeJBOD.cpp
index f202dda03a2..401822fc901 100644
--- a/src/Disks/VolumeJBOD.cpp
+++ b/src/Disks/VolumeJBOD.cpp
@@ -94,7 +94,7 @@ DiskPtr VolumeJBOD::getDisk(size_t /* index */) const
         case VolumeLoadBalancing::LEAST_USED:
         {
             std::lock_guard lock(mutex);
-            return disks_by_size.top();
+            return disks_by_size.top().disk;
         }
     }
     __builtin_unreachable();
@@ -128,10 +128,10 @@ ReservationPtr VolumeJBOD::reserve(UInt64 bytes)
         {
             std::lock_guard lock(mutex);
 
-            DiskPtr disk = disks_by_size.top();
-            ReservationPtr reservation = disk->reserve(bytes);
-
+            DiskWithSize disk = disks_by_size.top();
             disks_by_size.pop();
+
+            ReservationPtr reservation = disk.reserve(bytes);
             disks_by_size.push(disk);
 
             return reservation;
diff --git a/src/Disks/VolumeJBOD.h b/src/Disks/VolumeJBOD.h
index 5c9d2921ca5..0fb80422d55 100644
--- a/src/Disks/VolumeJBOD.h
+++ b/src/Disks/VolumeJBOD.h
@@ -64,12 +64,27 @@ public:
     bool are_merges_avoided = true;
 
 private:
-    struct DiskBySize
+    struct DiskWithSize
     {
-        bool operator()(const DiskPtr & lhs, const DiskPtr & rhs) const
+        DiskPtr disk;
+        uint64_t free_size = 0;
+
+        DiskWithSize(DiskPtr disk_)
+            : disk(disk_)
+            , free_size(disk->getUnreservedSpace())
+        {}
+
+        bool operator<(const DiskWithSize & rhs) const
         {
-            /// TODO: avoid getAvailableSpace() calls
-            return lhs->getUnreservedSpace() < rhs->getUnreservedSpace();
+            return free_size < rhs.free_size;
+        }
+
+        ReservationPtr reserve(uint64_t bytes)
+        {
+            ReservationPtr reservation = disk->reserve(bytes);
+            if (reservation)
+                free_size -= bytes;
+            return reservation;
         }
     };
 
@@ -77,7 +92,7 @@ private:
     /// Index of last used disk, for load_balancing=round_robin
     mutable std::atomic<size_t> last_used = 0;
     /// Priority queue of disks sorted by size, for load_balancing=least_used
-    mutable std::priority_queue<DiskPtr, std::vector<DiskPtr>, DiskBySize> disks_by_size;
+    mutable std::priority_queue<DiskWithSize> disks_by_size;
 
     /// True if parts on this volume participate in merges according to START/STOP MERGES ON VOLUME.
     std::atomic<std::optional<bool>> are_merges_avoided_user_override{std::nullopt};