Fix max_caps may not been update after GC

dbms/src/Storages/Page/V3/BlobStore.cpp

@@ -192,18 +192,36 @@ PageEntriesEdit BlobStore::write(DB::WriteBatch & wb, const WriteLimiterPtr & wr
 
 void BlobStore::remove(const PageEntriesV3 & del_entries)
 {
+    std::set<BlobFileId> blob_updated;
     for (const auto & entry : del_entries)
     {
+        blob_updated.insert(entry.file_id);
         // External page size is 0
         if (entry.size == 0)
         {
             continue;
-            // throw Exception(fmt::format("Invaild entry. entry size 0. [id={}] [offset={}]",
-            //                             entry.file_id,
-            //                             entry.offset));
         }
         removePosFromStats(entry.file_id, entry.offset, entry.size);
     }
+
+    // After we remove postion of blob, we need recalculate the blob.
+    for (const auto & blob_id : blob_updated)
+    {
+        const auto & stat = blob_stats.blobIdToStat(blob_id,
+                                                    /*restore_if_not_exist*/ false,
+                                                    /*ignore_not_exist*/ true);
+
+        // Some of blob may been removed.
+        // So if we can't use id find blob, just ignore it.
+        if (stat)
+        {
+            {
+                auto lock = stat->lock();
+                stat->recalculateCapability();
+            }
+            LOG_FMT_TRACE(log, "Blob begin to recalculate capability [blob_id={}]", blob_id);
+        }
+    }
 }
 
 std::pair<BlobFileId, BlobFileOffset> BlobStore::getPosFromStats(size_t size)
@@ -747,6 +765,7 @@ void BlobStore::BlobStats::restore()
     for (const auto & stat : stats_map)
     {
         stat->recalculateSpaceMap();
+        stat->recalculateCapability();
         max_restored_file_id = std::max(stat->id, max_restored_file_id);
         existing_file_ids.insert(stat->id);
     }
@@ -832,7 +851,7 @@ void BlobStore::BlobStats::eraseStat(BlobFileId blob_file_id, const std::lock_gu
     eraseStat(std::move(stat), lock);
 }
 
-std::pair<BlobStatPtr, BlobFileId> BlobStore::BlobStats::chooseStat(size_t buf_size, UInt64 file_limit_size, const std::lock_guard<std::mutex> &)
+std::pair<BlobStatPtr, BlobFileId> BlobStore::BlobStats::chooseStat(size_t buf_size, UInt64 /*file_limit_size*/, const std::lock_guard<std::mutex> &)
 {
     BlobStatPtr stat_ptr = nullptr;
     double smallest_valid_rate = 2;
@@ -847,10 +866,10 @@ std::pair<BlobStatPtr, BlobFileId> BlobStore::BlobStats::chooseStat(size_t buf_s
     {
         if (!stat->isReadOnly()
             && stat->sm_max_caps >= buf_size
-            && stat->sm_total_size + buf_size < file_limit_size
+            // && stat->sm_total_size + buf_size < file_limit_size
             && stat->sm_valid_rate < smallest_valid_rate)
         {
-            smallest_valid_rate = stat->sm_valid_size;
+            smallest_valid_rate = stat->sm_valid_rate;
             stat_ptr = stat;
         }
     }
@@ -863,7 +882,7 @@ std::pair<BlobStatPtr, BlobFileId> BlobStore::BlobStats::chooseStat(size_t buf_s
     return std::make_pair(stat_ptr, INVALID_BLOBFILE_ID);
 }
 
-BlobStatPtr BlobStore::BlobStats::blobIdToStat(BlobFileId file_id, bool restore_if_not_exist)
+BlobStatPtr BlobStore::BlobStats::blobIdToStat(BlobFileId file_id, bool restore_if_not_exist, bool ignore_not_exist)
 {
     auto guard = lock();
     for (auto & stat : stats_map)
@@ -880,9 +899,14 @@ BlobStatPtr BlobStore::BlobStats::blobIdToStat(BlobFileId file_id, bool restore_
         return createStatNotCheckingRoll(file_id, guard);
     }
 
-    throw Exception(fmt::format("Can't find BlobStat with [blob_id={}]",
-                                file_id),
-                    ErrorCodes::LOGICAL_ERROR);
+    if (!ignore_not_exist)
+    {
+        throw Exception(fmt::format("Can't find BlobStat with [blob_id={}]",
+                                    file_id),
+                        ErrorCodes::LOGICAL_ERROR);
+    }
+
+    return nullptr;
 }
 
 /*********************
@@ -962,5 +986,10 @@ void BlobStore::BlobStats::BlobStat::recalculateSpaceMap()
     sm_valid_rate = valid_size * 1.0 / total_size;
 }
 
+void BlobStore::BlobStats::BlobStat::recalculateCapability()
+{
+    sm_max_caps = smap->updateAccurateMaxCapacity();
+}
+
 } // namespace PS::V3
 } // namespace DB

dbms/src/Storages/Page/V3/BlobStore.h

@@ -110,6 +110,14 @@ class BlobStore : private Allocator<false>
              * We still need to recalculate a `sm_total_size`/`sm_valid_size`/`sm_valid_rate`.
              */
             void recalculateSpaceMap();
+
+            /**
+             * After GC we removed out of date data.
+             * But we can't calculate the `sm_max_cap`.It can only be updated through spacemap.
+             * 
+             * Also after restore, we need call this method to recalculate the `sm_max_cap`.
+             */
+            void recalculateCapability();
         };
 
         using BlobStatPtr = std::shared_ptr<BlobStat>;
@@ -143,7 +151,7 @@ class BlobStore : private Allocator<false>
          */
         std::pair<BlobStatPtr, BlobFileId> chooseStat(size_t buf_size, UInt64 file_limit_size, const std::lock_guard<std::mutex> &);
 
-        BlobStatPtr blobIdToStat(BlobFileId file_id, bool restore_if_not_exist = false);
+        BlobStatPtr blobIdToStat(BlobFileId file_id, bool restore_if_not_exist = false, bool ignore_not_exist = false);
 
         std::list<BlobStatPtr> getStats() const
         {

dbms/src/Storages/Page/V3/spacemap/SpaceMap.h

@@ -82,6 +82,10 @@ class SpaceMap
      */
     virtual UInt64 getRightMargin() = 0;
 
+    /**
+     * Get the max of capability of space map.
+     */
+    virtual UInt64 updateAccurateMaxCapacity() = 0;
 
     /**
      * Return the size of file and the size contains valid data.

dbms/src/Storages/Page/V3/spacemap/SpaceMapRBTree.cpp

@@ -650,6 +650,34 @@ std::pair<UInt64, UInt64> RBTreeSpaceMap::searchInsertOffset(size_t size)
     return std::make_pair(offset, max_cap);
 }
 
+UInt64 RBTreeSpaceMap::updateAccurateMaxCapacity()
+{
+    struct rb_node * node = nullptr;
+    struct SmapRbEntry * entry;
+    UInt64 max_offset = 0;
+    UInt64 max_cap = 0;
+
+    node = rb_tree_first(&rb_tree->root);
+    if (node == nullptr)
+    {
+        return max_cap;
+    }
+
+    for (; node != nullptr; node = rb_tree_next(node))
+    {
+        entry = node_to_entry(node);
+        if (entry->count > max_cap)
+        {
+            max_offset = entry->start;
+            max_cap = entry->count;
+        }
+    }
+
+    biggest_range = max_offset;
+    biggest_cap = max_cap;
+    return max_cap;
+}
+
 std::pair<UInt64, UInt64> RBTreeSpaceMap::getSizes() const
 {
     struct rb_node * node = rb_tree_last(&rb_tree->root);

dbms/src/Storages/Page/V3/spacemap/SpaceMapRBTree.h

@@ -28,6 +28,8 @@ class RBTreeSpaceMap
 
     std::pair<UInt64, UInt64> searchInsertOffset(size_t size) override;
 
+    UInt64 updateAccurateMaxCapacity() override;
+
     std::pair<UInt64, UInt64> getSizes() const override;
 
     UInt64 getRightMargin() override;

dbms/src/Storages/Page/V3/spacemap/SpaceMapSTDMap.h

@@ -243,6 +243,25 @@ class STDMapSpaceMap
         return std::make_pair(offset, hint_biggest_cap);
     }
 
+    UInt64 updateAccurateMaxCapacity() override
+    {
+        UInt64 max_offset = 0;
+        UInt64 max_cap = 0;
+
+        for (const auto & [start, size] : free_map)
+        {
+            if (size > max_cap)
+            {
+                max_cap = size;
+                max_offset = start;
+            }
+        }
+        hint_biggest_offset = max_offset;
+        hint_biggest_cap = max_cap;
+
+        return max_cap;
+    }
+
     bool markFreeImpl(UInt64 offset, size_t length) override
     {
         auto it = free_map.find(offset);

dbms/src/Storages/Page/V3/tests/gtest_free_map.cpp

@@ -356,6 +356,25 @@ TEST_P(SpaceMapTest, TestGetSizes)
     }
 }
 
+
+TEST_P(SpaceMapTest, TestGetMaxCap)
+{
+    {
+        auto smap = SpaceMap::createSpaceMap(test_type, 0, 100);
+        ASSERT_TRUE(smap->markUsed(50, 10));
+        ASSERT_TRUE(smap->markUsed(80, 10));
+
+        ASSERT_EQ(smap->getMaxCapability(), 50);
+    }
+
+    {
+        auto smap = SpaceMap::createSpaceMap(test_type, 0, 100);
+        ASSERT_TRUE(smap->markUsed(0, 100));
+
+        ASSERT_EQ(smap->getMaxCapability(), 0);
+    }
+}
+
 INSTANTIATE_TEST_CASE_P(
     Type,
     SpaceMapTest,