[clangd] Cancel certain operations if the file changes before we start.

Summary:
Otherwise they can force us to build lots of snapshots that we don't need.
Particularly, try to do this for operations that are frequently
generated by editors without explicit user interaction, and where
editing the file makes the result less useful. (Code action
enumeration is a good example).

https://github.com/clangd/clangd/issues/298

This doesn't return the "right" LSP error code (ContentModified) to the client,
we need to teach the cancellation API to distinguish between different causes.

Reviewers: kadircet

Subscribers: ilya-biryukov, javed.absar, MaskRay, jkorous, arphaman, jfb, usaxena95, cfe-commits

Tags: #clang

Differential Revision: https://reviews.llvm.org/D75602
diff --git a/clang-tools-extra/clangd/Cancellation.cpp b/clang-tools-extra/clangd/Cancellation.cpp
index 6b0e21d..ccd27a1 100644
--- a/clang-tools-extra/clangd/Cancellation.cpp
+++ b/clang-tools-extra/clangd/Cancellation.cpp
@@ -13,20 +13,30 @@
 namespace clangd {
 
 char CancelledError::ID = 0;
-static Key<std::shared_ptr<std::atomic<bool>>> FlagKey;
+
+// We don't want a cancelable scope to "shadow" an enclosing one.
+struct CancelState {
+  std::shared_ptr<std::atomic<bool>> Cancelled;
+  const CancelState *Parent;
+};
+static Key<CancelState> StateKey;
 
 std::pair<Context, Canceler> cancelableTask() {
-  auto Flag = std::make_shared<std::atomic<bool>>();
+  CancelState State;
+  State.Cancelled = std::make_shared<std::atomic<bool>>();
+  State.Parent = Context::current().get(StateKey);
   return {
-      Context::current().derive(FlagKey, Flag),
-      [Flag] { *Flag = true; },
+      Context::current().derive(StateKey, State),
+      [Flag(State.Cancelled)] { *Flag = true; },
   };
 }
 
 bool isCancelled(const Context &Ctx) {
-  if (auto *Flag = Ctx.get(FlagKey))
-    return **Flag;
-  return false; // Not in scope of a task.
+  for (const CancelState *State = Ctx.get(StateKey); State != nullptr;
+       State = State->Parent)
+    if (State->Cancelled->load())
+      return true;
+  return false;
 }
 
 } // namespace clangd
diff --git a/clang-tools-extra/clangd/Cancellation.h b/clang-tools-extra/clangd/Cancellation.h
index 6872c44..bdb0bd4 100644
--- a/clang-tools-extra/clangd/Cancellation.h
+++ b/clang-tools-extra/clangd/Cancellation.h
@@ -76,6 +76,7 @@
 std::pair<Context, Canceler> cancelableTask();
 
 /// True if the current context is within a cancelable task which was cancelled.
+/// (If the context is within multiple nested tasks, true if any are cancelled).
 /// Always false if there is no active cancelable task.
 /// This isn't free (context lookup) - don't call it in a tight loop.
 bool isCancelled(const Context &Ctx = Context::current());
diff --git a/clang-tools-extra/clangd/ClangdServer.cpp b/clang-tools-extra/clangd/ClangdServer.cpp
index 5dd0032..cd833af 100644
--- a/clang-tools-extra/clangd/ClangdServer.cpp
+++ b/clang-tools-extra/clangd/ClangdServer.cpp
@@ -437,7 +437,8 @@
     CB(std::move(Res));
   };
 
