diff --git a/src/include/duckdb/storage/block_allocator.hpp b/src/include/duckdb/storage/block_allocator.hpp index 0830c2a89fde..9ca38edba03a 100644 --- a/src/include/duckdb/storage/block_allocator.hpp +++ b/src/include/duckdb/storage/block_allocator.hpp @@ -45,6 +45,10 @@ class BlockAllocator { bool SupportsFlush() const; void ThreadFlush(bool allocator_background_threads, idx_t threshold, idx_t thread_count) const; void FlushAll(optional_idx extra_memory = optional_idx()) const; + //! Memory freed since the last flush, which the allocator may still hold + idx_t GetDeallocatedSinceFlush() const; + //! Flush if at least threshold bytes were freed since the last flush, returns whether a flush happened + bool TryFlushDeallocated(idx_t threshold) const; private: bool IsActive() const; @@ -67,6 +71,8 @@ class BlockAllocator { const hugeint_t uuid; //! Fallback allocator Allocator &allocator; + //! Memory freed since the last flush + mutable atomic deallocated_since_flush {0}; //! Block size (power of two) const idx_t block_size; diff --git a/src/include/duckdb/storage/buffer/buffer_pool.hpp b/src/include/duckdb/storage/buffer/buffer_pool.hpp index a591da592761..6d928f110a99 100644 --- a/src/include/duckdb/storage/buffer/buffer_pool.hpp +++ b/src/include/duckdb/storage/buffer/buffer_pool.hpp @@ -56,6 +56,8 @@ class BufferPool { //! If bulk deallocation larger than this occurs, flush outstanding allocations void SetAllocatorBulkDeallocationFlushThreshold(idx_t threshold); idx_t GetAllocatorBulkDeallocationFlushThreshold(); + //! Number of times the allocator was flushed by the buffer pool + idx_t GetAllocatorFlushCount() const; void UpdateUsedMemory(MemoryTag tag, int64_t size); @@ -95,6 +97,13 @@ class BufferPool { //! Evict object cache entries if needed. EvictionResult EvictObjectCacheEntries(MemoryTag tag, idx_t extra_memory, idx_t memory_limit); + //! Return freed memory to the OS, optionally releasing up to extra_memory from the block allocator's pool + void FlushAllocator(optional_idx extra_memory = optional_idx()); + //! Flush the allocator if the memory freed since the last flush could push resident memory past the limit + void FlushOnBulkDeallocation(); + //! Freed memory that has to accumulate before the allocator is flushed + idx_t GetBulkDeallocationFlushThreshold() const; + //! Purge all blocks that haven't been pinned within the last N seconds idx_t PurgeAgedBlocks(uint32_t max_age_sec); idx_t PurgeAgedBlocksInternal(EvictionQueue &queue, uint32_t max_age_sec, int64_t now, int64_t limit); @@ -168,6 +177,10 @@ class BufferPool { atomic maximum_memory; //! If bulk deallocation larger than this occurs, flush outstanding allocations atomic allocator_bulk_deallocation_flush_threshold; + //! Number of allocator flushes + atomic allocator_flush_count {0}; + //! Freed memory is flushed once it exceeds this fraction of the memory limit + static constexpr idx_t BULK_DEALLOCATION_FLUSH_DIVISOR = 16; //! Record timestamps of buffer manager unpin() events. Usable by custom eviction policies. bool track_eviction_timestamps; //! Eviction queues diff --git a/src/storage/block_allocator.cpp b/src/storage/block_allocator.cpp index b9159da66e30..1dc69f36f28d 100644 --- a/src/storage/block_allocator.cpp +++ b/src/storage/block_allocator.cpp @@ -328,6 +328,7 @@ data_ptr_t BlockAllocator::AllocateData(const idx_t size) const { } void BlockAllocator::FreeData(const data_ptr_t pointer, const idx_t size) const { + deallocated_since_flush.fetch_add(size, std::memory_order_relaxed); if (!IsActive() || !IsInPool(pointer)) { return allocator.FreeData(pointer, size); } @@ -366,6 +367,7 @@ void BlockAllocator::ThreadFlush(bool allocator_background_threads, idx_t thresh } void BlockAllocator::FlushAll(const optional_idx extra_memory) const { + deallocated_since_flush.store(0, std::memory_order_relaxed); if (IsActive() && IsEnabled() && extra_memory.IsValid()) { FreeInternal(extra_memory.GetIndex()); } @@ -374,6 +376,26 @@ void BlockAllocator::FlushAll(const optional_idx extra_memory) const { } } +idx_t BlockAllocator::GetDeallocatedSinceFlush() const { + return deallocated_since_flush.load(std::memory_order_relaxed); +} + +bool BlockAllocator::TryFlushDeallocated(const idx_t threshold) const { + if (!SupportsFlush()) { + return false; + } + auto deallocated = deallocated_since_flush.load(std::memory_order_relaxed); + if (deallocated < threshold) { + return false; + } + // the thread that claims the accumulated amount performs the flush + if (!deallocated_since_flush.compare_exchange_strong(deallocated, 0, std::memory_order_relaxed)) { + return false; + } + FlushAll(deallocated); + return true; +} + void BlockAllocator::FreeInternal(const idx_t extra_memory) const { auto count = DivBlockSize(extra_memory); unsafe_vector to_free_buffer; diff --git a/src/storage/buffer/buffer_pool.cpp b/src/storage/buffer/buffer_pool.cpp index 729bb89f7fd5..f7dfca98e28c 100644 --- a/src/storage/buffer/buffer_pool.cpp +++ b/src/storage/buffer/buffer_pool.cpp @@ -331,7 +331,7 @@ BufferPool::EvictionResult BufferPool::EvictObjectCacheEntries(MemoryTag tag, id if (memory_usage.GetUsedMemory(MemoryUsageCaches::NO_FLUSH) <= memory_limit) { if (extra_memory > allocator_bulk_deallocation_flush_threshold) { - block_allocator.FlushAll(extra_memory); + FlushAllocator(extra_memory); } return {true, std::move(r)}; } @@ -352,7 +352,7 @@ BufferPool::EvictionResult BufferPool::EvictObjectCacheEntries(MemoryTag tag, id if (!success) { r.Resize(0); } else if (extra_memory > allocator_bulk_deallocation_flush_threshold) { - block_allocator.FlushAll(extra_memory); + FlushAllocator(extra_memory); } return {success, std::move(r)}; @@ -360,6 +360,7 @@ BufferPool::EvictionResult BufferPool::EvictObjectCacheEntries(MemoryTag tag, id BufferPool::EvictionResult BufferPool::EvictBlocks(MemoryTag tag, idx_t extra_memory, idx_t memory_limit, unique_ptr *buffer) { + FlushOnBulkDeallocation(); for (auto &queue : queues) { auto block_result = EvictBlocksInternal(*queue, tag, extra_memory, memory_limit, buffer); if (block_result.success) { @@ -379,7 +380,7 @@ BufferPool::EvictionResult BufferPool::EvictBlocksInternal(EvictionQueue &queue, if (memory_usage.GetUsedMemory(MemoryUsageCaches::NO_FLUSH) <= memory_limit) { if (extra_memory > allocator_bulk_deallocation_flush_threshold) { - block_allocator.FlushAll(extra_memory); + FlushAllocator(extra_memory); } return {true, std::move(r)}; } @@ -408,7 +409,7 @@ BufferPool::EvictionResult BufferPool::EvictBlocksInternal(EvictionQueue &queue, if (!found) { r.Resize(0); } else if (extra_memory > allocator_bulk_deallocation_flush_threshold) { - block_allocator.FlushAll(extra_memory); + FlushAllocator(extra_memory); } return {found, std::move(r)}; @@ -524,6 +525,36 @@ idx_t BufferPool::GetAllocatorBulkDeallocationFlushThreshold() { return allocator_bulk_deallocation_flush_threshold; } +idx_t BufferPool::GetAllocatorFlushCount() const { + return allocator_flush_count.load(std::memory_order_relaxed); +} + +idx_t BufferPool::GetBulkDeallocationFlushThreshold() const { + // freed memory that the allocator still holds must stay a small fraction of the memory limit + return maximum_memory.load(std::memory_order_relaxed) / BULK_DEALLOCATION_FLUSH_DIVISOR; +} + +void BufferPool::FlushAllocator(const optional_idx extra_memory) { + allocator_flush_count.fetch_add(1, std::memory_order_relaxed); + block_allocator.FlushAll(extra_memory); +} + +void BufferPool::FlushOnBulkDeallocation() { + const auto threshold = GetBulkDeallocationFlushThreshold(); + const auto deallocated = block_allocator.GetDeallocatedSinceFlush(); + if (deallocated < threshold) { + return; + } + // freed memory that the allocator still holds only matters once it could push us past the memory limit + const auto used_memory = memory_usage.GetUsedMemory(MemoryUsageCaches::NO_FLUSH); + if (used_memory + deallocated <= maximum_memory.load(std::memory_order_relaxed)) { + return; + } + if (block_allocator.TryFlushDeallocated(threshold)) { + allocator_flush_count.fetch_add(1, std::memory_order_relaxed); + } +} + BufferPool::MemoryUsage::MemoryUsage() { for (auto &v : memory_usage) { v = 0; diff --git a/src/storage/standard_buffer_manager.cpp b/src/storage/standard_buffer_manager.cpp index 55a09a556b6a..6c7b4aa67f45 100644 --- a/src/storage/standard_buffer_manager.cpp +++ b/src/storage/standard_buffer_manager.cpp @@ -705,7 +705,7 @@ data_ptr_t StandardBufferManager::BufferAllocatorAllocate(PrivateAllocatorData * StringUtil::BytesToHumanReadableString(size)); // We rely on manual tracking of this one. :( reservation.size = 0; - return Allocator::Get(data.manager.db).AllocateData(size); + return BlockAllocator::Get(data.manager.db).AllocateData(size); } void StandardBufferManager::BufferAllocatorFree(PrivateAllocatorData *private_data, data_ptr_t pointer, idx_t size) { @@ -713,7 +713,7 @@ void StandardBufferManager::BufferAllocatorFree(PrivateAllocatorData *private_da BufferPoolReservation r(MemoryTag::ALLOCATOR, data.manager.GetBufferPool()); r.size = size; r.Resize(0); - return Allocator::Get(data.manager.db).FreeData(pointer, size); + return BlockAllocator::Get(data.manager.db).FreeData(pointer, size); } data_ptr_t StandardBufferManager::BufferAllocatorRealloc(PrivateAllocatorData *private_data, data_ptr_t pointer, @@ -726,7 +726,7 @@ data_ptr_t StandardBufferManager::BufferAllocatorRealloc(PrivateAllocatorData *p r.size = old_size; r.Resize(size); r.size = 0; - return Allocator::Get(data.manager.db).ReallocateData(pointer, old_size, size); + return BlockAllocator::Get(data.manager.db).ReallocateData(pointer, old_size, size); } Allocator &BufferAllocator::Get(ClientContext &context) { diff --git a/test/api/CMakeLists.txt b/test/api/CMakeLists.txt index 7376b9d0ea2c..40592ee8f538 100644 --- a/test/api/CMakeLists.txt +++ b/test/api/CMakeLists.txt @@ -39,6 +39,7 @@ set(TEST_API_OBJECTS test_windows_unicode_path.cpp test_object_cache.cpp test_partial_executor.cpp + test_buffer_pool_allocator_flush.cpp test_buffer_pool_eviction.cpp test_block_allocator_tls.cpp) diff --git a/test/api/test_buffer_pool_allocator_flush.cpp b/test/api/test_buffer_pool_allocator_flush.cpp new file mode 100644 index 000000000000..ecce3417cfe4 --- /dev/null +++ b/test/api/test_buffer_pool_allocator_flush.cpp @@ -0,0 +1,62 @@ +#include "catch.hpp" +#include "duckdb/main/database.hpp" +#include "duckdb/storage/block_allocator.hpp" +#include "duckdb/storage/buffer/buffer_pool.hpp" +#include "duckdb/storage/buffer_manager.hpp" +#include "test_helpers.hpp" + +using namespace duckdb; // NOLINT + +TEST_CASE("Buffer pool flushes the allocator after bulk deallocation", "[storage][buffer_pool]") { + DBConfig config; + config.options.maximum_memory = 256ULL * 1024 * 1024; + // the block allocator pool supports flushing on every platform, unlike the fallback allocator + config.options.block_allocator_size = 64ULL * 1024 * 1024; + DuckDB db(nullptr, &config); + auto &db_instance = *db.instance; + if (!BlockAllocator::Get(db_instance).SupportsFlush()) { + return; + } + auto &buffer_manager = BufferManager::GetBufferManager(db_instance); + auto &pool = db_instance.GetBufferPool(); + const auto block_size = buffer_manager.GetBlockAllocSize(); + const auto blocks_per_mb = (1024 * 1024) / block_size; + + vector live; + auto hold = [&](idx_t mb) { + for (idx_t i = 0; i < mb * blocks_per_mb; i++) { + live.push_back(buffer_manager.Allocate(MemoryTag::EXTENSION, block_size)); + } + }; + auto allocate_and_free = [&](idx_t mb) { + vector handles; + for (idx_t i = 0; i < mb * blocks_per_mb; i++) { + handles.push_back(buffer_manager.Allocate(MemoryTag::EXTENSION, block_size)); + } + }; + + // freed memory is flushed once at least 16MB of it could push resident memory past the 256MB limit + const auto flushes = pool.GetAllocatorFlushCount(); + hold(200); + allocate_and_free(20); + allocate_and_free(20); + // 200MB live plus 40MB freed still fits in the limit + auto handle = buffer_manager.Allocate(MemoryTag::EXTENSION, block_size); + REQUIRE(pool.GetAllocatorFlushCount() == flushes); + + // live memory climbs to 220MB while 40MB is unflushed: the flush happens once they could exceed the limit + allocate_and_free(20); + handle = buffer_manager.Allocate(MemoryTag::EXTENSION, block_size); + REQUIRE(pool.GetAllocatorFlushCount() == flushes + 1); + // the 20MB freed after that flush fit in the limit again + handle = buffer_manager.Allocate(MemoryTag::EXTENSION, block_size); + REQUIRE(pool.GetAllocatorFlushCount() == flushes + 1); + + // operator allocations are counted as well + auto &allocator = buffer_manager.GetBufferAllocator(); + for (idx_t i = 0; i < 60; i++) { + auto data = allocator.Allocate(1024 * 1024); + } + handle = buffer_manager.Allocate(MemoryTag::EXTENSION, block_size); + REQUIRE(pool.GetAllocatorFlushCount() == flushes + 2); +}