Storages: Fix MVCC bitmap read bytes estimation (#10378) (#10386)

ref #10380

Signed-off-by: ti-chi-bot <ti-community-prow-bot@tidb.io>

Co-authored-by: jinhelin <linjinhe33@gmail.com>

Author: ti-chi-bot

dbms/src/Storages/DeltaMerge/Segment.cpp

@@ -932,6 +932,62 @@ SegmentSnapshotPtr Segment::createSnapshot(const DMContext & dm_context, bool fo
         Logger::get(dm_context.tracing_id));
 }
 
+std::vector<DMFilePackFilter> Segment::loadDMFilePackFilters(
+    const DMFiles & dmfiles,
+    const MinMaxIndexCachePtr & index_cache,
+    const RowKeyRanges & rowkey_ranges,
+    const RSOperatorPtr & rs_filter,
+    const FileProviderPtr & file_provider,
+    const ReadLimiterPtr & read_limiter,
+    const ScanContextPtr & scan_context,
+    const String & tracing_id,
+    const ReadTag & read_tag)
+{
+    std::vector<DMFilePackFilter> pack_filters;
+    pack_filters.reserve(dmfiles.size());
+    for (const auto & dmfile : dmfiles)
+    {
+        auto pack_filter = DMFilePackFilter::loadFrom(
+            dmfile,
+            index_cache,
+            /*set_cache_if_miss*/ true,
+            rowkey_ranges,
+            rs_filter,
+            /*read_packs*/ {},
+            file_provider,
+            read_limiter,
+            scan_context,
+            tracing_id,
+            read_tag);
+        pack_filters.emplace_back(std::move(pack_filter));
+    }
+    return pack_filters;
+}
+
+UInt64 Segment::estimatedBytesOfInternalColumns(
+    const SegmentSnapshotPtr & read_snap,
+    std::vector<DMFilePackFilter> & pack_filters,
+    UInt64 start_ts)
+{
+    // stable->getDMFiles() at least return one DMFile.
+    const auto & dmfiles = read_snap->stable->getDMFiles();
+    RUNTIME_CHECK(!dmfiles.empty());
+    auto handle_size = dmfiles.front()->getColumnStat(EXTRA_HANDLE_COLUMN_ID).avg_size;
+    if (handle_size == 0)
+        handle_size = sizeof(UInt64);
+    constexpr auto version_size = sizeof(UInt64);
+    constexpr auto delmark_size = sizeof(UInt8);
+
+    // For delta, rs_filter does not filter rows, so we need to read all rows.
+    const auto delta_read_rows = read_snap->delta->getRows();
+    // For Stable, rs_filter may filter rows.
+    const auto stable_read_rows = read_snap->stable->estimatedReadRows(
+        pack_filters,
+        start_ts,
+        /*use_delta_index*/ delta_read_rows != 0);
+    return (handle_size + version_size + delmark_size) * (delta_read_rows + stable_read_rows);
+}
+
 BlockInputStreamPtr Segment::getInputStream(
     const ReadMode & read_mode,
     const DMContext & dm_context,
@@ -949,6 +1005,33 @@ BlockInputStreamPtr Segment::getInputStream(
         expected_block_size,
         columns_to_read,
         segment_snap->stable->stable);
+
+    // Building MVCC bitmap can be a resource-intensive operation that cannot be paused midway.
+    // To prevent scenarios where multiple segments concurrently build MVCC bitmaps but exhaust
+    // RU during execution and causing RU consumption to exceed limitations, we need to first
+    // estimate the cost of building the MVCC bitmap, pre-consume the corresponding RU,
+    // and then proceed with the building task.
+    const auto & res_group_name = dm_context.scan_context->resource_group_name;
+    if (likely(read_mode == ReadMode::Bitmap && !res_group_name.empty()))
+    {
+        auto real_ranges = shrinkRowKeyRanges(read_ranges);
+        if (!real_ranges.empty())
+        {
+            auto pack_filters = loadDMFilePackFilters(
+                segment_snap->stable->getDMFiles(),
+                dm_context.global_context.getMinMaxIndexCache(),
+                real_ranges,
+                filter ? filter->rs_operator : EMPTY_RS_OPERATOR,
+                dm_context.global_context.getFileProvider(),
+                dm_context.global_context.getReadLimiter(),
+                dm_context.scan_context,
+                dm_context.tracing_id,
+                ReadTag::MVCC);
+            auto bytes = estimatedBytesOfInternalColumns(segment_snap, pack_filters, start_ts);
+            LocalAdmissionController::global_instance->consumeBytesResource(res_group_name, bytesToRU(bytes));
+        }
+    }
+
     switch (read_mode)
     {
     case ReadMode::Normal:

dbms/src/Storages/DeltaMerge/Segment.h

@@ -44,6 +44,8 @@ class RSOperator;
 using RSOperatorPtr = std::shared_ptr<RSOperator>;
 class PushDownFilter;
 using PushDownFilterPtr = std::shared_ptr<PushDownFilter>;
+class MinMaxIndexCache;
+using MinMaxIndexCachePtr = std::shared_ptr<MinMaxIndexCache>;
 
 enum class ReadMode;
 
@@ -68,13 +70,6 @@ struct SegmentSnapshot : private boost::noncopyable
 
     bool isForUpdate() const { return delta->isForUpdate(); }
 
-    UInt64 estimatedBytesOfInternalColumns() const
-    {
-        // TODO: how about cluster index?
-        // handle + version + flag
-        return (sizeof(Int64) + sizeof(UInt64) + sizeof(UInt8)) * getRows();
-    }
-
     String detailInfo() const;
 };
 
@@ -730,6 +725,7 @@ class Segment
         bool relevant_place) const;
 
     static bool useCleanRead(const SegmentSnapshotPtr & segment_snap, const ColumnDefines & columns_to_read);
