[clangd] Cleanup: make diagnostics callbacks from TUScheduler non-racy

[ilya-biryukov]

Summary

Previously, removeDoc followed by an addDoc to TUScheduler resulted in
racy diagnostic responses, i.e. the old dianostics could be delivered
to the client after the new ones by TUScheduler.

To workaround this, we tracked a version number in ClangdServer and
discarded stale diagnostics. After this commit, the TUScheduler will
stop delivering diagnostics for removed files and the workaround in
ClangdServer is not required anymore.

[ilya-biryukov]

clangd/TUScheduler.cpp

241

Note to myself: should reuse Done instead by making turning it into atomic<bool>

[sammccall]

clangd/TUScheduler.cpp

432–436

as discussed offline, this doesn't guarantee we're not going to send OnUpdated calls for destroyed FileData objects. Intuitively because ShuttingDown may be set e.g. while building the AST, more subtly because atomic only prevents reordering with respect to other atomic operations.

I believe we need to hold a lock:

{
  mutex_lock Lock(SendDiagnosticsMu);
  if (SendDiagnostics)
    OnUpdated(...);
}

and set SendDiagnostics to false under the same lock.

438

This patch prevents the AST from being indexed after removal - for the same reasons, right?

But we don't prevent onPreambleAST from being called...
Not sure if it makes sense to fix this, but deserves at least a comment I think.

527

I think we need to add public docs to TUScheduler::remove and ClangdServer::removeDocument indicating that diagnostics may not be delivered after documents are removed (regardless of WantDiagnostics).

[ilya-biryukov]

• Ensure consistency for diagnostic callbacks with a critical section.

[ilya-biryukov]

• Added comments

[ilya-biryukov]

clangd/TUScheduler.cpp

432–436

Done. I've made the change before looking at the comment, so using ReportDiagnostics name instead of SendDiagnostics. Happy to change if you like the latter more.

438

Added a comment.

[sammccall]

clangd/ClangdServer.h

132

I'm not sure we actually want to *guarantee* this in the API: this implementation serializes removeDocument and the diagnostics callback, but only because it was the most convenient.

I'd suggest just "pending diagnostics for closed files may not be delivered". (and similarly for TUScheduler).

Up to you though.

clangd/TUScheduler.cpp

256

I think the second sentence is a bit obvious and doesn't get into the subtlety here. Maybe:

Used to prevent remove document + leading to out-of-order diagnostics:
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.

537

Done is covered by RequestsCV, and ReportDiagnostics is not, it seems a little weird to have the diags critical section in between here.

