[llvm/CAS] UnifiedOnDiskCache::getStorageSize() should include the size of the chained DB

llvm/include/llvm/CAS/UnifiedOnDiskCache.h

@@ -129,6 +129,9 @@ class UnifiedOnDiskCache {
   Expected<std::optional<ObjectID>>
   faultInFromUpstreamKV(ArrayRef<uint8_t> Key);
 
+  /// \returns the storage size of the primary directory.
+  uint64_t getPrimaryStorageSize() const;
+
   std::string RootPath;
   std::atomic<uint64_t> SizeLimit;
 

llvm/lib/CAS/UnifiedOnDiskCache.cpp

@@ -251,6 +251,15 @@ void UnifiedOnDiskCache::setSizeLimit(std::optional<uint64_t> SizeLimit) {
 }
 
 uint64_t UnifiedOnDiskCache::getStorageSize() const {
+  uint64_t TotalSize = getPrimaryStorageSize();
+  if (UpstreamGraphDB)
+    TotalSize += UpstreamGraphDB->getStorageSize();
+  if (UpstreamKVDB)
+    TotalSize += UpstreamKVDB->getStorageSize();
+  return TotalSize;
+}
+
+uint64_t UnifiedOnDiskCache::getPrimaryStorageSize() const {
   return PrimaryGraphDB->getStorageSize() + PrimaryKVDB->getStorageSize();
 }
 
@@ -268,7 +277,7 @@ bool UnifiedOnDiskCache::hasExceededSizeLimit() const {
   // the primary has reached its own limit. Essentially in such situation we
   // prefer reclaiming the storage later in order to have more consistent cache
   // hits behavior.
-  return (CurSizeLimit / 2) < getStorageSize();
+  return (CurSizeLimit / 2) < getPrimaryStorageSize();
 }
 
 Error UnifiedOnDiskCache::close(bool CheckSizeLimit) {

llvm/unittests/CAS/UnifiedOnDiskCacheTest.cpp

@@ -151,6 +151,7 @@ TEST(UnifiedOnDiskCacheTest, Basic) {
   while (!UniDB->hasExceededSizeLimit()) {
     storeBigObject(Index++);
   }
+  PrevStoreSize = UniDB->getStorageSize();
   ASSERT_THAT_ERROR(UniDB->close(), Succeeded());
   EXPECT_TRUE(UniDB->needsGarbaseCollection());