diff --git a/velox/common/memory/HashStringAllocator.cpp b/velox/common/memory/HashStringAllocator.cpp index f18df41b13f12..6496733b2e007 100644 --- a/velox/common/memory/HashStringAllocator.cpp +++ b/velox/common/memory/HashStringAllocator.cpp @@ -83,6 +83,8 @@ std::string HashStringAllocator::Header::toString() { HashStringAllocator::~HashStringAllocator() { clear(); + VELOX_CHECK_EQ(cumulativeBytes_, 0); + VELOX_CHECK_EQ(sizeFromPool_, 0); } void HashStringAllocator::clear() { @@ -99,6 +101,41 @@ void HashStringAllocator::clear() { for (auto i = 0; i < kNumFreeLists; ++i) { new (&free_[i]) CompactDoubleList(); } + + static const auto kHugePageSize = memory::AllocationTraits::kHugePageSize; + for (auto i = 0; i < pool_.numRanges(); ++i) { + auto topRange = pool_.rangeAt(i); + auto topRangeSize = topRange.size(); + if (topRangeSize >= kHugePageSize) { + VELOX_CHECK_EQ(0, topRangeSize % kHugePageSize); + } + + for (int64_t subRangeStart = 0; subRangeStart < topRangeSize; + subRangeStart += kHugePageSize) { + auto range = folly::Range( + topRange.data() + subRangeStart, + std::min(topRangeSize, kHugePageSize)); + const auto size = range.size() - simd::kPadding; + auto* end = reinterpret_cast(range.data() + size); + auto* header = reinterpret_cast(range.data()); + while (header != end) { + VELOX_CHECK_GE(reinterpret_cast(header), range.data()); + VELOX_CHECK_LT( + reinterpret_cast(header), reinterpret_cast(end)); + VELOX_CHECK_LE( + reinterpret_cast(header->end()), + reinterpret_cast(end)); + + if (header->isFree()) { + // No-op + } else { + // Continued block & Non-free block. + cumulativeBytes_ -= header->size() + sizeof(Header); + } + header = reinterpret_cast(header->end()); + } + } + } pool_.clear(); } @@ -286,7 +323,7 @@ void HashStringAllocator::newRange( // The last bytes of the last range are no longer payload. So do not count // them in size and do not overwrite them if overwriting the multi-range // entry. Set position at the new end. - lastRange->size -= sizeof(void*); + lastRange->size -= Header::kContinuedPtrSize; lastRange->position = std::min(lastRange->size, lastRange->position); } *range = ByteRange{ @@ -412,7 +449,7 @@ HashStringAllocator::Header* HashStringAllocator::allocateFromFreeList( if (next != nullptr) { next->clearPreviousFree(); } - cumulativeBytes_ += found->size(); + cumulativeBytes_ += found->size() + sizeof(Header); if (isFinalSize) { freeRestOfBlock(found, preferredSize); } @@ -434,7 +471,7 @@ void HashStringAllocator::free(Header* header) { } else { VELOX_CHECK(!headerToFree->isFree()); freeBytes_ += headerToFree->size() + sizeof(Header); - cumulativeBytes_ -= headerToFree->size(); + cumulativeBytes_ -= headerToFree->size() + sizeof(Header); Header* next = headerToFree->next(); if (next != nullptr) { VELOX_CHECK(!next->isPreviousFree()); @@ -581,7 +618,7 @@ inline bool HashStringAllocator::storeStringFast( reinterpret_cast(freeHeader->begin())); header->setSize(roundedBytes); freeBytes_ -= spaceTaken; - cumulativeBytes_ += roundedBytes; + cumulativeBytes_ += roundedBytes + sizeof(Header); } else { header = allocateFromFreeList(roundedBytes, true, true, kNumFreeLists - 1); @@ -703,7 +740,7 @@ int64_t HashStringAllocator::checkConsistency() const { // header is readable and not free. auto* continued = header->nextContinued(); VELOX_CHECK(!continued->isFree()); - allocatedBytes += header->size() - sizeof(void*); + allocatedBytes += header->size() - Header::kContinuedPtrSize; } else { allocatedBytes += header->size(); } diff --git a/velox/common/memory/HashStringAllocator.h b/velox/common/memory/HashStringAllocator.h index 163de0da5513d..58721cfb94bde 100644 --- a/velox/common/memory/HashStringAllocator.h +++ b/velox/common/memory/HashStringAllocator.h @@ -305,7 +305,8 @@ class HashStringAllocator : public StreamArena { /// the pointer because in the worst case we would have one allocation that /// chains many small free blocks together via kContinued. uint64_t freeSpace() const { - int64_t minFree = freeBytes_ - numFree_ * (sizeof(Header) + sizeof(void*)); + int64_t minFree = + freeBytes_ - numFree_ * (sizeof(Header) + Header::kContinuedPtrSize); VELOX_CHECK_GE(minFree, 0, "Guaranteed free space cannot be negative"); return minFree; } diff --git a/velox/common/memory/tests/HashStringAllocatorTest.cpp b/velox/common/memory/tests/HashStringAllocatorTest.cpp index e8541f24d77dc..3db3d6ea91e7e 100644 --- a/velox/common/memory/tests/HashStringAllocatorTest.cpp +++ b/velox/common/memory/tests/HashStringAllocatorTest.cpp @@ -99,6 +99,36 @@ class HashStringAllocatorTest : public testing::Test { folly::Random::DefaultGenerator rng_; }; +TEST_F(HashStringAllocatorTest, multipleFree) { + ASSERT_NO_THROW(allocator_->toString()); + + auto h1 = allocate(123); + ASSERT_EQ(h1->toString(), "size: 123"); + + allocator_->free(h1); + // Running free() multiple times on the same memory block should result in an + // error. + VELOX_ASSERT_THROW(allocator_->free(h1), ""); +} + +TEST_F(HashStringAllocatorTest, multipleFreeAncCheckCumulativeBytes) { + ASSERT_NO_THROW(allocator_->toString()); + + auto h1 = allocate(123); + auto h2 = allocate(456); + auto h3 = allocate(789); + + ASSERT_EQ(h1->toString(), "size: 123"); + ASSERT_EQ(h2->toString(), "size: 456"); + ASSERT_EQ(h3->toString(), "size: 789"); + + allocator_->free(h3); + allocator_->free(h2); + allocator_->free(h1); + // After all blocks execute free(), the allocated bytes should be equal to 0. + ASSERT_EQ(allocator_->cumulativeBytes(), 0); +} + TEST_F(HashStringAllocatorTest, headerToString) { ASSERT_NO_THROW(allocator_->toString());