[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

242

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

[sammccall]

clangd/TUScheduler.cpp

436–440

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.

442

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.

533

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

436–440

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.

442

Added a comment.

[sammccall]

clangd/ClangdServer.h

128–129

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

257

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.

543

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

128–129

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.
+  /// resources associated with it. Pending diagnostics for closed files may not
+  /// be delivered, even if requested with WantDiags::Auto or WantDiags::Yes.
   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.h

public:
   /// If diagnostics are requested (Yes), and the context is cancelled before
   /// they are prepared, they may be skipped if eventual-consistency permits it
   /// (i.e. WantDiagnostics is downgraded to Auto).
   /// FIXME(ibiryukov): remove the callback from this function.
   void update(PathRef File, ParseInputs Inputs, WantDiagnostics WD,
               llvm::unique_function<void(std::vector<Diag>)> OnUpdated);
 
   /// Remove \p File from the list of tracked files and schedule removal of its
-  /// resources.
+  /// resources. Pending diagnostics for closed files may not be delivered, even
+  /// if requested with WantDiags::Auto or WantDiags::Yes.
   void remove(PathRef File);
 
   /// Schedule an async task with no dependencies.
   void run(llvm::StringRef Name, llvm::unique_function<void()> Action);
 
   /// Schedule an async read of the AST. \p Action will be called when AST is
   /// ready. The AST passed to \p Action refers to the version of \p File
   /// tracked at the time of the call, even if new updates are received before

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;
+  // 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.
+  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) {
         return;
       }
     }
 
     // We only need to build the AST if diagnostics were requested.
     if (WantDiags == WantDiagnostics::No)
       return;
 
+    {
+      std::lock_guard<std::mutex> Lock(DiagsMu);
+      // No need to rebuild the AST if we won't send the diagnotics. However,
+      // note that we don't prevent preamble rebuilds.
+      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(DiagsMu);
+    ReportDiagnostics = false;
+  }
+  {
     std::lock_guard<std::mutex> Lock(Mutex);
     assert(!Done && "stop() called twice");
     Done = true;
   }
   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)
std::vector<std::string> DiagsSeen, ReadsSeen, ReadsCanceled;
     Update("U1")();
     Read("R1")();
     Update("U2")();
     Read("R2A")();
     Read("R2B");
     Update("U3");
     Read("R3")();
     Proceed.notify();
+
+    ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   }
   EXPECT_THAT(DiagsSeen, ElementsAre("U2", "U3"))
       << "U1 and all dependent reads were cancelled. "
          "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"))
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);
 }