[CAS] Cleanup chaining for UnifiedOnDiskCache

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
-}
+}