kaka11chen commented on code in PR #37257: URL: https://github.com/apache/doris/pull/37257#discussion_r1688445072
########## be/src/vec/common/allocator.h: ########## @@ -68,6 +69,143 @@ static constexpr size_t MALLOC_MIN_ALIGNMENT = 8; // is always a multiple of sixteen. (https://www.gnu.org/software/libc/manual/html_node/Aligned-Memory-Blocks.html) static constexpr int ALLOCATOR_ALIGNMENT_16 = 16; +class DefaultMemoryAllocator { +public: + static void* malloc(size_t size) __THROW { return std::malloc(size); } + + static void* calloc(size_t n, size_t size) __THROW { return std::calloc(n, size); } + + static constexpr bool need_record_actual_size() { return false; } + + static int posix_memalign(void** ptr, size_t alignment, size_t size) __THROW { + return ::posix_memalign(ptr, alignment, size); + } + + static void* realloc(void* ptr, size_t size) __THROW { return std::realloc(ptr, size); } + + static void free(void* p) __THROW { std::free(p); } + + static void release_unused() {} + +private: +}; + +/** It would be better to put these Memory Allocators where they are used, such as in the orc memory pool and arrow memory pool. + * But currently allocators use templates in .cpp instead of all in .h, so they can only be placed here. + */ +class ORCMemoryAllocator { +public: + static void* malloc(size_t size) __THROW { return reinterpret_cast<char*>(std::malloc(size)); } + + static void* calloc(size_t n, size_t size) __THROW { return std::calloc(n, size); } + + static constexpr bool need_record_actual_size() { return true; } + + static size_t allocated_size(void* ptr) { return malloc_usable_size(ptr); } + + static int posix_memalign(void** ptr, size_t alignment, size_t size) __THROW { + return ::posix_memalign(ptr, alignment, size); + } + + static void* realloc(void* ptr, size_t size) __THROW { + LOG(FATAL) << "__builtin_unreachable"; + __builtin_unreachable(); + } + + static void free(void* p) __THROW { std::free(p); } + + static void release_unused() {} +}; + +class RecordSizeMemoryAllocator { +public: + static void* malloc(size_t size) __THROW { + void* p = std::malloc(size); + if (p) { + std::lock_guard<std::mutex> lock(_mutex); + _allocated_sizes[p] = size; + } + return p; + } + + static void* calloc(size_t n, size_t size) __THROW { + void* p = std::calloc(n, size); + if (p) { + std::lock_guard<std::mutex> lock(_mutex); + _allocated_sizes[p] = n * size; + } + return p; + } + + static constexpr bool need_record_actual_size() { return false; } + + static size_t allocated_size(void* ptr) { + std::lock_guard<std::mutex> lock(_mutex); + auto it = _allocated_sizes.find(ptr); + if (it != _allocated_sizes.end()) { + return it->second; + } + return 0; + } + + static int posix_memalign(void** ptr, size_t alignment, size_t size) __THROW { + int ret = ::posix_memalign(ptr, alignment, size); + if (ret == 0 && *ptr) { + std::lock_guard<std::mutex> lock(_mutex); + _allocated_sizes[*ptr] = size; + } + return ret; + } + + static void* realloc(void* ptr, size_t size) __THROW { + std::lock_guard<std::mutex> lock(_mutex); + void* p = std::realloc(ptr, size); + if (p) { + _allocated_sizes.erase(ptr); + _allocated_sizes[p] = size; + } + return p; + } + + static void free(void* p) __THROW { + if (p) { + std::lock_guard<std::mutex> lock(_mutex); + _allocated_sizes.erase(p); + std::free(p); + } + } + + static void release_unused() {} + +private: + static std::unordered_map<void*, size_t> _allocated_sizes; + static std::mutex _mutex; +}; + +#if defined(USE_JEMALLOC) +#include <jemalloc/jemalloc.h> +class ArrowJemallocMemoryAllocator { +public: + static void* malloc(size_t size) __THROW { return je_malloc(size); } + + static void* calloc(size_t n, size_t size) __THROW { return je_calloc(n, size); } + + static constexpr bool need_record_actual_size() { return false; } + + static int posix_memalign(void** ptr, size_t alignment, size_t size) __THROW { + return je_posix_memalign(ptr, alignment, size); + } + + static void* realloc(void* ptr, size_t size) __THROW { return je_realloc(ptr, size); } + + static void free(void* p) __THROW { je_free(p); } + + static void release_unused() { Review Comment: Will change to `DefaultMemoryAllocator`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org For additional commands, e-mail: commits-h...@doris.apache.org