-  WorkScheduler.runWithAST("EnumerateTweaks", File, std::move(Action));
+  WorkScheduler.runWithAST("EnumerateTweaks", File, std::move(Action),
+                           TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::applyTweak(PathRef File, Range Sel, StringRef TweakID,
@@ -556,7 +557,8 @@
         CB(clangd::findDocumentHighlights(InpAST->AST, Pos));
       };
 
-  WorkScheduler.runWithAST("Highlights", File, std::move(Action));
+  WorkScheduler.runWithAST("Highlights", File, std::move(Action),
+                           TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::findHover(PathRef File, Position Pos,
@@ -570,7 +572,8 @@
     CB(clangd::getHover(InpAST->AST, Pos, std::move(Style), Index));
   };
 
-  WorkScheduler.runWithAST("Hover", File, std::move(Action));
+  WorkScheduler.runWithAST("Hover", File, std::move(Action),
+                           TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::typeHierarchy(PathRef File, Position Pos, int Resolve,
@@ -618,7 +621,8 @@
           return CB(InpAST.takeError());
         CB(clangd::getDocumentSymbols(InpAST->AST));
       };
-  WorkScheduler.runWithAST("documentSymbols", File, std::move(Action));
+  WorkScheduler.runWithAST("documentSymbols", File, std::move(Action),
+                           TUScheduler::InvalidateOnUpdate);
 }
 
 void ClangdServer::findReferences(PathRef File, Position Pos, uint32_t Limit,
@@ -664,7 +668,8 @@
           return CB(InpAST.takeError());
         CB(clangd::getDocumentLinks(InpAST->AST));
       };
-  WorkScheduler.runWithAST("DocumentLinks", File, std::move(Action));
+  WorkScheduler.runWithAST("DocumentLinks", File, std::move(Action),
+                           TUScheduler::InvalidateOnUpdate);
 }
 
 std::vector<std::pair<Path, std::size_t>>
diff --git a/clang-tools-extra/clangd/JSONTransport.cpp b/clang-tools-extra/clangd/JSONTransport.cpp
index 43be9f7..d82dca1 100644
--- a/clang-tools-extra/clangd/JSONTransport.cpp
+++ b/clang-tools-extra/clangd/JSONTransport.cpp
@@ -19,6 +19,8 @@
 llvm::json::Object encodeError(llvm::Error E) {
   std::string Message;
   ErrorCode Code = ErrorCode::UnknownErrorCode;
+  // FIXME: encode cancellation errors using RequestCancelled or ContentModified
+  // as appropriate.
   if (llvm::Error Unhandled = llvm::handleErrors(
           std::move(E), [&](const LSPError &L) -> llvm::Error {
             Message = L.Message;
diff --git a/clang-tools-extra/clangd/TUScheduler.cpp b/clang-tools-extra/clangd/TUScheduler.cpp
index 3f3162a..714bc72 100644
--- a/clang-tools-extra/clangd/TUScheduler.cpp
+++ b/clang-tools-extra/clangd/TUScheduler.cpp
@@ -184,7 +184,8 @@
   void update(ParseInputs Inputs, WantDiagnostics);
   void
   runWithAST(llvm::StringRef Name,
-             llvm::unique_function<void(llvm::Expected<InputsAndAST>)> Action);
+             llvm::unique_function<void(llvm::Expected<InputsAndAST>)> Action,
+             TUScheduler::ASTActionInvalidation);
   bool blockUntilIdle(Deadline Timeout) const;
 
   std::shared_ptr<const PreambleData> getPossiblyStalePreamble() const;
@@ -214,7 +215,8 @@
   void stop();
   /// Adds a new task to the end of the request queue.
   void startTask(llvm::StringRef Name, llvm::unique_function<void()> Task,
-                 llvm::Optional<WantDiagnostics> UpdateType);
+                 llvm::Optional<WantDiagnostics> UpdateType,
+                 TUScheduler::ASTActionInvalidation);
   /// Updates the TUStatus and emits it. Only called in the worker thread.
   void emitTUStatus(TUAction FAction,
                     const TUStatus::BuildDetails *Detail = nullptr);
@@ -237,6 +239,8 @@
     steady_clock::time_point AddTime;
     Context Ctx;
     llvm::Optional<WantDiagnostics> UpdateType;
+    TUScheduler::ASTActionInvalidation InvalidationPolicy;
+    Canceler Invalidate;
   };
 
   /// Handles retention of ASTs.
@@ -542,12 +546,13 @@
     // Stash the AST in the cache for further use.
     IdleASTs.put(this, std::move(*AST));
   };
-  startTask(TaskName, std::move(Task), WantDiags);
+  startTask(TaskName, std::move(Task), WantDiags, TUScheduler::NoInvalidation);
 }
 
 void ASTWorker::runWithAST(
     llvm::StringRef Name,
-    llvm::unique_function<void(llvm::Expected<InputsAndAST>)> Action) {
+    llvm::unique_function<void(llvm::Expected<InputsAndAST>)> Action,
+    TUScheduler::ASTActionInvalidation Invalidation) {
   auto Task = [=, Action = std::move(Action)]() mutable {
     if (isCancelled())
       return Action(llvm::make_error<CancelledError>());
@@ -576,7 +581,7 @@
           "invalid AST", llvm::errc::invalid_argument));
     Action(InputsAndAST{*CurrentInputs, **AST});
   };
-  startTask(Name, std::move(Task), /*UpdateType=*/None);
+  startTask(Name, std::move(Task), /*UpdateType=*/None, Invalidation);
 }
 
 std::shared_ptr<const PreambleData>
@@ -606,7 +611,9 @@
                           },
                           "GetPreamble", steady_clock::now(),
                           Context::current().clone(),
-                          /*UpdateType=*/None});
+                          /*UpdateType=*/None,
+                          /*InvalidationPolicy=*/TUScheduler::NoInvalidation,
+                          /*Invalidate=*/nullptr});
   Lock.unlock();
   RequestsCV.notify_all();
 }
