Merge pull request #9441 from benlangmuir/no-cancel-4

[cursor-info] Add a way to opt out of automatic request cancellation
diff --git a/test/SourceKit/CursorInfo/cursor_no_cancel.swift b/test/SourceKit/CursorInfo/cursor_no_cancel.swift
new file mode 100644
index 0000000..19a86fb
--- /dev/null
+++ b/test/SourceKit/CursorInfo/cursor_no_cancel.swift
@@ -0,0 +1,46 @@
+func myFunc() {
+  _ = 1
+}
+
+// Perform 8 concurrent cursor infos, which is often enough to cause
+// contention.  We disable printing the requests to minimize delay.
+
+// RUN: %sourcekitd-test \
+// RUN:      -async -dont-print-request -cancel-on-subsequent-request=0 -req=cursor -pos=1:6 %s -- %s \
+// RUN:   == -async -dont-print-request -cancel-on-subsequent-request=0 -req=cursor -pos=1:6 %s -- %s \
+// RUN:   == -async -dont-print-request -cancel-on-subsequent-request=0 -req=cursor -pos=1:6 %s -- %s \
+// RUN:   == -async -dont-print-request -cancel-on-subsequent-request=0 -req=cursor -pos=1:6 %s -- %s \
+// RUN:   == -async -dont-print-request -cancel-on-subsequent-request=0 -req=cursor -pos=1:6 %s -- %s \
+// RUN:   == -async -dont-print-request -cancel-on-subsequent-request=0 -req=cursor -pos=1:6 %s -- %s \
+// RUN:   == -async -dont-print-request -cancel-on-subsequent-request=0 -req=cursor -pos=1:6 %s -- %s \
+// RUN:   == -async -dont-print-request -cancel-on-subsequent-request=0 -req=cursor -pos=1:6 %s -- %s 2>&1 \
+// RUN:   | %FileCheck %s -implicit-check-not='Request Cancel'
+
+// CHECK: source.lang.swift.decl.function.free
+// CHECK: source.lang.swift.decl.function.free
+// CHECK: source.lang.swift.decl.function.free
+// CHECK: source.lang.swift.decl.function.free
+// CHECK: source.lang.swift.decl.function.free
+// CHECK: source.lang.swift.decl.function.free
+// CHECK: source.lang.swift.decl.function.free
+// CHECK: source.lang.swift.decl.function.free
+
+// RUN: %sourcekitd-test \
+// RUN:      -async -dont-print-request -cancel-on-subsequent-request=0 -req=range -pos=2:3 -length=5 %s -- %s \
+// RUN:   == -async -dont-print-request -cancel-on-subsequent-request=0 -req=range -pos=2:3 -length=5 %s -- %s \
+// RUN:   == -async -dont-print-request -cancel-on-subsequent-request=0 -req=range -pos=2:3 -length=5 %s -- %s \
+// RUN:   == -async -dont-print-request -cancel-on-subsequent-request=0 -req=range -pos=2:3 -length=5 %s -- %s \
+// RUN:   == -async -dont-print-request -cancel-on-subsequent-request=0 -req=range -pos=2:3 -length=5 %s -- %s \
+// RUN:   == -async -dont-print-request -cancel-on-subsequent-request=0 -req=range -pos=2:3 -length=5 %s -- %s \
+// RUN:   == -async -dont-print-request -cancel-on-subsequent-request=0 -req=range -pos=2:3 -length=5 %s -- %s \
+// RUN:   == -async -dont-print-request -cancel-on-subsequent-request=0 -req=range -pos=2:3 -length=5 %s -- %s 2>&1 \
+// RUN:   | %FileCheck %s -check-prefix=RANGE -implicit-check-not='Request Cancel'
+
+// RANGE: source.lang.swift.range.singleexpression
+// RANGE: source.lang.swift.range.singleexpression
+// RANGE: source.lang.swift.range.singleexpression
+// RANGE: source.lang.swift.range.singleexpression
+// RANGE: source.lang.swift.range.singleexpression
+// RANGE: source.lang.swift.range.singleexpression
+// RANGE: source.lang.swift.range.singleexpression
+// RANGE: source.lang.swift.range.singleexpression
diff --git a/tools/SourceKit/docs/Protocol.md b/tools/SourceKit/docs/Protocol.md
index 40b9afa..7dae01b 100644
--- a/tools/SourceKit/docs/Protocol.md
+++ b/tools/SourceKit/docs/Protocol.md
@@ -635,6 +635,10 @@
     [opt] <key.compilerargs>: [string*] // Array of zero or more strings for the compiler arguments,
                                         // e.g ["-sdk", "/path/to/sdk"]. If key.sourcefile is provided,
                                         // these must include the path to that file.
