[clangd] Respect task cancellation in TUScheduler.

[sammccall]

Summary

• Reads are never executed if canceled before ready-to run. In practice, we finalize cancelled reads eagerly and out-of-order.
• Cancelled reads don't prevent prior updates from being elided, as they don't actually depend on the result of the update.
• Updates are downgraded from WantDiagnostics::Yes to WantDiagnostics::Auto when cancelled, which allows them to be elided when all dependent reads are cancelled and there are subsequent writes. (e.g. when the queue is backed up with cancelled requests).

The queue operations aren't optimal (we scan the whole queue for cancelled
tasks every time the scheduler runs, and check cancellation twice in the end).
However I believe these costs are still trivial in practice (compared to any
AST operation) and the logic can be cleanly separated from the rest of the
scheduler.

[ilya-biryukov]

LGTM

clangd/TUScheduler.cpp

598

NIT: maybe also break on non-cancelled reads? Going past this point does not seem particularly useful, as we won't be able to cancel any pending updates before them anyway. And it does not seem like responding for a cancelled task fast is higher priority than delivering the results of that non-cancelled read.

clangd/Cancellation.h

 /// The returned Context defines the scope of the task.
 /// When the context is active, isCancelled() is false until the Canceler is
 /// invoked, and true afterwards.
 std::pair<Context, Canceler> cancelableTask();
 
 /// True if the current context is within a cancelable task which was 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();
+bool isCancelled(const Context &Ctx = Context::current());
 
 /// Conventional error when no result is returned due to cancellation.
 class CancelledError : public llvm::ErrorInfo<CancelledError> {
 public:
   static char ID;
 
   void log(llvm::raw_ostream &OS) const override {
     OS << "Task was cancelled.";

clangd/Cancellation.cpp

 std::pair<Context, Canceler> cancelableTask() {
   auto Flag = std::make_shared<std::atomic<bool>>();
   return {
       Context::current().derive(FlagKey, Flag),
       [Flag] { *Flag = true; },
   };
 }
 
-bool isCancelled() {
-  if (auto *Flag = Context::current().get(FlagKey))
+bool isCancelled(const Context &Ctx) {
+  if (auto *Flag = Ctx.get(FlagKey))
     return **Flag;
   return false; // Not in scope of a task.
 }
 
 } // namespace clangd
 } // namespace clang

clangd/TUScheduler.h

public:
 
   /// Returns a list of files with ASTs currently stored in memory. This method
   /// is not very reliable and is only used for test. E.g., the results will not
   /// contain files that currently run something over their AST.
   std::vector<Path> getFilesWithCachedAST() const;
 
   /// Schedule an update for \p File. Adds \p File to a list of tracked files if
   /// \p File was not part of it before.
+  /// If diagnostics are requested (Yes), and the context is cancelled before
+  /// they are prepared, they may be skipped if eventual-consistency permits it
+  /// (i.e. WantDiagnostics is downgraded to Auto).
   /// FIXME(ibiryukov): remove the callback from this function.
   void update(PathRef File, ParseInputs Inputs, WantDiagnostics WD,
               llvm::unique_function<void(std::vector<Diag>)> OnUpdated);
 
   /// Remove \p File from the list of tracked files and schedule removal of its
   /// resources.
   void remove(PathRef File);
 
   /// Schedule an async task with no dependencies.
   void run(llvm::StringRef Name, llvm::unique_function<void()> Action);
 
   /// 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
   /// \p Action is executed.
   /// 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.
   void runWithAST(llvm::StringRef Name, PathRef File,
                   Callback<InputsAndAST> Action);
 
   /// Controls whether preamble reads wait for the preamble to be up-to-date.
   enum PreambleConsistency {
     /// The preamble is generated from the current version of the file.
     /// If the content was recently updated, we will wait until we have a
     /// preamble that reflects that update.
     /// This is the slowest option, and may be delayed by other tasks.
     Consistent,
     /// The preamble may be generated from an older version of the file.
     /// Reading from locations in the preamble may cause files to be re-read.
     /// This gives callers two options:
     /// - validate that the preamble is still valid, and only use it if so
     /// - accept that the preamble contents may be outdated, and try to avoid
     ///   reading source code from headers.
     /// This is the fastest option, usually a preamble is available immediately.
     Stale,
   };
   /// Schedule an async read of the preamble.
   /// If there's no preamble yet (because the file was just opened), we'll wait
   /// for it to build. The result may be null if it fails to build or is empty.
   /// If an error occurs, it is forwarded to the \p Action callback.
+  /// Context cancellation is ignored and should be handled by the Action.
+  /// (In practice, the Action is almost always executed immediately).
   void runWithPreamble(llvm::StringRef Name, PathRef File,
                        PreambleConsistency Consistency,
                        Callback<InputsAndPreamble> Action);
 
   /// Wait until there are no scheduled or running tasks.
   /// Mostly useful for synchronizing tests.
   bool blockUntilIdle(Deadline D) const;
 

clangd/TUScheduler.cpp

 // - For each update we remember the pending inputs, but delay rebuild of the
 //   AST for some timeout.
 // - If subsequent updates come before rebuild was started, we replace the
 //   pending inputs and reset the timer.
 // - If any reads of the AST are scheduled, we start building the AST
 //   immediately.
 
 #include "TUScheduler.h"
+#include "Cancellation.h"
 #include "Logger.h"
 #include "Trace.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/PCHContainerOperations.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Path.h"
 #include <algorithm>
void ASTWorker::update(ParseInputs Inputs, WantDiagnostics WantDiags,
   };
 
   startTask("Update", Bind(Task, std::move(OnUpdated)), WantDiags);
 }
 
 void ASTWorker::runWithAST(
     StringRef Name, unique_function<void(Expected<InputsAndAST>)> Action) {
   auto Task = [=](decltype(Action) Action) {
+    if (isCancelled())
+      return Action(make_error<CancelledError>());
     Optional<std::unique_ptr<ParsedAST>> AST = IdleASTs.take(this);
     if (!AST) {
       std::unique_ptr<CompilerInvocation> Invocation =
           buildCompilerInvocation(FileInputs);
       // Try rebuilding the AST.
       Optional<ParsedAST> NewAST =
           Invocation
               ? buildAST(FileName,
while (true) {
     }
     RequestsCV.notify_all();
   }
 }
 
 Deadline ASTWorker::scheduleLocked() {
   if (Requests.empty())
     return Deadline::infinity(); // Wait for new requests.
+  // Handle cancelled requests first so the rest of the scheduler doesn't.
+  // Scanning the whole queue may be suboptimal, but it's much simpler.
+  for (auto I = Requests.begin(), E = Requests.end(); I != E; ++I) {
+    if (!isCancelled(I->Ctx))
+      continue;
+    // Cancelled reads are moved to the front of the queue and run immediately.
+    if (I->UpdateType == None) {
+      Request R = std::move(*I);
+      Requests.erase(I);
+      Requests.push_front(std::move(R));
+      return Deadline::zero();
+    }
+    // Cancelled updates are downgraded to auto-diagnostics, and may be elided.
+    if (I->UpdateType == WantDiagnostics::Yes)
+      I->UpdateType = WantDiagnostics::Auto;
+  }
+
   while (shouldSkipHeadLocked())
     Requests.pop_front();
   assert(!Requests.empty() && "skipped the whole queue");
   // Some updates aren't dead yet, but never end up being used.
   // e.g. the first keystroke is live until obsoleted by the second.
   // We debounce "maybe-unused" writes, sleeping 500ms in case they become dead.
   // But don't delay reads (including updates where diagnostics are needed).
   for (const auto &R : Requests)

unittests/clangd/TUSchedulerTests.cpp

S.runWithPreamble("ConsistentRead", Path, TUScheduler::Consistent,
                         EXPECT_THAT(includes(Pre->Preamble),
                                     ElementsAre("<B>"));
                         ++CallbackCount;
                       });
   }
   EXPECT_EQ(2, CallbackCount);
 }
 
