[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

279

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.

612

what if the currently-running task is a write?

721

maybe assert there's no running task here?

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

376

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

379

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

376

no it wasn't intentional just copy paste ..

clang-tools-extra/clangd/TUScheduler.cpp

private:
   /// 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 RunningTask = false;     /* 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() && !RunningTask &&
+         "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(
   // 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()) {
     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());
                           },
if (CanPublishResults) {
     Callbacks.onFileUpdated(FileName, Status);
   }
 }
 
 void ASTWorker::run() {
   while (true) {
     Request Req;
     {
       std::unique_lock<std::mutex> Lock(Mutex);
       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.
         }
Request Req;
                             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.
+      Requests.pop_front();
+      RunningTask = true;
     } // unlock Mutex
 
     {
       std::unique_lock<Semaphore> Lock(Barrier, std::try_to_lock);
       if (!Lock.owns_lock()) {
         emitTUStatus({TUAction::Queued, Req.Name});
         Lock.lock();
       }
       WithContext Guard(std::move(Req.Ctx));
       trace::Span Tracer(Req.Name);
       emitTUStatus({TUAction::RunningAction, Req.Name});
       Req.Action();
     }
 
     bool IsEmpty = false;
     {
       std::lock_guard<std::mutex> Lock(Mutex);
-      Requests.pop_front();
+      RunningTask = false;
       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() && !RunningTask; });
 }
 
 // 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

 }
 
 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;
+  updateWithDiags(S, Path, "b", WantDiagnostics::Auto,
+                  [&](std::vector<Diag>) { ++Builds; });
+  ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
+  S.runWithAST(
+      "invalidatable-but-running", Path,
+      [&](llvm::Expected<InputsAndAST> AST) {
         Start.wait();
-  });
+        ++Actions;
+        EXPECT_TRUE(bool(AST)) << "Shouldn't be invalidated, because running.";
+      },
+      TUScheduler::InvalidateOnUpdate);
   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));
S.runWithAST(
         ++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)";
+  EXPECT_EQ(5, 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;