[SourceKit] Add an option to cancel async requests

Author: ahoppen

test/SourceKit/Misc/cancellation.swift

@@ -0,0 +1,13 @@
+// Check that we can cancel requests.
+// We need to wait a little bit after request scheduling and cancellation to make sure we are not cancelling the request before it got scheduled.
+
+// RUN: not %sourcekitd-test -req=cursor -id=slow -async -pos=10:3 %s -- %s == \
+// RUN: -shell -- sleep 1 == \
+// RUN: -cancel=slow 2>&1 \
+// RUN: | %FileCheck %s
+
+func foo(x: Invalid1, y: Invalid2) {
+  x / y / x / y / x / y / x / y
+}
+
+// CHECK: error response (Request Cancelled)

tools/SourceKit/tools/sourcekitd-test/Options.td

@@ -16,6 +16,10 @@ def req : Separate<["-"], "req">,
   HelpText<"Request kind">;
 def req_EQ : Joined<["-"], "req=">, Alias<req>;
 
+def id : Separate<["-"], "id">,
+  HelpText<"If this is an async request, an ID that can be used to cancel the request">;
+def id_EQ : Joined<["-"], "id=">, Alias<id>;
+
 def offset : Separate<["-"], "offset">,
   HelpText<"byte offset">;
 def offset_EQ : Joined<["-"], "offset=">, Alias<offset>;
@@ -156,6 +160,10 @@ def module_cache_path_EQ : Joined<["-"], "module-cache-path=">, Alias<module_cac
 def shell: Flag<["-"], "shell">,
   HelpText<"Run shell command">;
 
+def cancel : Separate<["-"], "cancel">,
+  HelpText<"Cancel all async requests with the given ID">;
+def cancel_EQ : Joined<["-"], "cancel=">, Alias<cancel>;
+
 def help : Flag<["-", "--"], "help">,
   HelpText<"Display available options">;
 

tools/SourceKit/tools/sourcekitd-test/TestOptions.cpp

@@ -274,6 +274,10 @@ bool TestOptions::parseArgs(llvm::ArrayRef<const char *> Args) {
       ModuleGroupName = InputArg->getValue();
       break;
 
+    case OPT_id:
+      RequestId = InputArg->getValue();
+      break;
+
     case OPT_interested_usr:
       InterestedUSR = InputArg->getValue();
       break;
@@ -420,6 +424,10 @@ bool TestOptions::parseArgs(llvm::ArrayRef<const char *> Args) {
       ShellExecution = true;
       break;
 
+    case OPT_cancel:
+      CancelRequest = InputArg->getValue();
+      break;
+
     case OPT_disable_implicit_concurrency_module_import:
       DisableImplicitConcurrencyModuleImport = true;
       break;

tools/SourceKit/tools/sourcekitd-test/TestOptions.h

@@ -103,6 +103,10 @@ struct TestOptions {
   std::string ObjCName;
   std::string ObjCSelector;
   std::string Name;
+  /// An ID that can be used to cancel this request.
+  std::string RequestId;
+  /// If not empty, all requests with this ID should be cancelled.
+  std::string CancelRequest;
   bool CheckInterfaceIsASCII = false;
   bool UsedSema = false;
   bool PrintRequest = true;

tools/SourceKit/tools/sourcekitd-test/sourcekitd-test.cpp

@@ -143,6 +143,7 @@ struct AsyncResponseInfo {
   TestOptions options;
   std::string sourceFilename;
   std::unique_ptr<llvm::MemoryBuffer> sourceBuffer;
+  sourcekitd_request_handle_t requestHandle;
 };
 } // end anonymous namespace
 
@@ -480,6 +481,15 @@ static int handleTestInvocation(ArrayRef<const char *> Args,
   if (Opts.ShellExecution)
     return performShellExecution(Opts.CompilerArgs);
 
+  if (!Opts.CancelRequest.empty()) {
+    for (auto &asyncResponse : asyncResponses) {
+      if (asyncResponse.options.RequestId == Opts.CancelRequest) {
+        sourcekitd_cancel_request(asyncResponse.requestHandle);
+      }
+    }
+    return 0;
+  }
+
   assert(Opts.repeatRequest >= 1);
   for (unsigned i = 0; i < Opts.repeatRequest; ++i) {
     if (int ret = handleTestInvocation(Opts, InitOpts)) {
@@ -1196,11 +1206,12 @@ static int handleTestInvocation(TestOptions Opts, TestOptions &InitOpts) {
     if (Opts.PrintRequest)
       sourcekitd_request_description_dump(Req);
 
-    sourcekitd_send_request(Req, nullptr, ^(sourcekitd_response_t resp) {
-      auto &info = asyncResponses[respIndex];
-      info.response = resp;
-      info.semaphore.signal(); // Ready to be handled!
-    });
+    sourcekitd_send_request(Req, &asyncResponses[respIndex].requestHandle,
+                            ^(sourcekitd_response_t resp) {
+                              auto &info = asyncResponses[respIndex];
+                              info.response = resp;
+                              info.semaphore.signal(); // Ready to be handled!
+                            });
 
 #else
     llvm::report_fatal_error(