[CAS] Introduce changes in ObjectStore/ActionCache to better accomodate distributed caching

There are two main changes:

* Introduce `Globally` parameter to `ActionCache::get/put`. This is a hint to the underlying implementation
that the key association and the associated object is profitable to be globally distributed.
* `ObjectStore::loadIfExists` can be called asynchronously and can return `std::nullopt`. With support from the
underlying implementation, these changes allow downloading the outputs of the compilation concurrently and lazily
(e.g. only downloading the 'main' module file output). With `std::nullopt` the implementation can indicate that one of
the requested outputs was missing from the CAS.

`CompileJobCacheResult::forEachLoadedOutput` is introduced to enable concurrent downloading of outputs and
`ObjectStoreCachingOutputs` is enhanced to take advantage of it.

`libCASPluginTest` is enhanced to simulate "uploading"/"downloading" objects for testing purposes.

Make sure to handle multiple errors by joining them.

(cherry picked from commit d4a0c07)

Author: akyrtzi

clang/include/clang/Frontend/CASDependencyCollector.h

@@ -36,10 +36,10 @@ class CASDependencyCollector : public DependencyFileGenerator {
   /// \param Opts Output options. Only options that affect the output format of
   ///             a dependency file are signficant.
   /// \param CAS The CAS to read the result from.
-  /// \param DepsRef The dependencies.
+  /// \param Deps The dependencies.
   /// \param OS The output stream to write the dependency file to.
   static llvm::Error replay(const DependencyOutputOptions &Opts,
-                            cas::ObjectStore &CAS, cas::ObjectRef DepsRef,
+                            cas::ObjectStore &CAS, cas::ObjectProxy Deps,
                             llvm::raw_ostream &OS);
 
 private:

clang/include/clang/Frontend/CompileJobCacheResult.h

@@ -46,11 +46,21 @@ class CompileJobCacheResult : public ObjectProxy {
   llvm::Error
   forEachOutput(llvm::function_ref<llvm::Error(Output)> Callback) const;
 
+  /// Loads all outputs concurrently and passes the resulting \c ObjectProxy
+  /// objects to \p Callback. If there was an error during loading then the
+  /// callback will not be invoked.
+  llvm::Error forEachLoadedOutput(
+      llvm::function_ref<llvm::Error(Output, std::optional<ObjectProxy>)>
+          Callback);
+
   size_t getNumOutputs() const;
 
   /// Retrieves a specific output specified by \p Kind, if it exists.
   Optional<Output> getOutput(OutputKind Kind) const;
 
+  /// \returns a string for the given \p Kind.
+  static StringRef getOutputKindName(OutputKind Kind);
+
   /// Print this result to \p OS.
   llvm::Error print(llvm::raw_ostream &OS);
 

clang/lib/CAS/CASOptions.cpp

@@ -86,7 +86,7 @@ void CASOptions::initCache(DiagnosticsEngine &Diags) const {
   StringRef CASPath = Cache.Config.CASPath;
 
   if (!PluginPath.empty()) {
-    std::pair<std::unique_ptr<ObjectStore>, std::unique_ptr<ActionCache>> DBs;
+    std::pair<std::shared_ptr<ObjectStore>, std::shared_ptr<ActionCache>> DBs;
     if (llvm::Error E =
             createPluginCASDatabases(PluginPath, CASPath, PluginOptions)
                 .moveInto(DBs)) {

clang/lib/Frontend/CASDependencyCollector.cpp

@@ -31,12 +31,8 @@ CASDependencyCollector::CASDependencyCollector(
       CAS(CAS), Callback(std::move(Callback)) {}
 
 llvm::Error CASDependencyCollector::replay(const DependencyOutputOptions &Opts,
-                                           ObjectStore &CAS, ObjectRef DepsRef,
+                                           ObjectStore &CAS, ObjectProxy Refs,
                                            llvm::raw_ostream &OS) {
-  auto Refs = CAS.getProxy(DepsRef);
-  if (!Refs)
-    return Refs.takeError();
-
   CASDependencyCollector DC(Opts, CAS, nullptr);
 
   // Add the filenames from DependencyOutputOptions::ExtraDeps. These are kept
@@ -47,7 +43,7 @@ llvm::Error CASDependencyCollector::replay(const DependencyOutputOptions &Opts,
     DC.addDependency(Dep.first);
   }
 
-  auto Err = Refs->forEachReference([&](ObjectRef Ref) -> llvm::Error {
+  auto Err = Refs.forEachReference([&](ObjectRef Ref) -> llvm::Error {
     auto PathHandle = CAS.getProxy(Ref);
     if (!PathHandle)
       return PathHandle.takeError();

clang/lib/Frontend/CompileJobCacheResult.cpp

@@ -34,6 +34,43 @@ Error CompileJobCacheResult::forEachOutput(
   return Error::success();
 }
 
+Error CompileJobCacheResult::forEachLoadedOutput(
+    llvm::function_ref<Error(Output, std::optional<ObjectProxy>)> Callback) {
+  // Kick-off materialization for all outputs concurrently.
+  SmallVector<std::future<Expected<std::optional<ObjectProxy>>>, 4>
+      FutureOutputs;
+  size_t Count = getNumOutputs();
+  for (size_t I = 0; I < Count; ++I) {
+    ObjectRef Ref = getOutputObject(I);
+    FutureOutputs.push_back(getCAS().getProxyAsync(Ref));
+  }
+
+  // Make sure all the outputs have materialized.
+  std::optional<Error> OccurredError;
+  SmallVector<std::optional<ObjectProxy>, 4> Outputs;
+  for (auto &FutureOutput : FutureOutputs) {
+    auto Obj = FutureOutput.get();
+    if (!Obj) {
+      if (!OccurredError)
+        OccurredError = Obj.takeError();
+      else
+        OccurredError =
+            llvm::joinErrors(std::move(*OccurredError), Obj.takeError());
+      continue;
+    }
+    Outputs.push_back(*Obj);
+  }
+  if (OccurredError)
+    return std::move(*OccurredError);
+
+  // Pass the loaded outputs.
+  for (size_t I = 0; I < Count; ++I) {
+    if (auto Err = Callback({getOutputObject(I), getOutputKind(I)}, Outputs[I]))
+      return Err;
+  }
+  return Error::success();
+}
+
 Optional<CompileJobCacheResult::Output>
 CompileJobCacheResult::getOutput(OutputKind Kind) const {
   size_t Count = getNumOutputs();
@@ -45,25 +82,21 @@ CompileJobCacheResult::getOutput(OutputKind Kind) const {
   return None;
 }
 
-static void printOutputKind(llvm::raw_ostream &OS,
-                            CompileJobCacheResult::OutputKind Kind) {
+StringRef CompileJobCacheResult::getOutputKindName(OutputKind Kind) {
   switch (Kind) {
-  case CompileJobCacheResult::OutputKind::MainOutput:
-    OS << "main   ";
-    break;
-  case CompileJobCacheResult::OutputKind::Dependencies:
-    OS << "deps   ";
-    break;
-  case CompileJobCacheResult::OutputKind::SerializedDiagnostics:
-    OS << "diags  ";
-    break;
+  case OutputKind::MainOutput:
+    return "main";
+  case OutputKind::SerializedDiagnostics:
+    return "deps";
+  case OutputKind::Dependencies:
+    return "diags";
   }
 }
 
 Error CompileJobCacheResult::print(llvm::raw_ostream &OS) {
   return forEachOutput([&](Output O) -> Error {
-    printOutputKind(OS, O.Kind);
-    OS << ' ' << getCAS().getID(O.Object) << '\n';
+    OS << getOutputKindName(O.Kind) << "    " << getCAS().getID(O.Object)
+       << '\n';
     return Error::success();
   });
 }

clang/lib/Frontend/CompilerInstance.cpp

@@ -2401,6 +2401,10 @@ static bool addCachedModuleFileToInMemoryCache(
       Result->getOutput(cas::CompileJobCacheResult::OutputKind::MainOutput);
   if (!Output)
     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
+  // better network utilization.
   auto OutputProxy = CAS.getProxy(Output->Object);
   if (!OutputProxy) {
     Diags.Report(diag::err_cas_cannot_get_module_cache_key)

clang/test/CAS/plugin-cas.c

@@ -2,19 +2,47 @@
 
 // RUN: rm -rf %t && mkdir -p %t
 
-// RUN: %clang -cc1depscan -o %t/t.rsp -fdepscan=inline -cc1-args \
+// RUN: %clang -cc1depscan -o %t/t1.rsp -fdepscan=inline -cc1-args \
 // RUN:   -cc1 %s -fcas-path %t/cas \
 // RUN:   -fcas-plugin-path %llvmshlibdir/libCASPluginTest%pluginext \
-// RUN:   -fcas-plugin-option first-prefix=myfirst- -fcas-plugin-option second-prefix=mysecond-
-// RUN: %clang @%t/t.rsp -emit-obj -o %t/t1.o -Rcompile-job-cache 2>&1 | FileCheck %s --check-prefix=CACHE-MISS
-// RUN: %clang @%t/t.rsp -emit-obj -o %t/t2.o -Rcompile-job-cache 2>&1 | FileCheck %s --check-prefix=CACHE-HIT
+// RUN:   -fcas-plugin-option first-prefix=myfirst- -fcas-plugin-option second-prefix=mysecond- \
+// RUN:   -fcas-plugin-option upstream-path=%t/cas-upstream
+// RUN: %clang @%t/t1.rsp -emit-obj -o %t/t1.o -Rcompile-job-cache 2>&1 | FileCheck %s --check-prefix=CACHE-MISS
+
+// Clear the CAS and check the outputs can still be "downloaded" from upstream.
+// RUN: rm -rf %t/cas
+// RUN: %clang -cc1depscan -o %t/t2.rsp -fdepscan=inline -cc1-args \
+// RUN:   -cc1 %s -fcas-path %t/cas \
+// RUN:   -fcas-plugin-path %llvmshlibdir/libCASPluginTest%pluginext \
+// RUN:   -fcas-plugin-option first-prefix=myfirst- -fcas-plugin-option second-prefix=mysecond- \
+// RUN:   -fcas-plugin-option upstream-path=%t/cas-upstream
+// RUN: %clang @%t/t2.rsp -emit-obj -o %t/t2.o -Rcompile-job-cache 2>&1 | FileCheck %s --check-prefix=CACHE-HIT
 // RUN: diff %t/t1.o %t/t2.o
 
+// Check that it's a cache miss if outputs are not found in the upstream CAS.
+// RUN: rm -rf %t/cas
+// RUN: %clang -cc1depscan -o %t/t3.rsp -fdepscan=inline -cc1-args \
+// RUN:   -cc1 %s -fcas-path %t/cas \
+// RUN:   -fcas-plugin-path %llvmshlibdir/libCASPluginTest%pluginext \
+// RUN:   -fcas-plugin-option first-prefix=myfirst- -fcas-plugin-option second-prefix=mysecond- \
+// RUN:   -fcas-plugin-option upstream-path=%t/cas-upstream \
+// RUN:   -fcas-plugin-option simulate-missing-objects
+// RUN: %clang @%t/t3.rsp -emit-obj -o %t/t.o -Rcompile-job-cache 2>&1 | FileCheck %s --check-prefix=CACHE-NOTFOUND
+
 // CACHE-MISS: remark: compile job cache miss for 'myfirst-mysecond-
 // CACHE-MISS: warning: some warning
 
-// CACHE-HIT: remark: compile job cache hit for 'myfirst-mysecond-
-// CACHE-HIT: warning: some warning
+// Check that outputs are downloaded concurrently.
+// CACHE-HIT:      load_object_async downstream begin:
+// CACHE-HIT-NEXT: load_object_async downstream begin:
+// CACHE-HIT-NEXT: load_object_async downstream end:
+// CACHE-HIT-NEXT: load_object_async downstream end:
+// CACHE-HIT-NEXT: remark: compile job cache hit for 'myfirst-mysecond-
+// CACHE-HIT-NEXT: warning: some warning
+
+// CACHE-NOTFOUND: remark: compile job cache backend did not find output 'main' for key
+// CACHE-NOTFOUND: remark: compile job cache miss
+// CACHE-NOTFOUND: warning: some warning
 
 // RUN: not %clang -cc1depscan -o %t/t.rsp -fdepscan=inline -cc1-args \
 // RUN:   -cc1 %s -fcas-path %t/cas \

clang/tools/driver/cc1_main.cpp

@@ -273,8 +273,9 @@ class ObjectStoreCachingOutputs : public CachingOutputs {
   /// Replay a cache hit.
   ///
   /// Return status if should exit immediately, otherwise None.
-  Optional<int> replayCachedResult(llvm::cas::ObjectRef ResultID,
-                                   bool JustComputedResult);
+  std::optional<int> replayCachedResult(const llvm::cas::CASID &ResultCacheKey,
+                                        llvm::cas::ObjectRef ResultID,
+                                        bool JustComputedResult);
 
   std::shared_ptr<llvm::cas::ObjectStore> CAS;
   std::shared_ptr<llvm::cas::ActionCache> Cache;
@@ -561,7 +562,8 @@ Expected<bool> ObjectStoreCachingOutputs::tryReplayCachedResult(
       Diags.Report(diag::remark_compile_job_cache_timing_backend_key_query)
           << llvm::format("%.6fs", Seconds);
     });
-    if (Error E = Cache->get(ResultCacheKey).moveInto(Result))
+    if (Error E =
+            Cache->get(ResultCacheKey, /*Globally=*/true).moveInto(Result))
       return std::move(E);
   }
 
@@ -583,11 +585,11 @@ Expected<bool> ObjectStoreCachingOutputs::tryReplayCachedResult(
     return false;
   }
 
-  Diags.Report(diag::remark_compile_job_cache_hit)
-      << ResultCacheKey.toString() << CAS->getID(*ResultRef).toString();
-  Optional<int> Status =
-      replayCachedResult(*ResultRef, /*JustComputedResult=*/false);
-  assert(Status && "Expected a status for a cache hit");
+  // \c replayCachedResult emits remarks for a cache hit or miss.
+  std::optional<int> Status = replayCachedResult(ResultCacheKey, *ResultRef,
+                                                 /*JustComputedResult=*/false);
+  if (!Status)
+    return false; // cache miss.
   assert(*Status == 0 && "Expected success status for a cache hit");
   return true;
 }
@@ -770,22 +772,23 @@ Error ObjectStoreCachingOutputs::finishComputedResult(
       Diags.Report(diag::remark_compile_job_cache_timing_backend_key_update)
           << llvm::format("%.6fs", Seconds);
     });
-    if (llvm::Error E = Cache->put(ResultCacheKey, CAS->getID(*Result)))
+    if (llvm::Error E =
+            Cache->put(ResultCacheKey, CAS->getID(*Result), /*Globally=*/true))
       return E;
   }
 
   // Replay / decanonicalize as necessary.
-  Optional<int> Status = replayCachedResult(*Result,
-                                            /*JustComputedResult=*/true);
+  std::optional<int> Status = replayCachedResult(ResultCacheKey, *Result,
+                                                 /*JustComputedResult=*/true);
   (void)Status;
   assert(Status == None);
   return Error::success();
 }
 
 /// Replay a result after a cache hit.
-Optional<int>
-ObjectStoreCachingOutputs::replayCachedResult(llvm::cas::ObjectRef ResultID,
-                                              bool JustComputedResult) {
+std::optional<int> ObjectStoreCachingOutputs::replayCachedResult(
+    const llvm::cas::CASID &ResultCacheKey, llvm::cas::ObjectRef ResultID,
+    bool JustComputedResult) {
   if (JustComputedResult)
     return None;
 
@@ -795,19 +798,30 @@ ObjectStoreCachingOutputs::replayCachedResult(llvm::cas::ObjectRef ResultID,
   if (Error E = Schema.load(ResultID).moveInto(Result))
     llvm::report_fatal_error(std::move(E));
 
-  auto Err = Result->forEachOutput([&](clang::cas::CompileJobCacheResult::Output
-                                           O) -> Error {
+  DiagnosticsEngine &Diags = Clang.getDiagnostics();
+  bool HasMissingOutput = false;
+  std::optional<llvm::cas::ObjectProxy> SerialDiags;
+
+  auto processOutput = [&](clang::cas::CompileJobCacheResult::Output O,
+                           std::optional<llvm::cas::ObjectProxy> Obj) -> Error {
+    if (!Obj.has_value()) {
+      Diags.Report(diag::remark_compile_job_cache_backend_output_not_found)
+          << clang::cas::CompileJobCacheResult::getOutputKindName(O.Kind)
+          << ResultCacheKey.toString() << CAS->getID(O.Object).toString();
+      HasMissingOutput = true;
+      return Error::success();
+    }
+    if (HasMissingOutput)
+      return Error::success();
+
     if (O.Kind == OutputKind::SerializedDiagnostics) {
-      Optional<llvm::cas::ObjectProxy> DiagsObj;
-      if (Error E = CAS->getProxy(O.Object).moveInto(DiagsObj))
-        return E;
-      return replayCachedDiagnostics(DiagsObj->getData());
+      SerialDiags = Obj;
+      return Error::success();
     }
 
     std::string Path = std::string(getPathForOutputKind(O.Kind));
     if (Path.empty())
-      // The output may be always generated but not needed with this invocation,
-      // like the serialized diagnostics file.
+      // The output may be always generated but not needed with this invocation.
       return Error::success(); // continue
 
     // Always create parent directory of outputs, since it is hard to precisely
@@ -823,14 +837,11 @@ ObjectStoreCachingOutputs::replayCachedResult(llvm::cas::ObjectRef ResultID,
     if (O.Kind == OutputKind::Dependencies) {
       llvm::raw_svector_ostream OS(ContentsStorage);
       if (auto E = CASDependencyCollector::replay(
-              Clang.getDependencyOutputOpts(), *CAS, O.Object, OS))
+              Clang.getDependencyOutputOpts(), *CAS, *Obj, OS))
         return E;
       Contents = ContentsStorage;
     } else {
-      Optional<llvm::cas::ObjectProxy> Bytes;
-      if (Error E = CAS->getProxy(O.Object).moveInto(Bytes))
-        return E;
-      Contents = Bytes->getData();
+      Contents = Obj->getData();
     }
 
     std::unique_ptr<llvm::FileOutputBuffer> Output;
@@ -839,12 +850,29 @@ ObjectStoreCachingOutputs::replayCachedResult(llvm::cas::ObjectRef ResultID,
       return E;
     llvm::copy(*Contents, Output->getBufferStart());
     return Output->commit();
-  });
+  };
 
   // FIXME: Stop calling report_fatal_error().
-  if (Err)
+  if (auto Err = Result->forEachLoadedOutput(processOutput))
     llvm::report_fatal_error(std::move(Err));
 
+  if (HasMissingOutput) {
+    Diags.Report(diag::remark_compile_job_cache_miss)
+        << ResultCacheKey.toString();
+    return std::nullopt;
+  }
+
+  if (!JustComputedResult) {
+    Diags.Report(diag::remark_compile_job_cache_hit)
+        << ResultCacheKey.toString() << CAS->getID(ResultID).toString();
+
+    if (SerialDiags) {
+      // FIXME: Stop calling report_fatal_error().
+      if (Error E = replayCachedDiagnostics(SerialDiags->getData()))
+        llvm::report_fatal_error(std::move(E));
+    }
+  }
+
   if (JustComputedResult)
     return None;
   return 0;