[SourceKit] Introduce key for cancelling on close

Defaults to `true`.

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/SwiftEditor.cpp

@@ -2447,26 +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) {
-    // Cancel any in-flight builds for the given AST.
+  // 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 ?
+  // Then remove the cached AST if we've been asked to do so.
+  if (RemoveCache)
+    Removed->removeCachedAST();
 }
 
-
 //===----------------------------------------------------------------------===//
 // EditorReplaceText
 //===----------------------------------------------------------------------===//

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h

@@ -615,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/CloseTest.cpp

@@ -105,8 +105,8 @@ class CloseTest : public ::testing::Test {
     getLang().editorOpen(DocName, Buf.get(), Consumer, Args, std::nullopt);
   }
 
-  void close(const char *DocName, bool removeCache) {
-    getLang().editorClose(DocName, removeCache);
+  void close(const char *DocName, bool CancelBuilds, bool RemoveCache) {
+    getLang().editorClose(DocName, CancelBuilds, RemoveCache);
   }
 
   void getDiagnosticsAsync(
@@ -171,7 +171,7 @@ TEST_F(CloseTest, Cancel) {
 
     getCompileTracker().waitForBuildToStart();
 
-    close(DocName, RemoveCache);
+    close(DocName, /*CancelBuilds*/ true, RemoveCache);
 
     bool Expired = BuildResultSema.wait(30 * 1000);
     if (Expired)

unittests/SourceKit/SwiftLang/EditingTest.cpp

@@ -166,7 +166,8 @@ class EditTest : public ::testing::Test {
   }
 
   void close(const char *DocName) {
-    getLang().editorClose(DocName, /*removeCache=*/false);
+    getLang().editorClose(DocName, /*CancelBuilds*/ true,
+                          /*RemoveCache*/ false);
   }
 
   void replaceText(StringRef DocName, unsigned Offset, unsigned Length,

utils/gyb_sourcekit_support/UIDs.py

@@ -116,6 +116,7 @@ def __init__(self, internal_name, external_name):
     KEY('Introduced', 'key.introduced'),
     KEY('Deprecated', 'key.deprecated'),
     KEY('Obsoleted', 'key.obsoleted'),
+    KEY('CancelBuilds', 'key.cancel_builds'),
     KEY('RemoveCache', 'key.removecache'),
     KEY('TypeUsr', 'key.typeusr'),
     KEY('ContainerTypeUsr', 'key.containertypeusr'),