@@ -650,7 +657,8 @@
 
 void ASTWorker::startTask(llvm::StringRef Name,
                           llvm::unique_function<void()> Task,
-                          llvm::Optional<WantDiagnostics> UpdateType) {
+                          llvm::Optional<WantDiagnostics> UpdateType,
+                          TUScheduler::ASTActionInvalidation Invalidation) {
   if (RunSync) {
     assert(!Done && "running a task after stop()");
     trace::Span Tracer(Name + ":" + llvm::sys::path::filename(FileName));
@@ -661,9 +669,26 @@
   {
     std::lock_guard<std::mutex> Lock(Mutex);
     assert(!Done && "running a task after stop()");
-    Requests.push_back(
-        {std::move(Task), std::string(Name), steady_clock::now(),
-         Context::current().derive(kFileBeingProcessed, FileName), UpdateType});
+    // Cancel any requests invalidated by this request.
+    if (UpdateType) {
+      for (auto &R : llvm::reverse(Requests)) {
+        if (R.InvalidationPolicy == TUScheduler::InvalidateOnUpdate)
+          R.Invalidate();
+        if (R.UpdateType)
+          break; // Older requests were already invalidated by the older update.
+      }
+    }
+
+    // Allow this request to be cancelled if invalidated.
+    Context Ctx = Context::current().derive(kFileBeingProcessed, FileName);
+    Canceler Invalidate = nullptr;
+    if (Invalidation) {
+      WithContext WC(std::move(Ctx));
+      std::tie(Ctx, Invalidate) = cancelableTask();
+    }
+    Requests.push_back({std::move(Task), std::string(Name), steady_clock::now(),
+                        std::move(Ctx), UpdateType, Invalidation,
+                        std::move(Invalidate)});
   }
   RequestsCV.notify_all();
 }
@@ -942,7 +967,8 @@
 
 void TUScheduler::runWithAST(
     llvm::StringRef Name, PathRef File,
-    llvm::unique_function<void(llvm::Expected<InputsAndAST>)> Action) {
+    llvm::unique_function<void(llvm::Expected<InputsAndAST>)> Action,
+    TUScheduler::ASTActionInvalidation Invalidation) {
   auto It = Files.find(File);
   if (It == Files.end()) {
     Action(llvm::make_error<LSPError>(
@@ -950,7 +976,7 @@
     return;
   }
 
-  It->second->Worker->runWithAST(Name, std::move(Action));
+  It->second->Worker->runWithAST(Name, std::move(Action), Invalidation);
 }
 
 void TUScheduler::runWithPreamble(llvm::StringRef Name, PathRef File,
@@ -1067,5 +1093,15 @@
   return P;
 }
 
+void TUScheduler::InvalidatedError::log(llvm::raw_ostream &OS) const {
+  switch (Policy) {
+  case InvalidateOnUpdate:
+    OS << "Task was cancelled due to a subsequent change to the file.";
+    break;
+  case NoInvalidation:
+    llvm_unreachable("Invalidated for no reason?");
+  }
+}
+
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/clangd/TUScheduler.h b/clang-tools-extra/clangd/TUScheduler.h
index 948fde7..074404c 100644
--- a/clang-tools-extra/clangd/TUScheduler.h
+++ b/clang-tools-extra/clangd/TUScheduler.h
@@ -219,6 +219,29 @@
   /// Schedule an async task with no dependencies.
   void run(llvm::StringRef Name, llvm::unique_function<void()> Action);
 
+  /// Defines how a runWithAST action is implicitly cancelled by other actions.
+  enum ASTActionInvalidation {
+    /// The request will run unless explicitly cancelled.
+    NoInvalidation,
+    /// The request will be implicitly cancelled by a subsequent update().
+    /// (Only if the request was not yet cancelled).
+    /// Useful for requests that are generated by clients, without any explicit
+    /// user action. These can otherwise e.g. force every version to be built.
+    InvalidateOnUpdate,
+  };
+
+  /// Error to return when an ASTActionInvalidation policy fires.
+  class InvalidatedError : public llvm::ErrorInfo<InvalidatedError> {
+  public:
+    static char ID;
+    ASTActionInvalidation Policy;
+
+    void log(llvm::raw_ostream &OS) const override;
+    std::error_code convertToErrorCode() const override {
+      return std::make_error_code(std::errc::interrupted);
+    }
+  };
+
   /// Schedule an async read of the AST. \p Action will be called when AST is
   /// ready. The AST passed to \p Action refers to the version of \p File
   /// tracked at the time of the call, even if new updates are received before
@@ -226,9 +249,11 @@
   /// If an error occurs during processing, it is forwarded to the \p Action
   /// callback.
   /// If the context is cancelled before the AST is ready, the callback will
-  /// receive a CancelledError.
+  /// receive a CancelledError. If the invalidation policy is triggered, the
+  /// callback will receive an InvalidatedError.
   void runWithAST(llvm::StringRef Name, PathRef File,
-                  Callback<InputsAndAST> Action);
+                  Callback<InputsAndAST> Action,
+                  ASTActionInvalidation = NoInvalidation);
 
   /// Controls whether preamble reads wait for the preamble to be up-to-date.
   enum PreambleConsistency {
diff --git a/clang-tools-extra/clangd/unittests/CancellationTests.cpp b/clang-tools-extra/clangd/unittests/CancellationTests.cpp
index 611ce07..3fe7190 100644
--- a/clang-tools-extra/clangd/unittests/CancellationTests.cpp
+++ b/clang-tools-extra/clangd/unittests/CancellationTests.cpp
@@ -44,6 +44,30 @@
   Task.second();
 }
 
+struct NestedTasks {
+  std::pair<Context, Canceler> Outer, Inner;
+  NestedTasks() {
+    Outer = cancelableTask();
+    {
+      WithContext WithOuter(Outer.first.clone());
+      Inner = cancelableTask();
+    }
+  }
+};
+
+TEST(CancellationTest, Nested) {
+  // Cancelling inner task works but leaves outer task unaffected.
+  NestedTasks CancelInner;
+  CancelInner.Inner.second();
+  EXPECT_TRUE(isCancelled(CancelInner.Inner.first));
+  EXPECT_FALSE(isCancelled(CancelInner.Outer.first));
+  // Cancellation of outer task is inherited by inner task.
+  NestedTasks CancelOuter;
+  CancelOuter.Outer.second();
+  EXPECT_TRUE(isCancelled(CancelOuter.Inner.first));
+  EXPECT_TRUE(isCancelled(CancelOuter.Outer.first));
+}
+
 TEST(CancellationTest, AsynCancellationTest) {
   std::atomic<bool> HasCancelled(false);
   Notification Cancelled;
diff --git a/clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp b/clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
index 9e5952f..365e534 100644
--- a/clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
+++ b/clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "Annotations.h"
+#include "Cancellation.h"
 #include "Context.h"
 #include "Diagnostics.h"
 #include "Matchers.h"
@@ -346,6 +347,63 @@
       << "All reads other than R2B were cancelled";
 }
 
+TEST_F(TUSchedulerTests, Invalidation) {
+  auto Path = testPath("foo.cpp");
+  TUScheduler S(CDB, optsForTest(), captureDiags());
+  std::atomic<int> Builds(0), Actions(0);
+
+  Notification Start;
+  updateWithDiags(S, Path, "a", WantDiagnostics::Yes, [&](std::vector<Diag>) {
+    ++Builds;
+    Start.wait();
+  });
+  S.runWithAST(
+      "invalidatable", Path,
+      [&](llvm::Expected<InputsAndAST> AST) {
+        ++Actions;
+        EXPECT_FALSE(bool(AST));
+        llvm::Error E = AST.takeError();
+        EXPECT_TRUE(E.isA<CancelledError>());
+        consumeError(std::move(E));
+      },
+      TUScheduler::InvalidateOnUpdate);
+  S.runWithAST(
+      "not-invalidatable", Path,
+      [&](llvm::Expected<InputsAndAST> AST) {
+        ++Actions;
+        EXPECT_TRUE(bool(AST));
+      },
+      TUScheduler::NoInvalidation);
+  updateWithDiags(S, Path, "b", WantDiagnostics::Auto, [&](std::vector<Diag>) {
+    ++Builds;
+    ADD_FAILURE() << "Shouldn't build, all dependents invalidated";
+  });
+  S.runWithAST(
+      "invalidatable", Path,
+      [&](llvm::Expected<InputsAndAST> AST) {
+        ++Actions;
+        EXPECT_FALSE(bool(AST));
+        llvm::Error E = AST.takeError();
+        EXPECT_TRUE(E.isA<CancelledError>());
+        consumeError(std::move(E));
+      },
+      TUScheduler::InvalidateOnUpdate);
+  updateWithDiags(S, Path, "c", WantDiagnostics::Auto,
+                  [&](std::vector<Diag>) { ++Builds; });
+  S.runWithAST(
+      "invalidatable", Path,
+      [&](llvm::Expected<InputsAndAST> AST) {
+        ++Actions;
+        EXPECT_TRUE(bool(AST)) << "Shouldn't be invalidated, no update follows";
+      },
+      TUScheduler::InvalidateOnUpdate);
+  Start.notify();
+  ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
+
+  EXPECT_EQ(2, Builds.load()) << "Middle build should be skipped";
+  EXPECT_EQ(4, Actions.load()) << "All actions should run (some with error)";
+}
+
 TEST_F(TUSchedulerTests, ManyUpdates) {
   const int FilesCount = 3;
   const int UpdatesPerFile = 10;