[clang][cas] Implement cancellation functionality for the asynchronous CAS query APIs

(cherry picked from commit b88f023)

Author: akyrtzi

clang/test/CAS/libclang-replay-job.c

@@ -7,7 +7,7 @@
 // RUN: clang-scan-deps -compilation-database %t/cdb.json \
 // RUN:   -format experimental-include-tree-full \
 // RUN:   -cas-path %t/cas -fcas-plugin-path %llvmshlibdir/libCASPluginTest%pluginext \
-// RUN:   -fcas-plugin-option upstream-path=%t/cas-upstream -fcas-plugin-option no-logging \
+// RUN:   -fcas-plugin-option no-logging \
 // RUN:   > %t/deps.json
 
 // RUN: %deps-to-rsp %t/deps.json --tu-index 0 > %t/cc1.rsp
@@ -24,21 +24,30 @@
 // RUN:   -e "s/^.*hit for '//" \
 // RUN:   -e "s/' .*$//" > %t/cache-key
 
+// RUN: c-index-test core -upload-cached-job -cas-path %t/cas @%t/cache-key -test-cas-cancellation \
+// RUN:   -fcas-plugin-path %llvmshlibdir/libCASPluginTest%pluginext \
+// RUN:   -fcas-plugin-option upstream-path=%t/cas-upstream \
+// RUN:   2>&1 | FileCheck %s --check-prefix=UPLOAD-CANCEL
+// UPLOAD-CANCEL: actioncache_put_for_digest_async cancelled
+
 // Delete the "local" cache and use the "upstream" one to re-materialize the outputs locally.
 // RUN: rm -rf %t/cas
 
 // Re-run the scan to populate the include-tree in the cas
 // RUN: clang-scan-deps -compilation-database %t/cdb.json \
 // RUN:   -format experimental-include-tree-full \
 // RUN:   -cas-path %t/cas -fcas-plugin-path %llvmshlibdir/libCASPluginTest%pluginext \
-// RUN:   -fcas-plugin-option upstream-path=%t/cas-upstream -fcas-plugin-option no-logging \
+// RUN:   -fcas-plugin-option no-logging \
 // RUN:   > %t/deps2.json
 // RUN: diff -u %t/deps.json %t/deps2.json
 
 
-// RUN: c-index-test core -materialize-cached-job -cas-path %t/cas @%t/cache-key \
+// RUN: c-index-test core -materialize-cached-job -cas-path %t/cas @%t/cache-key -test-cas-cancellation \
 // RUN:   -fcas-plugin-path %llvmshlibdir/libCASPluginTest%pluginext \
-// RUN:   -fcas-plugin-option upstream-path=%t/cas-upstream -fcas-plugin-option no-logging
+// RUN:   -fcas-plugin-option upstream-path=%t/cas-upstream \
+// RUN:   2>&1 | FileCheck %s --check-prefix=MATERIALIZE-CANCEL
+// MATERIALIZE-CANCEL: actioncache_get_for_digest_async cancelled
+// MATERIALIZE-CANCEL: load_object_async cancelled
 
 // RUN: c-index-test core -replay-cached-job -cas-path %t/cas @%t/cache-key \
 // RUN:   -fcas-plugin-path %llvmshlibdir/libCASPluginTest%pluginext \

clang/tools/c-index-test/core_main.cpp