(I don't have an opinion about the relative order of the sections, but notify_all should immediately follow the Mutex section for clarity, I think)

[ilya-biryukov]

• Address the comments

clangd/ClangdServer.h

132

Yeah, it might be a bit stronger than what we guaranteed before (which was "you won't get stale diagnostics for a file that you reopen")

clangd/ClangdServer.h

public:
   /// Add a \p File to the list of tracked C++ files or update the contents if
   /// \p File is already tracked. Also schedules parsing of the AST for it on a
   /// separate thread. When the parsing is complete, DiagConsumer passed in
   /// constructor will receive onDiagnosticsReady callback.
   void addDocument(PathRef File, StringRef Contents,
                    WantDiagnostics WD = WantDiagnostics::Auto);
 
   /// Remove \p File from list of tracked files, schedule a request to free
   /// resources associated with it.
   void removeDocument(PathRef File);
 
   /// Run code completion for \p File at \p Pos.
   /// Request is processed asynchronously.
   ///
   /// This method should only be called for currently tracked files. However, it
   /// is safe to call removeDocument for \p File after this method returns, even
   /// while returned future is not yet ready.

 
 private:
   /// FIXME: This stats several files to find a .clang-format file. I/O can be
   /// slow. Think of a way to cache this.
   llvm::Expected<tooling::Replacements>
   formatCode(llvm::StringRef Code, PathRef File,
              ArrayRef<tooling::Range> Ranges);
 
-  typedef uint64_t DocVersion;
-
-  void consumeDiagnostics(PathRef File, DocVersion Version,
-                          std::vector<Diag> Diags);
-
   tooling::CompileCommand getCompileCommand(PathRef File);
 
   const GlobalCompilationDatabase &CDB;
   DiagnosticsConsumer &DiagConsumer;
   const FileSystemProvider &FSProvider;
 
-  /// Used to synchronize diagnostic responses for added and removed files.
-  llvm::StringMap<DocVersion> InternalVersion;
-
   Path ResourceDir;
   // The index used to look up symbols. This could be:
   //   - null (all index functionality is optional)
   //   - the dynamic index owned by ClangdServer (DynamicIdx)
   //   - the static index passed to the constructor
   //   - a merged view of a static and dynamic index (MergedIndex)
   const SymbolIndex *Index;
   // If present, an index of symbols in open files. Read via *Index.
   std::unique_ptr<FileIndex> DynamicIdx;
   // If present, storage for the merged static/dynamic index. Read via *Index.
   std::unique_ptr<SymbolIndex> MergedIdx;
 
   // GUARDED_BY(CachedCompletionFuzzyFindRequestMutex)
   llvm::StringMap<llvm::Optional<FuzzyFindRequest>>
       CachedCompletionFuzzyFindRequestByFile;
   mutable std::mutex CachedCompletionFuzzyFindRequestMutex;
 
   llvm::Optional<std::string> WorkspaceRoot;
   std::shared_ptr<PCHContainerOperations> PCHs;
-  /// Used to serialize diagnostic callbacks.
-  /// FIXME(ibiryukov): get rid of an extra map and put all version counters
-  /// into CppFile.
-  std::mutex DiagnosticsMutex;
-  /// Maps from a filename to the latest version of reported diagnostics.
-  llvm::StringMap<DocVersion> ReportedDiagnosticVersions;
   // WorkScheduler has to be the last member, because its destructor has to be
   // called before all other members to stop the worker thread that references
   // ClangdServer.
   TUScheduler WorkScheduler;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clangd/ClangdServer.cpp

ClangdServer::ClangdServer(const GlobalCompilationDatabase &CDB,
   else if (Opts.StaticIndex)
     Index = Opts.StaticIndex;
   else
     Index = nullptr;
 }
 
 void ClangdServer::addDocument(PathRef File, StringRef Contents,
                                WantDiagnostics WantDiags) {
-  DocVersion Version = ++InternalVersion[File];
   ParseInputs Inputs = {getCompileCommand(File), FSProvider.getFileSystem(),
                         Contents.str()};
-
   Path FileStr = File.str();
   WorkScheduler.update(File, std::move(Inputs), WantDiags,
-                       [this, FileStr, Version](std::vector<Diag> Diags) {
-                         consumeDiagnostics(FileStr, Version, std::move(Diags));
+                       [this, FileStr](std::vector<Diag> Diags) {
+                         DiagConsumer.onDiagnosticsReady(FileStr,
+                                                         std::move(Diags));
                        });
 }
 
 void ClangdServer::removeDocument(PathRef File) {
-  ++InternalVersion[File];
   WorkScheduler.remove(File);
 }
 
 void ClangdServer::codeComplete(PathRef File, Position Pos,
                                 const clangd::CodeCompleteOptions &Opts,
                                 Callback<CodeCompleteResult> CB) {
   // Copy completion options for passing them to async task handler.
   auto CodeCompleteOpts = Opts;
auto Action = [Pos](Callback<Optional<Hover>> CB,
     if (!InpAST)
       return CB(InpAST.takeError());
     CB(clangd::getHover(InpAST->AST, Pos));
   };
 
   WorkScheduler.runWithAST("Hover", File, Bind(Action, std::move(CB)));
 }
 
-void ClangdServer::consumeDiagnostics(PathRef File, DocVersion Version,
-                                      std::vector<Diag> Diags) {
-  // We need to serialize access to resulting diagnostics to avoid calling
-  // `onDiagnosticsReady` in the wrong order.
-  std::lock_guard<std::mutex> DiagsLock(DiagnosticsMutex);
-  DocVersion &LastReportedDiagsVersion = ReportedDiagnosticVersions[File];
-
-  // FIXME(ibiryukov): get rid of '<' comparison here. In the current
-  // implementation diagnostics will not be reported after version counters'
-  // overflow. This should not happen in practice, since DocVersion is a
-  // 64-bit unsigned integer.
-  if (Version < LastReportedDiagsVersion)
-    return;
-  LastReportedDiagsVersion = Version;
-
-  DiagConsumer.onDiagnosticsReady(File, std::move(Diags));
-}
-
 tooling::CompileCommand ClangdServer::getCompileCommand(PathRef File) {
   trace::Span Span("GetCompileCommand");
   Optional<tooling::CompileCommand> C = CDB.getCompileCommand(File);
   if (!C) // FIXME: Suppress diagnostics? Let the user know?
     C = CDB.getFallbackCommand(File);
 
   // Inject the resource dir.
   // FIXME: Don't overwrite it if it's already there.

clangd/TUScheduler.cpp

private:
   /// Callback, invoked when preamble or main file AST is built.
   ParsingCallbacks &Callbacks;
   /// Helper class required to build the ASTs.
   const std::shared_ptr<PCHContainerOperations> PCHs;
 
   Semaphore &Barrier;
   /// Inputs, corresponding to the current state of AST.
   ParseInputs FileInputs;
   /// Whether the diagnostics for the current FileInputs were reported to the
   /// users before.
   bool DiagsWereReported = false;
   /// Size of the last AST
   /// Guards members used by both TUScheduler and the worker thread.
   mutable std::mutex 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;
+  /// Guards a critical section for running the diagnostics callbacks. 
+  std::mutex DiagsMu;
+  /// When set to false, no diagnostics will be reported. Used to ensure we
+  /// don't report any diagnostics after a shutdown of the worker is requested.
+  bool ReportDiagnostics = true; /* GUARDED_BY(DiagMu) */
 };
 
 /// 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;
if (!CanReuseAST) {
         // 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);
         return;
       }
     }
 
-    // We only need to build the AST if diagnostics were requested.
+    // We only need to build the AST if diagnostics were requested and the
+    // callers are still interested in this ASTWorker (it's not shutting down).
     if (WantDiags == WantDiagnostics::No)
       return;
 
+    {
+      std::lock_guard<std::mutex> Lock(DiagsMu);
+      if (!ReportDiagnostics)
+        return;
+    }
+
     // Get the AST for diagnostics.
     Optional<std::unique_ptr<ParsedAST>> AST = IdleASTs.take(this);
     if (!AST) {
       Optional<ParsedAST> NewAST =
           buildAST(FileName, std::move(Invocation), Inputs, NewPreamble, PCHs);
       AST = NewAST ? llvm::make_unique<ParsedAST>(std::move(*NewAST)) : nullptr;
     }
     // 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)
           OnUpdated((*AST)->getDiagnostics());
+      }
       trace::Span Span("Running main AST callback");
       Callbacks.onMainAST(FileName, **AST);
       DiagsWereReported = true;
     }
     // Stash the AST in the cache for further use.
     IdleASTs.put(this, std::move(*AST));
   };
 
   startTask("Update", Bind(Task, std::move(OnUpdated)), WantDiags);
 }
