[clangd] Provide a way to publish highlightings in non-racy manner

[ilya-biryukov]

Summary

By exposing a callback that can guard code publishing results of
'onMainAST' callback in the same manner we guard diagnostics.

[sammccall]

I think this is the right design. As mentioned offline, I think we can now move the onDiagnostics call out from TUScheduler into ClangdServer (and remove onDiagnostics from the callbacks interface).
This is a better layer because the diagnostics callback model is really about how LSP treats diagnostics, and TUScheduler shouldn't have to care much (maybe one day this will even live in ClangdLSPServer).

clang-tools-extra/clangd/TUScheduler.h

104

I'd consider having the typedef be the full using PublishFn = llvm::function_ref<void(llvm::function_ref<void()>)>

That way the signature of onMainAST is simpler, and realistically people are going to understand how to call Publish by example.

116

I think we need to explain the API before apologising for it (explaining why it's this shape etc). And I think this is weird enough that an example would help.

Something like:

When information about the file (diagnostics, syntax highlighting) is published to clients,
this should be wrapped in Publish, e.g.
  void onMainAST(...) {
    Highlights = computeHighlights();
    Publish([&] { notifyHighlights(Path, Highlights); });
  }
This guarantees that clients will see results in the correct sequence if the file is concurrently closed
and/or reopened. (The lambda passed to Publish() may never run in this case).

[ilya-biryukov]

• Use the same mechanism for diagnostics
• Change typedef to function<void(function<void()>)>
• Update a comment
• s/PublishResults/PublishFn
• Reformat

[ilya-biryukov]

• Update usage of DiagsMu in a comment

[ilya-biryukov]

• Remove a leftover comment from the previous version

[ilya-biryukov]

Thanks for all the suggestions. This is ready for the next round now.

[sammccall]

Awesome, I think this is in better shape than before highlighting was added.

clang-tools-extra/clangd/TUScheduler.h

102

uber-nit: newline before this, as currently it's grouped with onPreambleAST rather than onMainAST

[ilya-biryukov]

• Group PublishFn with onMainAST

clang-tools-extra/clangd/ClangdServer.cpp

 //===--- ClangdServer.cpp - Main clangd server code --------------*- 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 "ClangdServer.h"
 #include "ClangdUnit.h"
 #include "CodeComplete.h"
 #include "FindSymbols.h"
 #include "Format.h"
 #include "FormattedString.h"
 #include "Headers.h"
 #include "Protocol.h"
+#include "SemanticHighlighting.h"
 #include "SourceCode.h"
 #include "TUScheduler.h"
 #include "Trace.h"
 #include "XRefs.h"
 #include "index/CanonicalIncludes.h"
 #include "index/FileIndex.h"
 #include "index/Merge.h"
 #include "refactor/Rename.h"
 #include "refactor/Tweak.h"
 #include "clang/Format/Format.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/Optional.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
 #include <future>
struct UpdateIndexCallbacks : public ParsingCallbacks {
 
   void onPreambleAST(PathRef Path, ASTContext &Ctx,
                      std::shared_ptr<clang::Preprocessor> PP,
                      const CanonicalIncludes &CanonIncludes) override {
     if (FIndex)
       FIndex->updatePreamble(Path, Ctx, std::move(PP), CanonIncludes);
   }
 
-  void onMainAST(PathRef Path, ParsedAST &AST) override {
+  void onMainAST(PathRef Path, ParsedAST &AST, PublishFn Publish) override {
     if (FIndex)
       FIndex->updateMain(Path, AST);
+
+    std::vector<Diag> Diagnostics = AST.getDiagnostics();
+    std::vector<HighlightingToken> Highlightings;
     if (SemanticHighlighting)
-      DiagConsumer.onHighlightingsReady(Path, getSemanticHighlightings(AST));
-  }
+      Highlightings = getSemanticHighlightings(AST);
 
-  void onDiagnostics(PathRef File, std::vector<Diag> Diags) override {
-    DiagConsumer.onDiagnosticsReady(File, std::move(Diags));
+    Publish([&]() {
+      DiagConsumer.onDiagnosticsReady(Path, std::move(Diagnostics));
+      if (SemanticHighlighting)
+        DiagConsumer.onHighlightingsReady(Path, std::move(Highlightings));
+    });
   }
 
   void onFileUpdated(PathRef File, const TUStatus &Status) override {
     DiagConsumer.onFileUpdated(File, Status);
   }
 
 private:
   FileIndex *FIndex;

clang-tools-extra/clangd/TUScheduler.h

 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_TUSCHEDULER_H
 
 #include "ClangdUnit.h"
 #include "Function.h"
 #include "GlobalCompilationDatabase.h"
 #include "Threading.h"
 #include "index/CanonicalIncludes.h"
 #include "llvm/ADT/Optional.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include <future>
 
 namespace clang {
 namespace clangd {
 
 /// Returns a number of a default async threads to use for TUScheduler.
public:
   virtual ~ParsingCallbacks() = default;
 
   /// Called on the AST that was built for emitting the preamble. The built AST
   /// contains only AST nodes from the #include directives at the start of the
   /// file. AST node in the current file should be observed on onMainAST call.
   virtual void onPreambleAST(PathRef Path, ASTContext &Ctx,
                              std::shared_ptr<clang::Preprocessor> PP,
                              const CanonicalIncludes &) {}
+  /// The argument function is run under the critical section guarding against
+  /// races when closing the files.
+  using PublishFn = llvm::function_ref<void(llvm::function_ref<void()>)>;
+
   /// Called on the AST built for the file itself. Note that preamble AST nodes
   /// are not deserialized and should be processed in the onPreambleAST call
   /// instead.
   /// The \p AST always contains all AST nodes for the main file itself, and
   /// only a portion of the AST nodes deserialized from the preamble. Note that
   /// some nodes from the preamble may have been deserialized and may also be
   /// accessed from the main file AST, e.g. redecls of functions from preamble,
   /// etc. Clients are expected to process only the AST nodes from the main file
   /// in this callback (obtained via ParsedAST::getLocalTopLevelDecls) to obtain
   /// optimal performance.
-  virtual void onMainAST(PathRef Path, ParsedAST &AST) {}
-
-  /// Called whenever the diagnostics for \p File are produced.
-  virtual void onDiagnostics(PathRef File, std::vector<Diag> Diags) {}
+  ///
+  /// When information about the file (diagnostics, syntax highlighting) is
+  /// published to clients, this should be wrapped in Publish, e.g.
+  ///   void onMainAST(...) {
+  ///     Highlights = computeHighlights();
+  ///     Publish([&] { notifyHighlights(Path, Highlights); });
+  ///   }
+  /// This guarantees that clients will see results in the correct sequence if
+  /// the file is concurrently closed and/or reopened. (The lambda passed to
+  /// Publish() may never run in this case).
+  virtual void onMainAST(PathRef Path, ParsedAST &AST, PublishFn Publish) {}
 
   /// Called whenever the TU status is updated.
   virtual void onFileUpdated(PathRef File, const TUStatus &Status) {}
 };
 
 /// Handles running tasks for ClangdServer and managing the resources (e.g.,
 /// preambles and ASTs) for opened files.
 /// TUScheduler is not thread-safe, only one thread should be providing updates

clang-tools-extra/clangd/TUScheduler.cpp

private:
   /// Whether to keep the built preambles in memory or on disk.
   const bool StorePreambleInMemory;
   /// Callback invoked when preamble or main file AST is built.
   ParsingCallbacks &Callbacks;
   /// Only accessed by the worker thread.
   TUStatus Status;
 
   Semaphore &Barrier;
-  /// Whether the diagnostics for the current FileInputs were reported to the
-  /// users before.
-  bool DiagsWereReported = false;
+  /// Whether the 'onMainAST' callback ran for the current FileInputs.
+  bool RanASTCallback = false;
   /// Guards members used by both TUScheduler and the worker thread.
   mutable std::mutex Mutex;
   /// File inputs, currently being used by the worker.
   /// Inputs are written and read by the worker thread, compile command can also
   /// be consumed by clients of ASTWorker.
   std::shared_ptr<const ParseInputs> FileInputs;         /* GUARDED_BY(Mutex) */
   std::shared_ptr<const PreambleData> LastBuiltPreamble; /* 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) */
   mutable std::condition_variable RequestsCV;
-  // FIXME: rename it to better fix the current usage, we also use it to guard
-  // emitting TUStatus.
-  /// Guards a critical section for running the diagnostics callbacks.
-  std::mutex DiagsMu;
-  // Used to prevent remove document + leading to out-of-order diagnostics:
+  /// 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
   // don't. When the old handle is destroyed, the old worker will stop reporting
-  // diagnostics.
-  bool ReportDiagnostics = true; /* GUARDED_BY(DiagsMu) */
+  // any results to the user.
+  bool CanPublishResults = true; /* GUARDED_BY(PublishMu) */
 };
 
 /// A smart-pointer-like class that points to an active ASTWorker.
 /// In destructor, signals to the underlying ASTWorker that no new requests will
 /// be sent and the processing loop may exit (after running all pending
 /// requests).
 class ASTWorkerHandle {
   friend class ASTWorker;
else
       Inputs.CompileCommand = CDB.getFallbackCommand(FileName);
     auto PrevInputs = getCurrentFileInputs();
     // Will be used to check if we can avoid rebuilding the AST.
     bool InputsAreTheSame =
         std::tie(PrevInputs->CompileCommand, PrevInputs->Contents) ==
         std::tie(Inputs.CompileCommand, Inputs.Contents);
 
     tooling::CompileCommand OldCommand = PrevInputs->CompileCommand;
-    bool PrevDiagsWereReported = DiagsWereReported;
+    bool RanCallbackForPrevInputs = RanASTCallback;
     {
       std::lock_guard<std::mutex> Lock(Mutex);
       FileInputs = std::make_shared<ParseInputs>(Inputs);
     }
-    DiagsWereReported = false;
+    RanASTCallback = false;
     emitTUStatus({TUAction::BuildingPreamble, TaskName});
     log("Updating file {0} with command {1}\n[{2}]\n{3}", FileName,
         Inputs.CompileCommand.Heuristic,
         Inputs.CompileCommand.Directory,
         llvm::join(Inputs.CompileCommand.CommandLine, " "));
     // Rebuild the preamble and the AST.
     std::unique_ptr<CompilerInvocation> Invocation =
         buildCompilerInvocation(Inputs);
auto Task = [=]() mutable {
     // to the old preamble, so it can be freed if there are no other references
     // to it.
     OldPreamble.reset();
     PreambleWasBuilt.notify();
     emitTUStatus({TUAction::BuildingFile, TaskName});
     if (!CanReuseAST) {
       IdleASTs.take(this); // Remove the old AST if it's still in cache.
     } else {
-      // Since we don't need to rebuild the AST, we might've already reported
-      // the diagnostics for it.
-      if (PrevDiagsWereReported) {
-        DiagsWereReported = true;
+      // We don't need to rebuild the AST, check if we need to run the callback.
+      if (RanCallbackForPrevInputs) {
+        RanASTCallback = true;
         // Take a shortcut and don't report the diagnostics, since they should
         // not changed. All the clients should handle the lack of OnUpdated()
         // call anyway to handle empty result from buildAST.
         // FIXME(ibiryukov): the AST could actually change if non-preamble
         // includes changed, but we choose to ignore it.
         // FIXME(ibiryukov): should we refresh the cache in IdleASTs for the
         // current file at this point?
         log("Skipping rebuild of the AST for {0}, inputs are the same.",
             FileName);
         TUStatus::BuildDetails Details;
         Details.ReuseAST = true;
         emitTUStatus({TUAction::BuildingFile, TaskName}, &Details);
         return;
       }
     }
 
     // We only need to build the AST if diagnostics were requested.
     if (WantDiags == WantDiagnostics::No)
       return;
 
     {
-      std::lock_guard<std::mutex> Lock(DiagsMu);
+      std::lock_guard<std::mutex> Lock(PublishMu);
       // No need to rebuild the AST if we won't send the diagnotics. However,
       // note that we don't prevent preamble rebuilds.
-      if (!ReportDiagnostics)
+      if (!CanPublishResults)
         return;
     }
 
     // Get the AST for diagnostics.
     llvm::Optional<std::unique_ptr<ParsedAST>> AST = IdleASTs.take(this);
     if (!AST) {
       llvm::Optional<ParsedAST> NewAST =
           buildAST(FileName, std::move(Invocation), Inputs, NewPreamble);
if (!AST) {
       Details.ReuseAST = true;
       emitTUStatus({TUAction::BuildingFile, TaskName}, &Details);
     }
     // We want to report the diagnostics even if this update was cancelled.
     // It seems more useful than making the clients wait indefinitely if they
     // spam us with updates.
     // Note *AST can still be null if buildAST fails.
     if (*AST) {
-      {
-        std::lock_guard<std::mutex> Lock(DiagsMu);
-        if (ReportDiagnostics)
-          Callbacks.onDiagnostics(FileName, (*AST)->getDiagnostics());
-      }
       trace::Span Span("Running main AST callback");
-      Callbacks.onMainAST(FileName, **AST);
-      DiagsWereReported = true;
+      auto RunPublish = [&](llvm::function_ref<void()> Publish) {
+        // Ensure we only publish results from the worker if the file was not
+        // removed, making sure there are not race conditions.
+        std::lock_guard<std::mutex> Lock(PublishMu);
+        if (CanPublishResults)
+          Publish();
+      };
+
+      Callbacks.onMainAST(FileName, **AST, RunPublish);
+      RanASTCallback = true;
     }
     // Stash the AST in the cache for further use.
     IdleASTs.put(this, std::move(*AST));
   };
   startTask(TaskName, std::move(Task), WantDiags);
 }
 
 void ASTWorker::runWithAST(
if (auto Preamble = getPossiblyStalePreamble())
     Result += Preamble->Preamble.getSize();
   return Result;
 }
 
 bool ASTWorker::isASTCached() const { return IdleASTs.getUsedBytes(this) != 0; }
 
 void ASTWorker::stop() {
   {
-    std::lock_guard<std::mutex> Lock(DiagsMu);
-    ReportDiagnostics = false;
+    std::lock_guard<std::mutex> Lock(PublishMu);
+    CanPublishResults = false;
   }
   {
     std::lock_guard<std::mutex> Lock(Mutex);
     assert(!Done && "stop() called twice");
     Done = true;
   }
   RequestsCV.notify_all();
 }
void ASTWorker::startTask(llvm::StringRef Name,
   RequestsCV.notify_all();
 }
 
 void ASTWorker::emitTUStatus(TUAction Action,
                              const TUStatus::BuildDetails *Details) {
   Status.Action = std::move(Action);
   if (Details)
     Status.Details = *Details;
-  std::lock_guard<std::mutex> Lock(DiagsMu);
+  std::lock_guard<std::mutex> Lock(PublishMu);
   // Do not emit TU statuses when the ASTWorker is shutting down.
-  if (ReportDiagnostics) {
+  if (CanPublishResults) {
     Callbacks.onFileUpdated(FileName, Status);
   }
 }
 
 void ASTWorker::run() {
   while (true) {
     Request Req;
     {

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 "ClangdUnit.h"
 #include "Context.h"
+#include "Diagnostics.h"
 #include "Matchers.h"
+#include "Path.h"
 #include "TUScheduler.h"
 #include "TestFS.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <algorithm>
 #include <utility>
 
 namespace clang {
 namespace clangd {
protected:
 
   static Key<llvm::unique_function<void(PathRef File, std::vector<Diag>)>>
       DiagsCallbackKey;
 
   /// A diagnostics callback that should be passed to TUScheduler when it's used
   /// in updateWithDiags.
   static std::unique_ptr<ParsingCallbacks> captureDiags() {
     class CaptureDiags : public ParsingCallbacks {
-      void onDiagnostics(PathRef File, std::vector<Diag> Diags) override {
+      void onMainAST(PathRef File, ParsedAST &AST, PublishFn Publish) override {
+        auto Diags = AST.getDiagnostics();
         auto D = Context::current().get(DiagsCallbackKey);
         if (!D)
           return;
-        const_cast<llvm::unique_function<void(PathRef, std::vector<Diag>)> &> (
-            *D)(File, Diags);
+
+        Publish([&]() {
+          const_cast<
+              llvm::unique_function<void(PathRef, std::vector<Diag>)> &> (*D)(
+              File, std::move(Diags));
+        });
       }
     };
     return llvm::make_unique<CaptureDiags>();
   }
 
   /// Schedule an update and call \p CB with the diagnostics it produces, if
   /// any. The TUScheduler should be created with captureDiags as a
   /// DiagsCallback for this to work.
TEST_F(TUSchedulerTests, MissingFiles) {
 
   auto Missing = testPath("missing.cpp");
   Files[Missing] = "";
 
   EXPECT_EQ(S.getContents(Added), "");
   S.update(Added, getInputs(Added, "x"), WantDiagnostics::No);
   EXPECT_EQ(S.getContents(Added), "x");
 
-  // Assert each operation for missing file is an error (even if it's available
-  // in VFS).
+  // Assert each operation for missing file is an error (even if it's
+  // available in VFS).
   S.runWithAST("", Missing,
                [&](Expected<InputsAndAST> AST) { EXPECT_ERROR(AST); });
   S.runWithPreamble(
       "", Missing, TUScheduler::Stale,
       [&](Expected<InputsAndPreamble> Preamble) { EXPECT_ERROR(Preamble); });
   // remove() shouldn't crash on missing files.
   S.remove(Missing);
 
// Run TUScheduler and collect some stats.
 
     StringRef Contents1 = R"cpp(int a;)cpp";
     StringRef Contents2 = R"cpp(int main() { return 1; })cpp";
     StringRef Contents3 = R"cpp(int a; int b; int sum() { return a + b; })cpp";
 
     StringRef AllContents[] = {Contents1, Contents2, Contents3};
     const int AllContentsSize = 3;
 
-    // Scheduler may run tasks asynchronously, but should propagate the context.
-    // We stash a nonce in the context, and verify it in the task.
+    // Scheduler may run tasks asynchronously, but should propagate the
+    // context. We stash a nonce in the context, and verify it in the task.
     static Key<int> NonceKey;
     int Nonce = 0;
 
     for (int FileI = 0; FileI < FilesCount; ++FileI) {
       for (int UpdateI = 0; UpdateI < UpdatesPerFile; ++UpdateI) {
         auto Contents = AllContents[(FileI + UpdateI) % AllContentsSize];
 
         auto File = Files[FileI];
TEST_F(TUSchedulerTests, EvictedAST) {
 
   // Build one file in advance. We will not access it later, so it will be the
   // one that the cache will evict.
   updateWithCallback(S, Foo, SourceContents, WantDiagnostics::Yes,
                      [&BuiltASTCounter]() { ++BuiltASTCounter; });
   ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   ASSERT_EQ(BuiltASTCounter.load(), 1);
 
-  // Build two more files. Since we can retain only 2 ASTs, these should be the
-  // ones we see in the cache later.
+  // Build two more files. Since we can retain only 2 ASTs, these should be
+  // the ones we see in the cache later.
   updateWithCallback(S, Bar, SourceContents, WantDiagnostics::Yes,
                      [&BuiltASTCounter]() { ++BuiltASTCounter; });
   updateWithCallback(S, Baz, SourceContents, WantDiagnostics::Yes,
                      [&BuiltASTCounter]() { ++BuiltASTCounter; });
   ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   ASSERT_EQ(BuiltASTCounter.load(), 3);
 
   // Check only the last two ASTs are retained.