Storages: Track memory for PageDirectory

dbms/src/Common/TiFlashMetrics.h

@@ -212,6 +212,11 @@ static_assert(RAFT_REGION_BIG_WRITE_THRES * 4 < RAFT_REGION_BIG_WRITE_MAX, "Inva
       "The freshness of tiflash data with tikv data",                                                                               \
       Histogram,                                                                                                                    \
       F(type_syncing_data_freshness, {{"type", "data_freshness"}}, ExpBuckets{0.001, 2, 20}))                                       \
+    M(tiflash_memory_usage_by_class,                                                                                                \
+      "TiFlash memory consumes by class",                                                                                           \
+      Gauge,                                                                                                                        \
+      F(type_uni_page_ids, {"type", "uni_page_ids"}),                                                                               \
+      F(type_versioned_entry_or_delete, {"type", "versioned_entry_or_delete"}))                                                     \
     M(tiflash_storage_read_tasks_count, "Total number of storage engine read tasks", Counter)                                       \
     M(tiflash_storage_command_count,                                                                                                \
       "Total number of storage's command, such as delete range / shutdown /startup",                                                \

dbms/src/Interpreters/AsynchronousMetrics.cpp

@@ -32,6 +32,7 @@
 #include <Storages/Page/FileUsage.h>
 #include <Storages/Page/PageConstants.h>
 #include <Storages/Page/PageStorage.h>
+#include <Storages/Page/PageStorageMemorySummary.h>
 #include <Storages/Page/V3/Universal/UniversalPageStorageService.h>
 #include <Storages/StorageDeltaMerge.h>
 #include <common/config_common.h>
@@ -300,6 +301,7 @@ void AsynchronousMetrics::update()
         set("LogNums", usage.total_log_file_num);
         set("LogDiskBytes", usage.total_log_disk_size);
         set("PagesInMem", usage.num_pages);
+        set("VersionedEntries", DB::PS::PageStorageMemorySummary::num_versioned_entry_or_delete.load());
     }
 
     if (context.getSharedContextDisagg()->isDisaggregatedStorageMode())

dbms/src/Storages/KVStore/FFI/JointThreadAllocInfo.cpp

@@ -17,6 +17,7 @@
 #include <Common/setThreadName.h>
 #include <Storages/KVStore/FFI/JointThreadAllocInfo.h>
 #include <Storages/KVStore/FFI/ProxyFFI.h>
+#include <Storages/Page/V3/PageDirectory.h>
 
 #include <magic_enum.hpp>
 #include <mutex>
@@ -47,6 +48,7 @@ void JointThreadInfoJeallocMap::recordThreadAllocInfo()
 {
     recordThreadAllocInfoForProxy();
     recordThreadAllocInfoForStorage();
+    recordClassdAlloc();
 }
 
 JointThreadInfoJeallocMap::~JointThreadInfoJeallocMap()
@@ -269,4 +271,13 @@ void JointThreadInfoJeallocMap::accessStorageMap(std::function<void(const AllocM
     f(storage_map);
 }
 
+void JointThreadInfoJeallocMap::recordClassdAlloc()
+{
+    GET_METRIC(tiflash_memory_usage_by_class, type_uni_page_ids)
+        .Set(PS::PageStorageMemorySummary::mem_sum_uni_page_ids.load());
+    GET_METRIC(tiflash_memory_usage_by_class, type_versioned_entry_or_delete)
+        .Set(PS::PageStorageMemorySummary::mem_sum_versioned_entry_or_delete.load());
+}
+
+
 } // namespace DB

dbms/src/Storages/KVStore/FFI/JointThreadAllocInfo.h

@@ -120,6 +120,8 @@ class JointThreadInfoJeallocMap
         uint64_t value,
         char aggregate_delimer);
 
+    static void recordClassdAlloc();
+
 private:
     mutable std::shared_mutex memory_allocation_mut;
     AllocMap proxy_map;

dbms/src/Storages/KVStore/FFI/SSTReader.h

@@ -158,7 +158,12 @@ class MultiSSTReader : public SSTReader
             {
                 // We don't drop if mono is the last instance for safety,
                 // and it will be dropped as MultiSSTReader is dropped.
-                LOG_INFO(log, "Open sst file {}", buffToStrView(args[current].path));
+                LOG_INFO(
+                    log,
+                    "Open sst file {}, range={}, current={}",
+                    buffToStrView(args[current].path),
+                    range->toDebugString(),
+                    current);
                 mono = initer(proxy_helper, args[current], range, split_id);
             }
         }

