[CAS/OnDiskKeyValueDB] Factor out the underlying implementation of OnDiskActionCache

Introduce `OnDiskKeyValueDB` as an on-disk key-value database that is independent of a particular hashing algorithm and data values.
`OnDiskActionCache` is implemented as a wrapper of `OnDiskKeyValueDB` with `BLAKE3` keys and values.

Author: akyrtzi

llvm/include/llvm/CAS/OnDiskKeyValueDB.h

@@ -0,0 +1,65 @@
+//===- OnDiskKeyValueDB.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.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CAS_ONDISKKEYVALUEDB_H
+#define LLVM_CAS_ONDISKKEYVALUEDB_H
+
+#include "llvm/CAS/OnDiskHashMappedTrie.h"
+
+namespace llvm::cas::ondisk {
+
+/// 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.
+/// * The value of a key cannot be changed once it is set.
+/// * The value buffers returned from a key lookup have 8-byte alignment.
+class OnDiskKeyValueDB {
+public:
+  /// Associate a value with a key.
+  ///
+  /// \param Key the hash bytes for the key
+  /// \param Value the value bytes, same size as \p ValueSize parameter of
+  /// \p open call.
+  ///
+  /// \returns the value associated with the \p Key. It may be different than
+  /// \p Value if another value is already associated with this key.
+  Expected<ArrayRef<char>> put(ArrayRef<uint8_t> Key, ArrayRef<char> Value);
+
+  /// \returns the value associated with the \p Key, or \p std::nullopt if the
+  /// key does not exist.
+  Expected<std::optional<ArrayRef<char>>> get(ArrayRef<uint8_t> Key);
+
+  /// \returns Total size of stored data.
+  size_t getStorageSize() const {
+    return Cache.size();
+  }
+
+  /// Open the on-disk store from a directory.
+  ///
+  /// \param Path directory for the on-disk store. The directory will be created
+  /// if it doesn't exist.
+  /// \param HashName Identifier name for the hashing algorithm that is going to
+  /// be used.
+  /// \param KeySize Size for the key hash bytes.
+  /// \param ValueName Identifier name for the values.
+  /// \param ValueSize Size for the value bytes.
+  static Expected<std::unique_ptr<OnDiskKeyValueDB>>
+  open(StringRef Path, StringRef HashName, unsigned KeySize,
+       StringRef ValueName, size_t ValueSize);
+
+private:
+  OnDiskKeyValueDB(size_t ValueSize, OnDiskHashMappedTrie Cache)
+      : ValueSize(ValueSize), Cache(std::move(Cache)) {}
+
+  const size_t ValueSize;
+  OnDiskHashMappedTrie Cache;
+};
+
+} // namespace llvm::cas::ondisk
+
+#endif // LLVM_CAS_ONDISKKEYVALUEDB_H

llvm/lib/CAS/ActionCaches.cpp

@@ -11,6 +11,7 @@
 #include "llvm/CAS/HashMappedTrie.h"
 #include "llvm/CAS/ObjectStore.h"
 #include "llvm/CAS/OnDiskHashMappedTrie.h"
+#include "llvm/CAS/OnDiskKeyValueDB.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/BLAKE3.h"
@@ -61,19 +62,10 @@ class OnDiskActionCache final : public ActionCache {
 
 private:
   static StringRef getHashName() { return "BLAKE3"; }
-  static StringRef getActionCacheTableName() {
-    static const std::string Name =
-        ("llvm.actioncache[" + getHashName() + "->" + getHashName() + "]")
-            .str();
-    return Name;
-  }
-  static constexpr StringLiteral ActionCacheFile = "actions";
-  static constexpr StringLiteral FilePrefix = "v1.";
-
-  OnDiskActionCache(StringRef RootPath, OnDiskHashMappedTrie ActionCache);
-
-  std::string Path;
-  OnDiskHashMappedTrie Cache;
+
+  OnDiskActionCache(std::unique_ptr<ondisk::OnDiskKeyValueDB> DB);
+
+  std::unique_ptr<ondisk::OnDiskKeyValueDB> DB;
   using DataT = CacheEntry<sizeof(HashType)>;
 };
 } // end namespace
@@ -137,62 +129,46 @@ std::unique_ptr<ActionCache> createInMemoryActionCache() {
 } // namespace cas
 } // namespace llvm
 
-constexpr StringLiteral OnDiskActionCache::ActionCacheFile;
-constexpr StringLiteral OnDiskActionCache::FilePrefix;
-
-OnDiskActionCache::OnDiskActionCache(StringRef Path, OnDiskHashMappedTrie Cache)
+OnDiskActionCache::OnDiskActionCache(
+    std::unique_ptr<ondisk::OnDiskKeyValueDB> DB)
     : ActionCache(builtin::BuiltinCASContext::getDefaultContext()),
