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->recalculateCapacity();
+            }
+            LOG_FMT_TRACE(log, "Blob begin to recalculate capability [blob_id={}]", blob_id);
+        }
+    }
 }
 
 std::pair<BlobFileId, BlobFileOffset> BlobStore::getPosFromStats(size_t size)
@@ -213,7 +231,7 @@ std::pair<BlobFileId, BlobFileOffset> BlobStore::getPosFromStats(size_t size)
     auto lock_stat = [size, this, &stat]() -> std::lock_guard<std::mutex> {
         auto lock_stats = blob_stats.lock();
         BlobFileId blob_file_id = INVALID_BLOBFILE_ID;
-        std::tie(stat, blob_file_id) = blob_stats.chooseStat(size, config.file_limit_size, lock_stats);
+        std::tie(stat, blob_file_id) = blob_stats.chooseStat(size, lock_stats);
 
         // No valid stat for puting data with `size`, create a new one
         if (stat == nullptr)
@@ -832,7 +850,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, const std::lock_guard<std::mutex> &)
 {
     BlobStatPtr stat_ptr = nullptr;
     double smallest_valid_rate = 2;
@@ -847,10 +865,9 @@ 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_valid_rate < smallest_valid_rate)
         {
-            smallest_valid_rate = stat->sm_valid_size;
+            smallest_valid_rate = stat->sm_valid_rate;
             stat_ptr = stat;
         }
     }
@@ -863,7 +880,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 +897,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;
 }
 
 /*********************
@@ -960,6 +982,12 @@ void BlobStore::BlobStats::BlobStat::recalculateSpaceMap()
     sm_total_size = total_size;
     sm_valid_size = valid_size;
     sm_valid_rate = valid_size * 1.0 / total_size;
+    recalculateCapacity();
+}
+
+void BlobStore::BlobStats::BlobStat::recalculateCapacity()
+{
+    sm_max_caps = smap->updateAccurateMaxCapacity();
 }
 
 } // namespace PS::V3

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

@@ -110,6 +110,12 @@ class BlobStore : private Allocator<false>
              * We still need to recalculate a `sm_total_size`/`sm_valid_size`/`sm_valid_rate`.
              */
             void recalculateSpaceMap();
+
+            /**
+             * The `sm_max_cap` is not accurate after GC removes out-of-date data, or after restoring from disk.
+             * Caller should call this function to update the `sm_max_cap` so that we can reuse the space in this BlobStat.
+             */
+            void recalculateCapacity();
         };
 
         using BlobStatPtr = std::shared_ptr<BlobStat>;
@@ -141,9 +147,9 @@ class BlobStore : private Allocator<false>
          * The `INVALID_BLOBFILE_ID` means that you don't need create a new `BlobFile`.
          * 
          */
-        std::pair<BlobStatPtr, BlobFileId> chooseStat(size_t buf_size, UInt64 file_limit_size, const std::lock_guard<std::mutex> &);
+        std::pair<BlobStatPtr, BlobFileId> chooseStat(size_t buf_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 accurate max capacity of the 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_blob_store.cpp

@@ -123,17 +123,17 @@ TEST_F(BlobStoreStatsTest, testStat)
 
     BlobStats stats(logger, config);
 
-    std::tie(stat, blob_file_id) = stats.chooseStat(10, BLOBFILE_LIMIT_SIZE, stats.lock());
+    std::tie(stat, blob_file_id) = stats.chooseStat(10, stats.lock());
     ASSERT_EQ(blob_file_id, 1);
     ASSERT_FALSE(stat);
 
     // still 0
-    std::tie(stat, blob_file_id) = stats.chooseStat(10, BLOBFILE_LIMIT_SIZE, stats.lock());
+    std::tie(stat, blob_file_id) = stats.chooseStat(10, stats.lock());
     ASSERT_EQ(blob_file_id, 1);
     ASSERT_FALSE(stat);
 
     stats.createStat(0, stats.lock());
-    std::tie(stat, blob_file_id) = stats.chooseStat(10, BLOBFILE_LIMIT_SIZE, stats.lock());
+    std::tie(stat, blob_file_id) = stats.chooseStat(10, stats.lock());
     ASSERT_EQ(blob_file_id, INVALID_BLOBFILE_ID);
     ASSERT_TRUE(stat);
 
@@ -210,7 +210,7 @@ TEST_F(BlobStoreStatsTest, testFullStats)
     ASSERT_LE(stat->sm_valid_rate, 1);
 
     // Won't choose full one
-    std::tie(stat, blob_file_id) = stats.chooseStat(100, BLOBFILE_LIMIT_SIZE, stats.lock());
+    std::tie(stat, blob_file_id) = stats.chooseStat(100, stats.lock());
     ASSERT_EQ(blob_file_id, 2);
     ASSERT_FALSE(stat);
 
@@ -228,7 +228,7 @@ TEST_F(BlobStoreStatsTest, testFullStats)
 
     // Then choose stat , it should return the stat id 3
     // Stat which id is 2 is full.
-    std::tie(stat, blob_file_id) = stats.chooseStat(100, BLOBFILE_LIMIT_SIZE, stats.lock());
+    std::tie(stat, blob_file_id) = stats.chooseStat(100, stats.lock());
     ASSERT_EQ(blob_file_id, 3);
     ASSERT_FALSE(stat);
 }

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->updateAccurateMaxCapacity(), 50);
+    }
+
+    {
+        auto smap = SpaceMap::createSpaceMap(test_type, 0, 100);
+        ASSERT_TRUE(smap->markUsed(0, 100));
+
+        ASSERT_EQ(smap->updateAccurateMaxCapacity(), 0);
+    }
+}
+
 INSTANTIATE_TEST_CASE_P(
     Type,
     SpaceMapTest,