dbms/src/Storages/KVStore/MultiRaft/PrehandleSnapshot.cpp

@@ -509,10 +509,11 @@ void executeParallelTransform(
         split_key_count);
     LOG_INFO(
         log,
-        "Parallel prehandling for single big region, range={}, split keys={}, region_id={}",
+        "Parallel prehandling for single big region, range={}, split keys={}, region_id={}, snaps={}",
         new_region->getRange()->toDebugString(),
         split_key_count,
-        new_region->id());
+        new_region->id(),
+        snaps.len);
     Stopwatch watch;
     // Make sure the queue is bigger than `split_key_count`, otherwise `addTask` may fail.
     auto async_tasks = SingleSnapshotAsyncTasks(split_key_count, split_key_count, split_key_count + 5);
@@ -714,9 +715,10 @@ PrehandleResult KVStore::preHandleSSTsToDTFiles(
             {
                 LOG_INFO(
                     log,
-                    "Single threaded prehandling for single region, range={} region_id={}",
+                    "Single threaded prehandling for single region, range={} region_id={} snaps={}",
                     new_region->getRange()->toDebugString(),
-                    new_region->id());
+                    new_region->id(),
+                    snaps.len);
                 std::tie(result, prehandle_result) = executeTransform(log, prehandle_ctx, new_region, sst_stream);
             }
             else

dbms/src/Storages/KVStore/tests/gtest_new_kvstore.cpp

@@ -20,6 +20,11 @@
 #include <Storages/KVStore/Region.h>
 #include <Storages/KVStore/Utils/AsyncTasks.h>
 #include <Storages/KVStore/tests/region_kvstore_test.h>
+#include <Storages/Page/V3/PageDefines.h>
+#include <Storages/Page/V3/PageDirectory.h>
+#include <Storages/Page/V3/PageDirectoryFactory.h>
+#include <Storages/Page/V3/PageEntriesEdit.h>
+#include <Storages/Page/V3/Universal/UniversalPageIdFormatImpl.h>
 #include <Storages/RegionQueryInfo.h>
 #include <TiDB/Schema/SchemaSyncService.h>
 #include <TiDB/Schema/TiDBSchemaManager.h>

dbms/src/Storages/Page/PageStorageMemorySummary.h

@@ -0,0 +1,28 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <atomic>
+
+namespace DB::PS
+{
+struct PageStorageMemorySummary
+{
+    static inline std::atomic_int64_t mem_sum_uni_page_ids{0};
+    static inline std::atomic_int64_t mem_sum_versioned_entry_or_delete{0};
+    static inline std::atomic_int64_t num_versioned_entry_or_delete{0};
+};
+
+} // namespace DB::PS

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

@@ -222,31 +222,47 @@ struct EntryOrDelete
     MultiVersionRefCount being_ref_count;
     std::optional<PageEntryV3> entry;
 
-    static EntryOrDelete newDelete()
+    EntryOrDelete(const EntryOrDelete & other)
+        : being_ref_count(other.being_ref_count)
+        , entry(other.entry)
     {
-        return EntryOrDelete{
-            .entry = std::nullopt,
-        };
-    };
-    static EntryOrDelete newNormalEntry(const PageEntryV3 & entry)
+        PageStorageMemorySummary::num_versioned_entry_or_delete.fetch_add(1);
+        if (entry)
+            PageStorageMemorySummary::mem_sum_versioned_entry_or_delete.fetch_add(sizeof(PageEntryV3));
+    }
+    EntryOrDelete() { PageStorageMemorySummary::num_versioned_entry_or_delete.fetch_add(1); }
+    EntryOrDelete(std::optional<PageEntryV3> entry_)
+        : entry(std::move(entry_))
+    {
+        PageStorageMemorySummary::num_versioned_entry_or_delete.fetch_add(1);
+        if (entry)
+            PageStorageMemorySummary::mem_sum_versioned_entry_or_delete.fetch_add(sizeof(PageEntryV3));
+    }
+    EntryOrDelete(MultiVersionRefCount being_ref_count_, std::optional<PageEntryV3> entry_)
+        : being_ref_count(being_ref_count_)
+        , entry(std::move(entry_))
     {
-        return EntryOrDelete{
-            .entry = entry,
-        };
+        PageStorageMemorySummary::num_versioned_entry_or_delete.fetch_add(1);
+        if (entry)
+            PageStorageMemorySummary::mem_sum_versioned_entry_or_delete.fetch_add(sizeof(PageEntryV3));
     }
