[SourceKit] Cancel in-flight builds on editor.close

When closing a document, cancel any in-flight
builds happening for it.

rdar://127126348

tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp

@@ -451,6 +451,9 @@ public:
   /// consumer, removes it from the \c Consumers severed by this build operation
   /// and, if no consumers are left, cancels the AST build of this operation.
   void requestConsumerCancellation(SwiftASTConsumerRef Consumer);
+
+  /// Cancels all consumers for the given operation.
+  void cancelAllConsumers();
 };
 
 using ASTBuildOperationRef = std::shared_ptr<ASTBuildOperation>;
@@ -517,6 +520,9 @@ public:
                        IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
                        SwiftASTManagerRef Mgr);
 
+  /// Cancel all currently running build operations.
+  void cancelAllBuilds();
+
   size_t getMemoryCost() const {
     size_t Cost = sizeof(*this);
     for (auto &BuildOp : BuildOperations) {
@@ -848,6 +854,14 @@ void SwiftASTManager::removeCachedAST(SwiftInvocationRef Invok) {
   Impl.ASTCache.remove(Invok->Impl.Key);
 }
 
+void SwiftASTManager::cancelBuildsForCachedAST(SwiftInvocationRef Invok) {
+  auto Result = Impl.getASTProducer(Invok);
+  if (!Result)
+    return;
+
+  (*Result)->cancelAllBuilds();
+}
+
 ASTProducerRef SwiftASTManager::Implementation::getOrCreateASTProducer(
     SwiftInvocationRef InvokRef) {
   llvm::sys::ScopedLock L(CacheMtx);
@@ -1006,6 +1020,24 @@ void ASTBuildOperation::requestConsumerCancellation(
   });
 }
 
+void ASTBuildOperation::cancelAllConsumers() {
+  if (isFinished())
+    return;
+
+  llvm::sys::ScopedLock L(ConsumersAndResultMtx);
+  CancellationFlag->store(true, std::memory_order_relaxed);
+
+  // Take the consumers, and notify them of the cancellation.
+  decltype(this->Consumers) Consumers;
+  std::swap(Consumers, this->Consumers);
+
+  ASTManager->Impl.ConsumerNotificationQueue.dispatch(
+      [Consumers = std::move(Consumers)] {
+        for (auto &Consumer : Consumers)
+          Consumer->cancelled();
+      });
+}
+
 static void collectModuleDependencies(ModuleDecl *TopMod,
     llvm::SmallPtrSetImpl<ModuleDecl *> &Visited,
     SmallVectorImpl<std::string> &Filenames) {
@@ -1330,6 +1362,15 @@ ASTBuildOperationRef ASTProducer::getBuildOperationForConsumer(
   return LatestUsableOp;
 }
 
+void ASTProducer::cancelAllBuilds() {
+  // Cancel all build operations, cleanup will happen when each operation
+  // terminates.
+  BuildOperationsQueue.dispatch([This = shared_from_this()] {
+    for (auto &BuildOp : This->BuildOperations)
+      BuildOp->cancelAllConsumers();
+  });
+}
+
 void ASTProducer::enqueueConsumer(
     SwiftASTConsumerRef Consumer,
     IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,

tools/SourceKit/lib/SwiftLang/SwiftASTManager.h

@@ -313,6 +313,7 @@ public:
                                       bool AllowInputs = true);
 
   void removeCachedAST(SwiftInvocationRef Invok);
+  void cancelBuildsForCachedAST(SwiftInvocationRef Invok);
 
   struct Implementation;
   Implementation &Impl;

tools/SourceKit/lib/SwiftLang/SwiftEditor.cpp

@@ -707,6 +707,13 @@ public:
     }
   }
 
+  void cancelBuildsForCachedAST() {
+    if (!InvokRef)
+      return;
+    if (auto ASTMgr = this->ASTMgr.lock())
+      ASTMgr->cancelBuildsForCachedAST(InvokRef);
+  }
+
 private:
   std::vector<SwiftSemanticToken> takeSemanticTokens(
       ImmutableTextSnapshotRef NewSnapshot);
@@ -2174,6 +2181,10 @@ void SwiftEditorDocument::removeCachedAST() {
   Impl.SemanticInfo->removeCachedAST();
 }
 
