Merge pull request #60573 from ahoppen/pr/re-enable-cancellation

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

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