Add memory usage statistics for slabs and allocation classes

cachelib/allocator/Cache.h

@@ -100,6 +100,9 @@ class CacheBase {
   // @param poolId   the pool id
   virtual PoolStats getPoolStats(PoolId poolId) const = 0;
 
+  virtual AllocationClassBaseStat getAllocationClassStats(TierId, PoolId pid, ClassId cid)
+      const = 0;
+
   // @param poolId   the pool id
   virtual AllSlabReleaseEvents getAllSlabReleaseEvents(PoolId poolId) const = 0;
 

cachelib/allocator/CacheAllocator-inl.h

@@ -2506,6 +2506,44 @@ PoolStats CacheAllocator<CacheTrait>::getPoolStats(PoolId poolId) const {
   return ret;
 }
 
+template <typename CacheTrait>
+double CacheAllocator<CacheTrait>::slabsApproxFreePercentage(TierId tid) const
+{
+  return allocator_[tid]->approxFreeSlabsPercentage();
+}
+
+template <typename CacheTrait>
+AllocationClassBaseStat CacheAllocator<CacheTrait>::getAllocationClassStats(
+  TierId tid, PoolId pid, ClassId cid) const {
+  const auto &ac = allocator_[tid]->getPool(pid).getAllocationClass(cid);
+
+  AllocationClassBaseStat stats{};
+  stats.allocSize = ac.getAllocSize();
+  stats.memorySize = ac.getNumSlabs() * Slab::kSize;
+
+  if (slabsApproxFreePercentage(tid) > 0.0) {
+    auto totalMemory = MemoryAllocator::getMemorySize(memoryTierSize(tid));
+    auto freeMemory = static_cast<double>(totalMemory) * slabsApproxFreePercentage(tid) / 100.0;
+
+    // amount of free memory which has the same ratio to entire free memory as
+    // this allocation class memory size has to used memory
+    auto scaledFreeMemory = static_cast<size_t>(freeMemory * stats.memorySize / totalMemory);
+
+    auto acAllocatedMemory = (100.0 - ac.approxFreePercentage()) / 100.0 * ac.getNumSlabs() * Slab::kSize;
+    auto acMaxAvailableMemory = ac.getNumSlabs() * Slab::kSize + scaledFreeMemory;
+
+    if (acMaxAvailableMemory == 0) {
+      stats.approxFreePercent = 100.0;
+    } else {
+      stats.approxFreePercent = 100.0 - 100.0 * acAllocatedMemory / acMaxAvailableMemory;
+    }
+  } else {
+    stats.approxFreePercent = ac.approxFreePercentage();
+  }
+
+  return stats;
+}
+
 template <typename CacheTrait>
 PoolEvictionAgeStats CacheAllocator<CacheTrait>::getPoolEvictionAgeStats(
     PoolId pid, unsigned int slabProjectionLength) const {
@@ -3613,6 +3651,10 @@ CacheMemoryStats CacheAllocator<CacheTrait>::getCacheMemoryStats() const {
   size_t compactCacheSize = std::accumulate(
       ccCachePoolIds.begin(), ccCachePoolIds.end(), 0ULL, addSize);
 
+  std::vector<double> slabsApproxFreePercentages;
+  for (TierId tid = 0; tid < numTiers_; tid++)
+    slabsApproxFreePercentages.push_back(slabsApproxFreePercentage(tid));
+
   return CacheMemoryStats{totalCacheSize,
                           regularCacheSize,
                           compactCacheSize,
@@ -3621,7 +3663,8 @@ CacheMemoryStats CacheAllocator<CacheTrait>::getCacheMemoryStats() const {
                           allocator_[currentTier()]->getUnreservedMemorySize(),
                           nvmCache_ ? nvmCache_->getSize() : 0,
                           util::getMemAvailable(),
-                          util::getRSSBytes()};
+                          util::getRSSBytes(),
+                          slabsApproxFreePercentages};
 }
 
 template <typename CacheTrait>

cachelib/allocator/CacheAllocator.h

@@ -1064,6 +1064,10 @@ class CacheAllocator : public CacheBase {
   // return cache's memory usage stats.
   CacheMemoryStats getCacheMemoryStats() const override final;
 
+  // return basic stats for Allocation Class
+  AllocationClassBaseStat getAllocationClassStats(TierId tid, PoolId pid, ClassId cid)
+      const override final;
+
   // return the nvm cache stats map
   std::unordered_map<std::string, double> getNvmCacheStatsMap()
       const override final;
@@ -1208,6 +1212,8 @@ class CacheAllocator : public CacheBase {
 #pragma GCC diagnostic pop
 
  private:
+  double slabsApproxFreePercentage(TierId tid) const;
+
   // wrapper around Item's refcount and active handle tracking
   FOLLY_ALWAYS_INLINE void incRef(Item& it);
   FOLLY_ALWAYS_INLINE RefcountWithFlags::Value decRef(Item& it);

cachelib/allocator/CacheStats.h

@@ -98,6 +98,17 @@ struct MMContainerStat {
   uint64_t numTailAccesses;
 };
 
+struct AllocationClassBaseStat {
+  // size of allocation class
+  size_t allocSize{0};
+
+  // size of memory assigned to this allocation class
+  size_t memorySize{0};
+
+  // percent of free memory in this class
+  double approxFreePercent{0.0};
+};
+
 // cache related stats for a given allocation class.
 struct CacheStat {
   // allocation size for this container.
@@ -521,6 +532,9 @@ struct CacheMemoryStats {
 
   // rss size of the process
   size_t memRssSize{0};
+
+  // percentage of free slabs
+  std::vector<double> slabsApproxFreePercentages{0.0};
 };
 
 // Stats for compact cache

cachelib/allocator/memory/AllocationClass.cpp

@@ -51,6 +51,7 @@ AllocationClass::AllocationClass(ClassId classId,
       allocationSize_(allocSize),
       slabAlloc_(s),
       freedAllocations_{slabAlloc_.createSingleTierPtrCompressor<FreeAlloc>()} {
+  curAllocatedSlabs_ = allocatedSlabs_.size();
   checkState();
 }
 
@@ -87,6 +88,12 @@ void AllocationClass::checkState() const {
         "Current allocation slab {} is not in allocated slabs list",
         currSlab_));
   }
+
+  if (curAllocatedSlabs_ != allocatedSlabs_.size()) {
+    throw std::invalid_argument(folly::sformat(
+      "Mismatch in allocated slabs numbers"
+    ));
+  }
 }
 
 // TODO(stuclar): Add poolId to the metadata to be serialized when cache shuts
@@ -116,10 +123,12 @@ AllocationClass::AllocationClass(
     freeSlabs_.push_back(slabAlloc_.getSlabForIdx(freeSlabIdx));
   }
 
+  curAllocatedSlabs_ = allocatedSlabs_.size();
   checkState();
 }
 
 void AllocationClass::addSlabLocked(Slab* slab) {
+  curAllocatedSlabs_.fetch_add(1, std::memory_order_relaxed);
   canAllocate_ = true;
   auto header = slabAlloc_.getSlabHeader(slab);
   header->classId = classId_;
@@ -168,6 +177,7 @@ void* AllocationClass::allocateLocked() {
   }
 
   XDCHECK(canAllocate_);
+  curAllocatedSize_.fetch_add(getAllocSize(), std::memory_order_relaxed);
 
   // grab from the free list if possible.
   if (!freedAllocations_.empty()) {
@@ -270,6 +280,7 @@ SlabReleaseContext AllocationClass::startSlabRelease(
                          slab, getId()));
     }
     *allocIt = allocatedSlabs_.back();
+    curAllocatedSlabs_.fetch_sub(1, std::memory_order_relaxed);
     allocatedSlabs_.pop_back();
 
     // if slab is being carved currently, then update slabReleaseAllocMap
@@ -510,6 +521,7 @@ void AllocationClass::abortSlabRelease(const SlabReleaseContext& context) {
     }
     slabReleaseAllocMap_.erase(slabPtrVal);
     allocatedSlabs_.push_back(const_cast<Slab*>(slab));
+    curAllocatedSlabs_.fetch_add(1, std::memory_order_relaxed);
     // restore the classId and allocSize
     header->classId = classId_;
     header->allocSize = allocationSize_;
@@ -660,6 +672,8 @@ void AllocationClass::free(void* memory) {
     freedAllocations_.insert(*reinterpret_cast<FreeAlloc*>(memory));
     canAllocate_ = true;
   });
+
+  curAllocatedSize_.fetch_sub(getAllocSize(), std::memory_order_relaxed);
 }
 
 serialization::AllocationClassObject AllocationClass::saveState() const {
@@ -722,3 +736,12 @@ std::vector<bool>& AllocationClass::getSlabReleaseAllocMapLocked(
   const auto slabPtrVal = getSlabPtrValue(slab);
   return slabReleaseAllocMap_.at(slabPtrVal);
 }
+
+double AllocationClass::approxFreePercentage() const {
+  if (getNumSlabs() == 0) {
+    return 100.0;
+  }
+
+  return 100.0 - 100.0 * static_cast<double>(curAllocatedSize_.load(std::memory_order_relaxed)) /
+    static_cast<double>(getNumSlabs() * Slab::kSize);
+}

cachelib/allocator/memory/AllocationClass.h

@@ -90,10 +90,7 @@ class AllocationClass {
 
   // total number of slabs under this AllocationClass.
   unsigned int getNumSlabs() const {
-    return lock_->lock_combine([this]() {
-      return static_cast<unsigned int>(freeSlabs_.size() +
-                                       allocatedSlabs_.size());
-    });
+    return curAllocatedSlabs_.load(std::memory_order_relaxed);
   }
 
   // fetch stats about this allocation class.
@@ -309,6 +306,9 @@ class AllocationClass {
   // @throw std::logic_error if the object state can not be serialized
   serialization::AllocationClassObject saveState() const;
 
+  // approximate percent of free memory inside this allocation class
+  double approxFreePercentage() const;
+
  private:
   // check if the state of the AllocationClass is valid and if not, throws an
   // std::invalid_argument exception. This is intended for use in
@@ -468,6 +468,12 @@ class AllocationClass {
 
   std::atomic<int64_t> activeReleases_{0};
 
+  // amount of memory currently allocated by this AC
+  std::atomic<size_t> curAllocatedSize_{0};
+
+  // total number of slabs under this AllocationClass.
+  std::atomic<size_t> curAllocatedSlabs_{0};
+
   // stores the list of outstanding allocations for a given slab. This is
   // created when we start a slab release process and if there are any active
   // allocaitons need to be marked as free.

cachelib/allocator/memory/MemoryAllocator.h

@@ -416,6 +416,14 @@ class MemoryAllocator {
     return memoryPoolManager_.getPoolIds();
   }
 
+  double approxFreeSlabsPercentage() const {
+    if (slabAllocator_.getNumUsableAndAdvisedSlabs() == 0)
+      return 100.0;
+
+    return 100.0 - 100.0 * static_cast<double>(slabAllocator_.approxNumSlabsAllocated()) /
+     slabAllocator_.getNumUsableAndAdvisedSlabs();
+  }
+
   // fetches the memory pool for the id if one exists. This is purely to get
   // information out of the pool.
   //

cachelib/allocator/memory/SlabAllocator.cpp

@@ -359,6 +359,8 @@ Slab* SlabAllocator::makeNewSlab(PoolId id) {
     return nullptr;
   }
 
+  numSlabsAllocated_.fetch_add(1, std::memory_order_relaxed);
+
   memoryPoolSize_[id] += sizeof(Slab);
   // initialize the header for the slab.
   initializeHeader(slab, id);
@@ -374,6 +376,8 @@ void SlabAllocator::freeSlab(Slab* slab) {
   }
 
   memoryPoolSize_[header->poolId] -= sizeof(Slab);
+  numSlabsAllocated_.fetch_sub(1, std::memory_order_relaxed);
+
   // grab the lock
   LockHolder l(lock_);
   freeSlabs_.push_back(slab);

cachelib/allocator/memory/SlabAllocator.h

@@ -323,7 +323,13 @@ class SlabAllocator {
                                    memorySize_);
   }
 
- private:
+  size_t approxNumSlabsAllocated() const {
+    return numSlabsAllocated_.load(std::memory_order_relaxed);
+  }
+
+private:
+  std::atomic<size_t> numSlabsAllocated_{0};
+
   // null Slab* presenttation. With 4M Slab size, a valid slab index would never
   // reach 2^16 - 1;
   static constexpr SlabIdx kNullSlabIdx = std::numeric_limits<SlabIdx>::max();

cachelib/allocator/tests/CacheBaseTest.cpp

@@ -33,6 +33,11 @@ class CacheBaseTest : public CacheBase, public SlabAllocatorTestBase {
   const std::string getCacheName() const override { return cacheName; }
   const MemoryPool& getPool(PoolId) const override { return memoryPool_; }
   PoolStats getPoolStats(PoolId) const override { return PoolStats(); }
+  AllocationClassBaseStat getAllocationClassStats(TierId tid,
+                                                  PoolId,
+                                                  ClassId) const {
+    return AllocationClassBaseStat();
+  };
   AllSlabReleaseEvents getAllSlabReleaseEvents(PoolId) const override {
     return AllSlabReleaseEvents{};
   }

cachelib/cachebench/cache/Cache-inl.h

@@ -493,16 +493,28 @@ bool Cache<Allocator>::checkGet(ValueTracker::Index opId,
 
 template <typename Allocator>
 Stats Cache<Allocator>::getStats() const {
-  PoolStats aggregate = cache_->getPoolStats(pools_[0]);
+  PoolStats aggregate = cache_->getPoolStats(0);
   for (size_t pid = 1; pid < pools_.size(); pid++) {
     aggregate += cache_->getPoolStats(static_cast<PoolId>(pid));
   }
 
+  std::map<TierId, std::map<PoolId, std::map<ClassId, AllocationClassBaseStat>>> allocationClassStats{};
+
+  for (size_t pid = 0; pid < pools_.size(); pid++) {
+    auto cids = cache_->getPoolStats(static_cast<PoolId>(pid)).getClassIds();
+    for (TierId tid = 0; tid < cache_->getNumTiers(); tid++) {
+      for (auto cid : cids)
+        allocationClassStats[tid][pid][cid] = cache_->getAllocationClassStats(tid, pid, cid);
+    }
+  }
+
   const auto cacheStats = cache_->getGlobalCacheStats();
   const auto rebalanceStats = cache_->getSlabReleaseStats();
   const auto navyStats = cache_->getNvmCacheStatsMap();
 
   Stats ret;
+  ret.slabsApproxFreePercentages = cache_->getCacheMemoryStats().slabsApproxFreePercentages;
+  ret.allocationClassStats = allocationClassStats;
   ret.numEvictions = aggregate.numEvictions();
   ret.numItems = aggregate.numItems();
   ret.allocAttempts = cacheStats.allocAttempts;

cachelib/cachebench/cache/Cache.cpp

@@ -22,6 +22,10 @@ DEFINE_bool(report_api_latency,
             false,
             "Enable reporting cache API latency tracking");
 
+DEFINE_bool(report_memory_usage_stats,
+            false,
+            "Enable reporting statistics for each allocation class");
+
 namespace facebook {
 namespace cachelib {
 namespace cachebench {} // namespace cachebench

cachelib/cachebench/cache/Cache.h

@@ -33,6 +33,7 @@
 #include "cachelib/cachebench/util/CacheConfig.h"
 
 DECLARE_bool(report_api_latency);
+DECLARE_bool(report_memory_usage_stats);
 
 namespace facebook {
 namespace cachelib {
@@ -249,6 +250,10 @@ class Cache {
   // return the stats for the pool.
   PoolStats getPoolStats(PoolId pid) const { return cache_->getPoolStats(pid); }
 
+  AllocationClassBaseStat getAllocationClassStats(TierId tid, PoolId pid, ClassId cid) const {
+    return cache_->getAllocationClassStats(tid, pid, cid);
+  }
+
   // return the total number of inconsistent operations detected since start.
   unsigned int getInconsistencyCount() const {
     return inconsistencyCount_.load(std::memory_order_relaxed);