[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

@@ -5,6 +5,11 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
+///
+/// \file
+/// This file contains the declaration of the ObjectStore class.
+///
+//===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CAS_OBJECTSTORE_H
 #define LLVM_CAS_OBJECTSTORE_H
@@ -246,7 +251,7 @@ class ObjectStore {
   /// Set the size for limiting growth of on-disk storage. This has an effect
   /// for when the instance is closed.
   ///
-  /// Implementations may be not have this implemented.
+  /// Implementations may leave this unimplemented.
   virtual Error setSizeLimit(std::optional<uint64_t> SizeLimit) {
     return Error::success();
   }
@@ -262,7 +267,7 @@ class ObjectStore {
   /// Prune local storage to reduce its size according to the desired size
   /// limit. Pruning can happen concurrently with other operations.
   ///
-  /// Implementations may be not have this implemented.
+  /// Implementations may leave this unimplemented.
   virtual Error pruneStorageData() { return Error::success(); }
 
   /// Validate the whole node tree.
@@ -291,13 +296,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); }
@@ -352,12 +353,13 @@ class ObjectProxy {
   ObjectHandle H;
 };
 
+/// Create an in memory CAS.
 std::unique_ptr<ObjectStore> createInMemoryCAS();
 
 /// \returns true if \c LLVM_ENABLE_ONDISK_CAS configuration was enabled.
 bool isOnDiskCASEnabled();
 
-/// Gets or creates a persistent on-disk path at \p Path.
+/// Create a persistent on-disk path at \p Path.
 Expected<std::unique_ptr<ObjectStore>> createOnDiskCAS(const Twine &Path);
 
 /// Set \p Path to a reasonable default on-disk path for a persistent CAS for

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.
@@ -33,7 +33,7 @@ class OnDiskKeyValueDB;
 /// Usage patterns should be that an instance of \p UnifiedOnDiskCache is open
 /// for a limited period of time, e.g. for the duration of a build operation.
 /// For long-living processes that need periodic access to a
-/// \p UnifiedOnDiskCache, the client should device a scheme where access is
+/// \p UnifiedOnDiskCache, the client should devise a scheme where access is
 /// performed within some defined period. For example, if a service is designed
 /// to continuously wait for requests that access a \p UnifiedOnDiskCache, it
 /// could keep the instance alive while new requests are coming in but close it
@@ -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

@@ -13,15 +13,13 @@
 #include "BuiltinCAS.h"
 #include "llvm/ADT/TrieRawHashMap.h"
 #include "llvm/CAS/ActionCache.h"
-#include "llvm/CAS/ObjectStore.h"
 #include "llvm/CAS/OnDiskCASLogger.h"
-#include "llvm/CAS/OnDiskGraphDB.h"
 #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 +65,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 +86,8 @@ class OnDiskActionCache final : public ActionCache {
   using DataT = CacheEntry<sizeof(HashType)>;
 };
 
+/// Builtin unified ActionCache that wraps around UnifiedOnDiskCache to provide
+/// access to its ActionCache.
 class UnifiedOnDiskActionCache final : public ActionCache {
 public:
   Error putImpl(ArrayRef<uint8_t> ActionKey, const CASID &Result,
@@ -118,7 +119,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 +171,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 +221,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 +237,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
-}
+}