Merge pull request swiftlang#73323 from hamishknight/cancel-on-close

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

Author: hamishknight

tools/SourceKit/include/SourceKit/Core/LangSupport.h

@@ -1105,7 +1105,8 @@ class LangSupport {
                                  SourceKitCancellationToken CancellationToken,
                                  std::shared_ptr<EditorConsumer> Consumer) = 0;
 
-  virtual void editorClose(StringRef Name, bool RemoveCache) = 0;
+  virtual void editorClose(StringRef Name, bool CancelBuilds,
+                           bool RemoveCache) = 0;
 
   virtual void editorReplaceText(StringRef Name, llvm::MemoryBuffer *Buf,
                                  unsigned Offset, unsigned Length,

tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp

@@ -451,6 +451,9 @@ class ASTBuildOperation
   /// 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 @@ class ASTProducer : public std::enable_shared_from_this<ASTProducer> {
                        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) {
@@ -631,7 +637,16 @@ struct SwiftASTManager::Implementation {
     });
   }
 
-  ASTProducerRef getASTProducer(SwiftInvocationRef InvokRef);
+  /// Retrieve the ASTProducer for a given invocation, creating one if needed.
+  ASTProducerRef getOrCreateASTProducer(SwiftInvocationRef InvokRef);
+
+  /// Retrieve the ASTProducer for a given invocation, returning \c nullopt if
+  /// not present.
+  std::optional<ASTProducerRef> getASTProducer(SwiftInvocationRef Invok);
+
+  /// Updates the cache entry to account for any changes to the ASTProducer
+  /// for the given invocation.
+  void updateASTProducer(SwiftInvocationRef Invok);
 
   FileContent
   getFileContent(StringRef FilePath, bool IsPrimary,
@@ -780,7 +795,7 @@ void SwiftASTManager::processASTAsync(
     const void *OncePerASTToken, SourceKitCancellationToken CancellationToken,
     llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> fileSystem) {
   assert(fileSystem);
-  ASTProducerRef Producer = Impl.getASTProducer(InvokRef);
+  ASTProducerRef Producer = Impl.getOrCreateASTProducer(InvokRef);
 
   Impl.cleanDeletedConsumers();
   {
@@ -816,12 +831,39 @@ void SwiftASTManager::processASTAsync(
       });
 }
 
+std::optional<ASTProducerRef>
+SwiftASTManager::Implementation::getASTProducer(SwiftInvocationRef Invok) {
+  llvm::sys::ScopedLock L(CacheMtx);
+  return ASTCache.get(Invok->Impl.Key);
+}
+
+void SwiftASTManager::Implementation::updateASTProducer(
+    SwiftInvocationRef Invok) {
+  llvm::sys::ScopedLock L(CacheMtx);
+
+  // Get and set the producer to update its cost in the cache. If we don't
+  // have a value, then this is a race where we've removed the cached AST, but
+  // still have a build waiting to complete after cancellation, we don't need
+  // to do anything in that case.
+  if (auto Producer = ASTCache.get(Invok->Impl.Key))
+    ASTCache.set(Invok->Impl.Key, *Producer);
+}
+
 void SwiftASTManager::removeCachedAST(SwiftInvocationRef Invok) {
+  llvm::sys::ScopedLock L(Impl.CacheMtx);
   Impl.ASTCache.remove(Invok->Impl.Key);
 }
 
-ASTProducerRef
-SwiftASTManager::Implementation::getASTProducer(SwiftInvocationRef InvokRef) {
+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);
   std::optional<ASTProducerRef> OptProducer = ASTCache.get(InvokRef->Impl.Key);
   if (OptProducer.has_value())
@@ -978,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) {
@@ -1302,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,
@@ -1342,7 +1411,7 @@ void ASTProducer::enqueueConsumer(
               [This]() { This->cleanBuildOperations(); });
           // Re-register the object with the cache to update its memory
           // cost.
-          Mgr->Impl.ASTCache.set(This->InvokRef->Impl.Key, This);
+          Mgr->Impl.updateASTProducer(This->InvokRef);
         }
       };
 

tools/SourceKit/lib/SwiftLang/SwiftASTManager.h

