diff --git a/bolt/dwio/common/BufferUtil.h b/bolt/dwio/common/BufferUtil.h index bec58c3fb..51f4e5a2c 100644 --- a/bolt/dwio/common/BufferUtil.h +++ b/bolt/dwio/common/BufferUtil.h @@ -40,6 +40,7 @@ inline void ensureCapacity( bolt::memory::MemoryPool* pool) { if (!data || !data->isMutable() || data->capacity() < BaseVector::byteSize(capacity)) { + data.reset(); data = AlignedBuffer::allocate(capacity, pool); } } diff --git a/bolt/dwio/common/tests/DataBufferTests.cpp b/bolt/dwio/common/tests/DataBufferTests.cpp index 23713f6b8..1d7de0de4 100644 --- a/bolt/dwio/common/tests/DataBufferTests.cpp +++ b/bolt/dwio/common/tests/DataBufferTests.cpp @@ -33,6 +33,7 @@ #include #include "bolt/common/base/tests/GTestUtils.h" #include "bolt/common/memory/Memory.h" +#include "bolt/dwio/common/BufferUtil.h" #include "bolt/dwio/common/DataBuffer.h" namespace bytedance { namespace bolt { @@ -188,6 +189,28 @@ TEST_F(DataBufferTest, Move) { } ASSERT_EQ(0, pool_->currentBytes()); } + +TEST_F(DataBufferTest, EnsureCapacityReleasesOldBufferBeforeAllocatingNewOne) { + constexpr size_t kInitialSize = 4 << 20; + const auto oldAllocationBytes = + pool_->preferredSize(kInitialSize + AlignedBuffer::kPaddedSize); + const auto oldCapacity = oldAllocationBytes - AlignedBuffer::kPaddedSize; + const auto requestedCapacity = oldCapacity + 1; + const auto newAllocationBytes = + pool_->preferredSize(requestedCapacity + AlignedBuffer::kPaddedSize); + + auto rootPool = memoryManager()->addRootPool( + "ensureCapacityReleaseBeforeAlloc", newAllocationBytes); + auto leafPool = rootPool->addLeafChild("leaf"); + BufferPtr buffer = + AlignedBuffer::allocate(kInitialSize, leafPool.get()); + ASSERT_EQ(buffer->capacity(), oldCapacity); + + dwio::common::ensureCapacity(buffer, requestedCapacity, leafPool.get()); + + EXPECT_GE(buffer->capacity(), requestedCapacity); + EXPECT_LE(leafPool->currentBytes(), newAllocationBytes); +} } // namespace common } // namespace dwio } // namespace bolt