+    [opt] <key.cancel_on_subsequent_request>: (int64) // Whether this request should be canceled if a
+                                        // new cursor-info request is made that uses the same AST.
+                                        // This behaviour is a workaround for not having first-class
+                                        // cancelation. For backwards compatibility, the default is 1.
 }
 ```
 
diff --git a/tools/SourceKit/include/SourceKit/Core/LangSupport.h b/tools/SourceKit/include/SourceKit/Core/LangSupport.h
index 295dbe8..995ee60 100644
--- a/tools/SourceKit/include/SourceKit/Core/LangSupport.h
+++ b/tools/SourceKit/include/SourceKit/Core/LangSupport.h
@@ -475,6 +475,7 @@
 
   virtual void getCursorInfo(StringRef Filename, unsigned Offset,
                              unsigned Length, bool Actionables,
+                             bool CancelOnSubsequentRequest,
                              ArrayRef<const char *> Args,
                           std::function<void(const CursorInfo &)> Receiver) = 0;
 
@@ -485,16 +486,19 @@
                 std::function<void(const NameTranslatingInfo &)> Receiver) = 0;
 
   virtual void getRangeInfo(StringRef Filename, unsigned Offset, unsigned Length,
+                            bool CancelOnSubsequentRequest,
                             ArrayRef<const char *> Args,
                             std::function<void(const RangeInfo&)> Receiver) = 0;
 
   virtual void
   getCursorInfoFromUSR(StringRef Filename, StringRef USR,
+                       bool CancelOnSubsequentRequest,
                        ArrayRef<const char *> Args,
                        std::function<void(const CursorInfo &)> Receiver) = 0;
 
   virtual void findRelatedIdentifiersInFile(StringRef Filename,
                                             unsigned Offset,
+                                            bool CancelOnSubsequentRequest,
                                             ArrayRef<const char *> Args,
                    std::function<void(const RelatedIdentsInfo &)> Receiver) = 0;
 
diff --git a/tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h b/tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h
index 55c7c31..e1b1176 100644
--- a/tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h
+++ b/tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h
@@ -389,6 +389,7 @@
 
   void getCursorInfo(StringRef Filename, unsigned Offset,
                      unsigned Length, bool Actionables,
+                     bool CancelOnSubsequentRequest,
                      ArrayRef<const char *> Args,
                      std::function<void(const CursorInfo &)> Receiver) override;
 
@@ -398,14 +399,16 @@
                    std::function<void(const NameTranslatingInfo &)> Receiver) override;
 
   void getRangeInfo(StringRef Filename, unsigned Offset, unsigned Length,
-                    ArrayRef<const char *> Args,
+                    bool CancelOnSubsequentRequest, ArrayRef<const char *> Args,
                     std::function<void(const RangeInfo&)> Receiver) override;
 
   void getCursorInfoFromUSR(
-      StringRef Filename, StringRef USR, ArrayRef<const char *> Args,
+      StringRef Filename, StringRef USR, bool CancelOnSubsequentRequest,
+      ArrayRef<const char *> Args,
       std::function<void(const CursorInfo &)> Receiver) override;
 
   void findRelatedIdentifiersInFile(StringRef Filename, unsigned Offset,
+                                    bool CancelOnSubsequentRequest,
                                     ArrayRef<const char *> Args,
               std::function<void(const RelatedIdentsInfo &)> Receiver) override;
 
diff --git a/tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp b/tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp
index e4f3c86..2bd5289 100644
--- a/tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp
+++ b/tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp
@@ -1031,6 +1031,8 @@
   SmallVector<ImmutableTextSnapshotRef, 4> PreviousASTSnaps;
 
 protected:
+  bool CancelOnSubsequentRequest;
+protected:
   ArrayRef<ImmutableTextSnapshotRef> getPreviousASTSnaps() {
     return llvm::makeArrayRef(PreviousASTSnaps);
   }
@@ -1038,9 +1040,10 @@
 public:
   CursorRangeInfoConsumer(StringRef InputFile, unsigned Offset, unsigned Length,
                           SwiftLangSupport &Lang, SwiftInvocationRef ASTInvok,
-                          bool TryExistingAST)
+                          bool TryExistingAST, bool CancelOnSubsequentRequest)
     : Lang(Lang), ASTInvok(ASTInvok),InputFile(InputFile), Offset(Offset),
-      Length(Length), TryExistingAST(TryExistingAST) { }
+      Length(Length), TryExistingAST(TryExistingAST),
+      CancelOnSubsequentRequest(CancelOnSubsequentRequest) {}
 
   bool canUseASTWithSnapshots(ArrayRef<ImmutableTextSnapshotRef> Snapshots) override {
     if (!TryExistingAST) {
@@ -1101,6 +1104,7 @@
                           unsigned Length, bool Actionables,
                           SwiftInvocationRef Invok,
                           bool TryExistingAST,
+                          bool CancelOnSubsequentRequest,
                           std::function<void(const CursorInfo &)> Receiver) {
   assert(Invok);
 
@@ -1114,9 +1118,11 @@
                        SwiftLangSupport &Lang,
                        SwiftInvocationRef ASTInvok,
                        bool TryExistingAST,
+                       bool CancelOnSubsequentRequest,
                        std::function<void(const CursorInfo &)> Receiver)
     : CursorRangeInfoConsumer(InputFile, Offset, Length, Lang, ASTInvok,
-                              TryExistingAST), Actionables(Actionables),
+                              TryExistingAST, CancelOnSubsequentRequest),
+      Actionables(Actionables),
       Receiver(std::move(Receiver)){ }
 
     void handlePrimaryAST(ASTUnitRef AstUnit) override {
@@ -1167,7 +1173,8 @@
           if (!getPreviousASTSnaps().empty()) {
             // Attempt again using the up-to-date AST.
             resolveCursor(Lang, InputFile, Offset, Length, Actionables, ASTInvok,
-                          /*TryExistingAST=*/false, Receiver);
+                          /*TryExistingAST=*/false, CancelOnSubsequentRequest,
+                          Receiver);
           } else {
             Receiver({});
           }
@@ -1198,11 +1205,16 @@
 
   auto Consumer = std::make_shared<CursorInfoConsumer>(
     InputFile, Offset, Length, Actionables, Lang, Invok, TryExistingAST,
-    Receiver);
+    CancelOnSubsequentRequest, Receiver);
+
   /// FIXME: When request cancellation is implemented and Xcode adopts it,
   /// don't use 'OncePerASTToken'.
   static const char OncePerASTToken = 0;
-  Lang.getASTManager().processASTAsync(Invok, std::move(Consumer), &OncePerASTToken);
+  static const char OncePerASTTokenWithActionables = 0;
+  const void *Once = nullptr;
+  if (CancelOnSubsequentRequest)
+    Once = Actionables ? &OncePerASTTokenWithActionables : &OncePerASTToken;
+  Lang.getASTManager().processASTAsync(Invok, std::move(Consumer), Once);
 }
 
 static void resolveName(SwiftLangSupport &Lang, StringRef InputFile,
@@ -1222,8 +1234,9 @@
                      bool TryExistingAST, NameTranslatingInfo Input,
                      std::function<void(const NameTranslatingInfo &)> Receiver)
     : CursorRangeInfoConsumer(InputFile, Offset, 0, Lang, ASTInvok,
-                              TryExistingAST), Input(std::move(Input)),
-      Receiver(std::move(Receiver)){ }
+                              TryExistingAST,
+                              /*CancelOnSubsequentRequest=*/false),
+      Input(std::move(Input)), Receiver(std::move(Receiver)){ }
 
     void handlePrimaryAST(ASTUnitRef AstUnit) override {
       auto &CompIns = AstUnit->getCompilerInstance();
@@ -1296,17 +1309,13 @@
   auto Consumer = std::make_shared<NameInfoConsumer>(
     InputFile, Offset, Lang, Invok, TryExistingAST, Input, Receiver);
 
-  /// FIXME: When request cancellation is implemented and Xcode adopts it,
-  /// don't use 'OncePerASTToken'.
-  static const char OncePerASTToken = 0;
-  Lang.getASTManager().processASTAsync(Invok, std::move(Consumer),
-                                       &OncePerASTToken);
+  Lang.getASTManager().processASTAsync(Invok, std::move(Consumer), nullptr);
 }
 
 static void resolveRange(SwiftLangSupport &Lang,
                           StringRef InputFile, unsigned Offset, unsigned Length,
                           SwiftInvocationRef Invok,
-                          bool TryExistingAST,
+                          bool TryExistingAST, bool CancelOnSubsequentRequest,
                           std::function<void(const RangeInfo&)> Receiver) {
   assert(Invok);
 
@@ -1316,10 +1325,11 @@
   public:
     RangeInfoConsumer(StringRef InputFile, unsigned Offset, unsigned Length,
                        SwiftLangSupport &Lang, SwiftInvocationRef ASTInvok,
-                       bool TryExistingAST,
+                       bool TryExistingAST, bool CancelOnSubsequentRequest,
                        std::function<void(const RangeInfo&)> Receiver)
     : CursorRangeInfoConsumer(InputFile, Offset, Length, Lang, ASTInvok,
-                              TryExistingAST), Receiver(std::move(Receiver)){ }
+                              TryExistingAST, CancelOnSubsequentRequest),
+      Receiver(std::move(Receiver)){ }
 
     void handlePrimaryAST(ASTUnitRef AstUnit) override {
       if (trace::enabled()) {
@@ -1352,7 +1362,8 @@
         if (!getPreviousASTSnaps().empty()) {
           // Attempt again using the up-to-date AST.
           resolveRange(Lang, InputFile, Offset, Length, ASTInvok,
-                      /*TryExistingAST=*/false, Receiver);
+                      /*TryExistingAST=*/false, CancelOnSubsequentRequest,
+                      Receiver);
         } else {
           Receiver(Result);
         }
@@ -1373,17 +1384,18 @@
   };
 
   auto Consumer = std::make_shared<RangeInfoConsumer>(
-    InputFile, Offset, Length, Lang, Invok, TryExistingAST, Receiver);
+    InputFile, Offset, Length, Lang, Invok, TryExistingAST,
+    CancelOnSubsequentRequest, Receiver);
   /// FIXME: When request cancellation is implemented and Xcode adopts it,
   /// don't use 'OncePerASTToken'.
   static const char OncePerASTToken = 0;
-  Lang.getASTManager().processASTAsync(Invok, std::move(Consumer),
-                                       &OncePerASTToken);
+  const void *Once = CancelOnSubsequentRequest ? &OncePerASTToken : nullptr;
+  Lang.getASTManager().processASTAsync(Invok, std::move(Consumer), Once);
 }
 
 void SwiftLangSupport::getCursorInfo(
     StringRef InputFile, unsigned Offset, unsigned Length, bool Actionables,
-    ArrayRef<const char *> Args,
+    bool CancelOnSubsequentRequest, ArrayRef<const char *> Args,
     std::function<void(const CursorInfo &)> Receiver) {
 
   if (auto IFaceGenRef = IFaceGenContexts.get(InputFile)) {
@@ -1433,12 +1445,12 @@
   }
 
   resolveCursor(*this, InputFile, Offset, Length, Actionables, Invok,
-                /*TryExistingAST=*/true, Receiver);
+                /*TryExistingAST=*/true, CancelOnSubsequentRequest, Receiver);
 }
 
 void SwiftLangSupport::
 getRangeInfo(StringRef InputFile, unsigned Offset, unsigned Length,
-             ArrayRef<const char *> Args,
+             bool CancelOnSubsequentRequest, ArrayRef<const char *> Args,
              std::function<void(const RangeInfo&)> Receiver) {
   if (IFaceGenContexts.get(InputFile)) {
     // FIXME: return range info for generated interfaces.
@@ -1458,7 +1470,7 @@
     return;
   }
   resolveRange(*this, InputFile, Offset, Length, Invok, /*TryExistingAST=*/true,
-                Receiver);
+               CancelOnSubsequentRequest, Receiver);
 }
 
 void SwiftLangSupport::
@@ -1517,6 +1529,7 @@
 static void
 resolveCursorFromUSR(SwiftLangSupport &Lang, StringRef InputFile, StringRef USR,
                      SwiftInvocationRef Invok, bool TryExistingAST,
+                     bool CancelOnSubsequentRequest,
                      std::function<void(const CursorInfo &)> Receiver) {
   assert(Invok);
 
@@ -1526,16 +1539,18 @@
     SwiftLangSupport &Lang;
     SwiftInvocationRef ASTInvok;
     const bool TryExistingAST;
+    bool CancelOnSubsequentRequest;
     std::function<void(const CursorInfo &)> Receiver;
     SmallVector<ImmutableTextSnapshotRef, 4> PreviousASTSnaps;
 
   public:
     CursorInfoConsumer(StringRef InputFile, StringRef USR,
                        SwiftLangSupport &Lang, SwiftInvocationRef ASTInvok,
-                       bool TryExistingAST,
+                       bool TryExistingAST, bool CancelOnSubsequentRequest,
                        std::function<void(const CursorInfo &)> Receiver)
         : InputFile(InputFile), USR(USR), Lang(Lang),
           ASTInvok(std::move(ASTInvok)), TryExistingAST(TryExistingAST),
+          CancelOnSubsequentRequest(CancelOnSubsequentRequest),
           Receiver(std::move(Receiver)) {}
 
     bool canUseASTWithSnapshots(
@@ -1607,7 +1622,8 @@
           if (!PreviousASTSnaps.empty()) {
             // Attempt again using the up-to-date AST.
             resolveCursorFromUSR(Lang, InputFile, USR, ASTInvok,
-                                 /*TryExistingAST=*/false, Receiver);
+                                 /*TryExistingAST=*/false,
+                                 CancelOnSubsequentRequest, Receiver);
           } else {
             Receiver({});
           }
@@ -1628,16 +1644,18 @@
   };
 
   auto Consumer = std::make_shared<CursorInfoConsumer>(
-      InputFile, USR, Lang, Invok, TryExistingAST, Receiver);
+      InputFile, USR, Lang, Invok, TryExistingAST, CancelOnSubsequentRequest,
+      Receiver);
   /// FIXME: When request cancellation is implemented and Xcode adopts it,
   /// don't use 'OncePerASTToken'.
   static const char OncePerASTToken = 0;
-  Lang.getASTManager().processASTAsync(Invok, std::move(Consumer),
-                                       &OncePerASTToken);
+  const void *Once = CancelOnSubsequentRequest ? &OncePerASTToken : nullptr;
+  Lang.getASTManager().processASTAsync(Invok, std::move(Consumer), Once);
 }
 
 void SwiftLangSupport::getCursorInfoFromUSR(
-    StringRef filename, StringRef USR, ArrayRef<const char *> args,
+    StringRef filename, StringRef USR, bool CancelOnSubsequentRequest,
+    ArrayRef<const char *> args,
     std::function<void(const CursorInfo &)> receiver) {
   if (auto IFaceGenRef = IFaceGenContexts.get(filename)) {
     LOG_WARN_FUNC("info from usr for generated interface not implemented yet");
@@ -1655,7 +1673,7 @@
   }
 
   resolveCursorFromUSR(*this, filename, USR, invok, /*TryExistingAST=*/true,
-                       receiver);
+                       CancelOnSubsequentRequest, receiver);
 }
 
 //===----------------------------------------------------------------------===//
@@ -1724,6 +1742,7 @@
 
 void SwiftLangSupport::findRelatedIdentifiersInFile(
     StringRef InputFile, unsigned Offset,
+    bool CancelOnSubsequentRequest,
     ArrayRef<const char *> Args,
     std::function<void(const RelatedIdentsInfo &)> Receiver) {
 
@@ -1820,5 +1839,6 @@
   /// FIXME: When request cancellation is implemented and Xcode adopts it,
   /// don't use 'OncePerASTToken'.
   static const char OncePerASTToken = 0;
-  ASTMgr->processASTAsync(Invok, std::move(Consumer), &OncePerASTToken);
+  const void *Once = CancelOnSubsequentRequest ? &OncePerASTToken : nullptr;
+  ASTMgr->processASTAsync(Invok, std::move(Consumer), Once);
 }
diff --git a/tools/SourceKit/tools/sourcekitd-test/Options.td b/tools/SourceKit/tools/sourcekitd-test/Options.td
index 7cb2fa0..c241581 100644
--- a/tools/SourceKit/tools/sourcekitd-test/Options.td
+++ b/tools/SourceKit/tools/sourcekitd-test/Options.td
@@ -106,6 +106,11 @@
 def objc_selector : Separate<["-"], "objc-selector">,
   HelpText<"Objective-C selector name to translate from">;
 
+def cancel_on_subsequent_request : Separate<["-"], "cancel-on-subsequent-request">,
+  HelpText<"Whether to cancel if there is a subsequent request using the same AST">;
+def cancel_on_subsequent_request_EQ : Joined<["-"], "cancel-on-subsequent-request=">,
+  Alias<cancel_on_subsequent_request>;
+
 def help : Flag<["-", "--"], "help">,
   HelpText<"Display available options">;
 
diff --git a/tools/SourceKit/tools/sourcekitd-test/TestOptions.cpp b/tools/SourceKit/tools/sourcekitd-test/TestOptions.cpp
index 3cfb04e..2cd97dd 100644
--- a/tools/SourceKit/tools/sourcekitd-test/TestOptions.cpp
+++ b/tools/SourceKit/tools/sourcekitd-test/TestOptions.cpp
@@ -292,6 +292,15 @@
       ObjCSelector = InputArg->getValue();
       break;
 
+    case OPT_cancel_on_subsequent_request:
+      unsigned Cancel;
+      if (StringRef(InputArg->getValue()).getAsInteger(10, Cancel)) {
+        llvm::errs() << "error: expected integer for 'cancel-on-subsequent-request'\n";
+        return true;
+      }
+      CancelOnSubsequentRequest = Cancel;
+      break;
+
     case OPT_UNKNOWN:
       llvm::errs() << "error: unknown argument: "
                    << InputArg->getAsString(ParsedArgs) << '\n'
diff --git a/tools/SourceKit/tools/sourcekitd-test/TestOptions.h b/tools/SourceKit/tools/sourcekitd-test/TestOptions.h
index 1d31a7c..fb06c47 100644
--- a/tools/SourceKit/tools/sourcekitd-test/TestOptions.h
+++ b/tools/SourceKit/tools/sourcekitd-test/TestOptions.h
@@ -90,6 +90,7 @@
   bool SynthesizedExtensions = false;
   bool CollectActionables = false;
   bool isAsyncRequest = false;
+  llvm::Optional<bool> CancelOnSubsequentRequest;
   bool parseArgs(llvm::ArrayRef<const char *> Args);
   void printHelp(bool ShowHidden) const;
 };
diff --git a/tools/SourceKit/tools/sourcekitd-test/sourcekitd-test.cpp b/tools/SourceKit/tools/sourcekitd-test/sourcekitd-test.cpp
index 08b3080..bc5aa9d 100644
--- a/tools/SourceKit/tools/sourcekitd-test/sourcekitd-test.cpp
+++ b/tools/SourceKit/tools/sourcekitd-test/sourcekitd-test.cpp
@@ -141,6 +141,7 @@
 static sourcekitd_uid_t KeySelectorPieces;
 static sourcekitd_uid_t KeyNameKind;
 static sourcekitd_uid_t KeySwiftVersion;
+static sourcekitd_uid_t KeyCancelOnSubsequentRequest;
 
 static sourcekitd_uid_t RequestProtocolVersion;
 static sourcekitd_uid_t RequestDemangle;
@@ -273,6 +274,7 @@
   KeyNameKind = sourcekitd_uid_get_from_cstr("key.namekind");
 
   KeySwiftVersion = sourcekitd_uid_get_from_cstr("key.swift_version");
+  KeyCancelOnSubsequentRequest = sourcekitd_uid_get_from_cstr("key.cancel_on_subsequent_request");
 
   SemaDiagnosticStage = sourcekitd_uid_get_from_cstr("source.diagnostic.stage.swift.sema");
 
@@ -828,6 +830,10 @@
     sourcekitd_request_dictionary_set_string(Req, KeyFilePath,
                                              Opts.HeaderPath.c_str());
   }
+  if (Opts.CancelOnSubsequentRequest.hasValue()) {
+    sourcekitd_request_dictionary_set_int64(Req, KeyCancelOnSubsequentRequest,
+                                            *Opts.CancelOnSubsequentRequest);
+  }
 
   if (Opts.SwiftVersion.hasValue()) {
     sourcekitd_request_dictionary_set_int64(Req, KeySwiftVersion,
diff --git a/tools/SourceKit/tools/sourcekitd/lib/API/DictionaryKeys.h b/tools/SourceKit/tools/sourcekitd/lib/API/DictionaryKeys.h
index 247f84b..7baf984 100644
--- a/tools/SourceKit/tools/sourcekitd/lib/API/DictionaryKeys.h
+++ b/tools/SourceKit/tools/sourcekitd/lib/API/DictionaryKeys.h
@@ -125,6 +125,7 @@
 extern SourceKit::UIdent KeyModuleGroups;
 
 extern SourceKit::UIdent KeyRangeContent;
+extern SourceKit::UIdent KeyCancelOnSubsequentRequest;
 
 extern SourceKit::UIdent KeyBaseName;
 extern SourceKit::UIdent KeyArgNames;
diff --git a/tools/SourceKit/tools/sourcekitd/lib/API/Requests.cpp b/tools/SourceKit/tools/sourcekitd/lib/API/Requests.cpp
index 80e41d6..a5571c4 100644
--- a/tools/SourceKit/tools/sourcekitd/lib/API/Requests.cpp
+++ b/tools/SourceKit/tools/sourcekitd/lib/API/Requests.cpp
@@ -192,6 +192,7 @@
 
 static void findRelatedIdents(StringRef Filename,
                               int64_t Offset,
+                              bool CancelOnSubsequentRequest,
                               ArrayRef<const char *> Args,
                               ResponseReceiver Rec);
 
@@ -758,6 +759,11 @@
   if (ReqUID == RequestCursorInfo) {
     LangSupport &Lang = getGlobalContext().getSwiftLangSupport();
 
+    // For backwards compatibility, the default is 1.
+    int64_t CancelOnSubsequentRequest = 1;
+    Req.getInt64(KeyCancelOnSubsequentRequest, CancelOnSubsequentRequest,
+                 /*isOptional=*/true);
+
     int64_t Offset;
     if (!Req.getInt64(KeyOffset, Offset, /*isOptional=*/false)) {
       int64_t Length = 0;
@@ -765,12 +771,12 @@
       int64_t Actionables = false;
       Req.getInt64(KeyActionable, Actionables, /*isOptional=*/true);
       return Lang.getCursorInfo(
-          *SourceFile, Offset, Length, Actionables, Args,
-          [Rec](const CursorInfo &Info) { reportCursorInfo(Info, Rec); });
+          *SourceFile, Offset, Length, Actionables, CancelOnSubsequentRequest,
+          Args, [Rec](const CursorInfo &Info) { reportCursorInfo(Info, Rec); });
     }
     if (auto USR = Req.getString(KeyUSR)) {
       return Lang.getCursorInfoFromUSR(
-          *SourceFile, *USR, Args,
+          *SourceFile, *USR, CancelOnSubsequentRequest, Args,
           [Rec](const CursorInfo &Info) { reportCursorInfo(Info, Rec); });
     }
 
@@ -782,9 +788,14 @@
     LangSupport &Lang = getGlobalContext().getSwiftLangSupport();
     int64_t Offset;
     int64_t Length;
+    // For backwards compatibility, the default is 1.
+    int64_t CancelOnSubsequentRequest = 1;
+    Req.getInt64(KeyCancelOnSubsequentRequest, CancelOnSubsequentRequest,
+                 /*isOptional=*/true);
     if (!Req.getInt64(KeyOffset, Offset, /*isOptional=*/false)) {
       if (!Req.getInt64(KeyLength, Length, /*isOptional=*/false)) {
-        return Lang.getRangeInfo(*SourceFile, Offset, Length, Args,
+        return Lang.getRangeInfo(*SourceFile, Offset, Length,
+                                 CancelOnSubsequentRequest, Args,
           [Rec](const RangeInfo &Info) { reportRangeInfo(Info, Rec); });
       }
     }
@@ -835,7 +846,14 @@
     int64_t Offset;
     if (Req.getInt64(KeyOffset, Offset, /*isOptional=*/false))
       return Rec(createErrorRequestInvalid("missing 'key.offset'"));
-    return findRelatedIdents(*SourceFile, Offset, Args, Rec);
+
+    // For backwards compatibility, the default is 1.
+    int64_t CancelOnSubsequentRequest = 1;
+    Req.getInt64(KeyCancelOnSubsequentRequest, CancelOnSubsequentRequest,
+                 /*isOptional=*/true);
+
+    return findRelatedIdents(*SourceFile, Offset, CancelOnSubsequentRequest,
+                             Args, Rec);
   }
 
   {
@@ -1514,11 +1532,12 @@
 
 static void findRelatedIdents(StringRef Filename,
                               int64_t Offset,
+                              bool CancelOnSubsequentRequest,
                               ArrayRef<const char *> Args,
                               ResponseReceiver Rec) {
   LangSupport &Lang = getGlobalContext().getSwiftLangSupport();
-  Lang.findRelatedIdentifiersInFile(Filename, Offset, Args,
-                                    [Rec](const RelatedIdentsInfo &Info) {
+  Lang.findRelatedIdentifiersInFile(Filename, Offset, CancelOnSubsequentRequest,
+                                    Args, [Rec](const RelatedIdentsInfo &Info) {
     if (Info.IsCancelled)
       return Rec(createErrorRequestCancelled());
 
diff --git a/tools/SourceKit/tools/sourcekitd/lib/API/sourcekitdAPI-Common.cpp b/tools/SourceKit/tools/sourcekitd/lib/API/sourcekitdAPI-Common.cpp
index e3f1cef..b5d25e1 100644
--- a/tools/SourceKit/tools/sourcekitd/lib/API/sourcekitdAPI-Common.cpp
+++ b/tools/SourceKit/tools/sourcekitd/lib/API/sourcekitdAPI-Common.cpp
@@ -124,6 +124,7 @@
 UIdent sourcekitd::KeyHide("key.hide");
 UIdent sourcekitd::KeySimplified("key.simplified");
 UIdent sourcekitd::KeyRangeContent("key.rangecontent");
+UIdent sourcekitd::KeyCancelOnSubsequentRequest("key.cancel_on_subsequent_request");
 
 UIdent sourcekitd::KeyIsDeprecated("key.is_deprecated");
 UIdent sourcekitd::KeyIsUnavailable("key.is_unavailable");
diff --git a/unittests/SourceKit/SwiftLang/CursorInfoTest.cpp b/unittests/SourceKit/SwiftLang/CursorInfoTest.cpp
index ea1af88..b76e7a8 100644
--- a/unittests/SourceKit/SwiftLang/CursorInfoTest.cpp
+++ b/unittests/SourceKit/SwiftLang/CursorInfoTest.cpp
@@ -134,7 +134,7 @@
     Semaphore sema(0);
 
     TestCursorInfo TestInfo;
-    getLang().getCursorInfo(DocName, Offset, 0, false, Args,
+    getLang().getCursorInfo(DocName, Offset, 0, false, false, Args,
       [&](const CursorInfo &Info) {
         TestInfo.Name = Info.Name;
         TestInfo.Typename = Info.TypeName;