+void SwiftEditorDocument::cancelBuildsForCachedAST() {
+  Impl.SemanticInfo->cancelBuildsForCachedAST();
+}
+
 void SwiftEditorDocument::applyFormatOptions(OptionsDictionary &FmtOptions) {
   static UIdent KeyUseTabs("key.editor.format.usetabs");
   static UIdent KeyIndentWidth("key.editor.format.indentwidth");
@@ -2444,8 +2455,14 @@ void SwiftLangSupport::editorClose(StringRef Name, bool RemoveCache) {
     IFaceGenContexts.remove(Name);
   }
 
-  if (Removed && RemoveCache)
+  if (Removed) {
+    // Cancel any in-flight builds for the given AST.
+    Removed->cancelBuildsForCachedAST();
+
+    // Then remove the cached AST if we've been asked to do so.
+    if (RemoveCache)
       Removed->removeCachedAST();
+  }
   // FIXME: Report error if Name did not apply to anything ?
 }
 

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h

@@ -107,6 +107,7 @@ public:
   void updateSemaInfo(SourceKitCancellationToken CancellationToken);
 
   void removeCachedAST();
+  void cancelBuildsForCachedAST();
 
   ImmutableTextSnapshotRef getLatestSnapshot() const;
 

unittests/SourceKit/SwiftLang/CMakeLists.txt

@@ -1,6 +1,7 @@
 if(NOT SWIFT_HOST_VARIANT MATCHES "${SWIFT_DARWIN_EMBEDDED_VARIANTS}")
   add_swift_unittest(SourceKitSwiftLangTests
     CursorInfoTest.cpp
+    CloseTest.cpp
     EditingTest.cpp
     )
   target_link_libraries(SourceKitSwiftLangTests PRIVATE SourceKitSwiftLang)

unittests/SourceKit/SwiftLang/CloseTest.cpp

