Update MemTableSet for allocated_bytes

JaySon-Huang · JaySon-Huang · commit c638bc7f7d31 · 2025-07-04T15:16:50.000+08:00
Signed-off-by: JaySon-Huang &lt;tshent@qq.com&gt;
diff --git a/dbms/src/Common/CurrentMetrics.cpp b/dbms/src/Common/CurrentMetrics.cpp
@@ -88,7 +88,10 @@
     M(ConnectionPoolSize)                       \
     M(MemoryTrackingQueryStorageTask)           \
     M(MemoryTrackingFetchPages)                 \
-    M(MemoryTrackingSharedColumnData)
+    M(MemoryTrackingSharedColumnData)           \
+    M(NumSegments)                              \
+    M(NumDeltaCache)                            \
+    M(BytesDeltaCache)
 
 namespace CurrentMetrics
 {
diff --git a/dbms/src/Common/CurrentMetrics.h b/dbms/src/Common/CurrentMetrics.h
@@ -18,7 +18,6 @@
 #include <stddef.h>
 
 #include <atomic>
-#include <cstdint>
 #include <utility>
 
 /** Allows to count number of simultaneously happening processes or current value of some metric.
diff --git a/dbms/src/DataStreams/AsynchronousBlockInputStream.h b/dbms/src/DataStreams/AsynchronousBlockInputStream.h
@@ -14,7 +14,6 @@
 
 #pragma once
 
-#include <Common/CurrentMetrics.h>
 #include <Common/MemoryTracker.h>
 #include <Common/wrapInvocable.h>
 #include <DataStreams/IProfilingBlockInputStream.h>
diff --git a/dbms/src/DataStreams/ParallelInputsProcessor.h b/dbms/src/DataStreams/ParallelInputsProcessor.h
@@ -14,7 +14,6 @@
 
 #pragma once
 
-#include <Common/CurrentMetrics.h>
 #include <Common/Logger.h>
 #include <Common/MPMCQueue.h>
 #include <Common/MemoryTracker.h>
@@ -24,10 +23,6 @@
 #include <common/logger_useful.h>
 
 #include <atomic>
-#include <list>
-#include <mutex>
-#include <queue>
-#include <thread>
 
 
 /** Allows to process multiple block input streams (sources) in parallel, using specified number of threads.
diff --git a/dbms/src/Dictionaries/CacheDictionary.h b/dbms/src/Dictionaries/CacheDictionary.h
@@ -16,14 +16,12 @@
 
 #include <Columns/ColumnString.h>
 #include <Common/ArenaWithFreeLists.h>
-#include <Common/CurrentMetrics.h>
 #include <Dictionaries/DictionaryStructure.h>
 #include <Dictionaries/IDictionary.h>
 #include <Dictionaries/IDictionarySource.h>
 
 #include <atomic>
 #include <chrono>
-#include <cmath>
 #include <ext/bit_cast.h>
 #include <map>
 #include <pcg_random.hpp>
diff --git a/dbms/src/Interpreters/ProcessList.h b/dbms/src/Interpreters/ProcessList.h
@@ -14,7 +14,6 @@
 
 #pragma once
 
-#include <Common/CurrentMetrics.h>
 #include <Common/MemoryTracker.h>
 #include <Common/Stopwatch.h>
 #include <Common/Throttler.h>
diff --git a/dbms/src/Interpreters/QueryPriorities.h b/dbms/src/Interpreters/QueryPriorities.h
@@ -14,7 +14,6 @@
 
 #pragma once
 
-#include <Common/CurrentMetrics.h>
 #include <Common/Stopwatch.h>
 
 #include <chrono>
diff --git a/dbms/src/Server/TCPHandler.h b/dbms/src/Server/TCPHandler.h
@@ -15,7 +15,6 @@
 #pragma once
 
 #include <Client/TimeoutSetter.h>
-#include <Common/CurrentMetrics.h>
 #include <Common/Stopwatch.h>
 #include <Common/getFQDNOrHostName.h>
 #include <Core/Protocol.h>
diff --git a/dbms/src/Storages/DeltaMerge/ColumnFile/ColumnFile.h b/dbms/src/Storages/DeltaMerge/ColumnFile/ColumnFile.h
@@ -138,7 +138,13 @@ class ColumnFile
     /// been persisted in the disk and their data will be immutable.
     virtual bool isAppendable() const { return false; }
     virtual void disableAppend() {}
-    virtual bool append(
+
+    struct AppendResult
+    {
+        bool success = false; // whether the append is successful
+        size_t new_alloc_bytes = 0; // the new allocated bytes after append
+    };
+    virtual AppendResult append(
         const DMContext & /*dm_context*/,
         const Block & /*data*/,
         size_t /*offset*/,
diff --git a/dbms/src/Storages/DeltaMerge/ColumnFile/ColumnFileInMemory.cpp b/dbms/src/Storages/DeltaMerge/ColumnFile/ColumnFileInMemory.cpp
@@ -65,24 +65,24 @@ ColumnFileReaderPtr ColumnFileInMemory::getReader(
     return std::make_shared<ColumnFileInMemoryReader>(*this, col_defs);
 }
 
-bool ColumnFileInMemory::append(
+ColumnFile::AppendResult ColumnFileInMemory::append(
     const DMContext & context,
     const Block & data,
     size_t offset,
     size_t limit,
     size_t data_bytes)
 {
     if (disable_append)
-        return false;
+        return AppendResult{false, 0};
 
     std::scoped_lock lock(cache->mutex);
     if (!isSameSchema(cache->block, data))
-        return false;
+        return AppendResult{false, 0};
 
     // check whether this instance overflows
     if (cache->block.rows() >= context.delta_cache_limit_rows
         || cache->block.bytes() >= context.delta_cache_limit_bytes)
-        return false;
+        return AppendResult{false, 0};
 
     size_t new_alloc_block_bytes = 0;
     size_t max_capacity = 0;
@@ -97,8 +97,8 @@ bool ColumnFileInMemory::append(
             // If the column is not enough, we reserve more space by 1.5 factor
             mutable_cache_col->reserveWithStrategy(mutable_cache_col->size() + limit, IColumn::ReserveStrategy::ScaleFactor1_5);
         }
-        max_capacity = std::max(max_capacity, mutable_cache_col->capacity());
         mutable_cache_col->insertRangeFrom(*col, offset, limit);
+        max_capacity = std::max(max_capacity, mutable_cache_col->capacity());
         new_alloc_block_bytes += mutable_cache_col->allocatedBytes() - alloc_bytes;
     }
 
@@ -113,7 +113,7 @@ bool ColumnFileInMemory::append(
         max_capacity,
         rows,
         bytes);
-    return true;
+    return AppendResult{true, new_alloc_block_bytes};
 }
 
 Block ColumnFileInMemory::readDataForFlush() const
diff --git a/dbms/src/Storages/DeltaMerge/ColumnFile/ColumnFileInMemory.h b/dbms/src/Storages/DeltaMerge/ColumnFile/ColumnFileInMemory.h
@@ -84,8 +84,13 @@ class ColumnFileInMemory : public ColumnFile
 
     bool isAppendable() const override { return !disable_append; }
     void disableAppend() override { disable_append = true; }
-    bool append(const DMContext & dm_context, const Block & data, size_t offset, size_t limit, size_t data_bytes)
-        override;
+
+    AppendResult append(
+        const DMContext & dm_context,
+        const Block & data,
+        size_t offset,
+        size_t limit,
+        size_t data_bytes) override;
 
     Block readDataForFlush() const;
 
diff --git a/dbms/src/Storages/DeltaMerge/Delta/ColumnFilePersistedSet.h b/dbms/src/Storages/DeltaMerge/Delta/ColumnFilePersistedSet.h
@@ -14,7 +14,6 @@
 
 #pragma once
 
-#include <Common/CurrentMetrics.h>
 #include <Common/Exception.h>
 #include <Core/Block.h>
 #include <IO/WriteHelpers.h>
diff --git a/dbms/src/Storages/DeltaMerge/Delta/MemTableSet.cpp b/dbms/src/Storages/DeltaMerge/Delta/MemTableSet.cpp
@@ -182,16 +182,16 @@ void MemTableSet::appendColumnFile(const ColumnFilePtr & column_file)
 void MemTableSet::appendToCache(DMContext & context, const Block & block, size_t offset, size_t limit)
 {
     // If the `column_files` is not empty, and the last `column_file` is a `ColumnInMemoryFile`, we will merge the newly block into the last `column_file`.
-    bool success = false;
+    ColumnFile::AppendResult append_res;
     size_t append_bytes = block.bytes(offset, limit);
     if (!column_files.empty())
     {
         auto & last_column_file = column_files.back();
         if (last_column_file->isAppendable())
-            success = last_column_file->append(context, block, offset, limit, append_bytes);
+            append_res = last_column_file->append(context, block, offset, limit, append_bytes);
     }
 
-    if (!success)
+    if (!append_res.success)
     {
         /// Otherwise, create a new `ColumnInMemoryFile` and write into it.
 
@@ -202,12 +202,13 @@ void MemTableSet::appendToCache(DMContext & context, const Block & block, size_t
         // Must append the empty `new_column_file` to `column_files` before appending data to it,
         // because `appendColumnFileInner` will update stats related to `column_files` but we will update stats relate to `new_column_file` here.
         appendColumnFileInner(new_column_file);
-        success = new_column_file->append(context, block, offset, limit, append_bytes);
-        if (unlikely(!success))
+        append_res = new_column_file->append(context, block, offset, limit, append_bytes);
+        if (unlikely(!append_res.success))
             throw Exception("Write to MemTableSet failed", ErrorCodes::LOGICAL_ERROR);
     }
     rows += limit;
     bytes += append_bytes;
+    allocated_bytes += append_res.new_alloc_bytes;
 }
 
 void MemTableSet::appendDeleteRange(const RowKeyRange & delete_range)
diff --git a/dbms/src/Storages/DeltaMerge/Segment.cpp b/dbms/src/Storages/DeltaMerge/Segment.cpp
@@ -128,6 +128,7 @@ extern const Metric DT_SnapshotOfPlaceIndex;
 extern const Metric DT_SnapshotOfReplayVersionChain;
 extern const Metric DT_SnapshotOfSegmentIngest;
 extern const Metric DT_SnapshotOfBitmapFilter;
+extern const Metric NumSegments;
 } // namespace CurrentMetrics
 
 namespace DB
@@ -295,7 +296,8 @@ Segment::Segment( //
     PageIdU64 next_segment_id_,
     const DeltaValueSpacePtr & delta_,
     const StableValueSpacePtr & stable_)
-    : epoch(epoch_)
+    : counter_holder(CurrentMetrics::NumSegments)
+    , epoch(epoch_)
     , rowkey_range(rowkey_range_)
     , is_common_handle(rowkey_range.is_common_handle)
     , rowkey_column_size(rowkey_range.rowkey_column_size)
diff --git a/dbms/src/Storages/DeltaMerge/Segment.h b/dbms/src/Storages/DeltaMerge/Segment.h
@@ -837,6 +837,8 @@ class Segment
 #else
 public:
 #endif
+
+    CurrentMetrics::Increment counter_holder;
     /// The version of this segment. After split / merge / mergeDelta / replaceData, epoch got increased by 1.
     const UInt64 epoch;
 
diff --git a/dbms/src/Storages/DeltaMerge/VersionChain/tests/mvcc_test_utils.h b/dbms/src/Storages/DeltaMerge/VersionChain/tests/mvcc_test_utils.h
@@ -14,7 +14,6 @@
 
 #pragma once
 
-#include <Common/CurrentMetrics.h>
 #include <Debug/TiFlashTestEnv.h>
 #include <Storages/DeltaMerge/DMContext.h>
 #include <Storages/DeltaMerge/Segment.h>
diff --git a/dbms/src/Storages/DeltaMerge/tests/gtest_segment.cpp b/dbms/src/Storages/DeltaMerge/tests/gtest_segment.cpp
@@ -666,7 +666,7 @@ try
     const size_t level = 1;
     for (size_t lvl = 0; lvl < level; ++lvl)
     {
-        size_t num_expected_segs = 1000;
+        size_t num_expected_segs = 500;
         // size_t num_expected_segs = 10;
         size_t progress_interval = 100;
         const auto lvl_beg_seg_id = segments.rbegin()->first;
