[SourceKit] Add a request tracker that manages cancellaiton

ahoppen · ahoppen · commit 772485def348 · 2021-11-10T22:11:02.000+01:00
Previously, `SwiftASTManager` and `SlowRequestSimulator` maintained their own list of in-progress cancellation tokens. With code completion cancellation coming up, there would need to be yet another place to track in-progress requests, so let’s centralize it.

While at it, also support cancelling requests before they are scheduled, eliminating the need for a `sleep` in a test case.

The current implementaiton leaks tiny amounts of memory if a request is cancelled after if finishes. I think this is fine because it is a pretty nieche case and the leaked memory is pretty small (a `std::map` entry pointing to a `std::function` + `bool`). Alternatively, we could require the client to always dispose of the cancellation token manually.
diff --git a/test/SourceKit/Misc/cancellation.swift b/test/SourceKit/Misc/cancellation.swift
@@ -1,8 +1,6 @@
 // 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=9:5 -simulate-long-request=5000 %s -- %s == \
-// RUN: -shell -- sleep 1 == \
+// RUN: not %sourcekitd-test -req=cursor -id=slow -async -pos=7:5 -simulate-long-request=5000 %s -- %s == \
 // RUN: -cancel=slow 2>&1 \
 // RUN: | %FileCheck %s
 
diff --git a/tools/SourceKit/include/SourceKit/Core/Context.h b/tools/SourceKit/include/SourceKit/Core/Context.h
@@ -54,43 +54,118 @@ class GlobalConfig {
   Settings::CompletionOptions getCompletionOpts() const;
 };
 
