From d11365da9cbae10f18e15e82695e2b769314d553 Mon Sep 17 00:00:00 2001
From: lihangyu <15605149486@163.com>
Date: Mon, 18 Dec 2023 10:31:16 +0800
Subject: [PATCH] [Fix](memtable) fix `shrink_memtable_by_agg` should also
 update `_row_in_blocks` (#28536)

Otherwise using the stale `_row_in_blocks` will result in heap-buffer-overflow

```
==2695213==ERROR: AddressSanitizer: heap-buffer-overflow on address 0x62900122e210 at pc 0x56524744aecf bp 0x7f62c595ef7
0 sp 0x7f62c595ef68
READ of size 8 at 0x62900122e210 thread T1627 (MemTableFlushTh)
    #0 0x56524744aece in doris::vectorized::ColumnVector<long>::insert_indices_from(doris::vectorized::IColumn const&, unsigned int const*, unsigned int const*) /mnt/disk2/lihangyu/doris/be/src/vec/columns/column_vector.cpp:378:33
    #1 0x5652472a7538 in doris::vectorized::ColumnNullable::insert_indices_from(doris::vectorized::IColumn const&, unsigned int const*, unsigned int const*) /mnt/disk2/lihangyu/doris/be/src/vec/columns/column_nullable.cpp:310:25
    #2 0x56524782a62a in doris::vectorized::MutableBlock::add_rows(doris::vectorized::Block const*, unsigned int const*, unsigned int const*) /mnt/disk2/lihangyu/doris/be/src/vec/core/block.cpp:961:14
    #3 0x565233f187ae in doris::MemTable::_put_into_output(doris::vectorized::Block&) /mnt/disk2/lihangyu/doris/be/src/olap/memtable.cpp:248:27
    #4 0x565233f1db66 in doris::MemTable::to_block() /mnt/disk2/lihangyu/doris/be/src/olap/memtable.cpp:496:13
    #5 0x565233efae60 in doris::FlushToken::_do_flush_memtable(doris::MemTable*, int, long*) /mnt/disk2/lihangyu/doris/be/src/olap/memtable_flush_executor.cpp:121:62
    #6 0x565233efc8d6 in doris::FlushToken::_flush_memtable(doris::MemTable*, int, long) /mnt/disk2/lihangyu/doris/be/src/olap/memtable_flush_executor.cpp:150:16
    #7 0x565233f0c5eb in doris::MemtableFlushTask::run() /mnt/disk2/lihangyu/doris/be/src/olap/memtable_flush_executor.cpp:58:23
```
---
 be/src/olap/memtable.cpp                      |  2 ++
 ...st_insert_with_aggregation_memtable.groovy | 22 ++++++++++++++++++-
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/be/src/olap/memtable.cpp b/be/src/olap/memtable.cpp
index b69364c3dbb24e..95966edf9dbdc5 100644
--- a/be/src/olap/memtable.cpp
+++ b/be/src/olap/memtable.cpp
@@ -451,6 +451,8 @@ void MemTable::_aggregate() {
         _output_mutable_block =
                 vectorized::MutableBlock::build_mutable_block(empty_input_block.get());
         _output_mutable_block.clear_column_data();
+        _row_in_blocks = temp_row_in_blocks;
+        _last_sorted_pos = _row_in_blocks.size();
     }
 }
 
diff --git a/regression-test/suites/load_p0/insert/test_insert_with_aggregation_memtable.groovy b/regression-test/suites/load_p0/insert/test_insert_with_aggregation_memtable.groovy
index c722905e622f0f..4e6d838a17e6fc 100644
--- a/regression-test/suites/load_p0/insert/test_insert_with_aggregation_memtable.groovy
+++ b/regression-test/suites/load_p0/insert/test_insert_with_aggregation_memtable.groovy
@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-suite("test_insert_with_aggregation_memtable") {
+suite("test_insert_with_aggregation_memtable", "nonConcurrent") {
     def backendId_to_backendIP = [:]
     def backendId_to_backendHttpPort = [:]
     def backendId_to_params = [string:[:]]
@@ -114,6 +114,26 @@ suite("test_insert_with_aggregation_memtable") {
     sql "sync"
     qt_sql "select * from ${testTable} order by id asc"
 
+    // test with mv
+    def table_name = "agg_shrink"
+    sql "DROP TABLE IF EXISTS ${table_name}"
+    sql """
+        CREATE TABLE IF NOT EXISTS ${table_name} (
+            k bigint,
+            v text 
+        )
+        DUPLICATE KEY(`k`)
+        DISTRIBUTED BY HASH(k) BUCKETS 4
+        properties("replication_num" = "1");
+    """
+    set_be_param("write_buffer_size_for_agg", "10240") // change it to 10KB
+    sql """INSERT INTO ${table_name} SELECT *, '{"k1":1, "k2": "hello world", "k3" : [1234], "k4" : 1.10000, "k5" : [[123]]}' FROM numbers("number" = "4096")"""
+    sql """INSERT INTO ${table_name} SELECT k, v from ${table_name}"""
+    sql """INSERT INTO ${table_name} SELECT k, v from ${table_name}"""
+    createMV("""create materialized view var_cnt as select k, count(k) from ${table_name} group by k""")    
+    sql """INSERT INTO ${table_name} SELECT k, v from ${table_name} limit 8101"""
+
     reset_be_param("enable_shrink_memory")
     reset_be_param("write_buffer_size_for_agg")
+
 }
\ No newline at end of file