std::size_t ASTWorker::getUsedBytes() const {
   std::size_t Result = IdleASTs.getUsedBytes(this);
   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(Mutex);
     assert(!Done && "stop() called twice");
     Done = true;
   }
+  {
+    std::lock_guard<std::mutex> Lock(DiagsMu);
+    ReportDiagnostics = false;
+  }
   RequestsCV.notify_all();
 }
 
 void ASTWorker::startTask(StringRef Name, unique_function<void()> Task,
                           Optional<WantDiagnostics> UpdateType) {
   if (RunSync) {
     assert(!Done && "running a task after stop()");
     trace::Span Tracer(Name + ":" + sys::path::filename(FileName));
     Task();

unittests/clangd/ClangdTests.cpp

for (unsigned I = 1; I <= RequestsCount; ++I) {
         unsigned RequestIndex = AsyncRequestIndexDist(RandGen);
         AsyncRequests[RequestIndex]();
       } else {
         // Issue a blocking request once in a while.
         auto RequestIndex = BlockingRequestIndexDist(RandGen);
         BlockingRequests[RequestIndex]();
       }
     }
-  } // Wait for ClangdServer to shutdown before proceeding.
+    ASSERT_TRUE(Server.blockUntilIdleForTest());
+  }
 
   // Check some invariants about the state of the program.
   std::vector<FileStat> Stats = DiagConsumer.takeFileStats();
   for (unsigned I = 0; I < FilesCount; ++I) {
     if (!ReqStats[I].FileIsRemoved) {
       ASSERT_EQ(Stats[I].HadErrorsInLastDiags,
                 ReqStats[I].LastContentsHadErrors);
     }

unittests/clangd/TUSchedulerTests.cpp

S.update(Path, getInputs(Path, "auto (clobbered)"), WantDiagnostics::Auto,
              });
     S.update(Path, getInputs(Path, "request no diags"), WantDiagnostics::No,
              [&](std::vector<Diag> Diags) {
                ADD_FAILURE() << "no diags should not be called back";
              });
     S.update(Path, getInputs(Path, "auto (produces)"), WantDiagnostics::Auto,
              [&](std::vector<Diag> Diags) { ++CallbackCount; });
     Ready.notify();
