Backport ClickHouse#92390 to 25.8: Fix possible crash in aggregate functions after MEMORY_LIMIT_EXCEEDED

Author: robot-clickhouse

src/Columns/ColumnAggregateFunction.cpp

@@ -11,8 +11,10 @@
 #include <IO/WriteBufferFromArena.h>
 #include <IO/WriteBufferFromString.h>
 #include <Processors/Transforms/ColumnGathererTransform.h>
+#include <base/defines.h>
 #include <Common/AlignedBuffer.h>
 #include <Common/Arena.h>
+#include <Common/FailPoint.h>
 #include <Common/FieldVisitorToString.h>
 #include <Common/HashTable/Hash.h>
 #include <Common/SipHash.h>
@@ -32,6 +34,12 @@ namespace ErrorCodes
     extern const int SIZES_OF_COLUMNS_DOESNT_MATCH;
     extern const int ILLEGAL_TYPE_OF_ARGUMENT;
     extern const int NOT_IMPLEMENTED;
+    extern const int MEMORY_LIMIT_EXCEEDED;
+}
+
+namespace FailPoints
+{
+    extern const char column_aggregate_function_ensureOwnership_exception[];
 }
 
 
@@ -226,29 +234,44 @@ void ColumnAggregateFunction::ensureOwnership()
         size_t size_of_state = func->sizeOfData();
         size_t align_of_state = func->alignOfData();
 
+        bool inject_memory_limit_exceeded = false;
+        /// Avoid checking failpoint in a loop, since it uses atomic's
+        fiu_do_on(FailPoints::column_aggregate_function_ensureOwnership_exception,
+        {
+            inject_memory_limit_exceeded = true;
+        });
+
+        Container new_data;
+        new_data.resize_exact(size);
+
         size_t rollback_pos = 0;
         try
         {
             for (size_t i = 0; i < size; ++i)
             {
-                ConstAggregateDataPtr old_place = data[i];
-                data[i] = arena.alignedAlloc(size_of_state, align_of_state);
-                func->create(data[i]);
+                new_data[i] = arena.alignedAlloc(size_of_state, align_of_state);
+                func->create(new_data[i]);
                 ++rollback_pos;
-                func->merge(data[i], old_place, &arena);
+
+                if (unlikely(inject_memory_limit_exceeded))
+                    throw Exception(ErrorCodes::MEMORY_LIMIT_EXCEEDED, "Failpoint triggered");
+
+                func->merge(new_data[i], data[i], &arena);
             }
         }
         catch (...)
         {
             /// If we failed to take ownership, destroy all temporary data.
-
             if (!func->hasTrivialDestructor())
+            {
                 for (size_t i = 0; i < rollback_pos; ++i)
-                    func->destroy(data[i]);
-
+                    func->destroy(new_data[i]);
+            }
             throw;
         }
 
+        data = std::move(new_data);
+
         /// Now we own all data.
         src.reset();
     }

src/Columns/tests/gtest_column_aggregate_function.cpp

@@ -0,0 +1,70 @@
+#include <base/types.h>
+#include <AggregateFunctions/DDSketch.h>
+#include <AggregateFunctions/AggregateFunctionFactory.h>
+#include <AggregateFunctions/IAggregateFunction.h>
+#include <Columns/ColumnAggregateFunction.h>
+#include <Columns/ColumnsNumber.h>
+#include <DataTypes/DataTypesNumber.h>
+#include <IO/ReadBufferFromMemory.h>
+#include <IO/ReadBufferFromString.h>
+#include <IO/WriteBufferFromVector.h>
+#include <Common/Arena.h>
+#include <Common/Base64.h>
+#include <Common/FailPoint.h>
+#include <Common/tests/gtest_global_register.h>
+
+#include <gtest/gtest.h>
+
+namespace DB::FailPoints
+{
+    extern const char column_aggregate_function_ensureOwnership_exception[];
+}
+
+TEST(ColumnAggregateFunction, EnsureOwnershipExceptionLeavesCorruptedState)
+{
+    tryRegisterAggregateFunctions();
+
+    using namespace DB;
+
+    // Create the aggregate function quantileDD with relative accuracy 0.01
+    AggregateFunctionFactory & factory = AggregateFunctionFactory::instance();
+    DataTypes argument_types = {std::make_shared<DataTypeFloat64>()};
+    Array params = {Field(0.01), Field(0.5)};
+    AggregateFunctionProperties properties;
+    auto aggregate_function = factory.get("quantileDD", NullsAction::EMPTY, argument_types, params, properties);
+
+    // Create a source column with some data
+    auto src_column = ColumnAggregateFunction::create(aggregate_function);
+    Arena arena_src;
+    auto data_column = ColumnFloat64::create();
+    data_column->insert(Field(1.0));
+    data_column->insert(Field(2.0));
+    data_column->insert(Field(3.0));
+    const IColumn * columns[1] = {data_column.get()};
+
+    for (size_t i = 0; i < 3; ++i)
+    {
+        src_column->insertDefault();
+        aggregate_function->add(src_column->getData()[i], columns, i, &arena_src);
+    }
+
+    // Create a view column from the source - this sets src pointer
+    auto view_column = src_column->cloneEmpty();
+    view_column->insertRangeFrom(*src_column, 0, 3);
+
+    // Enable failpoint that will trigger an exception during ensureOwnership
+    // This will happen after at least one state is created and destroyed
+    FailPointInjection::enableFailPoint(FailPoints::column_aggregate_function_ensureOwnership_exception);
+
+    // Try to insert - this will call ensureOwnership() which will throw
+    // After the exception, previously, data[] points to destroyed memory where mapping == nullptr
+    ASSERT_THROW({
+        view_column->insertDefault();
+    }, Exception);
+
+    // Disable failpoint
+    FailPointInjection::disableFailPoint(FailPoints::column_aggregate_function_ensureOwnership_exception);
+
+    /// Previously leads to a crash
+    view_column->insertDefault();
+}

src/Common/FailPoint.cpp

@@ -88,6 +88,7 @@ static struct InitFiu
     ONCE(execute_query_calling_empty_set_result_func_on_exception) \
     ONCE(receive_timeout_on_table_status_response) \
     ONCE(delta_kernel_fail_literal_visitor) \
+    ONCE(column_aggregate_function_ensureOwnership_exception) \
     REGULAR(keepermap_fail_drop_data) \
     REGULAR(lazy_pipe_fds_fail_close) \
     PAUSEABLE(infinite_sleep) \