[CAS] Cleanup chaining for UnifiedOnDiskCache

Previously, the chaining of KeyValueDB and OnDiskGraphDB is not
consistant. Some operations are implemented directly in the lowest
layer, some are in the UnifiedOnDiskCache layer, and some are in the
ActionCache/ObjectStore layer.

Now unifies all the chaining logics down into OnDiskGraphDB and
OnDiskKeyValueDB layer, with the exception of KeyValueDB chaining will
need the help of functions in UnifiedOnDiskCache layer.

This cleans up the interfaces for UnifiedOnDiskCache member functions
so it only contains database managment functions. Old functions like
`KVPut/Get` can be done directly via underlying database file with a
little bit of extra wrapper around it (see libCASPluginTest.dylib)
implementation for the simple wrapper needed.

Author: cachemeifyoucan

llvm/include/llvm/CAS/BuiltinUnifiedCASDatabases.h

@@ -1,4 +1,4 @@
-//===- BuiltinUnifiedCASDatabases.h -----------------------------*- C++ -*-===//
+//===----------------------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.

llvm/include/llvm/CAS/ObjectStore.h

@@ -291,13 +291,9 @@ class ObjectStore {
 /// Reference to an abstract hierarchical node, with data and references.
 /// Reference is passed by value and is expected to be valid as long as the \a
 /// ObjectStore is.
-///
-/// TODO: Expose \a ObjectStore::readData() and only call \a
-/// ObjectStore::getDataString() when asked.
 class ObjectProxy {
 public:
-  const ObjectStore &getCAS() const { return *CAS; }
-  ObjectStore &getCAS() { return *CAS; }
+  ObjectStore &getCAS() const { return *CAS; }
   CASID getID() const { return CAS->getID(Ref); }
   ObjectRef getRef() const { return Ref; }
   size_t getNumReferences() const { return CAS->getNumRefs(H); }

llvm/include/llvm/CAS/OnDiskGraphDB.h

@@ -341,13 +341,16 @@ class OnDiskGraphDB {
   /// \param HashByteSize Size for the object digest hash bytes.
   /// \param UpstreamDB Optional on-disk store to be used for faulting-in nodes
   /// if they don't exist in the primary store. The upstream store is only used
-  /// for reading nodes, new nodes are only written to the primary store.
+  /// for reading nodes, new nodes are only written to the primary store. User
+  /// need to make sure \p UpstreamDB outlives current instance of
+  /// OnDiskGraphDB and the common usage is to have an \p UnifiedOnDiskCache to
+  /// manage both.
   /// \param Policy If \p UpstreamDB is provided, controls how nodes are copied
   /// to primary store. This is recorded at creation time and subsequent opens
   /// need to pass the same policy otherwise the \p open will fail.
   static Expected<std::unique_ptr<OnDiskGraphDB>>
   open(StringRef Path, StringRef HashName, unsigned HashByteSize,
-       std::unique_ptr<OnDiskGraphDB> UpstreamDB = nullptr,
+       OnDiskGraphDB *UpstreamDB = nullptr,
        std::shared_ptr<OnDiskCASLogger> Logger = nullptr,
        FaultInPolicy Policy = FaultInPolicy::FullTree);
 
@@ -440,9 +443,8 @@ class OnDiskGraphDB {
 
   // Private constructor.
   OnDiskGraphDB(StringRef RootPath, OnDiskTrieRawHashMap Index,
-                OnDiskDataAllocator DataPool,
-                std::unique_ptr<OnDiskGraphDB> UpstreamDB, FaultInPolicy Policy,
-                std::shared_ptr<OnDiskCASLogger> Logger);
+                OnDiskDataAllocator DataPool, OnDiskGraphDB *UpstreamDB,
+                FaultInPolicy Policy, std::shared_ptr<OnDiskCASLogger> Logger);
 
   /// Mapping from hash to object reference.
   ///
@@ -461,7 +463,7 @@ class OnDiskGraphDB {
   std::string RootPath;
 
   /// Optional on-disk store to be used for faulting-in nodes.
-  std::unique_ptr<OnDiskGraphDB> UpstreamDB;
+  OnDiskGraphDB* UpstreamDB = nullptr;
 
   /// The policy used to fault in data from upstream.
   FaultInPolicy FIPolicy;

llvm/include/llvm/CAS/OnDiskKeyValueDB.h

@@ -19,6 +19,8 @@
 
 namespace llvm::cas::ondisk {
 
+class UnifiedOnDiskCache;
+
 /// An on-disk key-value data store with the following properties:
 /// * Keys are fixed length binary hashes with expected normal distribution.
 /// * Values are buffers of the same size, specified at creation time.
@@ -59,9 +61,13 @@ class OnDiskKeyValueDB {
   /// \param KeySize Size for the key hash bytes.
   /// \param ValueName Identifier name for the values.
   /// \param ValueSize Size for the value bytes.
+  /// \param UnifiedCache An optional UnifiedOnDiskCache that manages the size
+  /// and lifetime of the CAS instance and it must owns current initializing
+  /// KeyValueDB after initialized.
   static Expected<std::unique_ptr<OnDiskKeyValueDB>>
   open(StringRef Path, StringRef HashName, unsigned KeySize,
        StringRef ValueName, size_t ValueSize,
+       UnifiedOnDiskCache *UnifiedCache = nullptr,
        std::shared_ptr<OnDiskCASLogger> Logger = nullptr);
 
   using CheckValueT =
@@ -71,11 +77,14 @@ class OnDiskKeyValueDB {
   Error validate(CheckValueT CheckValue) const;
 
 private:
-  OnDiskKeyValueDB(size_t ValueSize, OnDiskTrieRawHashMap Cache)
-      : ValueSize(ValueSize), Cache(std::move(Cache)) {}
+  OnDiskKeyValueDB(size_t ValueSize, OnDiskTrieRawHashMap Cache,
+                   UnifiedOnDiskCache *UnifiedCache)
+      : ValueSize(ValueSize), Cache(std::move(Cache)),
+        UnifiedCache(UnifiedCache) {}
 
   const size_t ValueSize;
   OnDiskTrieRawHashMap Cache;
+  UnifiedOnDiskCache *UnifiedCache = nullptr;
 };
 
 } // namespace llvm::cas::ondisk

llvm/include/llvm/CAS/UnifiedOnDiskCache.h

@@ -1,4 +1,4 @@
-//===- UnifiedOnDiskCache.h -------------------------------------*- C++ -*-===//
+//===----------------------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -43,28 +43,8 @@ class UnifiedOnDiskCache {
   /// The \p OnDiskGraphDB instance for the open directory.
   OnDiskGraphDB &getGraphDB() { return *PrimaryGraphDB; }
 
-  /// Associate an \p ObjectID, of the \p OnDiskGraphDB instance, with a key.
-  ///
-  /// \param Key the hash bytes for the key.
-  /// \param Value the \p ObjectID value.
-  ///
-  /// \returns the \p ObjectID associated with the \p Key. It may be different
-  /// than \p Value if another value was already associated with this key.
-  Expected<ObjectID> KVPut(ArrayRef<uint8_t> Key, ObjectID Value);
-
-  /// Associate an \p ObjectID, of the \p OnDiskGraphDB instance, with a key.
-  /// An \p ObjectID as a key is equivalent to its digest bytes.
-  ///
-  /// \param Key the \p ObjectID for the key.
-  /// \param Value the \p ObjectID value.
-  ///
-  /// \returns the \p ObjectID associated with the \p Key. It may be different
-  /// than \p Value if another value was already associated with this key.
-  Expected<ObjectID> KVPut(ObjectID Key, ObjectID Value);
-
-  /// \returns the \p ObjectID, of the \p OnDiskGraphDB instance, associated
-  /// with the \p Key, or \p std::nullopt if the key does not exist.
-  Expected<std::optional<ObjectID>> KVGet(ArrayRef<uint8_t> Key);
+  /// The \p OnDiskGraphDB instance for the open directory.
+  OnDiskKeyValueDB &getKeyValueDB() { return *PrimaryKVDB; }
 
   /// Open a \p UnifiedOnDiskCache instance for a directory.
   ///
@@ -150,18 +130,23 @@ class UnifiedOnDiskCache {
   static Error collectGarbage(StringRef Path,
                               ondisk::OnDiskCASLogger *Logger = nullptr);
 
+  /// Remove unused data from the current UnifiedOnDiskCache.
   Error collectGarbage();
 
-  ~UnifiedOnDiskCache();
+  /// Helper function to convert the value stored in KeyValueDB and ObjectID.
+  static ObjectID getObjectIDFromValue(ArrayRef<char> Value);
 
-  Error validateActionCache();
+  using ValueBytes = std::array<char, sizeof(uint64_t)>;
+  static ValueBytes getValueFromObjectID(ObjectID ID);
 
-  OnDiskGraphDB *getUpstreamGraphDB() const { return UpstreamGraphDB; }
+  ~UnifiedOnDiskCache();
 
 private:
+  friend class OnDiskGraphDB;
+  friend class OnDiskKeyValueDB;
   UnifiedOnDiskCache();
 
-  Expected<std::optional<ObjectID>>
+  Expected<std::optional<ArrayRef<char>>>
   faultInFromUpstreamKV(ArrayRef<uint8_t> Key);
 
   /// \returns the storage size of the primary directory.
@@ -175,7 +160,7 @@ class UnifiedOnDiskCache {
   std::atomic<bool> NeedsGarbageCollection;
   std::string PrimaryDBDir;
 
-  OnDiskGraphDB *UpstreamGraphDB = nullptr;
+  std::unique_ptr<OnDiskGraphDB> UpstreamGraphDB;
   std::unique_ptr<OnDiskGraphDB> PrimaryGraphDB;
 
   std::unique_ptr<OnDiskKeyValueDB> UpstreamKVDB;

llvm/lib/CAS/ActionCaches.cpp

@@ -19,9 +19,9 @@
 #include "llvm/CAS/OnDiskKeyValueDB.h"
 #include "llvm/CAS/UnifiedOnDiskCache.h"
 #include "llvm/Config/llvm-config.h"
-#include "llvm/Support/Alignment.h"
 #include "llvm/Support/BLAKE3.h"
 #include "llvm/Support/Compiler.h"
+#include "llvm/Support/Errc.h"
 #include "llvm/Support/Path.h"
 
 #define DEBUG_TYPE "cas-action-caches"
@@ -67,6 +67,7 @@ class InMemoryActionCache final : public ActionCache {
   InMemoryCacheT Cache;
 };
 
+/// Builtin basic OnDiskActionCache that uses one underlying OnDiskKeyValueDB.
 class OnDiskActionCache final : public ActionCache {
 public:
   Error putImpl(ArrayRef<uint8_t> ActionKey, const CASID &Result,
@@ -87,6 +88,8 @@ class OnDiskActionCache final : public ActionCache {
   using DataT = CacheEntry<sizeof(HashType)>;
 };
 
+/// Builtin unified ActionCache that wraps around UnifiedOnDiskCache to privode
+/// access to its ActionCache.
 class UnifiedOnDiskActionCache final : public ActionCache {
 public:
   Error putImpl(ArrayRef<uint8_t> ActionKey, const CASID &Result,
@@ -118,7 +121,8 @@ static Error createResultCachePoisonedError(ArrayRef<uint8_t> KeyHash,
 }
 
 Expected<std::optional<CASID>>
-InMemoryActionCache::getImpl(ArrayRef<uint8_t> Key, bool /*CanBeDistributed*/) const {
+InMemoryActionCache::getImpl(ArrayRef<uint8_t> Key,
+                             bool /*CanBeDistributed*/) const {
   auto Result = Cache.find(Key);
   if (!Result)
     return std::nullopt;
@@ -169,17 +173,18 @@ OnDiskActionCache::create(StringRef AbsPath) {
           ondisk::OnDiskCASLogger::openIfEnabled(AbsPath).moveInto(Logger))
     return std::move(E);
   std::unique_ptr<ondisk::OnDiskKeyValueDB> DB;
-  if (Error E = ondisk::OnDiskKeyValueDB::open(AbsPath, getHashName(),
-                                               sizeof(HashType), getHashName(),
-                                               sizeof(DataT), std::move(Logger))
+  if (Error E = ondisk::OnDiskKeyValueDB::open(
+                    AbsPath, getHashName(), sizeof(HashType), getHashName(),
+                    sizeof(DataT), /*UnifiedCache=*/nullptr, std::move(Logger))
                     .moveInto(DB))
     return std::move(E);
   return std::unique_ptr<OnDiskActionCache>(
       new OnDiskActionCache(std::move(DB)));
 }
 
 Expected<std::optional<CASID>>
-OnDiskActionCache::getImpl(ArrayRef<uint8_t> Key, bool /*CanBeDistributed*/) const {
+OnDiskActionCache::getImpl(ArrayRef<uint8_t> Key,
+                           bool /*CanBeDistributed*/) const {
   std::optional<ArrayRef<char>> Val;
   if (Error E = DB->get(Key).moveInto(Val))
     return std::move(E);
@@ -218,13 +223,14 @@ UnifiedOnDiskActionCache::UnifiedOnDiskActionCache(
 Expected<std::optional<CASID>>
 UnifiedOnDiskActionCache::getImpl(ArrayRef<uint8_t> Key,
                                   bool /*CanBeDistributed*/) const {
-  std::optional<ondisk::ObjectID> Val;
-  if (Error E = UniDB->KVGet(Key).moveInto(Val))
+  std::optional<ArrayRef<char>> Val;
+  if (Error E = UniDB->getKeyValueDB().get(Key).moveInto(Val))
     return std::move(E);
   if (!Val)
     return std::nullopt;
+  auto ID = ondisk::UnifiedOnDiskCache::getObjectIDFromValue(*Val);
   return CASID::create(&getContext(),
-                       toStringRef(UniDB->getGraphDB().getDigest(*Val)));
+                       toStringRef(UniDB->getGraphDB().getDigest(ID)));
 }
 
 Error UnifiedOnDiskActionCache::putImpl(ArrayRef<uint8_t> Key,
@@ -233,20 +239,35 @@ Error UnifiedOnDiskActionCache::putImpl(ArrayRef<uint8_t> Key,
   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))
+
+  auto Value = ondisk::UnifiedOnDiskCache::getValueFromObjectID(*Expected);
+  std::optional<ArrayRef<char>> Observed;
+  if (Error E = UniDB->getKeyValueDB().put(Key, Value).moveInto(Observed))
     return E;
 
-  if (*Expected == Observed)
+  auto ObservedID = ondisk::UnifiedOnDiskCache::getObjectIDFromValue(*Observed);
+  if (*Expected == ObservedID)
     return Error::success();
 
   return createResultCachePoisonedError(
-      Key, getContext(), Result,
-      UniDB->getGraphDB().getDigest(*Observed));
+      Key, getContext(), Result, UniDB->getGraphDB().getDigest(ObservedID));
 }
 
 Error UnifiedOnDiskActionCache::validate() const {
-  return UniDB->validateActionCache();
+  auto ValidateRef = [](FileOffset Offset, ArrayRef<char> Value) -> Error {
+    auto ID = ondisk::UnifiedOnDiskCache::getObjectIDFromValue(Value);
+    auto formatError = [&](Twine Msg) {
+      return createStringError(
+          llvm::errc::illegal_byte_sequence,
+          "bad record at 0x" +
+              utohexstr((unsigned)Offset.get(), /*LowerCase=*/true) + ": " +
+              Msg.str());
+    };
+    if (ID.getOpaqueData() == 0)
+      return formatError("zero is not a valid ref");
+    return Error::success();
+  };
+  return UniDB->getKeyValueDB().validate(ValidateRef);
 }
 
 Expected<std::unique_ptr<ActionCache>>

llvm/lib/CAS/BuiltinUnifiedCASDatabases.cpp

@@ -1,4 +1,4 @@
-//===- BuiltinUnifiedCASDatabases.cpp ---------------------------*- C++ -*-===//
+//===----------------------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -35,4 +35,4 @@ Expected<ValidationResult> cas::validateOnDiskUnifiedCASDatabasesIfNeeded(
 #else
   return createStringError(inconvertibleErrorCode(), "OnDiskCache is disabled");
 #endif
-}
+}

llvm/lib/CAS/OnDiskCAS.cpp

@@ -43,8 +43,8 @@ class OnDiskCAS : public BuiltinCAS {
 
   static Expected<std::unique_ptr<OnDiskCAS>> open(StringRef Path);
 
-  OnDiskCAS(std::shared_ptr<ondisk::UnifiedOnDiskCache> UniDB_)
-      : UniDB(std::move(UniDB_)), DB(&UniDB->getGraphDB()) {}
+  OnDiskCAS(std::shared_ptr<ondisk::UnifiedOnDiskCache> UniDB)
+      : UnifiedDB(std::move(UniDB)), DB(&UnifiedDB->getGraphDB()) {}
 
 private:
   ObjectHandle convertHandle(ondisk::ObjectHandle Node) const {
@@ -78,11 +78,11 @@ class OnDiskCAS : public BuiltinCAS {
   Expected<std::optional<uint64_t>> getStorageSize() const final;
   Error pruneStorageData() final;
 
-  OnDiskCAS(std::unique_ptr<ondisk::OnDiskGraphDB> DB_)
-      : OwnedDB(std::move(DB_)), DB(OwnedDB.get()) {}
+  OnDiskCAS(std::unique_ptr<ondisk::OnDiskGraphDB> GraphDB)
+      : OwnedDB(std::move(GraphDB)), DB(OwnedDB.get()) {}
 
   std::unique_ptr<ondisk::OnDiskGraphDB> OwnedDB;
-  std::shared_ptr<ondisk::UnifiedOnDiskCache> UniDB;
+  std::shared_ptr<ondisk::UnifiedOnDiskCache> UnifiedDB;
   ondisk::OnDiskGraphDB *DB;
 };
 
@@ -99,8 +99,6 @@ Error OnDiskCAS::validate(bool CheckHash) const {
 
   if (auto E = DB->validate(CheckHash, Hasher))
     return E;
-  if (UniDB && UniDB->getUpstreamGraphDB())
-    return UniDB->getUpstreamGraphDB()->validate(CheckHash, Hasher);
 
   return Error::success();
 }
@@ -165,15 +163,15 @@ Error OnDiskCAS::forEachRef(ObjectHandle Node,
 }
 
 Error OnDiskCAS::setSizeLimit(std::optional<uint64_t> SizeLimit) {
-  UniDB->setSizeLimit(SizeLimit);
+  UnifiedDB->setSizeLimit(SizeLimit);
   return Error::success();
 }
 
 Expected<std::optional<uint64_t>> OnDiskCAS::getStorageSize() const {
-  return UniDB->getStorageSize();
+  return UnifiedDB->getStorageSize();
 }
 
-Error OnDiskCAS::pruneStorageData() { return UniDB->collectGarbage(); }
+Error OnDiskCAS::pruneStorageData() { return UnifiedDB->collectGarbage(); }
 
 Expected<std::unique_ptr<OnDiskCAS>> OnDiskCAS::open(StringRef AbsPath) {
   std::shared_ptr<ondisk::OnDiskCASLogger> Logger;