[CAS] Improve when MappedFileRegionBumpPtr is out of space

Fix a corner case for how MappedFileRegionBumpPtr returns the allocator
into a valid state after it runs out of space. It makes sure the
allocation that bump the allocator out of range needs to reset the
pointer into its original state.

Also propagate the failed allocation errors to make sure those errors
are properly handled without throwing fatal errors to crash the system.

Author: cachemeifyoucan

llvm/include/llvm/CAS/MappedFileRegionBumpPtr.h

@@ -63,11 +63,14 @@ class MappedFileRegionBumpPtr {
   }
 
   /// Allocate at least \p AllocSize. Rounds up to \a getAlign().
-  char *allocate(uint64_t AllocSize) {
-    return data() + allocateOffset(AllocSize);
+  Expected<char *> allocate(uint64_t AllocSize) {
+    auto Offset = allocateOffset(AllocSize);
+    if (LLVM_UNLIKELY(!Offset))
+      return Offset.takeError();
+    return data() + *Offset;
   }
   /// Allocate, returning the offset from \a data() instead of a pointer.
-  int64_t allocateOffset(uint64_t AllocSize);
+  Expected<int64_t> allocateOffset(uint64_t AllocSize);
 
   char *data() const { return Region.data(); }
   uint64_t size() const { return *BumpPtr; }

llvm/include/llvm/CAS/OnDiskGraphDB.h

@@ -271,7 +271,7 @@ class OnDiskGraphDB {
 
   /// Form a reference for the provided hash. The reference can be used as part
   /// of a CAS object even if it's not associated with an object yet.
-  ObjectID getReference(ArrayRef<uint8_t> Hash);
+  Expected<ObjectID> getReference(ArrayRef<uint8_t> Hash);
 
   /// Get an existing reference to the object \p Digest.
   ///
@@ -362,7 +362,7 @@ class OnDiskGraphDB {
   Error importFullTree(ObjectID PrimaryID, ObjectHandle UpstreamNode);
   Error importSingleNode(ObjectID PrimaryID, ObjectHandle UpstreamNode);
 
-  IndexProxy indexHash(ArrayRef<uint8_t> Hash);
+  Expected<IndexProxy> indexHash(ArrayRef<uint8_t> Hash);
 
   Error createStandaloneLeaf(IndexProxy &I, ArrayRef<char> Data);
 

llvm/include/llvm/CAS/OnDiskHashMappedTrie.h

@@ -15,6 +15,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/Support/Casting.h"
+#include "llvm/Support/Compiler.h"
 #include "llvm/Support/FileSystem.h"
 #include <atomic>
 #include <mutex>
@@ -253,23 +254,23 @@ class OnDiskHashMappedTrie {
   /// The in-memory \a HashMappedTrie uses LazyAtomicPointer to synchronize
   /// simultaneous writes, but that seems dangerous to use in a memory-mapped
   /// file in case a process crashes in the busy state.
-  pointer insertLazy(const_pointer Hint, ArrayRef<uint8_t> Hash,
-                     LazyInsertOnConstructCB OnConstruct = nullptr,
-                     LazyInsertOnLeakCB OnLeak = nullptr);
-  pointer insertLazy(ArrayRef<uint8_t> Hash,
-                     LazyInsertOnConstructCB OnConstruct = nullptr,
-                     LazyInsertOnLeakCB OnLeak = nullptr) {
+  Expected<pointer> insertLazy(const_pointer Hint, ArrayRef<uint8_t> Hash,
+                               LazyInsertOnConstructCB OnConstruct = nullptr,
+                               LazyInsertOnLeakCB OnLeak = nullptr);
+  Expected<pointer> insertLazy(ArrayRef<uint8_t> Hash,
+                               LazyInsertOnConstructCB OnConstruct = nullptr,
+                               LazyInsertOnLeakCB OnLeak = nullptr) {
     return insertLazy(const_pointer(), Hash, OnConstruct, OnLeak);
   }
 
-  pointer insert(const_pointer Hint, const ConstValueProxy &Value) {
+  Expected<pointer> insert(const_pointer Hint, const ConstValueProxy &Value) {
     return insertLazy(Hint, Value.Hash, [&](FileOffset, ValueProxy Allocated) {
       assert(Allocated.Hash == Value.Hash);
       assert(Allocated.Data.size() == Value.Data.size());
       llvm::copy(Value.Data, Allocated.Data.begin());
     });
   }
-  pointer insert(const ConstValueProxy &Value) {
+  Expected<pointer> insert(const ConstValueProxy &Value) {
     return insert(const_pointer(), Value);
   }
 
@@ -348,13 +349,15 @@ class OnDiskDataAllocator {
         const_cast<const OnDiskDataAllocator *>(this)->beginData(Offset));
   }
 
-  pointer allocate(size_t Size);
-  pointer save(ArrayRef<char> Data) {
-    pointer P = allocate(Data.size());
-    llvm::copy(Data, P->begin());
+  Expected<pointer> allocate(size_t Size);
+  Expected<pointer> save(ArrayRef<char> Data) {
+    auto P = allocate(Data.size());
+    if (LLVM_UNLIKELY(!P))
+      return P.takeError();
+    llvm::copy(Data, (*P)->begin());
     return P;
   }
-  pointer save(StringRef Data) {
+  Expected<pointer> save(StringRef Data) {
     return save(ArrayRef<char>(Data.begin(), Data.size()));
   }
 

llvm/lib/CAS/ActionCaches.cpp

@@ -17,6 +17,7 @@
 #include "llvm/Config/llvm-config.h"
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/BLAKE3.h"
+#include "llvm/Support/Compiler.h"
 #include "llvm/Support/Path.h"
 
 #define DEBUG_TYPE "action-caches"
@@ -212,13 +213,14 @@ UnifiedOnDiskActionCache::getImpl(ArrayRef<uint8_t> Key,
 Error UnifiedOnDiskActionCache::putImpl(ArrayRef<uint8_t> Key,
                                         const CASID &Result,
                                         bool /*Globally*/) {
-  ondisk::ObjectID Expected =
-      UniDB->getGraphDB().getReference(Result.getHash());
+  auto Expected = UniDB->getGraphDB().getReference(Result.getHash());
+  if (LLVM_UNLIKELY(!Expected))
+    return Expected.takeError();
   std::optional<ondisk::ObjectID> Observed;
-  if (Error E = UniDB->KVPut(Key, Expected).moveInto(Observed))
+  if (Error E = UniDB->KVPut(Key, *Expected).moveInto(Observed))
     return E;
 
-  if (Expected == Observed)
+  if (*Expected == Observed)
     return Error::success();
 
   return createResultCachePoisonedError(

llvm/lib/CAS/MappedFileRegionBumpPtr.cpp

@@ -220,15 +220,26 @@ void MappedFileRegionBumpPtr::initializeBumpPtr(int64_t BumpPtrOffset) {
            "Expected 0, or past the end of the BumpPtr itself");
 }
 
-int64_t MappedFileRegionBumpPtr::allocateOffset(uint64_t AllocSize) {
+static Error createAllocatorOutOfSpaceError() {
+  return createStringError(std::make_error_code(std::errc::not_enough_memory),
+                           "memory mapped file allocator is out of space");
+}
+
+Expected<int64_t> MappedFileRegionBumpPtr::allocateOffset(uint64_t AllocSize) {
   AllocSize = alignTo(AllocSize, getAlign());
   int64_t OldEnd = BumpPtr->fetch_add(AllocSize);
   int64_t NewEnd = OldEnd + AllocSize;
   if (LLVM_UNLIKELY(NewEnd > (int64_t)capacity())) {
-    // Try to return the allocation.
-    (void)BumpPtr->compare_exchange_strong(OldEnd, NewEnd);
-    report_fatal_error(
-        errorCodeToError(std::make_error_code(std::errc::not_enough_memory)));
+    // Return the allocation. If the start already passed the end, that means
+    // some other concurrent allocations already consumed all the capacity.
+    // There is no need to return the original value. If the start was not
+    // passed the end, current allocation certainly bumped it passed the end.
+    // All other allocation afterwards must have failed and current allocation
+    // is in charge of return the allocation back to a valid value.
+    if (OldEnd <= (int64_t)capacity())
+      (void)BumpPtr->exchange(OldEnd);
+
+    return createAllocatorOutOfSpaceError();
   }
   return OldEnd;
 }

llvm/lib/CAS/OnDiskCAS.cpp

@@ -9,6 +9,7 @@
 #include "BuiltinCAS.h"
 #include "llvm/CAS/OnDiskGraphDB.h"
 #include "llvm/CAS/UnifiedOnDiskCache.h"
+#include "llvm/Support/Compiler.h"
 #include "llvm/Support/Path.h"
 
 using namespace llvm;
@@ -125,10 +126,12 @@ Expected<ObjectRef> OnDiskCAS::storeImpl(ArrayRef<uint8_t> ComputedHash,
     IDs.push_back(convertRef(Ref));
   }
 
-  ondisk::ObjectID StoredID = DB->getReference(ComputedHash);
-  if (Error E = DB->store(StoredID, IDs, Data))
+  auto StoredID = DB->getReference(ComputedHash);
+  if (LLVM_UNLIKELY(!StoredID))
+    return StoredID.takeError();
+  if (Error E = DB->store(*StoredID, IDs, Data))
     return std::move(E);
-  return convertRef(StoredID);
+  return convertRef(*StoredID);
 }
 
 Error OnDiskCAS::forEachRef(ObjectHandle Node,

llvm/lib/CAS/OnDiskGraphDB.cpp

@@ -52,10 +52,13 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Alignment.h"
+#include "llvm/Support/Compiler.h"
 #include "llvm/Support/Errc.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Process.h"
+#include <optional>
 
 #define DEBUG_TYPE "on-disk-cas"
 
@@ -915,17 +918,21 @@ void OnDiskGraphDB::print(raw_ostream &OS) const {
   }
 }
 
-OnDiskGraphDB::IndexProxy OnDiskGraphDB::indexHash(ArrayRef<uint8_t> Hash) {
-  OnDiskHashMappedTrie::pointer P = Index.insertLazy(
+Expected<OnDiskGraphDB::IndexProxy>
+OnDiskGraphDB::indexHash(ArrayRef<uint8_t> Hash) {
+  auto P = Index.insertLazy(
       Hash, [](FileOffset TentativeOffset,
                OnDiskHashMappedTrie::ValueProxy TentativeValue) {
         assert(TentativeValue.Data.size() == sizeof(TrieRecord));
         assert(
             isAddrAligned(Align::Of<TrieRecord>(), TentativeValue.Data.data()));
         new (TentativeValue.Data.data()) TrieRecord();
       });
-  assert(P && "Expected insertion");
-  return getIndexProxyFromPointer(P);
+  if (LLVM_UNLIKELY(!P))
+    return P.takeError();
+
+  assert(*P && "Expected insertion");
+  return getIndexProxyFromPointer(*P);
 }
 
 OnDiskGraphDB::IndexProxy OnDiskGraphDB::getIndexProxyFromPointer(
@@ -937,9 +944,11 @@ OnDiskGraphDB::IndexProxy OnDiskGraphDB::getIndexProxyFromPointer(
                         reinterpret_cast<const TrieRecord *>(P->Data.data()))};
 }
 
-ObjectID OnDiskGraphDB::getReference(ArrayRef<uint8_t> Hash) {
-  IndexProxy I = indexHash(Hash);
-  return getExternalReference(I);
+Expected<ObjectID> OnDiskGraphDB::getReference(ArrayRef<uint8_t> Hash) {
+  auto I = indexHash(Hash);
+  if (LLVM_UNLIKELY(!I))
+    return I.takeError();
+  return getExternalReference(*I);
 }
 
 ObjectID OnDiskGraphDB::getExternalReference(const IndexProxy &I) {
@@ -954,10 +963,13 @@ OnDiskGraphDB::getExistingReference(ArrayRef<uint8_t> Digest) {
       return std::nullopt;
     std::optional<ObjectID> UpstreamID =
         UpstreamDB->getExistingReference(Digest);
-    if (!UpstreamID)
+    if (LLVM_UNLIKELY(!UpstreamID))
+      return std::nullopt;
+    auto Ref = expectedToOptional(indexHash(Digest));
+    if (!Ref)
       return std::nullopt;
     if (!I)
-      I.emplace(indexHash(Digest));
+      I.emplace(*Ref);
     return getExternalReference(*I);
   };
 
@@ -1246,14 +1258,16 @@ Error OnDiskGraphDB::store(ObjectID ID, ArrayRef<ObjectID> Refs,
   auto Alloc = [&](size_t Size) -> Expected<char *> {
     if (Size <= TrieRecord::MaxEmbeddedSize) {
       SK = TrieRecord::StorageKind::DataPool;
-      OnDiskDataAllocator::pointer P = DataPool.allocate(Size);
-      PoolOffset = P.getOffset();
+      auto P = DataPool.allocate(Size);
+      if (LLVM_UNLIKELY(!P))
+        return P.takeError();
+      PoolOffset = P->getOffset();
       LLVM_DEBUG({
         dbgs() << "pool-alloc addr=" << (void *)PoolOffset.get()
                << " size=" << Size
                << " end=" << (void *)(PoolOffset.get() + Size) << "\n";
       });
-      return P->data();
+      return (*P)->data();
     }
 
     SK = TrieRecord::StorageKind::Standalone;
@@ -1465,19 +1479,21 @@ Error OnDiskGraphDB::importFullTree(ObjectID PrimaryID,
     }
 
     ObjectID UpstreamID = *(Cur.RefI++);
-    ObjectID PrimaryID = getReference(UpstreamDB->getDigest(UpstreamID));
-    if (containsObject(PrimaryID, /*CheckUpstream=*/false)) {
+    auto PrimaryID = getReference(UpstreamDB->getDigest(UpstreamID));
+    if (LLVM_UNLIKELY(!PrimaryID))
+      return PrimaryID.takeError();
+    if (containsObject(*PrimaryID, /*CheckUpstream=*/false)) {
       // This \p ObjectID already exists in the primary. Either it was imported
       // via \p importFullTree or the client created it, in which case the
       // client takes responsibility for how it was formed.
-      enqueueNode(PrimaryID, std::nullopt);
+      enqueueNode(*PrimaryID, std::nullopt);
       continue;
     }
     Expected<std::optional<ObjectHandle>> UpstreamNode =
         UpstreamDB->load(UpstreamID);
     if (!UpstreamNode)
       return UpstreamNode.takeError();
-    enqueueNode(PrimaryID, *UpstreamNode);
+    enqueueNode(*PrimaryID, *UpstreamNode);
   }
 
   assert(PrimaryNodesStack.size() == 1);
@@ -1497,8 +1513,12 @@ Error OnDiskGraphDB::importSingleNode(ObjectID PrimaryID,
   auto UpstreamRefs = UpstreamDB->getObjectRefs(UpstreamNode);
   SmallVector<ObjectID, 64> Refs;
   Refs.reserve(std::distance(UpstreamRefs.begin(), UpstreamRefs.end()));
-  for (ObjectID UpstreamRef : UpstreamRefs)
-    Refs.push_back(getReference(UpstreamDB->getDigest(UpstreamRef)));
+  for (ObjectID UpstreamRef : UpstreamRefs) {
+    auto Ref = getReference(UpstreamDB->getDigest(UpstreamRef));
+    if (LLVM_UNLIKELY(!Ref))
+      return Ref.takeError();
+    Refs.push_back(*Ref);
+  }
 
   return store(PrimaryID, Refs, Data);
 }
@@ -1507,9 +1527,12 @@ Expected<std::optional<ObjectHandle>>
 OnDiskGraphDB::faultInFromUpstream(ObjectID PrimaryID) {
   assert(UpstreamDB);
 
-  ObjectID UpstreamID = UpstreamDB->getReference(getDigest(PrimaryID));
+  auto UpstreamID = UpstreamDB->getReference(getDigest(PrimaryID));
+  if (LLVM_UNLIKELY(!UpstreamID))
+    return UpstreamID.takeError();
+
   Expected<std::optional<ObjectHandle>> UpstreamNode =
-      UpstreamDB->load(UpstreamID);
+      UpstreamDB->load(*UpstreamID);
   if (!UpstreamNode)
     return UpstreamNode.takeError();
   if (!*UpstreamNode)