+TEST_F(TUSchedulerTests, Cancellation) {
+  // We have the following update/read sequence
+  //   U0
+  //   U1(WantDiags=Yes) <-- cancelled
+  //    R1               <-- cancelled
+  //   U2(WantDiags=Yes) <-- cancelled
+  //    R2A              <-- cancelled
+  //    R2B
+  //   U3(WantDiags=Yes)
+  //    R3               <-- cancelled
+  std::vector<std::string> DiagsSeen, ReadsSeen, ReadsCanceled;
+  {
+    TUScheduler S(
+        getDefaultAsyncThreadsCount(), /*StorePreamblesInMemory=*/true,
+        /*ASTCallbacks=*/nullptr,
+        /*UpdateDebounce=*/std::chrono::steady_clock::duration::zero(),
+        ASTRetentionPolicy());
+    auto Path = testPath("foo.cpp");
+    // Helper to schedule a named update and return a function to cancel it.
+    auto Update = [&](std::string ID) -> Canceler {
+      auto T = cancelableTask();
+      WithContext C(std::move(T.first));
+      S.update(Path, getInputs(Path, "//" + ID), WantDiagnostics::Yes,
+               [&, ID](std::vector<Diag> Diags) { DiagsSeen.push_back(ID); });
+      return std::move(T.second);
+    };
+    // Helper to schedule a named read and return a function to cancel it.
+    auto Read = [&](std::string ID) -> Canceler {
+      auto T = cancelableTask();
+      WithContext C(std::move(T.first));
+      S.runWithAST(ID, Path, [&, ID](llvm::Expected<InputsAndAST> E) {
+        if (auto Err = E.takeError()) {
+          if (Err.isA<CancelledError>()) {
+            ReadsCanceled.push_back(ID);
+            consumeError(std::move(Err));
+          } else {
+            ADD_FAILURE() << "Non-cancelled error for " << ID << ": "
+                          << llvm::toString(std::move(Err));
+          }
+        } else {
+          ReadsSeen.push_back(ID);
+        }
+      });
+      return std::move(T.second);
+    };
+
+    Notification Proceed; // Ensure we schedule everything.
+    S.update(Path, getInputs(Path, ""), WantDiagnostics::Yes,
+             [&](std::vector<Diag> Diags) { Proceed.wait(); });
+    // The second parens indicate cancellation, where present.
+    Update("U1")();
+    Read("R1")();
+    Update("U2")();
+    Read("R2A")();
+    Read("R2B");
+    Update("U3");
+    Read("R3")();
+    Proceed.notify();
+  }
+  EXPECT_THAT(DiagsSeen, ElementsAre("U2", "U3"))
+      << "U1 and all dependent reads were cancelled. "
+         "U2 has a dependent read R2A. "
+         "U3 was not cancelled.";
+  EXPECT_THAT(ReadsSeen, ElementsAre("R2B"))
+      << "All reads other than R2B were cancelled";
+  EXPECT_THAT(ReadsCanceled, ElementsAre("R1", "R2A", "R3"))
+      << "All reads other than R2B were cancelled";
+}
+
 TEST_F(TUSchedulerTests, ManyUpdates) {
   const int FilesCount = 3;
   const int UpdatesPerFile = 10;
 
   std::mutex Mut;
   int TotalASTReads = 0;
   int TotalPreambleReads = 0;
   int TotalUpdates = 0;