From fb163b55c2c4866ad6c1db548283523dcdc07216 Mon Sep 17 00:00:00 2001
From: "github-actions[bot]"
 <41898282+github-actions[bot]@users.noreply.github.com>
Date: Wed, 20 Nov 2024 21:07:43 +0800
Subject: [PATCH] branch-2.1: [Fix](merge-on-write) Fix
 `MergeIndexDeleteBitmapCalculator::calculate_one()` coredump #44284 (#44330)

Cherry-picked from #44284

Co-authored-by: bobhan1 <baohan@selectdb.com>
---
 be/src/olap/delete_bitmap_calculator.cpp | 70 +++++++++++++-----------
 1 file changed, 38 insertions(+), 32 deletions(-)

diff --git a/be/src/olap/delete_bitmap_calculator.cpp b/be/src/olap/delete_bitmap_calculator.cpp
index 6f6e0ec8889954..017e3cff3d0489 100644
--- a/be/src/olap/delete_bitmap_calculator.cpp
+++ b/be/src/olap/delete_bitmap_calculator.cpp
@@ -90,8 +90,10 @@ bool MergeIndexDeleteBitmapCalculatorContext::Comparator::operator()(
     // std::proiroty_queue is a max heap, and function should return the result of `lhs < rhs`
     // so if the result of the function is true, rhs will be popped before lhs
     Slice key1, key2;
-    RETURN_IF_ERROR(lhs->get_current_key(key1));
-    RETURN_IF_ERROR(rhs->get_current_key(key2));
+    // MergeIndexDeleteBitmapCalculatorContext::get_current_key may return non-OK status if encounter
+    // memory allocation failure, we can only throw exception here to propagate error in this situation
+    THROW_IF_ERROR(lhs->get_current_key(key1));
+    THROW_IF_ERROR(rhs->get_current_key(key2));
     if (_sequence_length == 0 && _rowid_length == 0) {
         auto cmp_result = key1.compare(key2);
         // when key1 is the same as key2,
@@ -135,28 +137,30 @@ Status MergeIndexDeleteBitmapCalculator::init(RowsetId rowset_id,
                                               std::vector<SegmentSharedPtr> const& segments,
                                               size_t seq_col_length, size_t rowdid_length,
                                               size_t max_batch_size) {
-    _rowset_id = rowset_id;
-    _seq_col_length = seq_col_length;
-    _rowid_length = rowdid_length;
-    _comparator =
-            MergeIndexDeleteBitmapCalculatorContext::Comparator(seq_col_length, _rowid_length);
-    _contexts.reserve(segments.size());
-    _heap = std::make_unique<Heap>(_comparator);
+    RETURN_IF_CATCH_EXCEPTION({
+        _rowset_id = rowset_id;
+        _seq_col_length = seq_col_length;
+        _rowid_length = rowdid_length;
+        _comparator =
+                MergeIndexDeleteBitmapCalculatorContext::Comparator(seq_col_length, _rowid_length);
+        _contexts.reserve(segments.size());
+        _heap = std::make_unique<Heap>(_comparator);
 
-    for (auto& segment : segments) {
-        RETURN_IF_ERROR(segment->load_index());
-        auto pk_idx = segment->get_primary_key_index();
-        std::unique_ptr<segment_v2::IndexedColumnIterator> index;
-        RETURN_IF_ERROR(pk_idx->new_iterator(&index));
-        auto index_type = vectorized::DataTypeFactory::instance().create_data_type(
-                pk_idx->type_info()->type(), 1, 0);
-        _contexts.emplace_back(std::move(index), index_type, segment->id(), pk_idx->num_rows());
-        _heap->push(&_contexts.back());
-    }
-    if (_rowid_length > 0) {
-        _rowid_coder = get_key_coder(
-                get_scalar_type_info<FieldType::OLAP_FIELD_TYPE_UNSIGNED_INT>()->type());
-    }
+        for (auto& segment : segments) {
+            RETURN_IF_ERROR(segment->load_index());
+            auto pk_idx = segment->get_primary_key_index();
+            std::unique_ptr<segment_v2::IndexedColumnIterator> index;
+            RETURN_IF_ERROR(pk_idx->new_iterator(&index));
+            auto index_type = vectorized::DataTypeFactory::instance().create_data_type(
+                    pk_idx->type_info()->type(), 1, 0);
+            _contexts.emplace_back(std::move(index), index_type, segment->id(), pk_idx->num_rows());
+            _heap->push(&_contexts.back());
+        }
+        if (_rowid_length > 0) {
+            _rowid_coder = get_key_coder(
+                    get_scalar_type_info<FieldType::OLAP_FIELD_TYPE_UNSIGNED_INT>()->type());
+        }
+    });
     return Status::OK();
 }
 
@@ -209,16 +213,18 @@ Status MergeIndexDeleteBitmapCalculator::calculate_one(RowLocation& loc) {
 }
 
 Status MergeIndexDeleteBitmapCalculator::calculate_all(DeleteBitmapPtr delete_bitmap) {
-    RowLocation loc;
-    while (true) {
-        auto st = calculate_one(loc);
-        if (st.is<ErrorCode::END_OF_FILE>()) {
-            break;
+    RETURN_IF_CATCH_EXCEPTION({
+        RowLocation loc;
+        while (true) {
+            auto st = calculate_one(loc);
+            if (st.is<ErrorCode::END_OF_FILE>()) {
+                break;
+            }
+            RETURN_IF_ERROR(st);
+            delete_bitmap->add({_rowset_id, loc.segment_id, DeleteBitmap::TEMP_VERSION_COMMON},
+                               loc.row_id);
         }
-        RETURN_IF_ERROR(st);
-        delete_bitmap->add({_rowset_id, loc.segment_id, DeleteBitmap::TEMP_VERSION_COMMON},
-                           loc.row_id);
-    }
+    });
     return Status::OK();
 }