Index: clangd/Cancellation.h =================================================================== --- clangd/Cancellation.h +++ clangd/Cancellation.h @@ -79,7 +79,7 @@ /// 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 { Index: clangd/Cancellation.cpp =================================================================== --- clangd/Cancellation.cpp +++ clangd/Cancellation.cpp @@ -24,8 +24,8 @@ }; } -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. } Index: clangd/TUScheduler.h =================================================================== --- clangd/TUScheduler.h +++ clangd/TUScheduler.h @@ -100,6 +100,9 @@ /// 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)> OnUpdated); @@ -117,6 +120,8 @@ /// \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 Action); @@ -140,6 +145,8 @@ /// 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 Action); Index: clangd/TUScheduler.cpp =================================================================== --- clangd/TUScheduler.cpp +++ clangd/TUScheduler.cpp @@ -43,6 +43,7 @@ // immediately. #include "TUScheduler.h" +#include "Cancellation.h" #include "Logger.h" #include "Trace.h" #include "clang/Frontend/CompilerInvocation.h" @@ -432,6 +433,8 @@ void ASTWorker::runWithAST( StringRef Name, unique_function)> Action) { auto Task = [=](decltype(Action) Action) { + if (isCancelled()) + return Action(make_error()); Optional> AST = IdleASTs.take(this); if (!AST) { std::unique_ptr Invocation = @@ -588,6 +591,26 @@ 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. + for (auto I = Requests.begin(), E = Requests.end(); I != E; ++I) { + if (!isCancelled(I->Ctx)) { + // Cancellations after the first read don't affect current scheduling. + if (I->UpdateType == None) + break; + 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"); Index: unittests/clangd/TUSchedulerTests.cpp =================================================================== --- unittests/clangd/TUSchedulerTests.cpp +++ unittests/clangd/TUSchedulerTests.cpp @@ -209,6 +209,75 @@ 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 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 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 E) { + if (auto Err = E.takeError()) { + if (Err.isA()) { + 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 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;