[CAS/ActionCache] Add asynchronous versions for the ActionCache APIs

(cherry picked from commit ab601d8)

Author: akyrtzi

clang/lib/Frontend/CompileJobCacheResult.cpp

@@ -41,7 +41,7 @@ Error CompileJobCacheResult::forEachLoadedOutput(
   size_t Count = getNumOutputs();
   for (size_t I = 0; I < Count; ++I) {
     ObjectRef Ref = getOutputObject(I);
-    FutureOutputs.push_back(getCAS().getProxyAsync(Ref));
+    FutureOutputs.push_back(getCAS().getProxyFuture(Ref));
   }
 
   // Make sure all the outputs have materialized.

clang/lib/Frontend/CompilerInstance.cpp

@@ -2469,7 +2469,7 @@ static bool addCachedModuleFileToInMemoryCache(
     llvm::report_fatal_error("missing main output");
   // FIXME: We wait to materialize each module file before proceeding, which
   // introduces latency for a network CAS. Instead we should collect all the
-  // module keys and materialize them concurrently using \c getProxyAsync, for
+  // module keys and materialize them concurrently using \c getProxyFuture, for
   // better network utilization.
   auto OutputProxy = CAS.getProxy(Output->Object);
   if (!OutputProxy) {

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

@@ -189,7 +189,6 @@ LLCAS_PUBLIC llcas_lookup_result_t llcas_cas_load_object(
  * Whether the call is asynchronous or not depends on the implementation.
  *
  * \param ctx_cb pointer to pass to the callback function.
- *
  */
 LLCAS_PUBLIC void llcas_cas_load_object_async(llcas_cas_t, llcas_objectid_t,
                                               void *ctx_cb,
@@ -257,6 +256,18 @@ LLCAS_PUBLIC llcas_lookup_result_t llcas_actioncache_get_for_digest(
     llcas_cas_t, llcas_digest_t key, llcas_objectid_t *p_value, bool globally,
     char **error);
 
+/**
+ * Like \c llcas_actioncache_get_for_digest but result is provided to a callback
+ * function. Whether the call is asynchronous or not depends on the
+ * implementation.
+ *
+ * \param ctx_cb pointer to pass to the callback function.
+ */
+LLCAS_PUBLIC void
+llcas_actioncache_get_for_digest_async(llcas_cas_t, llcas_digest_t key,
+                                       bool globally, void *ctx_cb,
+                                       llcas_actioncache_get_cb);
+
 /**
  * Associates a \c llcas_objectid_t \p value with a \p key. It is invalid to set
  * a different \p value to the same \p key.
@@ -274,6 +285,18 @@ LLCAS_PUBLIC bool llcas_actioncache_put_for_digest(llcas_cas_t,
                                                    llcas_objectid_t value,
                                                    bool globally, char **error);
 
+/**
+ * Like \c llcas_actioncache_put_for_digest but result is provided to a callback
+ * function. Whether the call is asynchronous or not depends on the
+ * implementation.
+ *
+ * \param ctx_cb pointer to pass to the callback function.
+ */
+LLCAS_PUBLIC void
+llcas_actioncache_put_for_digest_async(llcas_cas_t, llcas_digest_t key,
+                                       llcas_objectid_t value, bool globally,
+                                       void *ctx_cb, llcas_actioncache_put_cb);
+
 LLVM_C_EXTERN_C_END
 
 #endif /* LLVM_C_CAS_PLUGINAPI_FUNCTIONS_H */

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

@@ -94,4 +94,25 @@ typedef enum {
 typedef void (*llcas_cas_load_object_cb)(void *ctx, llcas_lookup_result_t,
                                          llcas_loaded_object_t, char *error);
 
+/**
+ * Callback for \c llcas_actioncache_get_for_digest_async.
+ *
+ * \param ctx pointer passed through from the
+ * \c llcas_actioncache_get_for_digest_async call.
+ * \param error message if an error occurred. If set, the memory it points to
+ * needs to be released via \c llcas_string_dispose.
+ */
+typedef void (*llcas_actioncache_get_cb)(void *ctx, llcas_lookup_result_t,
+                                         llcas_objectid_t, char *error);
+
+/**
+ * Callback for \c llcas_actioncache_put_for_digest_async.
+ *
+ * \param ctx pointer passed through from the
+ * \c llcas_actioncache_put_for_digest_async call.
+ * \param error message if an error occurred. If set, the memory it points to
+ * needs to be released via \c llcas_string_dispose.
+ */
+typedef void (*llcas_actioncache_put_cb)(void *ctx, bool failed, char *error);
+
 #endif /* LLVM_C_CAS_PLUGINAPI_TYPES_H */

llvm/include/llvm/CAS/ActionCache.h

@@ -9,11 +9,13 @@
 #ifndef LLVM_CAS_CASACTIONCACHE_H
 #define LLVM_CAS_CASACTIONCACHE_H
 
+#include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/CAS/CASID.h"
 #include "llvm/CAS/CASReference.h"
 #include "llvm/Support/Error.h"
+#include <future>
 
 namespace llvm::cas {
 
@@ -40,6 +42,20 @@ class CacheKey {
   std::string Key;
 };
 
+using AsyncCASIDValue = AsyncValue<CASID>;
+
+/// This is used to workaround the issue of MSVC needing default-constructible
+/// types for \c std::promise/future.
+struct AsyncErrorValue {
+  Error take() { return std::move(Value); }
+
+  AsyncErrorValue() : Value(Error::success()) {}
+  AsyncErrorValue(Error &&E) : Value(std::move(E)) {}
+
+private:
+  Error Value;
+};
+
 /// A cache from a key describing an action to the result of doing it.
 ///
 /// Actions are expected to be pure (collision is an error).
@@ -57,6 +73,18 @@ class ActionCache {
     return getImpl(arrayRefFromStringRef(ActionKey.getKey()), Globally);
   }
 
+  /// Asynchronous version of \c get.
+  std::future<AsyncCASIDValue> getFuture(const CacheKey &ActionKey,
+                                         bool Globally = false) const;
+
+  /// Asynchronous version of \c get.
+  void getAsync(
+      const CacheKey &ActionKey, bool Globally,
+      unique_function<void(Expected<std::optional<CASID>>)> Callback) const {
+    return getImplAsync(arrayRefFromStringRef(ActionKey.getKey()), Globally,
+                        std::move(Callback));
+  }
+
   /// Cache \p Result for the \p ActionKey computation.
   ///
   /// \param Globally if true it is a hint to the underlying implementation that
@@ -70,13 +98,35 @@ class ActionCache {
     return putImpl(arrayRefFromStringRef(ActionKey.getKey()), Result, Globally);
   }
 
+  /// Asynchronous version of \c put.
+  std::future<AsyncErrorValue> putFuture(const CacheKey &ActionKey,
+                                         const CASID &Result,
+                                         bool Globally = false);
+
+  /// Asynchronous version of \c put.
+  void putAsync(const CacheKey &ActionKey, const CASID &Result, bool Globally,
+                unique_function<void(Error)> Callback) {
+    assert(Result.getContext().getHashSchemaIdentifier() ==
+               getContext().getHashSchemaIdentifier() &&
+           "Hash schema mismatch");
+    return putImplAsync(arrayRefFromStringRef(ActionKey.getKey()), Result,
+                        Globally, std::move(Callback));
+  }
+
   virtual ~ActionCache() = default;
 
 protected:
   virtual Expected<Optional<CASID>> getImpl(ArrayRef<uint8_t> ResolvedKey,
                                             bool Globally) const = 0;
+  virtual void getImplAsync(
+      ArrayRef<uint8_t> ResolvedKey, bool Globally,
+      unique_function<void(Expected<std::optional<CASID>>)> Callback) const;
+
   virtual Error putImpl(ArrayRef<uint8_t> ResolvedKey, const CASID &Result,
                         bool Globally) = 0;
+  virtual void putImplAsync(ArrayRef<uint8_t> ResolvedKey, const CASID &Result,
+                            bool Globally,
+                            unique_function<void(Error)> Callback);
 
   ActionCache(const CASContext &Context) : Context(Context) {}
 

llvm/include/llvm/CAS/CASID.h

@@ -14,6 +14,7 @@
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Error.h"
 
 namespace llvm {
 
@@ -121,6 +122,21 @@ class CASID {
   SmallString<32> Hash;
 };
 
+/// This is used to workaround the issue of MSVC needing default-constructible
+/// types for \c std::promise/future.
+template <typename T> struct AsyncValue {
+  Expected<std::optional<T>> take() { return std::move(Value); }
+
+  AsyncValue() : Value(std::nullopt) {}
+  AsyncValue(Error &&E) : Value(std::move(E)) {}
+  AsyncValue(T &&V) : Value(std::move(V)) {}
+  AsyncValue(std::nullopt_t) : Value(std::nullopt) {}
+  AsyncValue(Expected<std::optional<T>> &&Obj) : Value(std::move(Obj)) {}
+
+private:
+  Expected<std::optional<T>> Value;
+};
+
 } // namespace cas
 
 template <> struct DenseMapInfo<cas::CASID> {

llvm/include/llvm/CAS/ObjectStore.h

@@ -26,10 +26,11 @@ template <typename T> class unique_function;
 
 namespace cas {
 
-struct AsyncProxyValue;
 class ObjectStore;
 class ObjectProxy;
 
+using AsyncProxyValue = AsyncValue<ObjectProxy>;
+
 /// Content-addressable storage for objects.
 ///
 /// Conceptually, objects are stored in a "unique set".
@@ -251,7 +252,12 @@ class ObjectStore {
   Expected<std::optional<ObjectProxy>> getProxyIfExists(ObjectRef Ref);
 
   /// Asynchronous version of \c getProxyIfExists.
-  std::future<AsyncProxyValue> getProxyAsync(ObjectRef Ref);
+  std::future<AsyncProxyValue> getProxyFuture(ObjectRef Ref);
+
+  /// Asynchronous version of \c getProxyIfExists using a callback.
+  void getProxyAsync(
+      ObjectRef Ref,
+      unique_function<void(Expected<std::optional<ObjectProxy>>)> Callback);
 
   /// Read the data from \p Data into \p OS.
   uint64_t readData(ObjectHandle Node, raw_ostream &OS, uint64_t Offset = 0,
@@ -347,20 +353,6 @@ class ObjectProxy {
   ObjectHandle H;
 };
 
-/// This is used to workaround the issue of MSVC needing default-constructible
-/// types for \c std::promise/future.
-struct AsyncProxyValue {
-  Expected<std::optional<ObjectProxy>> take() { return std::move(Value); }
-
-  AsyncProxyValue() : Value(std::nullopt) {}
-  AsyncProxyValue(Error &&E) : Value(std::move(E)) {}
-  AsyncProxyValue(ObjectProxy V) : Value(std::move(V)) {}
-  AsyncProxyValue(std::nullopt_t) : Value(std::nullopt) {}
-
-private:
-  Expected<std::optional<ObjectProxy>> Value;
-};
-
 std::unique_ptr<ObjectStore> createInMemoryCAS();
 
 /// \returns true if \c LLVM_ENABLE_ONDISK_CAS configuration was enabled.

llvm/lib/CAS/ActionCache.cpp

@@ -20,3 +20,46 @@ CacheKey::CacheKey(const ObjectProxy &Proxy)
     : CacheKey(Proxy.getCAS(), Proxy.getRef()) {}
 CacheKey::CacheKey(const ObjectStore &CAS, const ObjectRef &Ref)
     : Key(toStringRef(CAS.getID(Ref).getHash())) {}
+
+std::future<AsyncCASIDValue> ActionCache::getFuture(const CacheKey &ActionKey,
+                                                    bool Globally) const {
+  std::promise<AsyncCASIDValue> Promise;
+  auto Future = Promise.get_future();
+  getAsync(ActionKey, Globally,
+           [Promise =
+                std::move(Promise)](Expected<std::optional<CASID>> ID) mutable {
+             Promise.set_value(std::move(ID));
+           });
+  return Future;
+}
+
+std::future<AsyncErrorValue> ActionCache::putFuture(const CacheKey &ActionKey,
+                                                    const CASID &Result,
+                                                    bool Globally) {
+  std::promise<AsyncErrorValue> Promise;
+  auto Future = Promise.get_future();
+  putAsync(ActionKey, Result, Globally,
+           [Promise = std::move(Promise)](Error E) mutable {
+             Promise.set_value(std::move(E));
+           });
+  return Future;
+}
+
+void ActionCache::getImplAsync(
+    ArrayRef<uint8_t> ResolvedKey, bool Globally,
+    unique_function<void(Expected<std::optional<CASID>>)> Callback) const {
+  // The default implementation is synchronous.
+  Optional<CASID> Val;
+  if (Error E = getImpl(ResolvedKey, Globally).moveInto(Val))
+    return Callback(std::move(E));
+  if (Val)
+    return Callback(std::move(*Val));
+  return Callback(std::nullopt);
+}
+
+void ActionCache::putImplAsync(ArrayRef<uint8_t> ResolvedKey,
+                               const CASID &Result, bool Globally,
+                               unique_function<void(Error)> Callback) {
+  // The default implementation is synchronous.
+  return Callback(putImpl(ResolvedKey, Result, Globally));
+}

llvm/lib/CAS/ObjectStore.cpp

@@ -117,23 +117,33 @@ ObjectStore::getProxyIfExists(ObjectRef Ref) {
   return ObjectProxy::load(*this, Ref, *H);
 }
 
-std::future<AsyncProxyValue> ObjectStore::getProxyAsync(ObjectRef Ref) {
+std::future<AsyncProxyValue> ObjectStore::getProxyFuture(ObjectRef Ref) {
   std::promise<AsyncProxyValue> Promise;
   auto Future = Promise.get_future();
+  getProxyAsync(Ref, [Promise = std::move(Promise)](
+                         Expected<std::optional<ObjectProxy>> Obj) mutable {
+    Promise.set_value(std::move(Obj));
+  });
+  return Future;
+}
+
+void ObjectStore::getProxyAsync(
+    ObjectRef Ref,
+    unique_function<void(Expected<std::optional<ObjectProxy>>)> Callback) {
   // FIXME: there is potential for use-after-free for the 'this' pointer.
   // Either we should always allocate shared pointers for \c ObjectStore objects
   // and pass \c shared_from_this() or expect that the caller will not release
   // the \c ObjectStore before the callback returns.
-  loadIfExistsAsync(Ref, [this, Ref, Promise = std::move(Promise)](
-                             Expected<std::optional<ObjectHandle>> H) mutable {
-    if (!H)
-      Promise.set_value(H.takeError());
-    else if (!*H)
-      Promise.set_value(std::nullopt);
-    else
-      Promise.set_value(ObjectProxy::load(*this, Ref, **H));
-  });
-  return Future;
+  return loadIfExistsAsync(
+      Ref, [this, Ref, Callback = std::move(Callback)](
+               Expected<std::optional<ObjectHandle>> H) mutable {
+        if (!H)
+          Callback(H.takeError());
+        else if (!*H)
+          Callback(std::nullopt);
+        else
+          Callback(ObjectProxy::load(*this, Ref, **H));
+      });
 }
 
 Error ObjectStore::createUnknownObjectError(const CASID &ID) {

llvm/lib/CAS/PluginAPI.h

@@ -80,9 +80,18 @@ struct llcas_functions_t {
                                                       bool globally,
                                                       char **error);
 
+  void (*actioncache_get_for_digest_async)(llcas_cas_t, llcas_digest_t key,
+                                           bool globally, void *ctx_cb,
+                                           llcas_actioncache_get_cb);
+
   bool (*actioncache_put_for_digest)(llcas_cas_t, llcas_digest_t key,
                                      llcas_objectid_t value, bool globally,
                                      char **error);
+
+  void (*actioncache_put_for_digest_async)(llcas_cas_t, llcas_digest_t key,
+                                           llcas_objectid_t value,
+                                           bool globally, void *ctx_cb,
+                                           llcas_actioncache_put_cb);
 };
 
 #endif // LLVM_LIB_CAS_PLUGINAPI_H