+    ~EntryOrDelete()
+    {
+        PageStorageMemorySummary::num_versioned_entry_or_delete.fetch_sub(1);
+        if (entry)
+            PageStorageMemorySummary::mem_sum_versioned_entry_or_delete.fetch_sub(sizeof(PageEntryV3));
+    }
+
+    static EntryOrDelete newDelete() { return EntryOrDelete(std::nullopt); };
+    static EntryOrDelete newNormalEntry(const PageEntryV3 & entry) { return EntryOrDelete(entry); }
     static EntryOrDelete newReplacingEntry(const EntryOrDelete & ori_entry, const PageEntryV3 & entry)
     {
-        return EntryOrDelete{
-            .being_ref_count = ori_entry.being_ref_count,
-            .entry = entry,
-        };
+        return EntryOrDelete(ori_entry.being_ref_count, entry);
     }
 
     static EntryOrDelete newFromRestored(PageEntryV3 entry, const PageVersion & ver, Int64 being_ref_count)
     {
-        auto result = EntryOrDelete{
-            .entry = entry,
-        };
+        auto result = EntryOrDelete(std::move(entry));
         result.being_ref_count.restoreFrom(ver, being_ref_count);
         return result;
     }

dbms/src/Storages/Page/V3/PageDirectory/PageIdTrait.h

@@ -29,6 +29,7 @@ struct PageIdTrait
     static inline PageIdU64 getU64ID(const PageId & page_id) { return page_id.low; }
     static inline Prefix getPrefix(const PageId & page_id) { return page_id.high; }
     static inline PageIdU64 getPageMapKey(const PageId & page_id) { return page_id.low; }
+    static inline size_t getPageIDSize(const PageId & page_id) { return sizeof(page_id); }
 };
 } // namespace u128
 namespace universal
@@ -45,6 +46,8 @@ struct PageIdTrait
     static Prefix getPrefix(const PageId & page_id);
 
     static inline PageId getPageMapKey(const PageId & page_id) { return page_id; }
+
+    static inline size_t getPageIDSize(const PageId & page_id) { return page_id.size(); }
 };
 } // namespace universal
 } // namespace DB::PS::V3

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

@@ -15,6 +15,7 @@
 #pragma once
 
 #include <Common/Exception.h>
+#include <Storages/Page/PageStorageMemorySummary.h>
 #include <Storages/Page/V3/PageDefines.h>
 #include <Storages/Page/V3/PageEntryCheckpointInfo.h>
 #include <fmt/format.h>

dbms/src/Storages/Page/V3/Universal/UniversalPageId.h

@@ -18,26 +18,49 @@
 #include <IO/Buffer/WriteBufferFromString.h>
 #include <IO/WriteHelpers.h>
 #include <Storages/Page/PageDefinesBase.h>
+#include <Storages/Page/PageStorageMemorySummary.h>
 
 namespace DB
 {
 class UniversalPageId final
 {
 public:
-    UniversalPageId() = default;
+    UniversalPageId() { PS::PageStorageMemorySummary::mem_sum_uni_page_ids.fetch_add(id.size()); }
+    UniversalPageId(const UniversalPageId & other)
+        : id(other.id)
+    {
+        PS::PageStorageMemorySummary::mem_sum_uni_page_ids.fetch_add(id.size());
+    }
 
     UniversalPageId(String id_) // NOLINT(google-explicit-constructor)
         : id(std::move(id_))
-    {}
+    {
+        PS::PageStorageMemorySummary::mem_sum_uni_page_ids.fetch_add(id.size());
+    }
     UniversalPageId(const char * id_) // NOLINT(google-explicit-constructor)
         : id(id_)
-    {}
+    {
+        PS::PageStorageMemorySummary::mem_sum_uni_page_ids.fetch_add(id.size());
+    }
     UniversalPageId(const char * id_, size_t sz_)
         : id(id_, sz_)
-    {}
+    {
+        PS::PageStorageMemorySummary::mem_sum_uni_page_ids.fetch_add(id.size());
+    }
+
+    ~UniversalPageId() { PS::PageStorageMemorySummary::mem_sum_uni_page_ids.fetch_sub(id.size()); }
 
     UniversalPageId & operator=(String && id_) noexcept
     {
+        if (id.size() == id_.size()) {}
+        else if (id.size() > id_.size())
+        {
+            PS::PageStorageMemorySummary::mem_sum_uni_page_ids.fetch_sub(id.size() - id_.size());
+        }
+        else
+        {
+            PS::PageStorageMemorySummary::mem_sum_uni_page_ids.fetch_add(id_.size() - id.size());
+        }
         id.swap(id_);
         return *this;
     }