[llvm/CAS] Introduce ObjectStore::isMaterialized()

This returns true if the object is directly available from the local CAS, for
implementations that have this kind of distinction. For such implementations it's useful
to be able to identify whether objects need to be downloaded from the distributed cache.

(cherry picked from commit c69249f)

Author: akyrtzi

llvm/include/llvm-c/CAS/PluginAPI_functions.h

@@ -164,13 +164,18 @@ LLCAS_PUBLIC llcas_digest_t llcas_objectid_get_digest(llcas_cas_t,
 /**
  * Checks whether a \c llcas_objectid_t points to an existing object.
  *
+ * \param globally For CAS implementations that distinguish between local CAS
+ * and remote/distributed CAS, \p globally set to false indicates that the
+ * lookup will be restricted to the local CAS, returning "not found" even if the
+ * object might exist in the remote CAS.
  * \param error optional pointer to receive an error message if an error
  * occurred. If set, the memory it points to needs to be released via
  * \c llcas_string_dispose.
  * \returns one of \c llcas_lookup_result_t.
  */
 LLCAS_PUBLIC llcas_lookup_result_t llcas_cas_contains_object(llcas_cas_t,
                                                              llcas_objectid_t,
+                                                             bool globally,
                                                              char **error);
 
 /**

llvm/include/llvm/CAS/ObjectStore.h

@@ -152,6 +152,10 @@ class ObjectStore {
   /// Returns \c None if the object is not stored in this CAS.
   virtual Optional<ObjectRef> getReference(const CASID &ID) const = 0;
 
+  /// \returns true if the object is directly available from the local CAS, for
+  /// implementations that have this kind of distinction.
+  virtual Expected<bool> isMaterialized(ObjectRef Ref) const = 0;
+
   /// Validate the underlying object referred by CASID.
   virtual Error validate(const CASID &ID) = 0;
 

llvm/lib/CAS/InMemoryCAS.cpp

@@ -220,6 +220,8 @@ class InMemoryCAS : public BuiltinCAS {
     return None;
   }
 
+  Expected<bool> isMaterialized(ObjectRef Ref) const final { return true; }
+
   ArrayRef<char> getDataConst(ObjectHandle Node) const final {
     return cast<InMemoryObject>(asInMemoryObject(Node)).getData();
   }

llvm/lib/CAS/OnDiskCAS.cpp

@@ -29,6 +29,8 @@ class OnDiskCAS : public BuiltinCAS {
 
   Optional<ObjectRef> getReference(const CASID &ID) const final;
 
+  Expected<bool> isMaterialized(ObjectRef Ref) const final;
+
   ArrayRef<char> getDataConst(ObjectHandle Node) const final;
 
   void print(raw_ostream &OS) const final;
@@ -91,6 +93,10 @@ Optional<ObjectRef> OnDiskCAS::getReference(const CASID &ID) const {
   return convertRef(*ObjID);
 }
 
+Expected<bool> OnDiskCAS::isMaterialized(ObjectRef ExternalRef) const {
+  return DB->containsObject(convertRef(ExternalRef));
+}
+
 ArrayRef<char> OnDiskCAS::getDataConst(ObjectHandle Node) const {
   return DB->getObjectData(convertHandle(Node));
 }

llvm/lib/CAS/PluginAPI.h

@@ -48,7 +48,7 @@ struct llcas_functions_t {
   llcas_digest_t (*objectid_get_digest)(llcas_cas_t, llcas_objectid_t);
 
   llcas_lookup_result_t (*cas_contains_object)(llcas_cas_t, llcas_objectid_t,
-                                               char **error);
+                                               bool globally, char **error);
 
   llcas_lookup_result_t (*cas_load_object)(llcas_cas_t, llcas_objectid_t,
                                            llcas_loaded_object_t *,

llvm/lib/CAS/PluginCAS.cpp

@@ -130,6 +130,7 @@ class PluginObjectStore
                             ArrayRef<char> Data) final;
   CASID getID(ObjectRef Ref) const final;
   Optional<ObjectRef> getReference(const CASID &ID) const final;
+  Expected<bool> isMaterialized(ObjectRef Ref) const final;
   Expected<std::optional<ObjectHandle>> loadIfExists(ObjectRef Ref) final;
   void
   loadIfExistsAsync(ObjectRef Ref,
@@ -223,18 +224,22 @@ PluginObjectStore::getReference(const CASID &ID) const {
           Ctx->c_cas, llcas_digest_t{Hash.data(), Hash.size()}, &c_id, &c_err))
     report_fatal_error(Ctx->errorAndDispose(c_err));
 
-  llcas_lookup_result_t c_result =
-      Ctx->Functions.cas_contains_object(Ctx->c_cas, c_id, &c_err);
+  return ObjectRef::getFromInternalRef(*this, c_id.opaque);
+}
+
+Expected<bool> PluginObjectStore::isMaterialized(ObjectRef Ref) const {
+  llcas_objectid_t c_id{Ref.getInternalRef(*this)};
+  char *c_err = nullptr;
+  llcas_lookup_result_t c_result = Ctx->Functions.cas_contains_object(
+      Ctx->c_cas, c_id, /*globally=*/false, &c_err);
   switch (c_result) {
   case LLCAS_LOOKUP_RESULT_SUCCESS:
-    return ObjectRef::getFromInternalRef(*this, c_id.opaque);
+    return true;
   case LLCAS_LOOKUP_RESULT_NOTFOUND:
-    return std::nullopt;
+    return false;
   case LLCAS_LOOKUP_RESULT_ERROR:
-    report_fatal_error(Ctx->errorAndDispose(c_err));
+    return Ctx->errorAndDispose(c_err);
   }
