[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

105

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.

118

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

103

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

 #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,
+                 llvm::function_ref<void(PublishResults)> Publish) override {
     if (FIndex)
       FIndex->updateMain(Path, AST);
-    if (SemanticHighlighting)
-      DiagConsumer.onHighlightingsReady(Path, getSemanticHighlightings(AST));
+
+    if (!SemanticHighlighting)
+      return;
+    auto Highlightings = getSemanticHighlightings(AST);
+    Publish([&]() {
+      DiagConsumer.onHighlightingsReady(Path, std::move(Highlightings));
+    });
   }
 
   void onDiagnostics(PathRef File, std::vector<Diag> Diags) override {
     DiagConsumer.onDiagnosticsReady(File, std::move(Diags));
   }
 
   void onFileUpdated(PathRef File, const TUStatus &Status) override {
     DiagConsumer.onFileUpdated(File, Status);

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.
struct TUStatus {
   };
   /// Serialize this to an LSP file status item.
   FileStatus render(PathRef File) const;
 
   TUAction Action;
   BuildDetails Details;
 };
 
+
 class ParsingCallbacks {
 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 &) {}
+  /// A function that is run under the critical section guarding against races
+  /// when closing the files.
+  using PublishResults = 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) {}
+  ///
+  /// This callback is run on a worker thread and if we need to send results of
+  /// processing AST to the users, we might run into a race condition with file
+  /// removal. To avoid the race condition, send the results by running \p
+  /// Publish. That would ensure publishing the results is either sequenced
+  /// before removing the file or does not happen at all.
+  virtual void onMainAST(PathRef Path, ParsedAST &AST,
+                         llvm::function_ref<void(PublishResults)> Publish) {}
 
   /// Called whenever the diagnostics for \p File are produced.
   virtual void onDiagnostics(PathRef File, std::vector<Diag> Diags) {}
 
   /// Called whenever the TU status is updated.
   virtual void onFileUpdated(PathRef File, const TUStatus &Status) {}
 };
 

clang-tools-extra/clangd/TUScheduler.cpp

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) {
-      {
+      // Results from the worker thread should only be published if the file was
+      // not removed. This callback ensures the passed operation is not racing
+      // with file removal.
+      auto Publish = [&](ParsingCallbacks::PublishResults Action) {
         std::lock_guard<std::mutex> Lock(DiagsMu);
         if (ReportDiagnostics)
-          Callbacks.onDiagnostics(FileName, (*AST)->getDiagnostics());
-      }
+          Action();
+      };
+
+      // Publish diagnostics.
+      Publish([&]() {
+      });
+
       trace::Span Span("Running main AST callback");
-      Callbacks.onMainAST(FileName, **AST);
+      Callbacks.onMainAST(FileName, **AST, Publish);
       DiagsWereReported = true;
     }
     // Stash the AST in the cache for further use.
     IdleASTs.put(this, std::move(*AST));
   };
   startTask(TaskName, std::move(Task), WantDiags);
 }