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

[sammccall]

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.

[kadircet]

thanks, LGTM

clang-tools-extra/clangd/Cancellation.cpp

35–39

comment seems to be out-of-date.

Either not cancelled or not in a cancellable task ?

clang-tools-extra/clangd/JSONTransport.cpp

23

s/approprate/appropriate/

clang-tools-extra/clangd/TUScheduler.cpp

673

braces

682

this one is already being put into the context before pushing into requests.

clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp

364

nit: bool(AST) ? same in other places.

[kadircet]

clang-tools-extra/clangd/ClangdServer.cpp

58

nit: i would actually keep InvalidateOnUpdate on the callsites, but up to you.

also constexpr instead of static ?

[sammccall]

address comments

[thakis]

This broke building on my bots:

http://45.33.8.238/linux/11862/step_4.txt

ld.lld: error: undefined symbol: clang::clangd::TUScheduler::InvalidatedError::ID

Looks like that is indeed declared in this change, but it doesn't have a definition.

[sammccall]

In D75602#1906587, @thakis wrote:

This broke building on my bots:

http://45.33.8.238/linux/11862/step_4.txt

ld.lld: error: undefined symbol: clang::clangd::TUScheduler::InvalidatedError::ID

Looks like that is indeed declared in this change, but it doesn't have a definition.

Sorry about that, fixing. (I didn't see this on other linkers, maybe because the class itself wasn't used in this patch)

clang-tools-extra/clangd/Cancellation.cpp

35–39

Yeah, just removed the comment, this is now pretty obvious.

clang-tools-extra/clangd/Cancellation.h

 
 /// Defines a new task whose cancellation may be requested.
 /// 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.