+
+    ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   }
   EXPECT_EQ(2, CallbackCount);
 }
 
 TEST_F(TUSchedulerTests, Debounce) {
   std::atomic<int> CallbackCount(0);
   {
     TUScheduler S(getDefaultAsyncThreadsCount(),
                   /*StorePreamblesInMemory=*/true, /*ASTCallbacks=*/nullptr,
                   /*UpdateDebounce=*/std::chrono::seconds(1),
                   ASTRetentionPolicy());
     // FIXME: we could probably use timeouts lower than 1 second here.
     auto Path = testPath("foo.cpp");
     S.update(Path, getInputs(Path, "auto (debounced)"), WantDiagnostics::Auto,
              [&](std::vector<Diag> Diags) {
                ADD_FAILURE() << "auto should have been debounced and canceled";
              });
     std::this_thread::sleep_for(std::chrono::milliseconds(200));
     S.update(Path, getInputs(Path, "auto (timed out)"), WantDiagnostics::Auto,
              [&](std::vector<Diag> Diags) { ++CallbackCount; });
     std::this_thread::sleep_for(std::chrono::seconds(2));
     S.update(Path, getInputs(Path, "auto (shut down)"), WantDiagnostics::Auto,
              [&](std::vector<Diag> Diags) { ++CallbackCount; });
+
+    ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   }
   EXPECT_EQ(2, CallbackCount);
 }
 
 static std::vector<std::string> includes(const PreambleData *Preamble) {
   std::vector<std::string> Result;
   if (Preamble)
     for (const auto &Inclusion : Preamble->Includes.MainFileIncludes)
for (int FileI = 0; FileI < FilesCount; ++FileI) {
 
                 std::lock_guard<std::mutex> Lock(Mut);
                 ++TotalPreambleReads;
                 EXPECT_EQ(File, *TUScheduler::getFileBeingProcessedInContext());
               });
         }
       }
     }
+    ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   } // TUScheduler destructor waits for all operations to finish.
 
   std::lock_guard<std::mutex> Lock(Mut);
   EXPECT_EQ(TotalUpdates, FilesCount * UpdatesPerFile);
   EXPECT_EQ(TotalASTReads, FilesCount * UpdatesPerFile);
   EXPECT_EQ(TotalPreambleReads, FilesCount * UpdatesPerFile);
 }