[clangd] Fixed a data race.

[ilya-biryukov]

Summary

Calling addDocument after removeDocument could have resulted in an
invalid program state (AST and Preamble for the valid document could
have been incorrectly removed).
This commit also includes an improved CppFile::cancelRebuild
implementation that allows to cancel reparse without waiting for
ongoing rebuild to finish.

[klimek]

Tests?

[ilya-biryukov]

In D36397#833883, @klimek wrote:

Tests?

TSAN does not catch this (as it's a logical error) and it requires a rather bizarre timing of file reparses to trigger.
I couldn't come up with an example that would reproduce this.

[klimek]

In D36397#833890, @ilya-biryukov wrote:

In D36397#833883, @klimek wrote:

Tests?

TSAN does not catch this (as it's a logical error) and it requires a rather bizarre timing of file reparses to trigger.
I couldn't come up with an example that would reproduce this.

Yea, we'll probably want to add a "smash it hard with parallel" requests kind of test to catch these things. You're right, there is probably not a better solution for this specific bug.

[ilya-biryukov]

In D36397#833912, @klimek wrote:

Yea, we'll probably want to add a "smash it hard with parallel" requests kind of test to catch these things. You're right, there is probably not a better solution for this specific bug.

Added a relevant test to https://reviews.llvm.org/D36261. It actually found this problem (after repeatedly running it for some time).

[klimek]

clangd/ClangdUnit.cpp

714

Somewhat orthogonal question - why CppFile? We do support other languages in Clangd, right? (C / Obj-C)
I'd have expected this to be called ClangdUnit, especially given the file name.

735

Now I know why I don't like the shard_ptr, this is all pretty awkward.

738

Why do we need that instead of just calling the bound shared_ptr This?

[ilya-biryukov]

clangd/ClangdUnit.cpp

714

Just wanted to give it a different name when the class changed semantics. Didn't rename the file yet, though, because was afraid that git would not detect renames and we'll lose history.
ClangdUnit used to provide implementations for all features (codeComplete, etc.), CppFile manages AST and Preamble., i.e. doing a different thing.
But I agree, CppFile is a probably a bad name, but I would still go with something not ending with Unit to not bring up parallels with ASTUnit, that provides a very different interface.

How about ClangdFile, ClangdFileAST?

735

Do you mean shared_from_this() specifically?

738

This lambda is a deferred computation and might get executed on a different thread, so we want to keep a shared_ptr to our class around (weak_ptr would also work) to make sure this class does not get deleted before this lambda is called.
This local var is for capture in a lambda that is passed to wait in the code below (we could've captured That, but that would be an extra shared_ptr copy).

[ilya-biryukov]

A follow-up after discussion offline: we should make an attempt to simplify the threading code in ClangdUnit.h by moving the async worker loop of CppFile into the class itself, rather than relying on external scheduling to do the right thing.
@klimek, are there any other issues we need to address in this specific patch or are we good to go with submitting it?

[bkramer]

I think this can be committed now. It's still a bit awkward but we can address that later.

[ilya-biryukov]

Thanks. I'll make sure to address the awkwardness in further commits.

clangd/ClangdServer.h

public:
   std::string getDocument(PathRef File);
 
   /// Only for testing purposes.
   /// Waits until all requests to worker thread are finished and dumps AST for
   /// \p File. \p File must be in the list of added documents.
   std::string dumpAST(PathRef File);
 
 private:
+  std::future<void>
+  scheduleReparseAndDiags(PathRef File, VersionedDraft Contents,
+                          std::shared_ptr<CppFile> Resources,
+                          Tagged<IntrusiveRefCntPtr<vfs::FileSystem>> TaggedFS);
+
+  std::future<void> scheduleCancelRebuild(std::shared_ptr<CppFile> Resources);
+
   GlobalCompilationDatabase &CDB;
   DiagnosticsConsumer &DiagConsumer;
   FileSystemProvider &FSProvider;
   DraftStore DraftMgr;
   CppFileCollection Units;
   std::string ResourceDir;
   std::shared_ptr<PCHContainerOperations> PCHs;
   // WorkScheduler has to be the last member, because its destructor has to be

clangd/ClangdServer.cpp

: CDB(CDB), DiagConsumer(DiagConsumer), FSProvider(FSProvider),
       WorkScheduler(RunSynchronously) {}
 
 std::future<void> ClangdServer::addDocument(PathRef File, StringRef Contents) {
   DocVersion Version = DraftMgr.updateDraft(File, Contents);
 
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
   std::shared_ptr<CppFile> Resources =
       Units.getOrCreateFile(File, ResourceDir, CDB, PCHs, TaggedFS.Value);
-
-  std::future<llvm::Optional<std::vector<DiagWithFixIts>>> DeferredRebuild =
-      Resources->deferRebuild(Contents, TaggedFS.Value);
-  std::promise<void> DonePromise;
-  std::future<void> DoneFuture = DonePromise.get_future();
-
-  Path FileStr = File;
-  VFSTag Tag = TaggedFS.Tag;
-  auto ReparseAndPublishDiags =
-      [this, FileStr, Version,
-       Tag](std::future<llvm::Optional<std::vector<DiagWithFixIts>>>
-                DeferredRebuild,
-            std::promise<void> DonePromise) -> void {
-    FulfillPromiseGuard Guard(DonePromise);
-
-    auto CurrentVersion = DraftMgr.getVersion(FileStr);
-    if (CurrentVersion != Version)
-      return; // This request is outdated
-
-    auto Diags = DeferredRebuild.get();
-    if (!Diags)
-      return; // A new reparse was requested before this one completed.
-    DiagConsumer.onDiagnosticsReady(FileStr,
-                                    make_tagged(std::move(*Diags), Tag));
-  };
-
-  WorkScheduler.addToFront(std::move(ReparseAndPublishDiags),
-                           std::move(DeferredRebuild), std::move(DonePromise));
-  return DoneFuture;
+  return scheduleReparseAndDiags(File, VersionedDraft{Version, Contents.str()},
+                                 std::move(Resources), std::move(TaggedFS));
 }
 
 std::future<void> ClangdServer::removeDocument(PathRef File) {
-  auto Version = DraftMgr.removeDraft(File);
-  Path FileStr = File;
-
-  std::promise<void> DonePromise;
-  std::future<void> DoneFuture = DonePromise.get_future();
-
-  auto RemoveDocFromCollection = [this, FileStr,
-                                  Version](std::promise<void> DonePromise) {
-    FulfillPromiseGuard Guard(DonePromise);
-
-    if (Version != DraftMgr.getVersion(FileStr))
-      return; // This request is outdated, do nothing
-
-    std::shared_ptr<CppFile> File = Units.removeIfPresent(FileStr);
-    if (!File)
-      return;
-    // Cancel all ongoing rebuilds, so that we don't do extra work before
-    // deleting this file.
-    File->cancelRebuilds();
-  };
-  WorkScheduler.addToFront(std::move(RemoveDocFromCollection),
-                           std::move(DonePromise));
-  return DoneFuture;
+  DraftMgr.removeDraft(File);
+  std::shared_ptr<CppFile> Resources = Units.removeIfPresent(File);
+  return scheduleCancelRebuild(std::move(Resources));
 }
 
 std::future<void> ClangdServer::forceReparse(PathRef File) {
   // The addDocument schedules the reparse even if the contents of the file
   // never changed, so we just call it here.
   return addDocument(File, getDocument(File));
 }
 
Tagged<std::vector<Location>> ClangdServer::findDefinitions(PathRef File,
   std::vector<Location> Result;
   Resources->getAST().get()->runUnderLock([Pos, &Result](ParsedAST *AST) {
     if (!AST)
       return;
     Result = clangd::findDefinitions(*AST, Pos);
   });
   return make_tagged(std::move(Result), TaggedFS.Tag);
 }
+
+std::future<void> ClangdServer::scheduleReparseAndDiags(
+    PathRef File, VersionedDraft Contents, std::shared_ptr<CppFile> Resources,
+    Tagged<IntrusiveRefCntPtr<vfs::FileSystem>> TaggedFS) {
+
+  assert(Contents.Draft && "Draft must have contents");
+  std::future<llvm::Optional<std::vector<DiagWithFixIts>>> DeferredRebuild =
+      Resources->deferRebuild(*Contents.Draft, TaggedFS.Value);
+  std::promise<void> DonePromise;
+  std::future<void> DoneFuture = DonePromise.get_future();
+
+  DocVersion Version = Contents.Version;
+  Path FileStr = File;
+  VFSTag Tag = TaggedFS.Tag;
+  auto ReparseAndPublishDiags =
+      [this, FileStr, Version,
+       Tag](std::future<llvm::Optional<std::vector<DiagWithFixIts>>>
+                DeferredRebuild,
+            std::promise<void> DonePromise) -> void {
+    FulfillPromiseGuard Guard(DonePromise);
+
+    auto CurrentVersion = DraftMgr.getVersion(FileStr);
+    if (CurrentVersion != Version)
+      return; // This request is outdated
+
+    auto Diags = DeferredRebuild.get();
+    if (!Diags)
+      return; // A new reparse was requested before this one completed.
+    DiagConsumer.onDiagnosticsReady(FileStr,
+                                    make_tagged(std::move(*Diags), Tag));
+  };
+
+  WorkScheduler.addToFront(std::move(ReparseAndPublishDiags),
+                           std::move(DeferredRebuild), std::move(DonePromise));
+  return DoneFuture;
+}
+
+std::future<void>
+ClangdServer::scheduleCancelRebuild(std::shared_ptr<CppFile> Resources) {
+  std::promise<void> DonePromise;
+  std::future<void> DoneFuture = DonePromise.get_future();
+  if (!Resources) {
+    // No need to schedule any cleanup.
+    DonePromise.set_value();
+    return DoneFuture;
+  }
+
+  std::future<void> DeferredCancel = Resources->deferCancelRebuild();
+  auto CancelReparses = [Resources](std::promise<void> DonePromise,
+                                    std::future<void> DeferredCancel) {
+    FulfillPromiseGuard Guard(DonePromise);
+    DeferredCancel.get();
+  };
+  WorkScheduler.addToFront(std::move(CancelReparses), std::move(DonePromise),
+                           std::move(DeferredCancel));
+  return DoneFuture;
+}

clangd/ClangdUnit.h

 private:
   CppFile(PathRef FileName, tooling::CompileCommand Command,
           std::shared_ptr<PCHContainerOperations> PCHs);
 
 public:
   CppFile(CppFile const &) = delete;
   CppFile(CppFile &&) = delete;
 
-  /// Cancels all scheduled rebuilds and sets AST and Preamble to nulls.
+  /// Cancels a scheduled rebuild, if any, and sets AST and Preamble to nulls.
   /// If a rebuild is in progress, will wait for it to finish.
-  void cancelRebuilds();
+  void cancelRebuild();
+
+  /// Similar to deferRebuild, but sets both Preamble and AST to nulls instead
+  /// of doing an actual parsing. Returned future is a deferred computation that
+  /// will wait for any ongoing rebuilds to finish and actually set the AST and
+  /// Preamble to nulls. It can be run on a different thread.
+  /// This function is useful to cancel ongoing rebuilds, if any, before
+  /// removing CppFile.
+  std::future<void> deferCancelRebuild();
 
   /// Rebuild AST and Preamble synchronously on the calling thread.
   /// Returns a list of diagnostics or a llvm::None, if another rebuild was
   /// requested in parallel (effectively cancelling this rebuild) before
   /// diagnostics were produced.
   llvm::Optional<std::vector<DiagWithFixIts>>
   rebuild(StringRef NewContents, IntrusiveRefCntPtr<vfs::FileSystem> VFS);
 

clangd/ClangdUnit.cpp

CppFile::CppFile(PathRef FileName, tooling::CompileCommand Command,
   LatestAvailablePreamble = nullptr;
   PreamblePromise.set_value(nullptr);
   PreambleFuture = PreamblePromise.get_future();
 
   ASTPromise.set_value(std::make_shared<ParsedASTWrapper>(llvm::None));
   ASTFuture = ASTPromise.get_future();
 }
 
-void CppFile::cancelRebuilds() {
+void CppFile::cancelRebuild() { deferCancelRebuild().get(); }
+
+std::future<void> CppFile::deferCancelRebuild() {
   std::unique_lock<std::mutex> Lock(Mutex);
   // Cancel an ongoing rebuild, if any, and wait for it to finish.
-  ++this->RebuildCounter;
+  unsigned RequestRebuildCounter = ++this->RebuildCounter;
   // Rebuild asserts that futures aren't ready if rebuild is cancelled.
   // We want to keep this invariant.
   if (futureIsReady(PreambleFuture)) {
     PreamblePromise = std::promise<std::shared_ptr<const PreambleData>>();
     PreambleFuture = PreamblePromise.get_future();
   }
   if (futureIsReady(ASTFuture)) {
     ASTPromise = std::promise<std::shared_ptr<ParsedASTWrapper>>();
     ASTFuture = ASTPromise.get_future();
   }
-  // Now wait for rebuild to finish.
-  RebuildCond.wait(Lock, [this]() { return !this->RebuildInProgress; });
 
-  // Return empty results for futures.
-  PreamblePromise.set_value(nullptr);
-  ASTPromise.set_value(std::make_shared<ParsedASTWrapper>(llvm::None));
+  Lock.unlock();
+  // Notify about changes to RebuildCounter.
+  RebuildCond.notify_all();
+
+  std::shared_ptr<CppFile> That = shared_from_this();
+  return std::async(std::launch::deferred, [That, RequestRebuildCounter]() {
+    std::unique_lock<std::mutex> Lock(That->Mutex);
+    CppFile *This = &*That;
+    This->RebuildCond.wait(Lock, [This, RequestRebuildCounter]() {
+      return !This->RebuildInProgress ||
+             This->RebuildCounter != RequestRebuildCounter;
+    });
+
+    // This computation got cancelled itself, do nothing.
+    if (This->RebuildCounter != RequestRebuildCounter)
+      return;
+
+    // Set empty results for Promises.
+    That->PreamblePromise.set_value(nullptr);
+    That->ASTPromise.set_value(std::make_shared<ParsedASTWrapper>(llvm::None));
+  });
 }
 
 llvm::Optional<std::vector<DiagWithFixIts>>
 CppFile::rebuild(StringRef NewContents,
                  IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
   return deferRebuild(NewContents, std::move(VFS)).get();
 }
 
if (futureIsReady(this->PreambleFuture)) {
           std::promise<std::shared_ptr<const PreambleData>>();
       this->PreambleFuture = this->PreamblePromise.get_future();
     }
     if (futureIsReady(this->ASTFuture)) {
       this->ASTPromise = std::promise<std::shared_ptr<ParsedASTWrapper>>();
       this->ASTFuture = this->ASTPromise.get_future();
     }
   } // unlock Mutex.
+  // Notify about changes to RebuildCounter.
+  RebuildCond.notify_all();
 
   // A helper to function to finish the rebuild. May be run on a different
   // thread.
 
   // Don't let this CppFile die before rebuild is finished.
   std::shared_ptr<CppFile> That = shared_from_this();
   auto FinishRebuild = [OldPreamble, VFS, RequestRebuildCounter, PCHs,
                         That](std::string NewContents)

 CppFile::RebuildGuard::RebuildGuard(CppFile &File,
                                     unsigned RequestRebuildCounter)
     : File(File), RequestRebuildCounter(RequestRebuildCounter) {
   std::unique_lock<std::mutex> Lock(File.Mutex);
   WasCancelledBeforeConstruction = File.RebuildCounter != RequestRebuildCounter;
   if (WasCancelledBeforeConstruction)
     return;
 
-  File.RebuildCond.wait(Lock, [&File]() { return !File.RebuildInProgress; });
+  File.RebuildCond.wait(Lock, [&File, RequestRebuildCounter]() {
+    return !File.RebuildInProgress ||
+           File.RebuildCounter != RequestRebuildCounter;
+  });
 
   WasCancelledBeforeConstruction = File.RebuildCounter != RequestRebuildCounter;
   if (WasCancelledBeforeConstruction)
     return;
 
   File.RebuildInProgress = true;
 }