[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,