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