@@ -58,6 +58,7 @@ enum class ActionType {
   AggregateAsJSON,
   ScanDeps,
   ScanDepsByModuleName,
+  UploadCachedJob,
   MaterializeCachedJob,
   ReplayCachedJob,
   PruneCAS,
@@ -85,6 +86,8 @@ Action(cl::desc("Action:"), cl::init(ActionType::None),
                      "Get file dependencies"),
           clEnumValN(ActionType::ScanDepsByModuleName, "scan-deps-by-mod-name",
                      "Get file dependencies by module name alone"),
+          clEnumValN(ActionType::UploadCachedJob, "upload-cached-job",
+                     "Upload cached compilation data to upstream CAS"),
           clEnumValN(ActionType::MaterializeCachedJob, "materialize-cached-job",
                      "Materialize cached compilation data from upstream CAS"),
           clEnumValN(ActionType::ReplayCachedJob, "replay-cached-job",
@@ -154,6 +157,10 @@ static cl::list<std::string> CASPluginOpts("fcas-plugin-option",
                                            cl::desc("Plugin CAS Options"));
 static llvm::cl::opt<std::string>
     WorkingDir("working-dir", llvm::cl::desc("Path for working directory"));
+static cl::opt<bool> TestCASCancellation(
+    "test-cas-cancellation",
+    cl::desc(
+        "perform extra CAS API invocation and cancel it for testing purposes"));
 }
 } // anonymous namespace
 
@@ -865,41 +872,117 @@ static int scanDeps(ArrayRef<const char *> Args, std::string WorkingDirectory,
   return 1;
 }
 
