[CAS] Handle DataPool overflow in OnDiskGraphDB

Handle DataPool overflow in GraphDB:
* When the data pool can't hold more items, allocate the new entries
  on disk standalone (degraded performance and more disk space consumption)
* Teach UnifiedCAS to more aggressively chaining and reclaiming space if
  data pool is relatively full (currently hard code to 85%) but the CAS
  is not at the space limit for overall size.

rdar://139242950

Author: cachemeifyoucan

llvm/include/llvm/CAS/OnDiskGraphDB.h

@@ -302,6 +302,11 @@ class OnDiskGraphDB {
   /// large object if the process crashed right at the point of inserting it.
   size_t getStorageSize() const;
 
+  /// \returns The precentage of space utilization of hard space limits.
+  ///
+  /// Return value is an integer between 0 and 100 for percentage.
+  unsigned getHardStorageLimitUtilization() const;
+
   void print(raw_ostream &OS) const;
 
   /// How to fault-in nodes if an upstream database is used.

llvm/include/llvm/CAS/OnDiskHashMappedTrie.h

@@ -275,6 +275,7 @@ class OnDiskHashMappedTrie {
   }
 
   size_t size() const;
+  size_t capacity() const;
 
   /// Gets or creates a file at \p Path with a hash-mapped trie named \p
   /// TrieName. The hash size is \p NumHashBits (in bits) and the records store
@@ -366,6 +367,7 @@ class OnDiskDataAllocator {
   MutableArrayRef<uint8_t> getUserHeader();
 
   size_t size() const;
+  size_t capacity() const;
 
   static Expected<OnDiskDataAllocator>
   create(const Twine &Path, const Twine &TableName, uint64_t MaxFileSize,

llvm/include/llvm/CAS/OnDiskKeyValueDB.h

@@ -39,6 +39,13 @@ class OnDiskKeyValueDB {
     return Cache.size();
   }
 
+  /// \returns The precentage of space utilization of hard space limits.
+  ///
+  /// Return value is an integer between 0 and 100 for percentage.
+  unsigned getHardStorageLimitUtilization() const {
+    return Cache.size() * 100ULL / Cache.capacity();
+  }
+
   /// Open the on-disk store from a directory.
   ///
   /// \param Path directory for the on-disk store. The directory will be created

llvm/lib/CAS/OnDiskGraphDB.cpp

@@ -1259,12 +1259,33 @@ Error OnDiskGraphDB::store(ObjectID ID, ArrayRef<ObjectID> Refs,
   SmallString<256> Path;
   std::optional<MappedTempFile> File;
   std::optional<uint64_t> FileSize;
+  auto AllocStandaloneFile = [&](size_t Size) -> Expected<char *> {
+    getStandalonePath(TrieRecord::getStandaloneFileSuffix(
+                          TrieRecord::StorageKind::Standalone),
+                      I, Path);
+    if (Error E = createTempFile(Path, Size).moveInto(File))
+      return std::move(E);
+    assert(File->size() == Size);
+    FileSize = Size;
+    SK = TrieRecord::StorageKind::Standalone;
+    return File->data();
+  };
   auto Alloc = [&](size_t Size) -> Expected<char *> {
     if (Size <= TrieRecord::MaxEmbeddedSize) {
       SK = TrieRecord::StorageKind::DataPool;
       auto P = DataPool.allocate(Size);
-      if (LLVM_UNLIKELY(!P))
-        return P.takeError();
+      if (LLVM_UNLIKELY(!P)) {
+        char *NewAlloc = nullptr;
+        auto NewE = handleErrors(
+            P.takeError(), [&](std::unique_ptr<StringError> E) -> Error {
+              if (E->convertToErrorCode() == std::errc::not_enough_memory)
+                return AllocStandaloneFile(Size).moveInto(NewAlloc);
+              return Error(std::move(E));
+            });
+        if (!NewE)
+          return NewAlloc;
+        return std::move(NewE);
+      }
       PoolOffset = P->getOffset();
       LLVM_DEBUG({
         dbgs() << "pool-alloc addr=" << (void *)PoolOffset.get()
@@ -1273,15 +1294,9 @@ Error OnDiskGraphDB::store(ObjectID ID, ArrayRef<ObjectID> Refs,
       });
       return (*P)->data();
     }
-
-    SK = TrieRecord::StorageKind::Standalone;
-    getStandalonePath(TrieRecord::getStandaloneFileSuffix(SK), I, Path);
-    if (Error E = createTempFile(Path, Size).moveInto(File))
-      return std::move(E);
-    assert(File->size() == Size);
-    FileSize = Size;
-    return File->data();
+    return AllocStandaloneFile(Size);
   };
+
   DataRecordHandle Record;
   if (Error E =
           DataRecordHandle::createWithError(Alloc, Input).moveInto(Record))
@@ -1350,6 +1365,12 @@ size_t OnDiskGraphDB::getStorageSize() const {
   return Index.size() + DataPool.size() + getStandaloneStorageSize();
 }
 
+unsigned OnDiskGraphDB::getHardStorageLimitUtilization() const {
+  unsigned IndexPercent = Index.size() * 100ULL / Index.capacity();
+  unsigned DataPercent = DataPool.size() * 100ULL / DataPool.capacity();
+  return std::max(IndexPercent, DataPercent);
+}
+
 static bool useSmallMappedFiles(const Twine &P) {
   // macOS tmpfs does not support sparse tails.
 #if defined(__APPLE__) && __has_include(<sys/mount.h>)

llvm/lib/CAS/OnDiskHashMappedTrie.cpp

@@ -1063,6 +1063,9 @@ static Error checkTable(StringRef Label, size_t Expected, size_t Observed,
 }
 
 size_t OnDiskHashMappedTrie::size() const { return Impl->File.size(); }
+size_t OnDiskHashMappedTrie::capacity() const {
+  return Impl->File.getRegion().size();
+}
 
 Expected<OnDiskHashMappedTrie>
 OnDiskHashMappedTrie::create(const Twine &PathTwine, const Twine &TrieNameTwine,
@@ -1321,6 +1324,9 @@ MutableArrayRef<uint8_t> OnDiskDataAllocator::getUserHeader() {
 }
 
 size_t OnDiskDataAllocator::size() const { return Impl->File.size(); }
+size_t OnDiskDataAllocator::capacity() const {
+  return Impl->File.getRegion().size();
+}
 
 OnDiskDataAllocator::OnDiskDataAllocator(std::unique_ptr<ImplType> Impl)
     : Impl(std::move(Impl)) {}
@@ -1365,6 +1371,14 @@ size_t OnDiskHashMappedTrie::size() const {
   report_fatal_error("not supported");
 }
 
+size_t OnDiskHashMappedTrie::capacity() const {
+  report_fatal_error("not supported");
+}
+
+size_t OnDiskDataAllocator::capacity() const {
+  report_fatal_error("not supported");
+}
+
 struct OnDiskDataAllocator::ImplType {};
 
 Expected<OnDiskDataAllocator> OnDiskDataAllocator::create(

llvm/lib/CAS/UnifiedOnDiskCache.cpp

@@ -271,6 +271,14 @@ bool UnifiedOnDiskCache::hasExceededSizeLimit() const {
   uint64_t CurSizeLimit = SizeLimit;
   if (!CurSizeLimit)
     return false;
+
+  // If the hard limit is beyond 85%, declare above limit and request clean up.
+  unsigned CurrentPrecent =
+      std::max(PrimaryGraphDB->getHardStorageLimitUtilization(),
+               PrimaryKVDB->getHardStorageLimitUtilization());
+  if (CurrentPrecent > 85)
+    return true;
+
   // We allow each of the directories in the chain to reach up to half the
   // intended size limit. Check whether the primary directory has exceeded half
   // the limit or not, in order to decide whether we need to start a new chain.

llvm/unittests/CAS/OnDiskGraphDBTest.cpp

@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "CASTestConfig.h"
 #include "OnDiskCommonUtils.h"
 #include "llvm/Testing/Support/Error.h"
 #include "llvm/Testing/Support/SupportHelpers.h"
@@ -284,4 +285,29 @@ TEST(OnDiskGraphDBTest, FaultInPolicyConflict) {
                            OnDiskGraphDB::FaultInPolicy::SingleNode);
 }
 
+#if defined(EXPENSIVE_CHECKS)
+TEST(OnDiskGraphDBTest, SpaceLimit) {
+  setMaxOnDiskCASMappingSize();
+  unittest::TempDir Temp("ondiskcas", /*Unique=*/true);
+  std::unique_ptr<OnDiskGraphDB> DB;
+  ASSERT_THAT_ERROR(
+      OnDiskGraphDB::open(Temp.path(), "blake3", sizeof(HashType)).moveInto(DB),
+      Succeeded());
+
+  std::optional<ObjectID> ID;
+  std::string Data(500, '0');
+  auto storeSmallObject = [&]() {
+    SmallVector<ObjectID, 1> Refs;
+    if (ID)
+      Refs.push_back(*ID);
+    ASSERT_THAT_ERROR(store(*DB, Data, Refs).moveInto(ID), Succeeded());
+  };
+
+  // Insert enough small elements to overflow the data pool.
+  for (unsigned I = 0; I < 1024 * 256; ++I)
+    storeSmallObject();
+
+  EXPECT_GE(DB->getHardStorageLimitUtilization(), 99U);
+}
+#endif
 #endif // LLVM_ENABLE_ONDISK_CAS