Flush the allocator when freed memory could push RSS past the limit

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<idx_t> deallocated_since_flush {0};
 
 	//! Block size (power of two)
 	const idx_t block_size;

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<idx_t> maximum_memory;
 	//! If bulk deallocation larger than this occurs, flush outstanding allocations
 	atomic<idx_t> allocator_bulk_deallocation_flush_threshold;
+	//! Number of allocator flushes
+	atomic<idx_t> 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

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<uint32_t> to_free_buffer;

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,14 +352,15 @@ 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)};
 }
 
 BufferPool::EvictionResult BufferPool::EvictBlocks(MemoryTag tag, idx_t extra_memory, idx_t memory_limit,
                                                    unique_ptr<FileBuffer> *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;

src/storage/standard_buffer_manager.cpp

@@ -705,15 +705,15 @@ 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) {
 	auto &data = private_data->Cast<BufferAllocatorData>();
 	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) {

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)
 

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<BufferHandle> 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<BufferHandle> 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);
+}