[clangd] Fixed a data race.

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.

Reviewers: krasimir, bkramer, klimek

Reviewed By: bkramer

Subscribers: cfe-commits

Differential Revision: https://reviews.llvm.org/D36397

llvm-svn: 310818

4 files changed, 110 insertions, 62 deletions

diff --git a/clang-tools-extra/clangd/ClangdServer.cpp b/clang-tools-extra/clangd/ClangdServer.cpp
index eb4995be0cd..3b7c662e2e9 100644
--- a/clang-tools-extra/clangd/ClangdServer.cpp
+++ b/clang-tools-extra/clangd/ClangdServer.cpp
@@ -143,61 +143,14 @@ std::future<void> ClangdServer::addDocument(PathRef File, StringRef 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) {
@@ -306,3 +259,60 @@ Tagged<std::vector<Location>> ClangdServer::findDefinitions(PathRef File,
   });
   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;
+}
diff --git a/clang-tools-extra/clangd/ClangdServer.h b/clang-tools-extra/clangd/ClangdServer.h
index 00b55a66321..b5a75b3bd0e 100644
--- a/clang-tools-extra/clangd/ClangdServer.h
+++ b/clang-tools-extra/clangd/ClangdServer.h
@@ -233,6 +233,13 @@ public:
   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;
diff --git a/clang-tools-extra/clangd/ClangdUnit.cpp b/clang-tools-extra/clangd/ClangdUnit.cpp
index 5ec9045dad0..befdf75f807 100644
--- a/clang-tools-extra/clangd/ClangdUnit.cpp
+++ b/clang-tools-extra/clangd/ClangdUnit.cpp
@@ -711,10 +711,12 @@ CppFile::CppFile(PathRef FileName, tooling::CompileCommand Command,
   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)) {
@@ -725,12 +727,28 @@ void CppFile::cancelRebuilds() {
     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>>
@@ -767,6 +785,8 @@ CppFile::deferRebuild(StringRef NewContents,
       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.
@@ -916,7 +936,10 @@ CppFile::RebuildGuard::RebuildGuard(CppFile &File,
   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)
diff --git a/clang-tools-extra/clangd/ClangdUnit.h b/clang-tools-extra/clangd/ClangdUnit.h
index 8d7639e0594..9fa60b0f6e1 100644
--- a/clang-tools-extra/clangd/ClangdUnit.h
+++ b/clang-tools-extra/clangd/ClangdUnit.h
@@ -151,9 +151,17 @@ 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