[clangd] Use a separate RunningTask flag instead of leaving a broken request on top of the queue

[kadircet]

Summary

This helps us prevent races when scheduler (or any other thread) tries
to read a request while it's still running.

[sammccall]

clang-tools-extra/clangd/TUScheduler.cpp

278

I'd consider making this an llvm::Optional<Request> and using it as storage too. Even though it won't be accessed directly outside the main loop, I think it makes it easier to reason about the code: it emphasizes the parallel with Requests and minimizes the spooky-flags-at-a-distance effect.

611–614

what if the currently-running task is a write?

722–723

maybe assert there's no running task here?

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

403

does this deliberately match "b" below? if not, change back to something unique?
If so, probably needs a comment.

406

This looks racy to me. What guarantees that the worker thread pulls this item off the queue and into the "running" slot before the second updateWithDiags?

[kadircet]

• Address comments

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

403

no it wasn't intentional just copy paste ..

clang-tools-extra/clangd/TUScheduler.cpp

private:
   std::shared_ptr<const ParseInputs> FileInputs;         /* GUARDED_BY(Mutex) */
   std::shared_ptr<const PreambleData> LastBuiltPreamble; /* GUARDED_BY(Mutex) */
   /// Times of recent AST rebuilds, used for UpdateDebounce computation.
   llvm::SmallVector<DebouncePolicy::clock::duration, 8>
       RebuildTimes; /* GUARDED_BY(Mutex) */
   /// Becomes ready when the first preamble build finishes.
   Notification PreambleWasBuilt;
   /// Set to true to signal run() to finish processing.