+    // Shrink the read_ranges by the segment real rowkey_range
     RowKeyRanges shrinkRowKeyRanges(const RowKeyRanges & read_ranges) const;
     BitmapFilterPtr buildBitmapFilter(
         const DMContext & dm_context,
@@ -795,6 +791,21 @@ class Segment
         const ColumnDefines & read_columns,
         const StableValueSpacePtr & stable);
 
+    static std::vector<DMFilePackFilter> loadDMFilePackFilters(
+        const DMFiles & dmfiles,
+        const MinMaxIndexCachePtr & index_cache,
+        const RowKeyRanges & rowkey_ranges,
+        const RSOperatorPtr & rs_filter,
+        const FileProviderPtr & file_provider,
+        const ReadLimiterPtr & read_limiter,
+        const ScanContextPtr & scan_context,
+        const String & tracing_id,
+        const ReadTag & read_tag);
+    static UInt64 estimatedBytesOfInternalColumns(
+        const SegmentSnapshotPtr & read_snap,
+        std::vector<DMFilePackFilter> & pack_filters,
+        UInt64 start_ts);
+
 #ifndef DBMS_PUBLIC_GTEST
 private:
 #else

dbms/src/Storages/DeltaMerge/SegmentReadTaskPool.cpp

@@ -89,11 +89,6 @@ BlockInputStreamPtr SegmentReadTaskPool::buildInputStream(SegmentReadTaskPtr & t
 
     t->fetchPages();
 
-    if (likely(read_mode == ReadMode::Bitmap && !res_group_name.empty()))
-    {
-        auto bytes = t->read_snapshot->estimatedBytesOfInternalColumns();
-        LocalAdmissionController::global_instance->consumeBytesResource(res_group_name, bytesToRU(bytes));
-    }
     t->initInputStream(
         columns_to_read,
         start_ts,

dbms/src/Storages/DeltaMerge/StableValueSpace.cpp

@@ -633,6 +633,48 @@ StableValueSpace::Snapshot::getAtLeastRowsAndBytes(const DMContext & context, co
     return ret;
 }
 
+// Estimating the number of rows that need to be read to build the MVCC bitmap.
+// If a delta index is used, the relevant data must be read and merged with the delta.
+// For stable-only (use_delta_index is false), no merging with the delta is required,
+// so reading occurs only when filtering of the pack is necessary.
+UInt64 StableValueSpace::Snapshot::estimatedReadRows(
+    std::vector<DMFilePackFilter> & pack_filters,
+    UInt64 start_ts,
+    bool use_delta_index) const
+{
+    const auto & dmfiles = getDMFiles();
+    UInt64 rows = 0;
+    for (size_t i = 0; i < dmfiles.size(); ++i)
+    {
+        const auto & dmfile = dmfiles[i];
+        auto & pack_filter = pack_filters[i];
+        const auto & pack_res = pack_filter.getPackResConst();
+        const auto & handle_res = pack_filter.getHandleRes();
+        const auto & pack_stats = dmfile->getPackStats();
+        for (size_t pack_id = 0; pack_id < pack_stats.size(); ++pack_id)
+        {
+            const auto & pack_stat = pack_stats[pack_id];
+            if (!pack_res[pack_id].isUse())
+                continue;
+
+            if (use_delta_index)
+            {
+                rows += pack_stat.rows;
+            }
+            else if (
+                handle_res[pack_id] == RSResult::Some || pack_stat.not_clean > 0
+                || pack_filter.getMaxVersion(pack_id) > start_ts)
+            {
+                // `not_clean > 0` means there are more than one version for some rowkeys in this pack
+                // `pack.max_version > start_ts` means some rows will be filtered by MVCC reading
+                // We need to read this pack to do RowKey or MVCC filter.
+                rows += pack_stat.rows;
+            }
+        }
+    }
+    return rows;
+}
+
 static size_t defaultValueBytes(const Field & f)
 {
     switch (f.getType())

dbms/src/Storages/DeltaMerge/StableValueSpace.h

@@ -253,6 +253,9 @@ class StableValueSpace : public std::enable_shared_from_this<StableValueSpace>
          */
         AtLeastRowsAndBytesResult getAtLeastRowsAndBytes(const DMContext & context, const RowKeyRange & range) const;
 
+        UInt64 estimatedReadRows(std::vector<DMFilePackFilter> & pack_filters, UInt64 start_ts, bool use_delta_index)
+            const;
+
     private:
         LoggerPtr log;
     };

dbms/src/Storages/DeltaMerge/tests/gtest_segment_bitmap.cpp

@@ -16,6 +16,7 @@
 #include <DataStreams/OneBlockInputStream.h>
 #include <Interpreters/Context.h>
 #include <Storages/DeltaMerge/DeltaMergeStore.h>
+#include <Storages/DeltaMerge/File/DMFilePackFilter.h>
 #include <Storages/DeltaMerge/tests/gtest_segment_test_basic.h>
 #include <Storages/DeltaMerge/tests/gtest_segment_util.h>
 #include <TestUtils/FunctionTestUtils.h>
@@ -177,12 +178,34 @@ class SegmentBitmapFilterTest : public SegmentTestBasic
             ASSERT_TRUE(sequenceEqual(expected_handle.data(), handle->data(), test_case.expected_size));
         }
     }
