Add chain-aware getTotalAllocSize

Summary: Adds getTotalAllocSize which walks the chained item linked list and totals the allocation size of all chained items.

Reviewed By: rlyerly

Differential Revision: D110933811

fbshipit-source-id: 41a504a41e8fc16ceeab25de5f94079f714d4988

Author: IvanTopolcic

cachelib/allocator/CacheAllocator.h

@@ -1170,6 +1170,16 @@ class CacheAllocator : public CacheBase {
     return allocator_->getAllocInfo(memory);
   }
 
+  // total allocated size of the item and all of its chained items.
+  uint64_t getTotalAllocSize(const Item& item,
+                             folly::Range<ChainedItemIter> chainedItems) const {
+    uint64_t size = getAllocInfo(item.getMemory()).allocSize;
+    for (const auto& c : chainedItems) {
+      size += getAllocInfo(c.getMemory()).allocSize;
+    }
+    return size;
+  }
+
   // return the ids for the set of existing pools in this cache.
   std::set<PoolId> getPoolIds() const override final {
     return allocator_->getPoolIds();

cachelib/allocator/tests/AllocatorTypeTest.cpp

@@ -342,6 +342,10 @@ TYPED_TEST(BaseAllocatorTest, ChainedAllocsIteration) {
   this->testChainedAllocsIteration();
 }
 
+TYPED_TEST(BaseAllocatorTest, GetTotalAllocSize) {
+  this->testGetTotalAllocSize();
+}
+
 TYPED_TEST(BaseAllocatorTest, ReplaceChainedItem) {
   this->testReplaceChainedItem();
 }

cachelib/allocator/tests/BaseAllocatorTest.h

@@ -2501,6 +2501,38 @@ class BaseAllocatorTest : public AllocatorTest<AllocatorT> {
     }
   }
 
+  void testGetTotalAllocSize() {
+    typename AllocatorT::Config config;
+    config.configureChainedItems();
+    config.setCacheSize(10 * Slab::kSize);
+    AllocatorT alloc(config);
+    const size_t numBytes = alloc.getCacheMemoryStats().ramCacheSize;
+    auto poolId = alloc.addPool("foobar", numBytes);
+
+    auto parent = util::allocateAccessible(alloc, poolId, "parent", 100);
+    ASSERT_NE(nullptr, parent);
+    const int nChainedAllocs = 3;
+    for (int i = 0; i < nChainedAllocs; ++i) {
+      auto chained = alloc.allocateChainedItem(parent, 100 * (i + 1));
+      ASSERT_NE(nullptr, chained);
+      alloc.addChainedItem(parent, std::move(chained));
+    }
+
+    auto handle = alloc.find("parent");
+    auto chainedAllocs = alloc.viewAsChainedAllocs(handle);
+    auto chain = chainedAllocs.getChain();
+
+    // Independently sum the allocation-class size across the whole chain.
+    uint64_t expected = alloc.getAllocInfo(handle->getMemory()).allocSize;
+    for (const auto& c : chain) {
+      expected += alloc.getAllocInfo(c.getMemory()).allocSize;
+    }
+    // The chained parts must contribute beyond the parent's own block.
+    ASSERT_GT(expected, alloc.getAllocInfo(handle->getMemory()).allocSize);
+
+    ASSERT_EQ(expected, alloc.getTotalAllocSize(*handle, chain));
+  }
+
   // create a chain of allocations, replace the allocation and ensure that the
   // order is preserved (multithreaded version).
   void testChainedAllocsReplaceInChainMultithread() {