Fix HashStringAllocator::clear() and cumulativeBytes_

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<char*>(
+          topRange.data() + subRangeStart,
+          std::min<int64_t>(topRangeSize, kHugePageSize));
+      const auto size = range.size() - simd::kPadding;
+      auto* end = reinterpret_cast<Header*>(range.data() + size);
+      auto* header = reinterpret_cast<Header*>(range.data());
+      while (header != end) {
+        VELOX_CHECK_GE(reinterpret_cast<char*>(header), range.data());
+        VELOX_CHECK_LT(
+            reinterpret_cast<char*>(header), reinterpret_cast<char*>(end));
+        VELOX_CHECK_LE(
+            reinterpret_cast<char*>(header->end()),
+            reinterpret_cast<char*>(end));
+
+        if (header->isFree()) {
+          // No-op
+        } else {
+          // Continued block & Non-free block.
+          cumulativeBytes_ -= header->size() + sizeof(Header);
+        }
+        header = reinterpret_cast<Header*>(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<CompactDoubleList*>(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();
         }

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;
   }

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());