+
+    UInt64 estimatedBytesOfInternalColumns(UInt64 start_ts)
+    {
+        auto [seg, snap] = getSegmentForRead(SEG_ID);
+        auto pack_filters = Segment::loadDMFilePackFilters(
+            snap->stable->getDMFiles(),
+            dm_context->global_context.getMinMaxIndexCache(),
+            /*rowkey_ranges*/ {},
+            EMPTY_RS_OPERATOR,
+            dm_context->global_context.getFileProvider(),
+            dm_context->global_context.getReadLimiter(),
+            dm_context->scan_context,
+            dm_context->tracing_id,
+            ReadTag::MVCC);
+        return Segment::estimatedBytesOfInternalColumns(snap, pack_filters, start_ts);
+    }
 };
 
 TEST_F(SegmentBitmapFilterTest, InMemory1)
 try
 {
     runTestCase(TestCase("d_mem:[0, 1000)", 1000, "[0, 1000)", "[0, 1000)"));
+
+    // Delta only
+    auto bytes = estimatedBytesOfInternalColumns(std::numeric_limits<UInt64>::max());
+    ASSERT_EQ(bytes, 17 * 1000);
+    bytes = estimatedBytesOfInternalColumns(0);
+    ASSERT_EQ(bytes, 17 * 1000);
 }
 CATCH
 
@@ -272,6 +295,12 @@ TEST_F(SegmentBitmapFilterTest, Stable1)
 try
 {
     runTestCase(TestCase{"s:[0, 1024)", 1024, "[0, 1024)", "[0, 1024)"});
+
+    // Stable only
+    auto bytes = estimatedBytesOfInternalColumns(std::numeric_limits<UInt64>::max());
+    ASSERT_EQ(bytes, 0);
+    bytes = estimatedBytesOfInternalColumns(0);
+    ASSERT_EQ(bytes, 17 * 1024);
 }
 CATCH
 
@@ -302,6 +331,12 @@ try
         946,
         "[0, 128)|[1034, 1089)|[256, 298)|[1089, 1096)|[310, 1024)",
         "[0, 128)|[200, 255)|[256, 305)|[310, 1024)"});
+
+    // Delta + Stable
+    auto bytes = estimatedBytesOfInternalColumns(std::numeric_limits<UInt64>::max());
+    ASSERT_EQ(bytes, 17 * (1024 + 10 + 55 + 7));
+    bytes = estimatedBytesOfInternalColumns(0);
+    ASSERT_EQ(bytes, 17 * (1024 + 10 + 55 + 7));
 }
 CATCH