-class SlowRequestSimulator {
-  std::map<SourceKitCancellationToken, std::shared_ptr<Semaphore>>
-      InProgressRequests;
+/// Keeps track of all requests that are currently in progress and coordinates
+/// cancellation. All operations are no-ops if \c nullptr is used as a
+/// cancellation token.
+///
+/// Tracked requests will be removed from this object when the request returns
+/// a response (\c sourcekitd::handleRequest disposes the cancellation token).
+/// We leak memory if a request is cancelled after it has returned a result
+/// because we add an \c IsCancelled entry in that case but \c
+/// sourcekitd::handleRequest won't have a chance to dispose of the token.
+class RequestTracker {
+  struct TrackedRequest {
+    /// Whether the request has already been cancelled.
+    bool IsCancelled = false;
+    /// A handler that will be called as the request gets cancelled.
+    std::function<void(void)> CancellationHandler;
+  };
 
-  /// Mutex guarding \c InProgressRequests.
-  llvm::sys::Mutex InProgressRequestsMutex;
+  /// Once we have information about a request (either a cancellation handler
+  /// or information that it has been cancelled), it is added to this list.
+  std::map<SourceKitCancellationToken, TrackedRequest> Requests;
+
+  /// Guards the \c Requests variable.
+  llvm::sys::Mutex RequestsMtx;
+
+  /// Must only be called if \c RequestsMtx has been claimed.
+  /// Returns \c true if the request has already been cancelled. Requests that
+  /// are not tracked are not assumed to be cancelled.
+  bool isCancelledImpl(SourceKitCancellationToken CancellationToken) {
+    if (Requests.count(CancellationToken) == 0) {
+      return false;
+    } else {
+      return Requests[CancellationToken].IsCancelled;
+    }
+  }
 
 public:
+  /// Returns \c true if the request with the given \p CancellationToken has
+  /// already been cancelled.
+  bool isCancelled(SourceKitCancellationToken CancellationToken) {
+    if (!CancellationToken) {
+      return false;
+    }
+    llvm::sys::ScopedLock L(RequestsMtx);
+    return isCancelledImpl(CancellationToken);
+  }
+
+  /// Adds a \p CancellationHandler that will be called when the request
+  /// associated with the \p CancellationToken is cancelled.
+  /// If that request has already been cancelled when this method is called,
+  /// \p CancellationHandler will be called synchronously.
+  void setCancellationHandler(SourceKitCancellationToken CancellationToken,
+                              std::function<void(void)> CancellationHandler) {
+    if (!CancellationToken) {
+      return;
+    }
+    llvm::sys::ScopedLock L(RequestsMtx);
+    if (isCancelledImpl(CancellationToken)) {
+      if (CancellationHandler) {
+        CancellationHandler();
+      }
+    } else {
+      Requests[CancellationToken].CancellationHandler = CancellationHandler;
+    }
+  }
+
+  /// Cancel the request with the given \p CancellationToken. If a cancellation
+  /// handler is associated with it, call it.
+  void cancel(SourceKitCancellationToken CancellationToken) {
+    if (!CancellationToken) {
+      return;
+    }
+    llvm::sys::ScopedLock L(RequestsMtx);
+    Requests[CancellationToken].IsCancelled = true;
+    if (auto CancellationHandler =
+            Requests[CancellationToken].CancellationHandler) {
+      CancellationHandler();
+    }
+  }
+
+  /// Stop tracking the request with the given \p CancellationToken, freeing up
+  /// any memory needed for the tracking.
+  void stopTracking(SourceKitCancellationToken CancellationToken) {
+    if (!CancellationToken) {
+      return;
+    }
+    llvm::sys::ScopedLock L(RequestsMtx);
+    Requests.erase(CancellationToken);
+  }
+};
+
+class SlowRequestSimulator {
+  std::shared_ptr<RequestTracker> ReqTracker;
+
+public:
+  SlowRequestSimulator(std::shared_ptr<RequestTracker> ReqTracker)
+      : ReqTracker(ReqTracker) {}
+
   /// Simulate that a request takes \p DurationMs to execute. While waiting that
   /// duration, the request can be cancelled using the \p CancellationToken.
   /// Returns \c true if the request waited the required duration and \c false
   /// if it was cancelled.
   bool simulateLongRequest(int64_t DurationMs,
                            SourceKitCancellationToken CancellationToken) {
-    auto Sema = std::make_shared<Semaphore>(0);
-    {
-      llvm::sys::ScopedLock L(InProgressRequestsMutex);
-      InProgressRequests[CancellationToken] = Sema;
-    }
-    bool DidTimeOut = Sema->wait(DurationMs);
-    {
-      llvm::sys::ScopedLock L(InProgressRequestsMutex);
-      InProgressRequests[CancellationToken] = nullptr;
-    }
+    auto Sema = Semaphore(0);
+    ReqTracker->setCancellationHandler(CancellationToken,
+                                       [&] { Sema.signal(); });
+    bool DidTimeOut = Sema.wait(DurationMs);
+    ReqTracker->setCancellationHandler(CancellationToken, nullptr);
     // If we timed out, we waited the required duration. If we didn't time out,
     // the semaphore was cancelled.
     return DidTimeOut;
   }
-
-  /// Cancel a simulated long request. If the required wait duration already
-  /// elapsed, this is a no-op.
-  void cancel(SourceKitCancellationToken CancellationToken) {
-    llvm::sys::ScopedLock L(InProgressRequestsMutex);
-    if (auto InProgressSema = InProgressRequests[CancellationToken]) {
-      InProgressSema->signal();
-    }
-  }
 };
 
 class Context {
@@ -99,6 +174,7 @@ class Context {
   std::unique_ptr<LangSupport> SwiftLang;
   std::shared_ptr<NotificationCenter> NotificationCtr;
   std::shared_ptr<GlobalConfig> Config;
+  std::shared_ptr<RequestTracker> ReqTracker;
   std::shared_ptr<SlowRequestSimulator> SlowRequestSim;
 
 public:
@@ -122,6 +198,8 @@ class Context {
   std::shared_ptr<SlowRequestSimulator> getSlowRequestSimulator() {
     return SlowRequestSim;
   }
+
+  std::shared_ptr<RequestTracker> getRequestTracker() { return ReqTracker; }
 };
 
 } // namespace SourceKit
diff --git a/tools/SourceKit/include/SourceKit/Core/LangSupport.h b/tools/SourceKit/include/SourceKit/Core/LangSupport.h
@@ -745,8 +745,6 @@ class LangSupport {
 
   virtual void dependencyUpdated() {}
 
-  virtual void cancelRequest(SourceKitCancellationToken CancellationToken) = 0;
-
   virtual void indexSource(StringRef Filename,
                            IndexingConsumer &Consumer,
                            ArrayRef<const char *> Args) = 0;
diff --git a/tools/SourceKit/lib/Core/Context.cpp b/tools/SourceKit/lib/Core/Context.cpp
@@ -44,7 +44,8 @@ SourceKit::Context::Context(
       DiagnosticDocumentationPath(DiagnosticDocumentationPath),
       NotificationCtr(
           new NotificationCenter(shouldDispatchNotificationsOnMain)),
-      Config(new GlobalConfig()), SlowRequestSim(new SlowRequestSimulator()) {
+      Config(new GlobalConfig()), ReqTracker(new RequestTracker()),
+      SlowRequestSim(new SlowRequestSimulator(ReqTracker)) {
   // Should be called last after everything is initialized.
   SwiftLang = LangSupportFactoryFn(*this);
 }
diff --git a/tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp b/tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp
@@ -556,17 +556,19 @@ struct SwiftASTManager::Implementation {
   explicit Implementation(
       std::shared_ptr<SwiftEditorDocumentFileMap> EditorDocs,
       std::shared_ptr<GlobalConfig> Config,
-      std::shared_ptr<SwiftStatistics> Stats, StringRef RuntimeResourcePath,
+      std::shared_ptr<SwiftStatistics> Stats,
+      std::shared_ptr<RequestTracker> ReqTracker, StringRef RuntimeResourcePath,
       StringRef DiagnosticDocumentationPath)
       : EditorDocs(EditorDocs), Config(Config), Stats(Stats),
-        RuntimeResourcePath(RuntimeResourcePath),
+        ReqTracker(ReqTracker), RuntimeResourcePath(RuntimeResourcePath),
         DiagnosticDocumentationPath(DiagnosticDocumentationPath),
         SessionTimestamp(llvm::sys::toTimeT(std::chrono::system_clock::now())) {
   }
 
   std::shared_ptr<SwiftEditorDocumentFileMap> EditorDocs;
   std::shared_ptr<GlobalConfig> Config;
   std::shared_ptr<SwiftStatistics> Stats;
+  std::shared_ptr<RequestTracker> ReqTracker;
   std::string RuntimeResourcePath;
   std::string DiagnosticDocumentationPath;
   SourceManager SourceMgr;
@@ -590,9 +592,11 @@ struct SwiftASTManager::Implementation {
   struct ScheduledConsumer {
     SwiftASTConsumerWeakRef Consumer;
     const void *OncePerASTToken;
-    SourceKitCancellationToken CancellationToken;
   };
 
+  /// FIXME: Once we no longer support implicit cancellation using
+  /// OncePerASTToken, we can stop keeping track of ScheduledConsumers and
+  /// completely rely on RequestTracker for cancellation.
   llvm::sys::Mutex ScheduledConsumersMtx;
   std::vector<ScheduledConsumer> ScheduledConsumers;
 
@@ -632,9 +636,11 @@ struct SwiftASTManager::Implementation {
 SwiftASTManager::SwiftASTManager(
     std::shared_ptr<SwiftEditorDocumentFileMap> EditorDocs,
     std::shared_ptr<GlobalConfig> Config,
-    std::shared_ptr<SwiftStatistics> Stats, StringRef RuntimeResourcePath,
+    std::shared_ptr<SwiftStatistics> Stats,
+    std::shared_ptr<RequestTracker> ReqTracker, StringRef RuntimeResourcePath,
     StringRef DiagnosticDocumentationPath)
-    : Impl(*new Implementation(EditorDocs, Config, Stats, RuntimeResourcePath,
+    : Impl(*new Implementation(EditorDocs, Config, Stats, ReqTracker,
+                               RuntimeResourcePath,
                                DiagnosticDocumentationPath)) {}
 
 SwiftASTManager::~SwiftASTManager() {
@@ -766,30 +772,18 @@ void SwiftASTManager::processASTAsync(
         }
       }
     }
-    Impl.ScheduledConsumers.push_back(
-        {ASTConsumer, OncePerASTToken, CancellationToken});
+    Impl.ScheduledConsumers.push_back({ASTConsumer, OncePerASTToken});
   }
 
   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.
-      }
+  auto WeakConsumer = SwiftASTConsumerWeakRef(ASTConsumer);
+  Impl.ReqTracker->setCancellationHandler(CancellationToken, [WeakConsumer] {
+    if (auto Consumer = WeakConsumer.lock()) {
+      Consumer->requestCancellation();
     }
-  }
+  });
 }
 
 void SwiftASTManager::removeCachedAST(SwiftInvocationRef Invok) {
diff --git a/tools/SourceKit/lib/SwiftLang/SwiftASTManager.h b/tools/SourceKit/lib/SwiftLang/SwiftASTManager.h
@@ -74,6 +74,7 @@
 #ifndef LLVM_SOURCEKIT_LIB_SWIFTLANG_SWIFTASTMANAGER_H
 #define LLVM_SOURCEKIT_LIB_SWIFTLANG_SWIFTASTMANAGER_H
 
+#include "SourceKit/Core/Context.h"
 #include "SourceKit/Core/LLVM.h"
 #include "SourceKit/Support/CancellationToken.h"
 #include "SwiftInvocation.h"
@@ -223,6 +224,7 @@ class SwiftASTManager : public std::enable_shared_from_this<SwiftASTManager> {
   explicit SwiftASTManager(std::shared_ptr<SwiftEditorDocumentFileMap>,
                            std::shared_ptr<GlobalConfig> Config,
                            std::shared_ptr<SwiftStatistics> Stats,
+                           std::shared_ptr<RequestTracker> ReqTracker,
                            StringRef RuntimeResourcePath,
                            StringRef DiagnosticDocumentationPath);
   ~SwiftASTManager();
@@ -250,12 +252,6 @@ class SwiftASTManager : public std::enable_shared_from_this<SwiftASTManager> {
                   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);
 
diff --git a/tools/SourceKit/lib/SwiftLang/SwiftLangSupport.cpp b/tools/SourceKit/lib/SwiftLang/SwiftLangSupport.cpp
@@ -273,7 +273,7 @@ configureCompletionInstance(std::shared_ptr<CompletionInstance> CompletionInst,
 
 SwiftLangSupport::SwiftLangSupport(SourceKit::Context &SKCtx)
     : NotificationCtr(SKCtx.getNotificationCenter()),
-      CCCache(new SwiftCompletionCache) {
+      ReqTracker(SKCtx.getRequestTracker()), CCCache(new SwiftCompletionCache) {
   llvm::SmallString<128> LibPath(SKCtx.getRuntimeLibPath());
   llvm::sys::path::append(LibPath, "swift");
   RuntimeResourcePath = std::string(LibPath.str());
@@ -282,7 +282,7 @@ SwiftLangSupport::SwiftLangSupport(SourceKit::Context &SKCtx)
   Stats = std::make_shared<SwiftStatistics>();
   EditorDocuments = std::make_shared<SwiftEditorDocumentFileMap>();
   ASTMgr = std::make_shared<SwiftASTManager>(
-      EditorDocuments, SKCtx.getGlobalConfiguration(), Stats,
+      EditorDocuments, SKCtx.getGlobalConfiguration(), Stats, ReqTracker,
       RuntimeResourcePath, DiagnosticDocumentationPath);
 
   CompletionInst = std::make_shared<CompletionInstance>();
@@ -307,11 +307,6 @@ void SwiftLangSupport::dependencyUpdated() {
   CompletionInst->markCachedCompilerInstanceShouldBeInvalidated();
 }
 
-void SwiftLangSupport::cancelRequest(
-    SourceKitCancellationToken CancellationToken) {
-  getASTManager()->cancelASTConsumer(CancellationToken);
-}
-
 UIdent SwiftLangSupport::getUIDForDeclLanguage(const swift::Decl *D) {
   if (D->hasClangNode())
     return KindObjC;
diff --git a/tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h b/tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h
@@ -14,6 +14,7 @@
 #define LLVM_SOURCEKIT_LIB_SWIFTLANG_SWIFTLANGSUPPORT_H
 
 #include "CodeCompletion.h"
+#include "SourceKit/Core/Context.h"
 #include "SourceKit/Core/LangSupport.h"
 #include "SourceKit/Support/Concurrency.h"
 #include "SourceKit/Support/Statistic.h"
@@ -304,6 +305,7 @@ class SwiftLangSupport : public LangSupport {
   std::string DiagnosticDocumentationPath;
   std::shared_ptr<SwiftASTManager> ASTMgr;
   std::shared_ptr<SwiftEditorDocumentFileMap> EditorDocuments;
+  std::shared_ptr<RequestTracker> ReqTracker;
   SwiftInterfaceGenMap IFaceGenContexts;
   ThreadSafeRefCntPtr<SwiftCompletionCache> CCCache;
   ThreadSafeRefCntPtr<SwiftPopularAPI> PopularAPI;
@@ -487,8 +489,6 @@ class SwiftLangSupport : public LangSupport {
 
   void dependencyUpdated() override;
 
-  void cancelRequest(SourceKitCancellationToken CancellationToken) override;
-
   void indexSource(StringRef Filename, IndexingConsumer &Consumer,
                    ArrayRef<const char *> Args) override;
 
diff --git a/tools/SourceKit/tools/sourcekitd/include/sourcekitd/Internal.h b/tools/SourceKit/tools/sourcekitd/include/sourcekitd/Internal.h
@@ -180,6 +180,8 @@ void handleRequest(sourcekitd_object_t Request,
 
 void cancelRequest(SourceKitCancellationToken CancellationToken);
 
+void disposeCancellationToken(SourceKitCancellationToken CancellationToken);
+
 void printRequestObject(sourcekitd_object_t Obj, llvm::raw_ostream &OS);
 void printResponse(sourcekitd_response_t Resp, llvm::raw_ostream &OS);
 
diff --git a/tools/SourceKit/tools/sourcekitd/lib/API/Requests.cpp b/tools/SourceKit/tools/sourcekitd/lib/API/Requests.cpp
@@ -390,20 +390,27 @@ void sourcekitd::handleRequest(sourcekitd_object_t Req,
   }
 
   handleRequestImpl(
-      Req, CancellationToken, [Receiver](sourcekitd_response_t Resp) {
+      Req, CancellationToken,
+      [Receiver, CancellationToken](sourcekitd_response_t Resp) {
         LOG_SECTION("handleRequest-after", InfoHighPrio) {
           // Responses are big, print them out with info medium priority.
           if (Logger::isLoggingEnabledForLevel(Logger::Level::InfoMediumPrio))
             sourcekitd::printResponse(Resp, Log->getOS());
         }
 
+        sourcekitd::disposeCancellationToken(CancellationToken);
+
         Receiver(Resp);
       });
 }
 
 void sourcekitd::cancelRequest(SourceKitCancellationToken CancellationToken) {
-  getGlobalContext().getSlowRequestSimulator()->cancel(CancellationToken);
-  getGlobalContext().getSwiftLangSupport().cancelRequest(CancellationToken);
+  getGlobalContext().getRequestTracker()->cancel(CancellationToken);
+}
+
+void sourcekitd::disposeCancellationToken(
+    SourceKitCancellationToken CancellationToken) {
+  getGlobalContext().getRequestTracker()->stopTracking(CancellationToken);
 }
 
 static std::unique_ptr<llvm::MemoryBuffer> getInputBufForRequest(
diff --git a/unittests/SourceKit/SwiftLang/CursorInfoTest.cpp b/unittests/SourceKit/SwiftLang/CursorInfoTest.cpp
