[SourceKit] Allow explicit cancellation of requests with a cancellation token

test/SourceKit/Misc/cancellation.swift

@@ -0,0 +1,13 @@
+// Check that we can cancel requests.
+// We need to wait a little bit after request scheduling and cancellation to make sure we are not cancelling the request before it got scheduled.
+
+// RUN: not %sourcekitd-test -req=cursor -id=slow -async -pos=10:3 %s -- %s == \
+// RUN: -shell -- sleep 1 == \
+// RUN: -cancel=slow 2>&1 \
+// RUN: | %FileCheck %s
+
+func foo(x: Invalid1, y: Invalid2) {
+  x / y / x / y / x / y / x / y
+}
+
+// CHECK: error response (Request Cancelled)

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

@@ -14,13 +14,14 @@
 #define LLVM_SOURCEKIT_CORE_LANGSUPPORT_H
 
 #include "SourceKit/Core/LLVM.h"
+#include "SourceKit/Support/CancellationToken.h"
 #include "SourceKit/Support/UIdent.h"
-#include "llvm/Support/VersionTuple.h"
+#include "swift/AST/Type.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallString.h"
-#include "swift/AST/Type.h"
+#include "llvm/Support/VersionTuple.h"
 #include "llvm/Support/VirtualFileSystem.h"
 #include <functional>
 #include <memory>