-      Path(Path.str()), Cache(std::move(Cache)) {}
+      DB(std::move(DB)) {}
 
 Expected<std::unique_ptr<OnDiskActionCache>>
 OnDiskActionCache::create(StringRef AbsPath) {
-  if (std::error_code EC = sys::fs::create_directories(AbsPath))
-    return createFileError(AbsPath, EC);
-
-  SmallString<256> CachePath(AbsPath);
-  sys::path::append(CachePath, FilePrefix + ActionCacheFile);
-  constexpr uint64_t MB = 1024ull * 1024ull;
-  constexpr uint64_t GB = 1024ull * 1024ull * 1024ull;
-
-  Optional<OnDiskHashMappedTrie> ActionCache;
-  if (Error E = OnDiskHashMappedTrie::create(
-                    CachePath, getActionCacheTableName(), sizeof(HashType) * 8,
-                    /*DataSize=*/sizeof(DataT), /*MaxFileSize=*/GB,
-                    /*MinFileSize=*/MB)
-                    .moveInto(ActionCache))
+  std::unique_ptr<ondisk::OnDiskKeyValueDB> DB;
+  if (Error E = ondisk::OnDiskKeyValueDB::open(AbsPath, getHashName(),
+                                               sizeof(HashType), getHashName(),
+                                               sizeof(DataT))
+                    .moveInto(DB))
     return std::move(E);
-
   return std::unique_ptr<OnDiskActionCache>(
-      new OnDiskActionCache(AbsPath, std::move(*ActionCache)));
+      new OnDiskActionCache(std::move(DB)));
 }
 
 Expected<Optional<CASID>>
 OnDiskActionCache::getImpl(ArrayRef<uint8_t> Key) const {
-  // Check the result cache.
-  OnDiskHashMappedTrie::const_pointer ActionP = Cache.find(Key);
-  if (!ActionP)
-    return None;
-
-  const DataT *Output = reinterpret_cast<const DataT *>(ActionP->Data.data());
-  return CASID::create(&getContext(), toStringRef(Output->getValue()));
+  std::optional<ArrayRef<char>> Val;
+  if (Error E = DB->get(Key).moveInto(Val))
+    return std::move(E);
+  if (!Val)
+    return std::nullopt;
+  return CASID::create(&getContext(), toStringRef(*Val));
 }
 
 Error OnDiskActionCache::putImpl(ArrayRef<uint8_t> Key, const CASID &Result) {
-  DataT Expected(Result.getHash());
-  OnDiskHashMappedTrie::pointer ActionP = Cache.insertLazy(
-      Key, [&](FileOffset TentativeOffset,
-               OnDiskHashMappedTrie::ValueProxy TentativeValue) {
-        assert(TentativeValue.Data.size() == sizeof(DataT));
-        assert(isAddrAligned(Align::Of<DataT>(), TentativeValue.Data.data()));
-        new (TentativeValue.Data.data()) DataT{Expected};
-      });
-  const DataT *Observed = reinterpret_cast<const DataT *>(ActionP->Data.data());
-
-  if (Expected.getValue() == Observed->getValue())
+  auto ResultHash = Result.getHash();
+  ArrayRef Expected((const char *)ResultHash.data(), ResultHash.size());
+  ArrayRef<char> Observed;
+  if (Error E = DB->put(Key, Expected).moveInto(Observed))
+    return E;
+
+  if (Expected == Observed)
     return Error::success();
 
-  return createResultCachePoisonedError(hashToString(Key), getContext(), Result,
-                                        Observed->getValue());
+  return createResultCachePoisonedError(
+      hashToString(Key), getContext(), Result,
+      ArrayRef((const uint8_t *)Observed.data(), Observed.size()));
 }
 
 #if LLVM_ENABLE_ONDISK_CAS

llvm/lib/CAS/CMakeLists.txt