-
-  return ObjectRef::getFromInternalRef(*this, c_id.opaque);
 }
 
 Expected<std::optional<ObjectHandle>>

llvm/lib/RemoteCachingService/CAS/GRPCRelayCAS.cpp

@@ -116,6 +116,7 @@ class GRPCRelayCAS : public ObjectStore {
                                ArrayRef<char> Data) final;
   CASID getID(ObjectRef Ref) const final;
   Optional<ObjectRef> getReference(const CASID &ID) const final;
+  Expected<bool> isMaterialized(ObjectRef Ref) const final;
   Expected<std::optional<ObjectHandle>> loadIfExists(ObjectRef Ref) final;
   Error validate(const CASID &ID) final {
     // Not supported yet. Always return success.
@@ -303,6 +304,11 @@ Optional<ObjectRef> GRPCRelayCAS::getReference(const CASID &ID) const {
   return toReference(I);
 }
 
+Expected<bool> GRPCRelayCAS::isMaterialized(ObjectRef Ref) const {
+  auto &I = asInMemoryIndexValue(Ref);
+  return (bool)I.Data.load();
+}
+
 Expected<std::optional<ObjectHandle>>
 GRPCRelayCAS::loadIfExists(ObjectRef Ref) {
   auto &I = asInMemoryIndexValue(Ref);

llvm/tools/libCASPluginTest/libCASPluginTest.cpp

@@ -112,13 +112,18 @@ struct CASWrapper {
 
   std::mutex Lock{};
 
+  /// Check if the object is contained, in the "local" CAS only or "globally".
+  bool containsObject(ObjectID ID, bool Globally);
+
   /// Load the object, potentially "downloading" it from upstream.
   Expected<std::optional<ondisk::ObjectHandle>> loadObject(ObjectID ID);
 
-  /// "Uploads" a key the associated full node graph.
+  /// "Uploads" a key and the associated full node graph.
   Error upstreamKey(ArrayRef<uint8_t> Key, ObjectID Value);
-  /// "Downloads" the single root node that is associated with the key. The rest
-  /// of the nodes in the graph will be "downloaded" lazily as they are visited.
+
+  /// "Downloads" the ID associated with the key but not the node data. The node
+  /// itself and the rest of the nodes in the graph will be "downloaded" lazily
+  /// as they are visited.
   Expected<std::optional<ObjectID>> downstreamKey(ArrayRef<uint8_t> Key);
 
   /// Synchronized access to \c llvm::errs().
@@ -140,6 +145,17 @@ DEFINE_SIMPLE_CONVERSION_FUNCTIONS(CASWrapper, llcas_cas_t)
 
 } // namespace
 
+bool CASWrapper::containsObject(ObjectID ID, bool Globally) {
+  if (DB->getGraphDB().containsObject(ID))
+    return true;
+  if (!Globally || !UpstreamDB)
+    return false;
+
+  ObjectID UpstreamID =
+      UpstreamDB->getGraphDB().getReference(DB->getGraphDB().getDigest(ID));
+  return UpstreamDB->getGraphDB().containsObject(UpstreamID);
+}
+
 Expected<std::optional<ondisk::ObjectHandle>>
 CASWrapper::loadObject(ObjectID ID) {
   std::optional<ondisk::ObjectHandle> Obj;
@@ -231,16 +247,8 @@ CASWrapper::downstreamKey(ArrayRef<uint8_t> Key) {
   if (!UpstreamValue)
     return std::nullopt;
 
-  Expected<ObjectID> Value = downstreamNode(*UpstreamValue);
-  if (!Value)
-    return Value.takeError();
-  assert(DB->getGraphDB().getDigest(*Value) ==
-         UpstreamDB->getGraphDB().getDigest(*UpstreamValue));
-  Expected<ObjectID> PutValue = DB->KVPut(Key, *Value);
-  if (!PutValue)
-    return PutValue.takeError();
-  assert(*PutValue == *Value);
-  return PutValue;
+  return DB->getGraphDB().getReference(
+      UpstreamDB->getGraphDB().getDigest(*UpstreamValue));
 }
 
 llcas_cas_t llcas_cas_create(llcas_cas_options_t c_opts, char **error) {
@@ -329,11 +337,11 @@ llcas_digest_t llcas_objectid_get_digest(llcas_cas_t c_cas,
 
 llcas_lookup_result_t llcas_cas_contains_object(llcas_cas_t c_cas,
                                                 llcas_objectid_t c_id,
-                                                char **error) {
-  auto &CAS = unwrap(c_cas)->DB->getGraphDB();
+                                                bool globally, char **error) {
   ObjectID ID = ObjectID::fromOpaqueData(c_id.opaque);
-  return CAS.containsObject(ID) ? LLCAS_LOOKUP_RESULT_SUCCESS
-                                : LLCAS_LOOKUP_RESULT_NOTFOUND;
+  return unwrap(c_cas)->containsObject(ID, globally)
+             ? LLCAS_LOOKUP_RESULT_SUCCESS
+             : LLCAS_LOOKUP_RESULT_NOTFOUND;
 }
 
 llcas_lookup_result_t llcas_cas_load_object(llcas_cas_t c_cas,

llvm/unittests/CAS/CMakeLists.txt

@@ -29,9 +29,11 @@ add_llvm_unittest(CASTests
   OnDiskGraphDBTest.cpp
   OnDiskHashMappedTrieTest.cpp
   OnDiskKeyValueDBTest.cpp
+  PluginCASTest.cpp
   ThreadSafeAllocatorTest.cpp
   TreeSchemaTest.cpp
   UnifiedOnDiskCacheTest.cpp
   )
 
 target_link_libraries(CASTests PRIVATE LLVMTestingSupport)
+add_dependencies(CASTests CASPluginTest)

llvm/unittests/CAS/PluginCASTest.cpp

@@ -0,0 +1,106 @@
+//===- llvm/unittest/CAS/PluginCASTest.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/ActionCache.h"
+#include "llvm/CAS/ObjectStore.h"
+#include "llvm/Config/config.h"
+#include "llvm/Support/Path.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;
+
+// See llvm/utils/unittest/UnitTestMain/TestMain.cpp
+extern const char *TestMainArgv0;
+
+// Just a reachable symbol to ease resolving of the executable's path.
+static std::string TestStringArg1("plugincas-test-string-arg1");
+
+static std::string getCASPluginPath() {
+  std::string Executable =
+      sys::fs::getMainExecutable(TestMainArgv0, &TestStringArg1);
+  llvm::SmallString<256> PathBuf(sys::path::parent_path(
+      sys::path::parent_path(sys::path::parent_path(Executable))));
+  std::string LibName = "libCASPluginTest";
+  sys::path::append(PathBuf, "lib", LibName + LLVM_PLUGIN_EXT);
+  return std::string(PathBuf);
+}
+
+TEST(PluginCASTest, isMaterialized) {
+  unittest::TempDir Temp("plugin-cas", /*Unique=*/true);
+  std::string UpDir(Temp.path("up"));
+  std::string DownDir(Temp.path("down"));
+  std::pair<std::string, std::string> PluginOpts[] = {
+      {"upstream-path", std::string(UpDir)}};
+
+  {
+    std::optional<
+        std::pair<std::shared_ptr<ObjectStore>, std::shared_ptr<ActionCache>>>
+        DBs;
+    ASSERT_THAT_ERROR(
+        createPluginCASDatabases(getCASPluginPath(), DownDir, PluginOpts)
+            .moveInto(DBs),
+        Succeeded());
+    std::shared_ptr<ObjectStore> CAS;
+    std::shared_ptr<ActionCache> AC;
+    std::tie(CAS, AC) = std::move(*DBs);
+
+    Optional<CASID> ID1, ID2;
+    ASSERT_THAT_ERROR(CAS->createProxy(std::nullopt, "1").moveInto(ID1),
+                      Succeeded());
+    ASSERT_THAT_ERROR(CAS->createProxy(std::nullopt, "2").moveInto(ID2),
+                      Succeeded());
+    Optional<ObjectRef> ID2Ref = CAS->getReference(*ID2);
+    ASSERT_TRUE(ID2Ref);
+    bool IsMaterialized = false;
+    ASSERT_THAT_ERROR(CAS->isMaterialized(*ID2Ref).moveInto(IsMaterialized),
+                      Succeeded());
+    EXPECT_TRUE(IsMaterialized);
+    ASSERT_THAT_ERROR(AC->put(*ID1, *ID2, /*Globally=*/true), Succeeded());
+  }
+
+  // Clear "local" cache.
+  sys::fs::remove_directories(DownDir);
+
+  {
+    std::optional<
+        std::pair<std::shared_ptr<ObjectStore>, std::shared_ptr<ActionCache>>>
+        DBs;
+    ASSERT_THAT_ERROR(
+        createPluginCASDatabases(getCASPluginPath(), DownDir, PluginOpts)
+            .moveInto(DBs),
+        Succeeded());
+    std::shared_ptr<ObjectStore> CAS;
+    std::shared_ptr<ActionCache> AC;
+    std::tie(CAS, AC) = std::move(*DBs);
+
+    Optional<CASID> ID1, ID2;
+    ASSERT_THAT_ERROR(CAS->createProxy(std::nullopt, "1").moveInto(ID1),
+                      Succeeded());
+    ASSERT_THAT_ERROR(AC->get(*ID1, /*Globally=*/true).moveInto(ID2),
+                      Succeeded());
+    Optional<ObjectRef> ID2Ref = CAS->getReference(*ID2);
+    ASSERT_TRUE(ID2Ref);
+    bool IsMaterialized = false;
+    ASSERT_THAT_ERROR(CAS->isMaterialized(*ID2Ref).moveInto(IsMaterialized),
+                      Succeeded());
+    EXPECT_FALSE(IsMaterialized);
+
+    Optional<ObjectProxy> Obj;
+    ASSERT_THAT_ERROR(CAS->getProxy(*ID2Ref).moveInto(Obj), Succeeded());
+    ASSERT_THAT_ERROR(CAS->isMaterialized(*ID2Ref).moveInto(IsMaterialized),
+                      Succeeded());
+    EXPECT_TRUE(IsMaterialized);
+  }
+}
+
+#endif // LLVM_ENABLE_ONDISK_CAS