-static int materializeCachedJob(std::string CacheKey, CXCASDatabases DBs) {
-  struct CompResult {
-    CXCASCachedCompilation Comp = nullptr;
-    CXError Err = nullptr;
+static int uploadCachedJob(std::string CacheKey, CXCASDatabases DBs) {
+  CXError Err = nullptr;
+  CXCASCachedCompilation CComp = clang_experimental_cas_getCachedCompilation(
+      DBs, CacheKey.c_str(), /*Globally*/ false, &Err);
+  auto CleanupCachedComp = llvm::make_scope_exit(
+      [&] { clang_experimental_cas_CachedCompilation_dispose(CComp); });
+  if (!CComp) {
+    if (Err) {
+      llvm::errs() << clang_Error_getDescription(Err) << "\n";
+      clang_Error_dispose(Err);
+    } else {
+      llvm::errs() << "cache key was not found\n";
+    }
+    return 1;
+  }
+
+  /// \returns true of an error occurred.
+  auto invokeMakeGlobal = [&](bool Cancel) -> bool {
+    CXCASCancellationToken CancelToken = nullptr;
+    auto CleanupCancelTok = llvm::make_scope_exit([&] {
+      if (Cancel)
+        clang_experimental_cas_CancellationToken_dispose(CancelToken);
+    });
+
+    std::promise<CXError> CallPromise;
+    clang_experimental_cas_CachedCompilation_makeGlobal(
+        CComp, &CallPromise,
+        [](void *Ctx, CXError Err) {
+          static_cast<std::promise<CXError> *>(Ctx)->set_value(Err);
+        },
+        Cancel ? &CancelToken : nullptr);
+    if (Cancel) {
+      clang_experimental_cas_CancellationToken_cancel(CancelToken);
+    }
+    CXError CallRes = CallPromise.get_future().get();
+    if (CallRes) {
+      llvm::errs() << clang_Error_getDescription(CallRes) << "\n";
+      return true;
+    }
+    return false;
   };
-  std::promise<CompResult> CompPromise;
-  auto CompFuture = CompPromise.get_future();
-  struct CompCall {
-    std::promise<CompResult> Promise;
+
+  if (options::TestCASCancellation) {
+    // Cancel an invocation for testing purposes.
+    if (invokeMakeGlobal(/*Cancel=*/true))
+      return 1;
+  }
+  if (invokeMakeGlobal(/*Cancel=*/false))
+    return 1;
+
+  return 0;
+}
+
+static int materializeCachedJob(std::string CacheKey, CXCASDatabases DBs) {
+  /// \returns true of an error occurred.
+  auto invokeGetCachedCompilation =
+      [&](bool Cancel, CXCASCachedCompilation &OutComp) -> bool {
+    OutComp = nullptr;
+    CXCASCancellationToken CancelToken = nullptr;
+    auto CleanupCancelTok = llvm::make_scope_exit([&] {
+      if (Cancel)
+        clang_experimental_cas_CancellationToken_dispose(CancelToken);
+    });
+
+    struct CompResult {
+      CXCASCachedCompilation Comp = nullptr;
+      CXError Err = nullptr;
+    };
+    std::promise<CompResult> CompPromise;
+    auto CompFuture = CompPromise.get_future();
+    struct CompCall {
+      std::promise<CompResult> Promise;
+    };
+    CompCall *CallCtx = new CompCall{std::move(CompPromise)};
+    clang_experimental_cas_getCachedCompilation_async(
+        DBs, CacheKey.c_str(), /*Globally*/ true, CallCtx,
+        [](void *Ctx, CXCASCachedCompilation Comp, CXError Err) {
+          std::unique_ptr<CompCall> CallCtx(static_cast<CompCall *>(Ctx));
+          CallCtx->Promise.set_value(CompResult{Comp, Err});
+        },
+        Cancel ? &CancelToken : nullptr);
+    if (Cancel) {
+      clang_experimental_cas_CancellationToken_cancel(CancelToken);
+    }
+    CompResult Res = CompFuture.get();
+    OutComp = Res.Comp;
+    if (!OutComp && !Cancel) {
+      if (Res.Err) {
+        llvm::errs() << clang_Error_getDescription(Res.Err) << "\n";
+        clang_Error_dispose(Res.Err);
+      } else {
+        llvm::errs() << "cache key was not found\n";
+      }
+      return true;
+    }
+    return false;
   };
-  CompCall *CallCtx = new CompCall{std::move(CompPromise)};
-  clang_experimental_cas_getCachedCompilation_async(
-      DBs, CacheKey.c_str(), /*Globally*/ true, CallCtx,
-      [](void *Ctx, CXCASCachedCompilation Comp, CXError Err) {
-        std::unique_ptr<CompCall> CallCtx(static_cast<CompCall *>(Ctx));
-        CallCtx->Promise.set_value(CompResult{Comp, Err});
-      },
-      /*cancelToken*/ nullptr);
-  CompResult Res = CompFuture.get();
-  CXCASCachedCompilation CComp = Res.Comp;
 
+  CXCASCachedCompilation CComp = nullptr;
   auto CleanupCachedComp = llvm::make_scope_exit([&] {
     if (CComp)
       clang_experimental_cas_CachedCompilation_dispose(CComp);
-    if (Res.Err)
-      clang_Error_dispose(Res.Err);
   });
-  if (!CComp) {
-    if (Res.Err) {
-      llvm::errs() << clang_Error_getDescription(Res.Err) << "\n";
-    } else {
-      llvm::errs() << "cache key was not found\n";
-    }
-    return 1;
+
+  if (options::TestCASCancellation) {
+    // Cancel an invocation for testing purposes.
+    if (invokeGetCachedCompilation(/*Cancel=*/true, CComp))
+      return 1;
   }
+  if (invokeGetCachedCompilation(/*Cancel=*/false, CComp))
+    return 1;
 
   for (unsigned
            I = 0,
@@ -912,42 +995,63 @@ static int materializeCachedJob(std::string CacheKey, CXCASDatabases DBs) {
     auto CleanupOutputID =
         llvm::make_scope_exit([&] { clang_disposeString(OutputID); });
 
-    struct LoadResult {
-      CXCASObject Obj = nullptr;
-      CXError Err = nullptr;
-    };
-    std::promise<LoadResult> LoadPromise;
-    auto LoadFuture = LoadPromise.get_future();
-    struct LoadCall {
-      std::promise<LoadResult> Promise;
-    };
-    LoadCall *CallCtx = new LoadCall{std::move(LoadPromise)};
-    clang_experimental_cas_loadObjectByString_async(
-        DBs, clang_getCString(OutputID), CallCtx,
-        [](void *Ctx, CXCASObject Obj, CXError Err) {
-          std::unique_ptr<LoadCall> CallCtx(static_cast<LoadCall *>(Ctx));
-          CallCtx->Promise.set_value(LoadResult{Obj, Err});
-        },
-        /*cancelToken*/ nullptr);
+    /// \returns true of an error occurred.
+    auto invokeLoadObject = [&](bool Cancel, CXCASObject &OutObj) -> bool {
+      OutObj = nullptr;
+      CXCASCancellationToken CancelToken = nullptr;
+      auto CleanupCancelTok = llvm::make_scope_exit([&] {
+        if (Cancel)
+          clang_experimental_cas_CancellationToken_dispose(CancelToken);
+      });
 
-    LoadResult Res = LoadFuture.get();
-    CXCASObject CASObj = Res.Obj;
+      struct LoadResult {
+        CXCASObject Obj = nullptr;
+        CXError Err = nullptr;
+      };
+      std::promise<LoadResult> LoadPromise;
+      auto LoadFuture = LoadPromise.get_future();
+      struct LoadCall {
+        std::promise<LoadResult> Promise;
+      };
+      LoadCall *CallCtx = new LoadCall{std::move(LoadPromise)};
+      clang_experimental_cas_loadObjectByString_async(
+          DBs, clang_getCString(OutputID), CallCtx,
+          [](void *Ctx, CXCASObject Obj, CXError Err) {
+            std::unique_ptr<LoadCall> CallCtx(static_cast<LoadCall *>(Ctx));
+            CallCtx->Promise.set_value(LoadResult{Obj, Err});
+          },
+          Cancel ? &CancelToken : nullptr);
+      if (Cancel) {
+        clang_experimental_cas_CancellationToken_cancel(CancelToken);
+      }
+      LoadResult Res = LoadFuture.get();
+      OutObj = Res.Obj;
+      if (!OutObj && !Cancel) {
+        if (Res.Err) {
+          llvm::errs() << clang_Error_getDescription(Res.Err) << "\n";
+          clang_Error_dispose(Res.Err);
+        } else {
+          llvm::errs() << "cache key was not found\n";
+        }
+        return true;
+      }
+      return false;
+    };
 
+    CXCASObject CASObj = nullptr;
     auto CleanupLoadObj = llvm::make_scope_exit([&] {
       if (CASObj)
         clang_experimental_cas_CASObject_dispose(CASObj);
-      if (Res.Err)
-        clang_Error_dispose(Res.Err);
     });
 
-    if (!CASObj) {
-      if (Res.Err) {
-        llvm::errs() << clang_Error_getDescription(Res.Err) << "\n";
-      } else {
-        llvm::errs() << "compilation output ID was not found\n";
-      }
-      return 1;
+    if (options::TestCASCancellation) {
+      // Cancel an invocation for testing purposes.
+      if (invokeLoadObject(/*Cancel=*/true, CASObj))
+        return 1;
     }
+    if (invokeLoadObject(/*Cancel=*/false, CASObj))
+      return 1;
+
     if (!clang_experimental_cas_CachedCompilation_isOutputMaterialized(CComp,
                                                                        I))
       report_fatal_error("output was not materialized?");
@@ -1393,6 +1497,18 @@ int indextest_core_main(int argc, const char **argv) {
                     options::OutputDir, DBs, options::ModuleName);
   }
 
+  if (options::Action == ActionType::UploadCachedJob) {
+    if (options::InputFiles.empty()) {
+      errs() << "error: missing cache key\n";
+      return 1;
+    }
+    if (!DBs) {
+      errs() << "error: CAS was not configured\n";
+      return 1;
+    }
+    return uploadCachedJob(options::InputFiles[0], DBs);
+  }
+
   if (options::Action == ActionType::MaterializeCachedJob) {
     if (options::InputFiles.empty()) {
       errs() << "error: missing cache key\n";