@@ -21,6 +21,7 @@ add_llvm_component_library(LLVMCAS
   OnDiskCAS.cpp
   OnDiskGraphDB.cpp
   OnDiskHashMappedTrie.cpp
+  OnDiskKeyValueDB.cpp
   TreeSchema.cpp
   Utils.cpp
 

llvm/lib/CAS/OnDiskKeyValueDB.cpp

@@ -0,0 +1,70 @@
+//===- OnDiskKeyValueDB.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.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CAS/OnDiskKeyValueDB.h"
+#include "llvm/Support/Alignment.h"
+#include "llvm/Support/Errc.h"
+#include "llvm/Support/Path.h"
+
+using namespace llvm;
+using namespace llvm::cas;
+using namespace llvm::cas::ondisk;
+
+static constexpr StringLiteral ActionCacheFile = "actions";
+static constexpr StringLiteral FilePrefix = "v2.";
+
+Expected<ArrayRef<char>> OnDiskKeyValueDB::put(ArrayRef<uint8_t> Key,
+                                               ArrayRef<char> Value) {
+  if (LLVM_UNLIKELY(Value.size() != ValueSize))
+    return createStringError(errc::invalid_argument,
+                             "expected value size of " + itostr(ValueSize) +
+                                 ", got: " + itostr(Value.size()));
+  assert(Value.size() == ValueSize);
+  OnDiskHashMappedTrie::pointer ActionP = Cache.insertLazy(
+      Key, [&](FileOffset TentativeOffset,
+               OnDiskHashMappedTrie::ValueProxy TentativeValue) {
+        assert(TentativeValue.Data.size() == ValueSize);
+        llvm::copy(Value, TentativeValue.Data.data());
+      });
+  return ActionP->Data;
+}
+
+Expected<std::optional<ArrayRef<char>>>
+OnDiskKeyValueDB::get(ArrayRef<uint8_t> Key) {
+  // Check the result cache.
+  OnDiskHashMappedTrie::const_pointer ActionP = Cache.find(Key);
+  if (!ActionP)
+    return std::nullopt;
+  assert(isAddrAligned(Align(8), ActionP->Data.data()));
+  return ActionP->Data;
+}
+
+Expected<std::unique_ptr<OnDiskKeyValueDB>>
+OnDiskKeyValueDB::open(StringRef Path, StringRef HashName, unsigned KeySize,
+                       StringRef ValueName, size_t ValueSize) {
+  if (std::error_code EC = sys::fs::create_directories(Path))
+    return createFileError(Path, EC);
+
+  SmallString<256> CachePath(Path);
+  sys::path::append(CachePath, FilePrefix + ActionCacheFile);
+  constexpr uint64_t MB = 1024ull * 1024ull;
+  constexpr uint64_t GB = 1024ull * 1024ull * 1024ull;
+
+  std::optional<OnDiskHashMappedTrie> ActionCache;
+  if (Error E = OnDiskHashMappedTrie::create(
+                    CachePath,
+                    "llvm.actioncache[" + HashName + "->" + ValueName + "]",
+                    KeySize * 8,
+                    /*DataSize=*/ValueSize, /*MaxFileSize=*/GB,
+                    /*MinFileSize=*/MB)
+                    .moveInto(ActionCache))
+    return std::move(E);
+
+  return std::unique_ptr<OnDiskKeyValueDB>(
+      new OnDiskKeyValueDB(ValueSize, std::move(*ActionCache)));
+}

llvm/unittests/CAS/CMakeLists.txt

@@ -28,6 +28,7 @@ add_llvm_unittest(CASTests
   ObjectStoreTest.cpp
   OnDiskGraphDBTest.cpp
   OnDiskHashMappedTrieTest.cpp
+  OnDiskKeyValueDBTest.cpp
   ThreadSafeAllocatorTest.cpp
   TreeSchemaTest.cpp
   )

llvm/unittests/CAS/OnDiskKeyValueDBTest.cpp

@@ -0,0 +1,68 @@
+//===- llvm/unittest/CAS/OnDiskKeyValueDBTest.cpp -------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CAS/OnDiskKeyValueDB.h"
+#include "llvm/CAS/BuiltinObjectHasher.h"
+#include "llvm/Support/BLAKE3.h"
+#include "llvm/Testing/Support/Error.h"
+#include "llvm/Testing/Support/SupportHelpers.h"
+#include "gtest/gtest.h"
+
+#if LLVM_ENABLE_ONDISK_CAS
+
+using namespace llvm;
+using namespace llvm::cas;
+using namespace llvm::cas::ondisk;
+
+using HasherT = BLAKE3;
+using HashType = decltype(HasherT::hash(std::declval<ArrayRef<uint8_t> &>()));
+using ValueType = std::array<char, 20>;
+
+static HashType digest(StringRef Data) {
+  return HasherT::hash(arrayRefFromStringRef(Data));
+}
+
+static ValueType valueFromString(StringRef S) {
+  ValueType Val;
+  llvm::copy(S.substr(0, sizeof(Val)), Val.data());
+  return Val;
+}
+
+TEST(OnDiskKeyValueDBTest, Basic) {
+  unittest::TempDir Temp("ondiskkv", /*Unique=*/true);
+  std::unique_ptr<OnDiskKeyValueDB> DB;
+  ASSERT_THAT_ERROR(OnDiskKeyValueDB::open(Temp.path(), "blake3",
+                                           sizeof(HashType), "test",
+                                           sizeof(ValueType))
+                        .moveInto(DB),
+                    Succeeded());
+
+  {
+    std::optional<ArrayRef<char>> Val;
+    ASSERT_THAT_ERROR(DB->get(digest("hello")).moveInto(Val), Succeeded());
+    EXPECT_FALSE(Val.has_value());
+  }
+
+  ValueType ValW = valueFromString("world");
+  ArrayRef<char> Val;
+  ASSERT_THAT_ERROR(DB->put(digest("hello"), ValW).moveInto(Val), Succeeded());
+  EXPECT_EQ(Val, ArrayRef(ValW));
+  ASSERT_THAT_ERROR(
+      DB->put(digest("hello"), valueFromString("other")).moveInto(Val),
+      Succeeded());
+  EXPECT_EQ(Val, ArrayRef(ValW));
+
+  {
+    std::optional<ArrayRef<char>> Val;
+    ASSERT_THAT_ERROR(DB->get(digest("hello")).moveInto(Val), Succeeded());
+    EXPECT_TRUE(Val.has_value());
+    EXPECT_EQ(*Val, ArrayRef(ValW));
+  }
+}
+
+#endif // LLVM_ENABLE_ONDISK_CAS