+/// (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());
 
 /// Conventional error when no result is returned due to cancellation.
 class CancelledError : public llvm::ErrorInfo<CancelledError> {
 public:
   static char ID;

clang-tools-extra/clangd/Cancellation.cpp

 //===--- Cancellation.cpp -----------------------------------------*-C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "Cancellation.h"
 #include <atomic>
 
 namespace clang {
 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
 } // namespace clang

clang-tools-extra/clangd/ClangdServer.cpp

 #include <mutex>
 #include <type_traits>
 
 namespace clang {
 namespace clangd {
 namespace {
 
 // Update the FileIndex with new ASTs and plumb the diagnostics responses.
 struct UpdateIndexCallbacks : public ParsingCallbacks {
   UpdateIndexCallbacks(FileIndex *FIndex,
                        ClangdServer::Callbacks *ServerCallbacks,
                        bool SemanticHighlighting)
       : FIndex(FIndex), ServerCallbacks(ServerCallbacks),
         SemanticHighlighting(SemanticHighlighting) {}
 
   void onPreambleAST(PathRef Path, ASTContext &Ctx,
                      std::shared_ptr<clang::Preprocessor> PP,
for (const auto &Sel : *Selections) {
         Res.push_back({T->id(), T->title(), T->intent()});
         PreparedTweaks.insert(T->id());
       }
     }
 
     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,
                               Callback<Tweak::Effect> CB) {
   auto Action =
       [File = File.str(), Sel, TweakID = TweakID.str(), CB = std::move(CB),
        FS = FSProvider.getFileSystem()](Expected<InputsAndAST> InpAST) mutable {
         if (!InpAST)
void ClangdServer::findDocumentHighlights(
     PathRef File, Position Pos, Callback<std::vector<DocumentHighlight>> CB) {
   auto Action =
       [Pos, CB = std::move(CB)](llvm::Expected<InputsAndAST> InpAST) mutable {
         if (!InpAST)
           return CB(InpAST.takeError());
         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,
                              Callback<llvm::Optional<HoverInfo>> CB) {
   auto Action = [File = File.str(), Pos, CB = std::move(CB),
                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
     format::FormatStyle Style = getFormatStyleForFile(
         File, InpAST->Inputs.Contents, InpAST->Inputs.FS.get());
     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,
                                  TypeHierarchyDirection Direction,
                                  Callback<Optional<TypeHierarchyItem>> CB) {
   auto Action = [File = File.str(), Pos, Resolve, Direction, CB = std::move(CB),
                  this](Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)

 void ClangdServer::documentSymbols(llvm::StringRef File,
                                    Callback<std::vector<DocumentSymbol>> CB) {
   auto Action =
       [CB = std::move(CB)](llvm::Expected<InputsAndAST> InpAST) mutable {
         if (!InpAST)
           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,
                                   Callback<ReferencesResult> CB) {
   auto Action = [Pos, Limit, CB = std::move(CB),
                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());

 void ClangdServer::documentLinks(PathRef File,
                                  Callback<std::vector<DocumentLink>> CB) {
   auto Action =
       [CB = std::move(CB)](llvm::Expected<InputsAndAST> InpAST) mutable {
         if (!InpAST)
           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>>
 ClangdServer::getUsedBytesPerFile() const {
   return WorkScheduler.getUsedBytesPerFile();
 }
 
 LLVM_NODISCARD bool
 ClangdServer::blockUntilIdleForTest(llvm::Optional<double> TimeoutSeconds) {
   return WorkScheduler.blockUntilIdle(timeoutSeconds(TimeoutSeconds)) &&
          (!BackgroundIdx ||
           BackgroundIdx->blockUntilIdleForTest(TimeoutSeconds));
 }
 
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/JSONTransport.cpp

 
 namespace clang {
 namespace clangd {
 namespace {
 
 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;
             Code = L.Code;
             return llvm::Error::success();
           }))
     Message = llvm::toString(std::move(Unhandled));
 

clang-tools-extra/clangd/TUScheduler.h

public:
   void remove(PathRef File);
 
   /// Returns a snapshot of all file buffer contents, per last update().
   llvm::StringMap<std::string> getAllFileContents() const;
 
   /// 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
   /// \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.
+  /// 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 {
     /// 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,

clang-tools-extra/clangd/TUScheduler.cpp

create(PathRef FileName, const GlobalCompilationDatabase &CDB,
          TUScheduler::ASTCache &IdleASTs, AsyncTaskRunner *Tasks,
          Semaphore &Barrier, DebouncePolicy UpdateDebounce,
          bool StorePreamblesInMemory, ParsingCallbacks &Callbacks);
   ~ASTWorker();
 
   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;
 
   /// Obtain a preamble reflecting all updates so far. Threadsafe.
   /// It may be delivered immediately, or later on the worker thread.
   void getCurrentPreamble(
       llvm::unique_function<void(std::shared_ptr<const PreambleData>)>);
private:
   // Must be called exactly once on processing thread. Will return after
   // stop() is called on a separate thread and all pending requests are
   // processed.
   void run();
   /// Signal that run() should finish processing pending requests and exit.
   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);
 
   /// Determines the next action to perform.
   /// All actions that should never run are discarded.
   /// Returns a deadline for the next action. If it's expired, run now.
   /// scheduleLocked() is called again at the deadline, or if requests arrive.
   Deadline scheduleLocked();
   /// Should the first task in the queue be skipped instead of run?
   bool shouldSkipHeadLocked() const;
   /// This is private because `FileInputs.FS` is not thread-safe and thus not
   /// safe to share. Callers should make sure not to expose `FS` via a public
   /// interface.
   std::shared_ptr<const ParseInputs> getCurrentFileInputs() const;
 
   struct Request {
     llvm::unique_function<void()> Action;
     std::string Name;
     steady_clock::time_point AddTime;
     Context Ctx;
     llvm::Optional<WantDiagnostics> UpdateType;
+    TUScheduler::ASTActionInvalidation InvalidationPolicy;
+    Canceler Invalidate;
   };
 
   /// Handles retention of ASTs.
   TUScheduler::ASTCache &IdleASTs;
   const bool RunSync;
   /// Time to wait after an update to see whether another update obsoletes it.
   const DebouncePolicy UpdateDebounce;
   /// File that ASTWorker is responsible for.
if (*AST) {
       // line if there were any.
       // FIXME: we might have got more errors while trying to build the AST,
       //        surface them too.
       Callbacks.onFailedAST(FileName, CompilerInvocationDiags, RunPublish);
     }
     // 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>());
     llvm::Optional<std::unique_ptr<ParsedAST>> AST = IdleASTs.take(this);
     auto CurrentInputs = getCurrentFileInputs();
     if (!AST) {
       StoreDiags CompilerInvocationDiagConsumer;
       std::unique_ptr<CompilerInvocation> Invocation = buildCompilerInvocation(
auto Task = [=, Action = std::move(Action)]() mutable {
     auto _ = llvm::make_scope_exit(
         [&AST, this]() { IdleASTs.put(this, std::move(*AST)); });
     // Run the user-provided action.
     if (!*AST)
       return Action(llvm::make_error<llvm::StringError>(
           "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>
 ASTWorker::getPossiblyStalePreamble() const {
   std::lock_guard<std::mutex> Lock(Mutex);
   return LastBuiltPreamble;
 }
 
void ASTWorker::getCurrentPreamble(
   }
   assert(!RunSync && "Running synchronously, but queue is non-empty!");
   Requests.insert(LastUpdate.base(),
                   Request{[Callback = std::move(Callback), this]() mutable {
                             Callback(getPossiblyStalePreamble());
                           },
                           "GetPreamble", steady_clock::now(),
                           Context::current().clone(),
-                          /*UpdateType=*/None});
+                          /*UpdateType=*/None,
+                          /*InvalidationPolicy=*/TUScheduler::NoInvalidation,
+                          /*Invalidate=*/nullptr});
   Lock.unlock();
   RequestsCV.notify_all();
 }
 
 void ASTWorker::waitForFirstPreamble() const { PreambleWasBuilt.wait(); }
 
 std::shared_ptr<const ParseInputs> ASTWorker::getCurrentFileInputs() const {
   std::unique_lock<std::mutex> Lock(Mutex);
void ASTWorker::stop() {
     assert(!Done && "stop() called twice");
     Done = true;
   }
   RequestsCV.notify_all();
 }
 
 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));
     Task();
     return;
   }
 
   {
     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();
 }
 
 void ASTWorker::emitTUStatus(TUAction Action,
                              const TUStatus::BuildDetails *Details) {
   Status.Action = std::move(Action);
   if (Details)
PreambleTasks->runAsync(Name, [this, Ctx = Context::current().clone(),
     std::lock_guard<Semaphore> BarrierLock(Barrier);
     WithContext WC(std::move(Ctx));
     Action();
   });
 }
 
 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>(
         "trying to get AST for non-added document", ErrorCode::InvalidParams));
     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,
                                   PreambleConsistency Consistency,
                                   Callback<InputsAndPreamble> Action) {
   auto It = Files.find(File);
   if (It == Files.end()) {
     Action(llvm::make_error<LSPError>(

 }
 
 DebouncePolicy DebouncePolicy::fixed(clock::duration T) {
   DebouncePolicy P;
   P.Min = P.Max = T;
   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

clang-tools-extra/clangd/unittests/CancellationTests.cpp

auto Task = cancelableTask();
     EXPECT_FALSE(isCancelled());
     Task.second();
     EXPECT_TRUE(isCancelled());
   }
   // Still should be able to cancel without any problems.
   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;
   auto TaskToBeCancelled = [&](Context Ctx) {
     WithContext ContextGuard(std::move(Ctx));
     Cancelled.wait();
     HasCancelled = isCancelled();
   };

clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp

 //===-- TUSchedulerTests.cpp ------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "Annotations.h"
+#include "Cancellation.h"
 #include "Context.h"
 #include "Diagnostics.h"
 #include "Matchers.h"
 #include "ParsedAST.h"
 #include "Path.h"
 #include "Preamble.h"
 #include "TUScheduler.h"
 #include "TestFS.h"
EXPECT_THAT(DiagsSeen, ElementsAre("U2", "U3"))
          "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, 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;
 
   std::mutex Mut;
   int TotalASTReads = 0;
   int TotalPreambleReads = 0;
   int TotalUpdates = 0;