-  bool Done;                    /* GUARDED_BY(Mutex) */
-  std::deque<Request> Requests; /* GUARDED_BY(Mutex) */
+  bool Done;                              /* GUARDED_BY(Mutex) */
+  std::deque<Request> Requests;           /* GUARDED_BY(Mutex) */
+  llvm::Optional<Request> CurrentRequest; /* GUARDED_BY(Mutex) */
   mutable std::condition_variable RequestsCV;
   /// Guards the callback that publishes results of AST-related computations
   /// (diagnostics, highlightings) and file statuses.
   std::mutex PublishMu;
   // Used to prevent remove document + add document races that lead to
   // out-of-order callbacks for publishing results of onMainAST callback.
   //
   // The lifetime of the old/new ASTWorkers will overlap, but their handles

 }
 
 ASTWorker::~ASTWorker() {
   // Make sure we remove the cached AST, if any.
   IdleASTs.take(this);
 #ifndef NDEBUG
   std::lock_guard<std::mutex> Lock(Mutex);
   assert(Done && "handle was not destroyed");
-  assert(Requests.empty() && "unprocessed requests when destroying ASTWorker");
+  assert(Requests.empty() && !CurrentRequest &&
+         "unprocessed requests when destroying ASTWorker");
 #endif
 }
 
 void ASTWorker::update(ParseInputs Inputs, WantDiagnostics WantDiags) {
   std::string TaskName = llvm::formatv("Update ({0})", Inputs.Version);
   auto Task = [=]() mutable {
     auto RunPublish = [&](llvm::function_ref<void()> Publish) {
       // Ensure we only publish results from the worker if the file was not
void ASTWorker::getCurrentPreamble(
     llvm::unique_function<void(std::shared_ptr<const PreambleData>)> Callback) {
   // We could just call startTask() to throw the read on the queue, knowing
   // it will run after any updates. But we know this task is cheap, so to
   // improve latency we cheat: insert it on the queue after the last update.
   std::unique_lock<std::mutex> Lock(Mutex);
   auto LastUpdate =
       std::find_if(Requests.rbegin(), Requests.rend(),
                    [](const Request &R) { return R.UpdateType.hasValue(); });
-  // If there were no writes in the queue, the preamble is ready now.
-  if (LastUpdate == Requests.rend()) {
+  // If there were no writes in the queue, and CurrentRequest is not a write,
+  // the preamble is ready now.
+  if (LastUpdate == Requests.rend() &&
+      (!CurrentRequest || CurrentRequest->UpdateType.hasValue())) {
     Lock.unlock();
     return Callback(getPossiblyStalePreamble());
   }
   assert(!RunSync && "Running synchronously, but queue is non-empty!");
   Requests.insert(LastUpdate.base(),
                   Request{[Callback = std::move(Callback), this]() mutable {
                             Callback(getPossiblyStalePreamble());
                           },
void ASTWorker::emitTUStatus(TUAction Action,
   // Do not emit TU statuses when the ASTWorker is shutting down.
   if (CanPublishResults) {
     Callbacks.onFileUpdated(FileName, Status);
   }
 }
 
 void ASTWorker::run() {
   while (true) {
-    Request Req;
     {
       std::unique_lock<std::mutex> Lock(Mutex);
+      assert(!CurrentRequest && "A task is already running, multiple workers?");
       for (auto Wait = scheduleLocked(); !Wait.expired();
            Wait = scheduleLocked()) {
         if (Done) {
           if (Requests.empty())
             return;
           else     // Even though Done is set, finish pending requests.
             break; // However, skip delays to shutdown fast.
         }
 
         // Tracing: we have a next request, attribute this sleep to it.
         llvm::Optional<WithContext> Ctx;
         llvm::Optional<trace::Span> Tracer;
         if (!Requests.empty()) {
           Ctx.emplace(Requests.front().Ctx.clone());
           Tracer.emplace("Debounce");
           SPAN_ATTACH(*Tracer, "next_request", Requests.front().Name);
           if (!(Wait == Deadline::infinity())) {
-            emitTUStatus({TUAction::Queued, Req.Name});
+            emitTUStatus({TUAction::Queued, Requests.front().Name});
             SPAN_ATTACH(*Tracer, "sleep_ms",
                         std::chrono::duration_cast<std::chrono::milliseconds>(
                             Wait.time() - steady_clock::now())
                             .count());
           }
         }
 
         wait(Lock, RequestsCV, Wait);
       }
-      Req = std::move(Requests.front());
-      // Leave it on the queue for now, so waiters don't see an empty queue.
+      CurrentRequest = std::move(Requests.front());
+      Requests.pop_front();
     } // unlock Mutex
 
+    // It is safe to perform reads to CurrentRequest without holding the lock as
+    // only writer is also this thread.
     {
       std::unique_lock<Semaphore> Lock(Barrier, std::try_to_lock);
       if (!Lock.owns_lock()) {
-        emitTUStatus({TUAction::Queued, Req.Name});
+        emitTUStatus({TUAction::Queued, CurrentRequest->Name});
         Lock.lock();
       }
-      WithContext Guard(std::move(Req.Ctx));
-      trace::Span Tracer(Req.Name);
-      emitTUStatus({TUAction::RunningAction, Req.Name});
-      Req.Action();
+      WithContext Guard(std::move(CurrentRequest->Ctx));
+      trace::Span Tracer(CurrentRequest->Name);
+      emitTUStatus({TUAction::RunningAction, CurrentRequest->Name});
+      CurrentRequest->Action();
     }
 
     bool IsEmpty = false;
     {
       std::lock_guard<std::mutex> Lock(Mutex);
-      Requests.pop_front();
+      CurrentRequest.reset();
       IsEmpty = Requests.empty();
     }
     if (IsEmpty)
       emitTUStatus({TUAction::Idle, /*Name*/ ""});
     RequestsCV.notify_all();
   }
 }
 
for (; Next != Requests.end(); ++Next)
         return true; // Prefer later diagnostics.
     return false;
   }
   llvm_unreachable("Unknown WantDiagnostics");
 }
 
 bool ASTWorker::blockUntilIdle(Deadline Timeout) const {
   std::unique_lock<std::mutex> Lock(Mutex);
-  return wait(Lock, RequestsCV, Timeout, [&] { return Requests.empty(); });
+  return wait(Lock, RequestsCV, Timeout,
+              [&] { return Requests.empty() && !CurrentRequest; });
 }
 
 // Render a TUAction to a user-facing string representation.
 // TUAction represents clangd-internal states, we don't intend to expose them
 // to users (say C++ programmers) directly to avoid confusion, we use terms that
 // are familiar by C++ programmers.
 std::string renderTUAction(const TUAction &Action) {
   std::string Result;

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

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, InvalidationNoCrash) {
+  auto Path = testPath("foo.cpp");
+  TUScheduler S(CDB, optsForTest(), captureDiags());
+
+  Notification StartedRunning;
+  Notification ScheduledChange;
+  // We expect invalidation logic to not crash by trying to invalidate a running
+  // request.
+  S.update(Path, getInputs(Path, ""), WantDiagnostics::Auto);
+  ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
+  S.runWithAST(
+      "invalidatable-but-running", Path,
+      [&](llvm::Expected<InputsAndAST> AST) {
+        StartedRunning.notify();
+        ScheduledChange.wait();
+        ASSERT_TRUE(bool(AST));
+      },
+      TUScheduler::InvalidateOnUpdate);
+  StartedRunning.wait();
+  S.update(Path, getInputs(Path, ""), WantDiagnostics::Auto);
+  ScheduledChange.notify();
+  ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
+}
+
 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(