@@ -313,6 +313,7 @@ class SwiftASTManager : public std::enable_shared_from_this<SwiftASTManager> {
                                       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 @@ class SwiftDocumentSemanticInfo :
     }
   }
 
+  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");
@@ -2436,20 +2447,25 @@ void SwiftLangSupport::editorOpen(StringRef Name, llvm::MemoryBuffer *Buf,
 // EditorClose
 //===----------------------------------------------------------------------===//
 
-void SwiftLangSupport::editorClose(StringRef Name, bool RemoveCache) {
+void SwiftLangSupport::editorClose(StringRef Name, bool CancelBuilds,
+                                   bool RemoveCache) {
   auto Removed = EditorDocuments->remove(Name);
-  if (Removed) {
-    --Stats->numOpenDocs;
-  } else {
+  if (!Removed) {
+    // FIXME: Report error if Name did not apply to anything ?
     IFaceGenContexts.remove(Name);
+    return;
   }
+  --Stats->numOpenDocs;
 
-  if (Removed && RemoveCache)
+  // Cancel any in-flight builds for the given AST if needed.
+  if (CancelBuilds)
+    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 ?
 }
 
-
 //===----------------------------------------------------------------------===//
 // EditorReplaceText
 //===----------------------------------------------------------------------===//

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h

@@ -107,6 +107,7 @@ class SwiftEditorDocument :
   void updateSemaInfo(SourceKitCancellationToken CancellationToken);
 
   void removeCachedAST();
+  void cancelBuildsForCachedAST();
 
   ImmutableTextSnapshotRef getLatestSnapshot() const;
 
@@ -614,7 +615,8 @@ class SwiftLangSupport : public LangSupport {
       SourceKitCancellationToken CancellationToken,
       std::shared_ptr<EditorConsumer> Consumer) override;
 
-  void editorClose(StringRef Name, bool RemoveCache) override;
+  void editorClose(StringRef Name, bool CancelBuilds,
+                   bool RemoveCache) override;
 
   void editorReplaceText(StringRef Name, llvm::MemoryBuffer *Buf,
                          unsigned Offset, unsigned Length,

tools/SourceKit/tools/sourcekitd/lib/Service/Requests.cpp

@@ -288,8 +288,8 @@ static sourcekitd_response_t editorExtractTextFromComment(StringRef Source);
 
 static sourcekitd_response_t editorConvertMarkupToXML(StringRef Source);
 
-static sourcekitd_response_t
-editorClose(StringRef Name, bool RemoveCache);
+static sourcekitd_response_t editorClose(StringRef Name, bool CancelBuilds,
+                                         bool RemoveCache);
 
 static sourcekitd_response_t
 editorReplaceText(StringRef Name, llvm::MemoryBuffer *Buf, unsigned Offset,
@@ -831,10 +831,14 @@ handleRequestEditorClose(const RequestDict &Req,
     if (!Name.has_value())
       return Rec(createErrorRequestInvalid("missing 'key.name'"));
 
+    // Whether to cancel in-flight builds, default true.
+    int64_t CancelBuilds = true;
+    Req.getInt64(KeyCancelBuilds, CancelBuilds, /*isOptional=*/true);
+
     // Whether we remove the cached AST from libcache, by default, false.
     int64_t RemoveCache = false;
     Req.getInt64(KeyRemoveCache, RemoveCache, /*isOptional=*/true);
-    return Rec(editorClose(*Name, RemoveCache));
+    return Rec(editorClose(*Name, CancelBuilds, RemoveCache));
   }
 }
 
@@ -3656,11 +3660,11 @@ editorOpenHeaderInterface(StringRef Name, StringRef HeaderName,
   return EditC.createResponse();
 }
 
-static sourcekitd_response_t
-editorClose(StringRef Name, bool RemoveCache) {
+static sourcekitd_response_t editorClose(StringRef Name, bool CancelBuilds,
+                                         bool RemoveCache) {
   ResponseBuilder RespBuilder;
   LangSupport &Lang = getGlobalContext().getSwiftLangSupport();
-  Lang.editorClose(Name, RemoveCache);
+  Lang.editorClose(Name, CancelBuilds, RemoveCache);
   return RespBuilder.createResponse();
 }
 

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)