@@ -737,6 +738,8 @@ class LangSupport {
 
   virtual void dependencyUpdated() {}
 
+  virtual void cancelRequest(SourceKitCancellationToken CancellationToken) = 0;
+
   virtual void indexSource(StringRef Filename,
                            IndexingConsumer &Consumer,
                            ArrayRef<const char *> Args) = 0;
@@ -794,10 +797,11 @@ class LangSupport {
                                          bool SynthesizedExtensions,
                                          StringRef swiftVersion) = 0;
 
-  virtual void editorOpenSwiftSourceInterface(StringRef Name,
-                                              StringRef SourceName,
-                                              ArrayRef<const char *> Args,
-                                              std::shared_ptr<EditorConsumer> Consumer) = 0;
+  virtual void
+  editorOpenSwiftSourceInterface(StringRef Name, StringRef SourceName,
+                                 ArrayRef<const char *> Args,
+                                 SourceKitCancellationToken CancellationToken,
+                                 std::shared_ptr<EditorConsumer> Consumer) = 0;
 
   virtual void editorClose(StringRef Name, bool RemoveCache) = 0;
 
@@ -821,33 +825,37 @@ class LangSupport {
                                        unsigned Length,
                                        EditorConsumer &Consumer) = 0;
 
+  virtual void getCursorInfo(
+      StringRef Filename, unsigned Offset, unsigned Length, bool Actionables,
+      bool SymbolGraph, bool CancelOnSubsequentRequest,
+      ArrayRef<const char *> Args, Optional<VFSOptions> vfsOptions,
+      SourceKitCancellationToken CancellationToken,
+      std::function<void(const RequestResult<CursorInfoData> &)> Receiver) = 0;
+
   virtual void
-  getCursorInfo(StringRef Filename, unsigned Offset, unsigned Length,
-                bool Actionables, bool SymbolGraph,
-                bool CancelOnSubsequentRequest, ArrayRef<const char *> Args,
-                Optional<VFSOptions> vfsOptions,
-       std::function<void(const RequestResult<CursorInfoData> &)> Receiver) = 0;
-
-  virtual void getNameInfo(StringRef Filename, unsigned Offset,
-                           NameTranslatingInfo &Input,
-                           ArrayRef<const char *> Args,
-                std::function<void(const RequestResult<NameTranslatingInfo> &)> Receiver) = 0;
-
-  virtual void getRangeInfo(StringRef Filename, unsigned Offset, unsigned Length,
-                            bool CancelOnSubsequentRequest,
-                            ArrayRef<const char *> Args,
-                            std::function<void(const RequestResult<RangeInfo> &)> Receiver) = 0;
+  getNameInfo(StringRef Filename, unsigned Offset, NameTranslatingInfo &Input,
+              ArrayRef<const char *> Args,
+              SourceKitCancellationToken CancellationToken,
+              std::function<void(const RequestResult<NameTranslatingInfo> &)>
+                  Receiver) = 0;
+
+  virtual void getRangeInfo(
+      StringRef Filename, unsigned Offset, unsigned Length,
+      bool CancelOnSubsequentRequest, ArrayRef<const char *> Args,
+      SourceKitCancellationToken CancellationToken,
+      std::function<void(const RequestResult<RangeInfo> &)> Receiver) = 0;
 
   virtual void getCursorInfoFromUSR(
       StringRef Filename, StringRef USR, bool CancelOnSubsequentRequest,
       ArrayRef<const char *> Args, Optional<VFSOptions> vfsOptions,
+      SourceKitCancellationToken CancellationToken,
       std::function<void(const RequestResult<CursorInfoData> &)> Receiver) = 0;
 
-  virtual void findRelatedIdentifiersInFile(StringRef Filename,
-                                            unsigned Offset,
-                                            bool CancelOnSubsequentRequest,
-                                            ArrayRef<const char *> Args,
-                   std::function<void(const RequestResult<RelatedIdentsInfo> &)> Receiver) = 0;
+  virtual void findRelatedIdentifiersInFile(
+      StringRef Filename, unsigned Offset, bool CancelOnSubsequentRequest,
+      ArrayRef<const char *> Args, SourceKitCancellationToken CancellationToken,
+      std::function<void(const RequestResult<RelatedIdentsInfo> &)>
+          Receiver) = 0;
 
   virtual llvm::Optional<std::pair<unsigned, unsigned>>
       findUSRRange(StringRef DocumentName, StringRef USR) = 0;
@@ -872,25 +880,29 @@ class LangSupport {
   virtual void
   findLocalRenameRanges(StringRef Filename, unsigned Line, unsigned Column,
                         unsigned Length, ArrayRef<const char *> Args,
+                        SourceKitCancellationToken CancellationToken,
                         CategorizedRenameRangesReceiver Receiver) = 0;
 
-  virtual void semanticRefactoring(StringRef Filename, SemanticRefactoringInfo Info,
-                                   ArrayRef<const char*> Args,
+  virtual void semanticRefactoring(StringRef Filename,
+                                   SemanticRefactoringInfo Info,
+                                   ArrayRef<const char *> Args,
+                                   SourceKitCancellationToken CancellationToken,
                                    CategorizedEditsReceiver Receiver) = 0;
 
-  virtual void collectExpressionTypes(StringRef FileName,
-                                      ArrayRef<const char *> Args,
-                                      ArrayRef<const char *> ExpectedProtocols,
-                                      bool CanonicalType,
-                                      std::function<void(const
-                                          RequestResult<ExpressionTypesInFile> &)> Receiver) = 0;
+  virtual void collectExpressionTypes(
+      StringRef FileName, ArrayRef<const char *> Args,
+      ArrayRef<const char *> ExpectedProtocols, bool CanonicalType,
+      SourceKitCancellationToken CancellationToken,
+      std::function<void(const RequestResult<ExpressionTypesInFile> &)>
+          Receiver) = 0;
 
   /// Collects variable types for a range defined by `Offset` and `Length` in
   /// the source file. If `Offset` or `Length` are empty, variable types for
   /// the entire document are collected.
   virtual void collectVariableTypes(
       StringRef FileName, ArrayRef<const char *> Args,
       Optional<unsigned> Offset, Optional<unsigned> Length,
+      SourceKitCancellationToken CancellationToken,
       std::function<void(const RequestResult<VariableTypesInFile> &)>
           Receiver) = 0;
 

tools/SourceKit/include/SourceKit/Support/CancellationToken.h

@@ -0,0 +1,26 @@
+//===--- CancellationToken.h - ----------------------------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2021 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_SOURCEKIT_SUPPORT_CANCELLATION_TOKEN_H
+#define LLVM_SOURCEKIT_SUPPORT_CANCELLATION_TOKEN_H
+
+namespace SourceKit {
+
+/// A token that uniquely identifies a SourceKit request that's served by a
+/// \c SwiftASTConsumer. Used to cancel the request.
+/// If the cancellation token is \c nullptr, it means that cancellation is not
+/// supported.
+typedef const void *SourceKitCancellationToken;
+
+} // namespace SourceKit
+
+#endif

tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp

@@ -580,9 +580,17 @@ struct SwiftASTManager::Implementation {
   /// Since we only keep a reference to the consumers to cancel them, the
   /// reference to the consumer itself is weak - if it's already deallocated,
   /// there is no need to cancel it anymore.
+  /// The \c CancellationToken that allows cancellation of this consumer.
+  /// Multiple consumers might share the same \c CancellationToken if they were
+  /// created from the same SourceKit request. E.g. a \c CursorInfoConsumer
+  /// might schedule a second \c CursorInfoConsumer if it discovers that the AST
+  /// that was used to serve the first request is not up-to-date enough.
+  /// If \c CancellationToken is \c nullptr, the consumer can't be cancelled
+  /// using a cancellation token.
   struct ScheduledConsumer {
     SwiftASTConsumerWeakRef Consumer;
     const void *OncePerASTToken;
+    SourceKitCancellationToken CancellationToken;
   };
 
   llvm::sys::Mutex ScheduledConsumersMtx;
@@ -739,7 +747,7 @@ SwiftInvocationRef SwiftASTManager::getInvocation(
 
 void SwiftASTManager::processASTAsync(
     SwiftInvocationRef InvokRef, SwiftASTConsumerRef ASTConsumer,
-    const void *OncePerASTToken,
+    const void *OncePerASTToken, SourceKitCancellationToken CancellationToken,
     llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> fileSystem,
     ArrayRef<ImmutableTextSnapshotRef> Snapshots) {
   assert(fileSystem);
@@ -758,13 +766,32 @@ void SwiftASTManager::processASTAsync(
         }
       }
     }
-    Impl.ScheduledConsumers.push_back({ASTConsumer, OncePerASTToken});
+    Impl.ScheduledConsumers.push_back(
+        {ASTConsumer, OncePerASTToken, CancellationToken});
   }
 
   Producer->enqueueConsumer(ASTConsumer, fileSystem, Snapshots,
                             shared_from_this());
 }
 
+void SwiftASTManager::cancelASTConsumer(
+    SourceKitCancellationToken CancellationToken) {
+  if (!CancellationToken) {
+    return;
+  }
+  Impl.cleanDeletedConsumers();
+  llvm::sys::ScopedLock L(Impl.ScheduledConsumersMtx);
+  for (auto ScheduledConsumer : Impl.ScheduledConsumers) {
+    if (ScheduledConsumer.CancellationToken == CancellationToken) {
+      if (auto Consumer = ScheduledConsumer.Consumer.lock()) {
+        Consumer->requestCancellation();
+        // Multiple consumers might share the same cancellation token (see
+        // documentation on ScheduledConsumer), so we can't break here.
+      }
+    }
+  }
+}
+
 void SwiftASTManager::removeCachedAST(SwiftInvocationRef Invok) {
   Impl.ASTCache.remove(Invok->Impl.Key);
 }

tools/SourceKit/lib/SwiftLang/SwiftASTManager.h

@@ -75,6 +75,7 @@
 #define LLVM_SOURCEKIT_LIB_SWIFTLANG_SWIFTASTMANAGER_H
 
 #include "SourceKit/Core/LLVM.h"
+#include "SourceKit/Support/CancellationToken.h"
 #include "SwiftInvocation.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
@@ -244,10 +245,17 @@ class SwiftASTManager : public std::enable_shared_from_this<SwiftASTManager> {
   void
   processASTAsync(SwiftInvocationRef Invok, SwiftASTConsumerRef ASTConsumer,
                   const void *OncePerASTToken,
+                  SourceKitCancellationToken CancellationToken,
                   llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> fileSystem,
                   ArrayRef<ImmutableTextSnapshotRef> Snapshots =
                       ArrayRef<ImmutableTextSnapshotRef>());
 
+  /// Request the \c SwiftASTConsumer with the given \p CancellationToken to be
+  /// cancelled. If \p CancellationToken is \c nullptr or no consumer with the
+  /// given cancellation token exists (e.g. because the consumer already
+  /// finished), this is a no-op.
+  void cancelASTConsumer(SourceKitCancellationToken CancellationToken);
+
   std::unique_ptr<llvm::MemoryBuffer> getMemoryBuffer(StringRef Filename,
                                                       std::string &Error);
 

tools/SourceKit/lib/SwiftLang/SwiftDocSupport.cpp

@@ -1407,7 +1407,8 @@ void SwiftLangSupport::findRenameRanges(
 
 void SwiftLangSupport::findLocalRenameRanges(
     StringRef Filename, unsigned Line, unsigned Column, unsigned Length,
-    ArrayRef<const char *> Args, CategorizedRenameRangesReceiver Receiver) {
+    ArrayRef<const char *> Args, SourceKitCancellationToken CancellationToken,
+    CategorizedRenameRangesReceiver Receiver) {
   std::string Error;
   SwiftInvocationRef Invok = ASTMgr->getInvocation(Args, Filename, Error);
   if (!Invok) {
@@ -1447,6 +1448,7 @@ void SwiftLangSupport::findLocalRenameRanges(
   /// don't use 'OncePerASTToken'.
   static const char OncePerASTToken = 0;
   getASTManager()->processASTAsync(Invok, ASTConsumer, &OncePerASTToken,
+                                   CancellationToken,
                                    llvm::vfs::getRealFileSystem());
 }