[SourceKit] Re-enable cancellation of non-completion requests

This enables the ability to cancel requests, which aren’t code completion requests, again.

Previous crashes in SILGen are prevented by disabling cancellation during the SIL stages. Instead, we add dedicated cancellation checkpoints before and after SIL.

rdar://98390926

Author: ahoppen

test/SourceKit/Diagnostics/cancel_diags.swift

@@ -0,0 +1,7 @@
+// RUN: not %sourcekitd-test -req=diags %s -print-raw-response -id=diag -async -- %s == -cancel=diag 2>&1 | %FileCheck %s
+
+func foo(x: Invalid1, y: Invalid2) {
+    x / y / x / y / x / y / x / y
+}
+
+// CHECK: error response (Request Cancelled)

tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp

@@ -1069,6 +1069,7 @@ ASTUnitRef ASTBuildOperation::buildASTUnit(std::string &Error) {
     }
     return nullptr;
   }
+  CompIns.getASTContext().CancellationFlag = CancellationFlag;
   registerIDERequestFunctions(CompIns.getASTContext().evaluator);
   if (TracedOp.enabled()) {
     TracedOp.start(TraceInfo);
@@ -1112,9 +1113,19 @@ ASTUnitRef ASTBuildOperation::buildASTUnit(std::string &Error) {
     // don't block any other AST processing for the same SwiftInvocation.
 
     if (auto SF = CompIns.getPrimarySourceFile()) {
+      if (CancellationFlag->load(std::memory_order_relaxed)) {
+        return nullptr;
+      }
+      // Disable cancellation while performing SILGen. If the cancellation flag
+      // is set, type checking performed during SILGen checks the cancellation
+      // flag and might thus fail, which SILGen cannot handle.
+      llvm::SaveAndRestore<std::shared_ptr<std::atomic<bool>>> DisableCancellationDuringSILGen(CompIns.getASTContext().CancellationFlag, nullptr);
       SILOptions SILOpts = Invocation.getSILOptions();
       auto &TC = CompIns.getSILTypes();
       std::unique_ptr<SILModule> SILMod = performASTLowering(*SF, TC, SILOpts);
+      if (CancellationFlag->load(std::memory_order_relaxed)) {
+        return nullptr;
+      }
       runSILDiagnosticPasses(*SILMod);
     }
   }

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -1621,6 +1621,10 @@ static void computeDiagnostics(
       auto Diagnostics = DiagConsumer.getDiagnosticsForBuffer(BufferID);
       Receiver(RequestResult<DiagnosticsResult>::fromResult(Diagnostics));
     }
+
+    void cancelled() override {
+      Receiver(RequestResult<DiagnosticsResult>::cancelled());
+    }
   };
 
   auto Consumer = std::make_shared<DiagnosticsConsumer>(std::move(Receiver));

unittests/SourceKit/SwiftLang/CursorInfoTest.cpp

@@ -424,8 +424,7 @@ TEST_F(CursorInfoTest, CursorInfoMustWaitDueTokenRace) {
   EXPECT_EQ(strlen("fog"), Info.Length);
 }
 
-// Disabled until we re-enable cancellation (rdar://91251055)
-TEST_F(CursorInfoTest, DISABLED_CursorInfoCancelsPreviousRequest) {
+TEST_F(CursorInfoTest, CursorInfoCancelsPreviousRequest) {
   // TODO: This test case relies on the following snippet being slow to type 
   // check so that the first cursor info request takes longer to execute than it 
   // takes time to schedule the second request. If that is fixed, we need to 
@@ -475,8 +474,7 @@ TEST_F(CursorInfoTest, DISABLED_CursorInfoCancelsPreviousRequest) {
     llvm::report_fatal_error("Did not receive a resonse for the first request");
 }
 
-// Disabled until we re-enable cancellation (rdar://91251055)
-TEST_F(CursorInfoTest, DISABLED_CursorInfoCancellation) {
+TEST_F(CursorInfoTest, CursorInfoCancellation) {
   // TODO: This test case relies on the following snippet being slow to type
   // check so that the first cursor info request takes longer to execute than it
   // takes time to schedule the second request. If that is fixed, we need to