From e5663f98724bcea5e107df07cf3aa720df1854b8 Mon Sep 17 00:00:00 2001 From: Adonis Ling Date: Wed, 20 Jul 2022 16:37:27 +0800 Subject: [PATCH] [Bug](array-type) Fix the core dump caused by unaligned __int128 (#11020) Fix the core dump caused by unaligned __int128 and change DEFAULT_ALIGNMENT --- be/src/runtime/collection_value.cpp | 6 +++-- be/src/runtime/free_pool.hpp | 36 +++++++++++++++++++++++++++++ be/src/runtime/mem_pool.h | 6 +++-- be/src/udf/udf.cpp | 14 +++++++++++ be/src/udf/udf.h | 7 ++++++ be/src/util/array_parser.h | 2 +- be/test/util/array_parser_test.cpp | 17 ++++++++++++++ 7 files changed, 83 insertions(+), 5 deletions(-) diff --git a/be/src/runtime/collection_value.cpp b/be/src/runtime/collection_value.cpp index c354c9f1a2..fb9de9635c 100644 --- a/be/src/runtime/collection_value.cpp +++ b/be/src/runtime/collection_value.cpp @@ -480,13 +480,15 @@ Status CollectionValue::init_collection(CollectionValue* value, const AllocateMe Status CollectionValue::init_collection(MemPool* pool, uint64_t size, PrimitiveType child_type, CollectionValue* value) { return init_collection( - value, [pool](size_t size) { return pool->allocate(size); }, size, child_type); + value, [pool](size_t size) { return pool->allocate_aligned(size, 16); }, size, + child_type); } Status CollectionValue::init_collection(FunctionContext* context, uint64_t size, PrimitiveType child_type, CollectionValue* value) { return init_collection( - value, [context](size_t size) { return context->allocate(size); }, size, child_type); + value, [context](size_t size) { return context->aligned_allocate(16, size); }, size, + child_type); } CollectionValue CollectionValue::from_collection_val(const CollectionVal& val) { diff --git a/be/src/runtime/free_pool.hpp b/be/src/runtime/free_pool.hpp index 49b58d0ea2..db23d98fc6 100644 --- a/be/src/runtime/free_pool.hpp +++ b/be/src/runtime/free_pool.hpp @@ -82,6 +82,42 @@ public: return reinterpret_cast(allocation) + sizeof(FreeListNode); } + // Allocates a buffer of size. + uint8_t* aligned_allocate(int alignment, int64_t size) { + // The alignment should be a power of 2. + DCHECK(alignment > 0 && ((alignment - 1) & alignment) == 0); + + // This is the typical malloc behavior. NULL is reserved for failures. + if (size == 0) { + return reinterpret_cast(alignment); + } + + int padding = sizeof(FreeListNode) >= alignment ? 0 : (alignment - sizeof(FreeListNode)); + size += padding; + + // Do ceil(log_2(size)) + int free_list_idx = BitUtil::log2(size); + DCHECK_LT(free_list_idx, NUM_LISTS); + + FreeListNode* allocation = _lists[free_list_idx].next; + + if (allocation == nullptr) { + // There wasn't an existing allocation of the right size, allocate a new one. + size = 1L << free_list_idx; + allocation = reinterpret_cast( + _mem_pool->allocate_aligned(size + sizeof(FreeListNode), alignment)); + } else { + // Remove this allocation from the list. + _lists[free_list_idx].next = allocation->next; + } + + DCHECK(allocation != nullptr); + // Set the back node to point back to the list it came from so know where + // to add it on free(). + allocation->list = &_lists[free_list_idx]; + return reinterpret_cast(allocation) + sizeof(FreeListNode) + padding; + } + void free(uint8_t* ptr) { if (ptr == NULL || reinterpret_cast(ptr) == 0x1) { return; diff --git a/be/src/runtime/mem_pool.h b/be/src/runtime/mem_pool.h index 02842b9752..0523fd736e 100644 --- a/be/src/runtime/mem_pool.h +++ b/be/src/runtime/mem_pool.h @@ -104,7 +104,6 @@ public: /// of the current chunk. Creates a new chunk if there aren't any chunks /// with enough capacity. uint8_t* allocate(int64_t size, Status* rst = nullptr) { - // TODO: rethink if DEFAULT_ALIGNMENT should be changed, malloc is aligned by 16. return allocate(size, DEFAULT_ALIGNMENT, rst); } @@ -170,7 +169,10 @@ public: MemTracker* mem_tracker() { return _mem_tracker; } - static constexpr int DEFAULT_ALIGNMENT = 8; + // The memory for __int128 should be aligned to 16 bytes. + // By the way, in 64-bit system, the address of a block returned by malloc or realloc in GNU systems + // is always a multiple of sixteen. (https://www.gnu.org/software/libc/manual/html_node/Aligned-Memory-Blocks.html) + static constexpr int DEFAULT_ALIGNMENT = 16; #if (defined(__SANITIZE_ADDRESS__) || defined(ADDRESS_SANITIZER)) && !defined(BE_TEST) static constexpr int DEFAULT_PADDING_SIZE = 0x10; diff --git a/be/src/udf/udf.cpp b/be/src/udf/udf.cpp index 7dac939fd7..d73fc6d691 100644 --- a/be/src/udf/udf.cpp +++ b/be/src/udf/udf.cpp @@ -49,6 +49,9 @@ public: FreePool(MemPool*) {} uint8_t* allocate(int byte_size) { return reinterpret_cast(malloc(byte_size)); } + uint8_t* aligned_allocate(int alignment, int byte_size) { + return reinterpret_cast(aligned_alloc(alignment, byte_size)); + } uint8_t* reallocate(uint8_t* ptr, int byte_size) { return reinterpret_cast(realloc(ptr, byte_size)); @@ -271,6 +274,17 @@ uint8_t* FunctionContext::allocate(int byte_size) { return buffer; } +uint8_t* FunctionContext::aligned_allocate(int alignment, int byte_size) { + uint8_t* buffer = _impl->_pool->aligned_allocate(alignment, byte_size); + _impl->_allocations[buffer] = byte_size; + + if (_impl->_debug) { + memset(buffer, 0xff, byte_size); + } + + return buffer; +} + uint8_t* FunctionContext::reallocate(uint8_t* ptr, int byte_size) { _impl->_allocations.erase(ptr); ptr = _impl->_pool->reallocate(ptr, byte_size); diff --git a/be/src/udf/udf.h b/be/src/udf/udf.h index d0ce5eb264..b4916fe7b4 100644 --- a/be/src/udf/udf.h +++ b/be/src/udf/udf.h @@ -176,6 +176,13 @@ public: // in this object causing the query to fail. uint8_t* allocate(int byte_size); + // Allocate and align memory for UDAs. All UDA allocations should use this if possible instead of + // malloc/new. The UDA is responsible for calling Free() on all buffers returned + // by Allocate(). + // If this Allocate causes the memory limit to be exceeded, the error will be set + // in this object causing the query to fail. + uint8_t* aligned_allocate(int alignment, int byte_size); + // Reallocates 'ptr' to the new byte_size. If the currently underlying allocation // is big enough, the original ptr will be returned. If the allocation needs to // grow, a new allocation is made that is at least 'byte_size' and the contents diff --git a/be/src/util/array_parser.h b/be/src/util/array_parser.h index 7bcf122c34..bbdf2571f0 100644 --- a/be/src/util/array_parser.h +++ b/be/src/util/array_parser.h @@ -210,7 +210,7 @@ private: break; } case TYPE_DECIMALV2: { - *val = reinterpret_cast(context->allocate(sizeof(DecimalV2Val))); + *val = reinterpret_cast(context->aligned_allocate(16, sizeof(DecimalV2Val))); new (*val) DecimalV2Val(); if (iterator->IsNumber()) { diff --git a/be/test/util/array_parser_test.cpp b/be/test/util/array_parser_test.cpp index ba92b05020..e5e2564744 100644 --- a/be/test/util/array_parser_test.cpp +++ b/be/test/util/array_parser_test.cpp @@ -172,4 +172,21 @@ TEST(ArrayParserTest, TestDecimalArray) { value = {data, num_items, false, nullptr}; test_array_parser(column_pb, "[2147483647.5, \"34359738368.5\"]", value); } + +TEST(ArrayParserTest, TestFreePool) { + auto column_pb = create_column_pb("ARRAY", "DECIMAL"); + MemTracker tracker(1024 * 1024, "ArrayParserTest"); + MemPool mem_pool(&tracker); + FunctionContext context; + ArrayUtils::prepare_context(context, mem_pool, column_pb); + int alignment = 1; + for (int i = 1; i <= 4; ++i) { + alignment <<= 1; + auto* p = context.aligned_allocate(alignment, alignment); + EXPECT_TRUE(reinterpret_cast(p) % alignment == 0); + p = context.aligned_allocate(alignment, alignment); + EXPECT_TRUE(reinterpret_cast(p) % alignment == 0); + } +} + } // namespace doris