@@ -0,0 +1,180 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#include "NullEditorConsumer.h"
+#include "SourceKit/Core/Context.h"
+#include "SourceKit/Core/LangSupport.h"
+#include "SourceKit/Core/NotificationCenter.h"
+#include "SourceKit/Support/Concurrency.h"
+#include "SourceKit/SwiftLang/Factory.h"
+#include "swift/Basic/LLVMInitialize.h"
+#include "gtest/gtest.h"
+
+#include <chrono>
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+using namespace SourceKit;
+using namespace llvm;
+
+static StringRef getRuntimeLibPath() {
+  return sys::path::parent_path(SWIFTLIB_DIR);
+}
+
+static SmallString<128> getSwiftExecutablePath() {
+  SmallString<128> path = sys::path::parent_path(getRuntimeLibPath());
+  sys::path::append(path, "bin", "swift-frontend");
+  return path;
+}
+
+namespace {
+class CompileTrackingConsumer final : public trace::TraceConsumer {
+  std::mutex Mtx;
+  std::condition_variable CV;
+  bool HasStarted = false;
+
+public:
+  CompileTrackingConsumer() {}
+  CompileTrackingConsumer(const CompileTrackingConsumer &) = delete;
+
+  void operationStarted(uint64_t OpId, trace::OperationKind OpKind,
+                        const trace::SwiftInvocation &Inv,
+                        const trace::StringPairs &OpArgs) override {
+    std::unique_lock<std::mutex> lk(Mtx);
+    HasStarted = true;
+    CV.notify_all();
+  }
+
+  void waitForBuildToStart() {
+    std::unique_lock<std::mutex> lk(Mtx);
+    auto secondsToWait = std::chrono::seconds(20);
+    auto when = std::chrono::system_clock::now() + secondsToWait;
+    CV.wait_until(lk, when, [&]() { return HasStarted; });
+    HasStarted = false;
+  }
+
+  void operationFinished(uint64_t OpId, trace::OperationKind OpKind,
+                         ArrayRef<DiagnosticEntryInfo> Diagnostics) override {}
+
+  swift::OptionSet<trace::OperationKind> desiredOperations() override {
+    return trace::OperationKind::PerformSema;
+  }
+};
+
+class CloseTest : public ::testing::Test {
+  std::shared_ptr<SourceKit::Context> Ctx;
+  std::shared_ptr<CompileTrackingConsumer> CompileTracker;
+
+  NullEditorConsumer Consumer;
+
+public:
+  CloseTest() {
+    INITIALIZE_LLVM();
+    Ctx = std::make_shared<SourceKit::Context>(
+        getSwiftExecutablePath(), getRuntimeLibPath(),
+        /*diagnosticDocumentationPath*/ "", SourceKit::createSwiftLangSupport,
+        /*dispatchOnMain=*/false);
+  }
+
+  CompileTrackingConsumer &getCompileTracker() const { return *CompileTracker; }
+  LangSupport &getLang() { return Ctx->getSwiftLangSupport(); }
+
+  void SetUp() override {
+    CompileTracker = std::make_shared<CompileTrackingConsumer>();
+    trace::registerConsumer(CompileTracker.get());
+  }
+
+  void TearDown() override {
+    trace::unregisterConsumer(CompileTracker.get());
+    CompileTracker = nullptr;
+  }
+
+  void open(const char *DocName, StringRef Text, ArrayRef<const char *> CArgs) {
+    auto Args = makeArgs(DocName, CArgs);
+    auto Buf = MemoryBuffer::getMemBufferCopy(Text, DocName);
+    getLang().editorOpen(DocName, Buf.get(), Consumer, Args, std::nullopt);
+  }
+
+  void close(const char *DocName, bool removeCache) {
+    getLang().editorClose(DocName, removeCache);
+  }
+
+  void getDiagnosticsAsync(
+      const char *DocName, ArrayRef<const char *> CArgs,
+      llvm::function_ref<void(const RequestResult<DiagnosticsResult> &)>
+          callback) {
+    auto Args = makeArgs(DocName, CArgs);
+    getLang().getDiagnostics(DocName, Args, /*VFSOpts*/ std::nullopt,
+                             /*CancelToken*/ {}, callback);
+  }
+
+private:
+  std::vector<const char *> makeArgs(const char *DocName,
+                                     ArrayRef<const char *> CArgs) {
+    std::vector<const char *> Args = CArgs;
+    Args.push_back(DocName);
+    return Args;
+  }
+};
+
+} // end anonymous namespace
+
+static const char *getComplexSourceText() {
+  // best of luck, type-checker
+  return
+    "struct A: ExpressibleByIntegerLiteral { init(integerLiteral value: Int) {} }\n"
+    "struct B: ExpressibleByIntegerLiteral { init(integerLiteral value: Int) {} }\n"
+    "struct C: ExpressibleByIntegerLiteral { init(integerLiteral value: Int) {} }\n"
+
+    "func + (lhs: A, rhs: B) -> A { fatalError() }\n"
+    "func + (lhs: B, rhs: C) -> A { fatalError() }\n"
+    "func + (lhs: C, rhs: A) -> A { fatalError() }\n"
+
+    "func + (lhs: B, rhs: A) -> B { fatalError() }\n"
+    "func + (lhs: C, rhs: B) -> B { fatalError() }\n"
+    "func + (lhs: A, rhs: C) -> B { fatalError() }\n"
+
+    "func + (lhs: C, rhs: B) -> C { fatalError() }\n"
+    "func + (lhs: B, rhs: C) -> C { fatalError() }\n"
+    "func + (lhs: A, rhs: A) -> C { fatalError() }\n"
+
+    "let x: C = 1 + 2 + 3 + 4 + 5 + 6 + 7 + 8\n";
+}
+
+TEST_F(CloseTest, Cancel) {
+  const char *DocName = "test.swift";
+  auto *Contents = getComplexSourceText();
+  const char *Args[] = {"-parse-as-library"};
+
+  // Test twice with RemoveCache = false to test both the prior state of
+  // the ASTProducer being cached and not cached.
+  for (auto RemoveCache : {true, false, false}) {
+    open(DocName, Contents, Args);
+
+    Semaphore BuildResultSema(0);
+
+    getDiagnosticsAsync(DocName, Args,
+                        [&](const RequestResult<DiagnosticsResult> &Result) {
+      EXPECT_TRUE(Result.isCancelled());
+      BuildResultSema.signal();
+    });
+
+    getCompileTracker().waitForBuildToStart();
+
+    close(DocName, RemoveCache);
+
+    bool Expired = BuildResultSema.wait(30 * 1000);
+    if (Expired)
+      llvm::report_fatal_error("Did not receive a response for the request");
+  }
+}