[clangd] Rewrote AST and Preamble management.

[ilya-biryukov]

Summary

The new implementation allows code completion that never waits for AST.

[ilya-biryukov]

I also want to rename ClangdUnit and ClangdUnitStore accordingly, but will do that in a separate commit so that git-svn correctly detects the renames (i.e. don't want file contents changes).

[bkramer]

clangd/ClangdServer.h

113

Why the template? This is awkward now because it can only be called from ClangdServer.cpp

clangd/ClangdUnit.cpp

887–889

Move implementation into header

910

Move implementation into header.

clangd/ClangdUnit.h

104–108

Typo: Providse

clangd/ClangdUnitStore.cpp

24–25

Not introduced in this commit, but this is equivalent to OpenedFiles.erase(File) without the find and check.

[ilya-biryukov]

• Fixed a bug that caused CppFiles to be deleted while used.

[ilya-biryukov]

Addressed review comments.

• Moved implementations of template function to header.
• Fixed a typo.

[ilya-biryukov]

clangd/ClangdServer.h

113

Storing std::futures with std::async(std::launch::deferred, ... makes things easier than storing std::function, as that allows to move into the caller(std::function must be copyable).
But we want to store only properly created std::futures, so we create them on our own, hence the signature similar to std::async.

I actually think that the signature itself is not that bad, but the fact that definition is only available in .cpp file is wrong. I've moved definintions to header.

clangd/ClangdUnitStore.cpp

24–25

The code has changed a bit after initial submission. It returns removed element now and there seems to be no equivalent in StringMap for that.

[bkramer]

ship it!

clangd/ClangdServer.h

 
 /// Handles running WorkerRequests of ClangdServer on a separate threads.
 /// Currently runs only one worker thread.
 class ClangdScheduler {
 public:
   ClangdScheduler(bool RunSynchronously);
   ~ClangdScheduler();
 
-  /// Add \p Request to the start of the queue. \p Request will be run on a
-  /// separate worker thread.
-  /// \p Request is scheduled to be executed before all currently added
-  /// requests.
-  void addToFront(std::function<void()> Request);
-  /// Add \p Request to the end of the queue. \p Request will be run on a
-  /// separate worker thread.
-  /// \p Request is scheduled to be executed after all currently added
-  /// requests.
-  void addToEnd(std::function<void()> Request);
+  /// Add a new request to run function \p F with args \p As to the start of the
+  /// queue. The request will be run on a separate thread.
+  template <class Func, class... Args>
+  void addToFront(Func &&F, Args &&... As) {
+    if (RunSynchronously) {
+      std::forward<Func>(F)(std::forward<Args>(As)...);
+      return;
+    }
+
+    {
+      std::lock_guard<std::mutex> Lock(Mutex);
+      RequestQueue.push_front(std::async(std::launch::deferred,
+                                         std::forward<Func>(F),
+                                         std::forward<Args>(As)...));
+    }
+    RequestCV.notify_one();
+  }
+
+  /// Add a new request to run function \p F with args \p As to the end of the
+  /// queue. The request will be run on a separate thread.
+  template <class Func, class... Args> void addToEnd(Func &&F, Args &&... As) {
+    if (RunSynchronously) {
+      std::forward<Func>(F)(std::forward<Args>(As)...);
+      return;
+    }
+
+    {
+      std::lock_guard<std::mutex> Lock(Mutex);
+      RequestQueue.push_back(std::async(std::launch::deferred,
+                                        std::forward<Func>(F),
+                                        std::forward<Args>(As)...));
+    }
+    RequestCV.notify_one();
+  }
 
 private:
   bool RunSynchronously;
   std::mutex Mutex;
-  /// We run some tasks on a separate thread(parsing, ClangdUnit cleanup).
+  /// We run some tasks on a separate threads(parsing, CppFile cleanup).
   /// This thread looks into RequestQueue to find requests to handle and
   /// terminates when Done is set to true.
   std::thread Worker;
   /// Setting Done to true will make the worker thread terminate.
   bool Done = false;
   /// A queue of requests.
   /// FIXME(krasimir): code completion should always have priority over parsing
   /// for diagnostics.
-  std::deque<std::function<void()>> RequestQueue;
+  std::deque<std::future<void>> RequestQueue;
   /// Condition variable to wake up the worker thread.
   std::condition_variable RequestCV;
 };
 
 /// Provides API to manage ASTs for a collection of C++ files and request
 /// various language features(currently, only codeCompletion and async
 /// diagnostics for tracked files).
 class ClangdServer {
ClangdServer(GlobalCompilationDatabase &CDB,
                DiagnosticsConsumer &DiagConsumer,
                FileSystemProvider &FSProvider, bool RunSynchronously,
                llvm::Optional<StringRef> ResourceDir = llvm::None);
 
   /// 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);
+  /// \return A future that will become ready when the rebuild (including
+  /// diagnostics) is finished.
+  std::future<void> addDocument(PathRef File, StringRef Contents);
   /// Remove \p File from list of tracked files, schedule a request to free
   /// resources associated with it.
-  void removeDocument(PathRef File);
+  /// \return A future that will become ready the file is removed and all
+  /// associated reosources are freed.
+  std::future<void> removeDocument(PathRef File);
   /// Force \p File to be reparsed using the latest contents.
-  void forceReparse(PathRef File);
+  std::future<void> forceReparse(PathRef File);
 
   /// Run code completion for \p File at \p Pos. If \p OverridenContents is not
   /// None, they will used only for code completion, i.e. no diagnostics update
   /// will be scheduled and a draft for \p File will not be updated.
   /// If \p OverridenContents is None, contents of the current draft for \p File
   /// will be used.
   /// If \p UsedFS is non-null, it will be overwritten by vfs::FileSystem used
   /// for completion.
public:
   /// \p File. \p File must be in the list of added documents.
   std::string dumpAST(PathRef File);
 
 private:
   GlobalCompilationDatabase &CDB;
   DiagnosticsConsumer &DiagConsumer;
   FileSystemProvider &FSProvider;
   DraftStore DraftMgr;
-  ClangdUnitStore Units;
+  CppFileCollection Units;
   std::string ResourceDir;
   std::shared_ptr<PCHContainerOperations> PCHs;
   // 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
   ClangdScheduler WorkScheduler;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clangd/ClangdServer.cpp

 #include "llvm/Support/raw_ostream.h"
 #include <future>
 
 using namespace clang;
 using namespace clang::clangd;
 
 namespace {
 
+class FulfillPromiseGuard {
+public:
+  FulfillPromiseGuard(std::promise<void> &Promise) : Promise(Promise) {}
+
+  ~FulfillPromiseGuard() { Promise.set_value(); }
+
+private:
+  std::promise<void> &Promise;
+};
+
 std::vector<tooling::Replacement> formatCode(StringRef Code, StringRef Filename,
                                              ArrayRef<tooling::Range> Ranges) {
   // Call clang-format.
   // FIXME: Don't ignore style.
   format::FormatStyle Style = format::getLLVMStyle();
   auto Result = format::reformat(Style, Code, Ranges, Filename);
 
   return std::vector<tooling::Replacement>(Result.begin(), Result.end());
if (RunSynchronously) {
     // Don't start the worker thread if we're running synchronously
     return;
   }
 
   // Initialize Worker in ctor body, rather than init list to avoid potentially
   // using not-yet-initialized members
   Worker = std::thread([this]() {
     while (true) {
-      std::function<void()> Request;
+      std::future<void> Request;
 
       // Pick request from the queue
       {
         std::unique_lock<std::mutex> Lock(Mutex);
         // Wait for more requests.
         RequestCV.wait(Lock, [this] { return !RequestQueue.empty() || Done; });
         if (Done)
           return;
 
         assert(!RequestQueue.empty() && "RequestQueue was empty");
 
         // We process requests starting from the front of the queue. Users of
         // ClangdScheduler have a way to prioritise their requests by putting
         // them to the either side of the queue (using either addToEnd or
         // addToFront).
         Request = std::move(RequestQueue.front());
         RequestQueue.pop_front();
       } // unlock Mutex
 
-      Request();
+      Request.get();
     }
   });
 }
 
 ClangdScheduler::~ClangdScheduler() {
   if (RunSynchronously)
     return; // no worker thread is running in that case
 
   {
     std::lock_guard<std::mutex> Lock(Mutex);
     // Wake up the worker thread
     Done = true;
   } // unlock Mutex
   RequestCV.notify_one();
   Worker.join();
 }
 
-void ClangdScheduler::addToFront(std::function<void()> Request) {
-  if (RunSynchronously) {
-    Request();
-    return;
-  }
-
-  {
-    std::lock_guard<std::mutex> Lock(Mutex);
-    RequestQueue.push_front(Request);
-  }
-  RequestCV.notify_one();
-}
-
-void ClangdScheduler::addToEnd(std::function<void()> Request) {
-  if (RunSynchronously) {
-    Request();
-    return;
-  }
-
-  {
-    std::lock_guard<std::mutex> Lock(Mutex);
-    RequestQueue.push_back(Request);
-  }
-  RequestCV.notify_one();
-}
-
 ClangdServer::ClangdServer(GlobalCompilationDatabase &CDB,
                            DiagnosticsConsumer &DiagConsumer,
                            FileSystemProvider &FSProvider,
                            bool RunSynchronously,
                            llvm::Optional<StringRef> ResourceDir)
     : CDB(CDB), DiagConsumer(DiagConsumer), FSProvider(FSProvider),
       ResourceDir(ResourceDir ? ResourceDir->str() : getStandardResourceDir()),
       PCHs(std::make_shared<PCHContainerOperations>()),
       WorkScheduler(RunSynchronously) {}
 
-void ClangdServer::addDocument(PathRef File, StringRef Contents) {
+std::future<void> ClangdServer::addDocument(PathRef File, StringRef Contents) {
   DocVersion Version = DraftMgr.updateDraft(File, Contents);
-  Path FileStr = File;
-  WorkScheduler.addToFront([this, FileStr, Version]() {
-    auto FileContents = DraftMgr.getDraft(FileStr);
-    if (FileContents.Version != Version)
-      return; // This request is outdated, do nothing
 
-    assert(FileContents.Draft &&
-           "No contents inside a file that was scheduled for reparse");
-    auto TaggedFS = FSProvider.getTaggedFileSystem(FileStr);
-    Units.runOnUnit(
-        FileStr, *FileContents.Draft, ResourceDir, CDB, PCHs, TaggedFS.Value,
-        [&](ClangdUnit const &Unit) {
-          DiagConsumer.onDiagnosticsReady(
-              FileStr, make_tagged(Unit.getLocalDiagnostics(), TaggedFS.Tag));
-        });
-  });
+  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;
 }
 
-void ClangdServer::removeDocument(PathRef File) {
+std::future<void> ClangdServer::removeDocument(PathRef File) {
   auto Version = DraftMgr.removeDraft(File);
   Path FileStr = File;
-  WorkScheduler.addToFront([this, FileStr, Version]() {
+
+  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
 
-    Units.removeUnitIfPresent(FileStr);
-  });
+    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;
 }
 
-void ClangdServer::forceReparse(PathRef File) {
+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.
-  addDocument(File, getDocument(File));
+  return addDocument(File, getDocument(File));
 }
 
 Tagged<std::vector<CompletionItem>>
 ClangdServer::codeComplete(PathRef File, Position Pos,
                            llvm::Optional<StringRef> OverridenContents,
                            IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS) {
   std::string DraftStorage;
   if (!OverridenContents) {
     auto FileContents = DraftMgr.getDraft(File);
     assert(FileContents.Draft &&
            "codeComplete is called for non-added document");
 
     DraftStorage = std::move(*FileContents.Draft);
     OverridenContents = DraftStorage;
   }
 
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
   if (UsedFS)
     *UsedFS = TaggedFS.Value;
 
-  std::vector<CompletionItem> Result;
-  Units.runOnUnitWithoutReparse(File, *OverridenContents, ResourceDir, CDB,
-                                PCHs, TaggedFS.Value, [&](ClangdUnit &Unit) {
-                                  Result = Unit.codeComplete(
-                                      *OverridenContents, Pos, TaggedFS.Value);
-                                });
+  std::shared_ptr<CppFile> Resources = Units.getFile(File);
+  assert(Resources && "Calling completion on non-added file");
+
+  auto Preamble = Resources->getPossiblyStalePreamble();
+  std::vector<CompletionItem> Result =
+      clangd::codeComplete(File, Resources->getCompileCommand(),
+                           Preamble ? &Preamble->Preamble : nullptr,
+                           *OverridenContents, Pos, TaggedFS.Value, PCHs);
   return make_tagged(std::move(Result), TaggedFS.Tag);
 }
 
 std::vector<tooling::Replacement> ClangdServer::formatRange(PathRef File,
                                                             Range Rng) {
   std::string Code = getDocument(File);
 
   size_t Begin = positionToOffset(Code, Rng.start);

 
 std::string ClangdServer::getDocument(PathRef File) {
   auto draft = DraftMgr.getDraft(File);
   assert(draft.Draft && "File is not tracked, cannot get contents");
   return *draft.Draft;
 }
 
 std::string ClangdServer::dumpAST(PathRef File) {
-  std::promise<std::string> DumpPromise;
-  auto DumpFuture = DumpPromise.get_future();
-  auto Version = DraftMgr.getVersion(File);
-
-  WorkScheduler.addToEnd([this, &DumpPromise, File, Version]() {
-    assert(DraftMgr.getVersion(File) == Version && "Version has changed");
-    (void)Version;
+  std::shared_ptr<CppFile> Resources = Units.getFile(File);
+  assert(Resources && "dumpAST is called for non-added document");
 
-    Units.runOnExistingUnit(File, [&DumpPromise](ClangdUnit &Unit) {
-      std::string Result;
-
-      llvm::raw_string_ostream ResultOS(Result);
-      Unit.dumpAST(ResultOS);
-      ResultOS.flush();
-
-      DumpPromise.set_value(std::move(Result));
-    });
+  Resources->getAST().get().runUnderLock([&Result](ParsedAST *AST) {
+    if (AST) {
+      clangd::dumpAST(*AST, ResultOS);
+    } else {
+      ResultOS << "<no-ast>";
+    }
   });
-  return DumpFuture.get();
+  return Result;
 }
 
-Tagged<std::vector<Location>>
-ClangdServer::findDefinitions(PathRef File, Position Pos) {
+Tagged<std::vector<Location>> ClangdServer::findDefinitions(PathRef File,
+                                                            Position Pos) {
   auto FileContents = DraftMgr.getDraft(File);
-  assert(FileContents.Draft && "findDefinitions is called for non-added document");
+  assert(FileContents.Draft &&
+         "findDefinitions is called for non-added document");
 
-  std::vector<Location> Result;
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
-  Units.runOnUnit(File, *FileContents.Draft, ResourceDir, CDB, PCHs,
-      TaggedFS.Value, [&](ClangdUnit &Unit) {
-        Result = Unit.findDefinitions(Pos);
-      });
+
+  std::shared_ptr<CppFile> Resources = Units.getFile(File);
+  assert(Resources && "Calling findDefinitions on non-added 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);
 }

clangd/ClangdUnit.h

 
 #include "Path.h"
 #include "Protocol.h"
 #include "clang/Frontend/ASTUnit.h"
 #include "clang/Frontend/PrecompiledPreamble.h"
 #include "clang/Serialization/ASTBitCodes.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "clang/Tooling/Core/Replacement.h"
+#include <atomic>
+#include <future>
 #include <memory>
+#include <mutex>
 
 namespace llvm {
 class raw_ostream;
 }
 
 namespace clang {
 class ASTUnit;
 class PCHContainerOperations;

 namespace clangd {
 
 /// A diagnostic with its FixIts.
 struct DiagWithFixIts {
   clangd::Diagnostic Diag;
   llvm::SmallVector<tooling::Replacement, 1> FixIts;
 };
 
-/// Stores parsed C++ AST and provides implementations of all operations clangd
-/// would want to perform on parsed C++ files.
-class ClangdUnit {
-public:
-  ClangdUnit(PathRef FileName, StringRef Contents, StringRef ResourceDir,
-             std::shared_ptr<PCHContainerOperations> PCHs,
-             std::vector<tooling::CompileCommand> Commands,
-             IntrusiveRefCntPtr<vfs::FileSystem> VFS);
-
-  /// Reparse with new contents.
-  void reparse(StringRef Contents, IntrusiveRefCntPtr<vfs::FileSystem> VFS);
-
-  /// Get code completions at a specified \p Line and \p Column in \p File.
-  ///
-  /// This function is thread-safe and returns completion items that own the
-  /// data they contain.
-  std::vector<CompletionItem>
-  codeComplete(StringRef Contents, Position Pos,
-               IntrusiveRefCntPtr<vfs::FileSystem> VFS);
-  /// Get definition of symbol at a specified \p Line and \p Column in \p File.
-  std::vector<Location> findDefinitions(Position Pos);
-  /// Returns diagnostics and corresponding FixIts for each diagnostic that are
-  /// located in the current file.
-  std::vector<DiagWithFixIts> getLocalDiagnostics() const;
-
-  /// For testing/debugging purposes. Note that this method deserializes all
-  /// unserialized Decls, so use with care.
-  void dumpAST(llvm::raw_ostream &OS) const;
-
-private:
-  /// Stores and provides access to parsed AST.
-  class ParsedAST {
-  public:
-    /// Attempts to run Clang and store parsed AST. If \p Preamble is non-null
-    /// it is reused during parsing.
-    static llvm::Optional<ParsedAST>
-    Build(std::unique_ptr<clang::CompilerInvocation> CI,
-          const PrecompiledPreamble *Preamble,
-          ArrayRef<serialization::DeclID> PreambleDeclIDs,
-          std::unique_ptr<llvm::MemoryBuffer> Buffer,
-          std::shared_ptr<PCHContainerOperations> PCHs,
-          IntrusiveRefCntPtr<vfs::FileSystem> VFS);
 
-    ParsedAST(ParsedAST &&Other);
-    ParsedAST &operator=(ParsedAST &&Other);
 
-    ~ParsedAST();
 
-    ASTContext &getASTContext();
-    const ASTContext &getASTContext() const;
 
-    Preprocessor &getPreprocessor();
-    const Preprocessor &getPreprocessor() const;
 
-    /// This function returns all top-level decls, including those that come
-    /// from Preamble. Decls, coming from Preamble, have to be deserialized, so
-    /// this call might be expensive.
-    ArrayRef<const Decl *> getTopLevelDecls();
 
-    const std::vector<DiagWithFixIts> &getDiagnostics() const;
 
-  private:
-    ParsedAST(std::unique_ptr<CompilerInstance> Clang,
-              std::unique_ptr<FrontendAction> Action,
-              std::vector<const Decl *> TopLevelDecls,
-              std::vector<serialization::DeclID> PendingTopLevelDecls,
-              std::vector<DiagWithFixIts> Diags);
 
-  private:
-    void ensurePreambleDeclsDeserialized();
 
-    // We store an "incomplete" FrontendAction (i.e. no EndSourceFile was called
-    // on it) and CompilerInstance used to run it. That way we don't have to do
-    // complex memory management of all Clang structures on our own. (They are
-    // stored in CompilerInstance and cleaned up by
-    // FrontendAction.EndSourceFile).
-    std::unique_ptr<CompilerInstance> Clang;
-    std::unique_ptr<FrontendAction> Action;
 
-    // Data, stored after parsing.
-    std::vector<DiagWithFixIts> Diags;
-    std::vector<const Decl *> TopLevelDecls;
-    std::vector<serialization::DeclID> PendingTopLevelDecls;
-  };
 
-  // Store Preamble and all associated data
-  struct PreambleData {
-    PreambleData(PrecompiledPreamble Preamble,
-                 std::vector<serialization::DeclID> TopLevelDeclIDs,
-                 std::vector<DiagWithFixIts> Diags);
+// Provides thread-safe access to ParsedAST.
+class ParsedASTWrapper {
+public:
+  ParsedASTWrapper(ParsedASTWrapper &&Wrapper);
+  ParsedASTWrapper(llvm::Optional<ParsedAST> AST);
+
+  /// Runs \p F on wrapped ParsedAST under lock. Ensures it is not accessed
+  /// concurrently.
+  template <class Func> void runUnderLock(Func F) const {
+    std::lock_guard<std::mutex> Lock(Mutex);
+    F(AST ? AST.getPointer() : nullptr);
+  }
+
+private:
+  // This wrapper is used as an argument to std::shared_future (and it returns a
+  // const ref in get()), but we need to have non-const ref in order to
+  // implement some features.
+  mutable std::mutex Mutex;
+  mutable llvm::Optional<ParsedAST> AST;
+};
+
+// Stores Preamble and associated data.
 
-    PrecompiledPreamble Preamble;
-    std::vector<serialization::DeclID> TopLevelDeclIDs;
-    std::vector<DiagWithFixIts> Diags;
-  };
 
-  SourceLocation getBeginningOfIdentifier(const Position &Pos,
-                                          const FileEntry *FE) const;
+/// Manages resources, required by clangd. Allows to rebuild file with new
+/// contents, and provides AST and Preamble for it.
+class CppFile : public std::enable_shared_from_this<CppFile> {
+public:
+  // We only allow to create CppFile as shared_ptr, because a future returned by
+  // deferRebuild will hold references to it.
+  static std::shared_ptr<CppFile>
+  Create(PathRef FileName, tooling::CompileCommand Command,
+         std::shared_ptr<PCHContainerOperations> PCHs);
+
+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.
+  /// If a rebuild is in progress, will wait for it to finish.
+  void cancelRebuilds();
+
+  /// 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);
+
+  /// Schedule a rebuild and return a deferred computation that will finish the
+  /// rebuild, that can be called on a different thread.
+  /// After calling this method, resources, available via futures returned by
+  /// getPreamble() and getAST(), will be waiting for rebuild to finish. A
+  /// future fininshing rebuild, returned by this function, must be
+  /// computed(i.e. get() should be called on it) in order to make those
+  /// resources ready. If deferRebuild is called again before the rebuild is
+  /// finished (either because returned future had not been called or because it
+  /// had not returned yet), the previous rebuild request is cancelled and the
+  /// resource futures (returned by getPreamble() or getAST()) that were not
+  /// ready will be waiting for the last rebuild to finish instead.
+  /// The future to finish rebuild returns a list of diagnostics built during
+  /// reparse, or None, if another deferRebuild was called before this
+  /// rebuild was finished.
+  std::future<llvm::Optional<std::vector<DiagWithFixIts>>>
+  deferRebuild(StringRef NewContents, IntrusiveRefCntPtr<vfs::FileSystem> VFS);
+
+  /// Returns a future to get the most fresh PreambleData for a file. The
+  /// future will wait until the Preamble is rebuilt.
+  std::shared_future<std::shared_ptr<const PreambleData>> getPreamble() const;
+  /// Return some preamble for a file. It might be stale, but won't wait for
+  /// rebuild to finish.
+  std::shared_ptr<const PreambleData> getPossiblyStalePreamble() const;
+
+  /// Returns a future to get the most fresh AST for a file. Returned AST is
+  /// wrapped to prevent concurrent accesses.
+  std::shared_future<ParsedASTWrapper> getAST() const;
+
+  /// Get CompileCommand used to build this CppFile.
+  tooling::CompileCommand const &getCompileCommand() const;
+
+private:
+  /// A helper guard that manages the state of CppFile during rebuild.
+  class RebuildGuard {
+  public:
+    RebuildGuard(CppFile &File, unsigned RequestRebuildCounter);
+    ~RebuildGuard();
+
+    bool wasCancelledBeforeConstruction() const;
+
+  private:
+    CppFile &File;
+    unsigned RequestRebuildCounter;
+    bool WasCancelledBeforeConstruction;
+  };
 
   Path FileName;
   tooling::CompileCommand Command;
 
-  llvm::Optional<PreambleData> Preamble;
-  llvm::Optional<ParsedAST> Unit;
-
+  /// Mutex protects all fields, declared below it, FileName and Command are not
+  /// mutated.
+  mutable std::mutex Mutex;
+  /// A counter to cancel old rebuilds.
+  unsigned RebuildCounter;
+  /// Used to wait when rebuild is finished before starting another one.
+  bool RebuildInProgress;
+  /// Condition variable to indicate changes to RebuildInProgress.
+  std::condition_variable RebuildCond;
+
+  /// Promise and future for the latests AST. Fulfilled during rebuild.
+  std::promise<ParsedASTWrapper> ASTPromise;
+  std::shared_future<ParsedASTWrapper> ASTFuture;
+
+  /// Promise and future for the latests Preamble. Fulfilled during rebuild.
+  std::promise<std::shared_ptr<const PreambleData>> PreamblePromise;
+  std::shared_future<std::shared_ptr<const PreambleData>> PreambleFuture;
+  /// Latest preamble that was built. May be stale, but always available without
+  /// waiting for rebuild to finish.
+  std::shared_ptr<const PreambleData> LatestAvailablePreamble;
+  /// Utility class, required by clang.
   std::shared_ptr<PCHContainerOperations> PCHs;
 };
 
+/// Get code completions at a specified \p Pos in \p FileName.
+std::vector<CompletionItem>
+codeComplete(PathRef FileName, tooling::CompileCommand Command,
+             PrecompiledPreamble const *Preamble, StringRef Contents,
+             Position Pos, IntrusiveRefCntPtr<vfs::FileSystem> VFS,
+             std::shared_ptr<PCHContainerOperations> PCHs);
+
+/// Get definition of symbol at a specified \p Pos.
+std::vector<Location> findDefinitions(ParsedAST &AST, Position Pos);
+
+/// For testing/debugging purposes. Note that this method deserializes all
+/// unserialized Decls, so use with care.
+void dumpAST(ParsedAST &AST, llvm::raw_ostream &OS);
+
 } // namespace clangd
 } // namespace clang
 #endif

clangd/ClangdUnit.cpp

 #include "clang/Serialization/ASTWriter.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/CrashRecoveryContext.h"
 #include "llvm/Support/Format.h"
 
 #include <algorithm>
+#include <chrono>
 
 using namespace clang::clangd;
 using namespace clang;
 
 namespace {
 
 class DeclTrackingASTConsumer : public ASTConsumer {
 public:
std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
                                                  StringRef InFile) override {
     return llvm::make_unique<DeclTrackingASTConsumer>(/*ref*/ TopLevelDecls);
   }
 
 private:
   std::vector<const Decl *> TopLevelDecls;
 };
 
-class ClangdUnitPreambleCallbacks : public PreambleCallbacks {
+class CppFilePreambleCallbacks : public PreambleCallbacks {
 public:
   std::vector<serialization::DeclID> takeTopLevelDeclIDs() {
     return std::move(TopLevelDeclIDs);
   }
 
   void AfterPCHEmitted(ASTWriter &Writer) override {
     TopLevelDeclIDs.reserve(TopLevelDecls.size());
     for (Decl *D : TopLevelDecls) {
if (!Clang->hasTarget())
     return nullptr;
 
   // RemappedFileBuffers will handle the lifetime of the Buffer pointer,
   // release it.
   Buffer.release();
   return Clang;
 }
 
-} // namespace
-
-ClangdUnit::ClangdUnit(PathRef FileName, StringRef Contents,
-                       StringRef ResourceDir,
-                       std::shared_ptr<PCHContainerOperations> PCHs,
-                       std::vector<tooling::CompileCommand> Commands,
-                       IntrusiveRefCntPtr<vfs::FileSystem> VFS)
-    : FileName(FileName), PCHs(PCHs) {
-  assert(!Commands.empty() && "No compile commands provided");
-
-  // Inject the resource dir.
-  // FIXME: Don't overwrite it if it's already there.
-  Commands.front().CommandLine.push_back("-resource-dir=" +
-                                         std::string(ResourceDir));
-
-  Command = std::move(Commands.front());
-  reparse(Contents, VFS);
-}
-
-void ClangdUnit::reparse(StringRef Contents,
-                         IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
-  std::vector<const char *> ArgStrs;
-  for (const auto &S : Command.CommandLine)
-    ArgStrs.push_back(S.c_str());
-
-  VFS->setCurrentWorkingDirectory(Command.Directory);
-
-  std::unique_ptr<CompilerInvocation> CI;
-  {
-    // FIXME(ibiryukov): store diagnostics from CommandLine when we start
-    // reporting them.
-    EmptyDiagsConsumer CommandLineDiagsConsumer;
-    IntrusiveRefCntPtr<DiagnosticsEngine> CommandLineDiagsEngine =
-        CompilerInstance::createDiagnostics(new DiagnosticOptions,
-                                            &CommandLineDiagsConsumer, false);
-    CI = createCompilerInvocation(ArgStrs, CommandLineDiagsEngine, VFS);
-  }
-  assert(CI && "Couldn't create CompilerInvocation");
-
-  std::unique_ptr<llvm::MemoryBuffer> ContentsBuffer =
-      llvm::MemoryBuffer::getMemBufferCopy(Contents, FileName);
+template <class T> bool futureIsReady(std::shared_future<T> const &Future) {
+  return Future.wait_for(std::chrono::seconds(0)) == std::future_status::ready;
 
-  // Rebuild the preamble if it is missing or can not be reused.
-  auto Bounds =
-      ComputePreambleBounds(*CI->getLangOpts(), ContentsBuffer.get(), 0);
-  if (!Preamble || !Preamble->Preamble.CanReuse(*CI, ContentsBuffer.get(),
-                                                Bounds, VFS.get())) {
-    std::vector<DiagWithFixIts> PreambleDiags;
-    StoreDiagsConsumer PreambleDiagnosticsConsumer(/*ref*/ PreambleDiags);
-    IntrusiveRefCntPtr<DiagnosticsEngine> PreambleDiagsEngine =
-        CompilerInstance::createDiagnostics(
-            &CI->getDiagnosticOpts(), &PreambleDiagnosticsConsumer, false);
-    ClangdUnitPreambleCallbacks SerializedDeclsCollector;
-    auto BuiltPreamble = PrecompiledPreamble::Build(
-        *CI, ContentsBuffer.get(), Bounds, *PreambleDiagsEngine, VFS, PCHs,
-        SerializedDeclsCollector);
-    if (BuiltPreamble)
-      Preamble = PreambleData(std::move(*BuiltPreamble),
-                              SerializedDeclsCollector.takeTopLevelDeclIDs(),
-                              std::move(PreambleDiags));
-  }
-  Unit = ParsedAST::Build(
-      std::move(CI), Preamble ? &Preamble->Preamble : nullptr,
-      Preamble ? llvm::makeArrayRef(Preamble->TopLevelDeclIDs) : llvm::None,
-      std::move(ContentsBuffer), PCHs, VFS);
-}
+} // namespace
 
 namespace {
 
 CompletionItemKind getKind(CXCursorKind K) {
   switch (K) {
   case CXCursor_MacroInstantiation:
   case CXCursor_MacroDefinition:
     return CompletionItemKind::Text;
public:
 
   GlobalCodeCompletionAllocator &getAllocator() override { return *Allocator; }
 
   CodeCompletionTUInfo &getCodeCompletionTUInfo() override { return CCTUInfo; }
 };
 } // namespace
 
 std::vector<CompletionItem>
-ClangdUnit::codeComplete(StringRef Contents, Position Pos,
-                         IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
+clangd::codeComplete(PathRef FileName, tooling::CompileCommand Command,
+                     PrecompiledPreamble const *Preamble, StringRef Contents,
+                     Position Pos, IntrusiveRefCntPtr<vfs::FileSystem> VFS,
+                     std::shared_ptr<PCHContainerOperations> PCHs) {
   std::vector<const char *> ArgStrs;
   for (const auto &S : Command.CommandLine)
     ArgStrs.push_back(S.c_str());
 
   VFS->setCurrentWorkingDirectory(Command.Directory);
 
   std::unique_ptr<CompilerInvocation> CI;
   EmptyDiagsConsumer DummyDiagsConsumer;
   {
     IntrusiveRefCntPtr<DiagnosticsEngine> CommandLineDiagsEngine =
         CompilerInstance::createDiagnostics(new DiagnosticOptions,
                                             &DummyDiagsConsumer, false);
     CI = createCompilerInvocation(ArgStrs, CommandLineDiagsEngine, VFS);
   }
   assert(CI && "Couldn't create CompilerInvocation");
 
   std::unique_ptr<llvm::MemoryBuffer> ContentsBuffer =
       llvm::MemoryBuffer::getMemBufferCopy(Contents, FileName);
 
   // Attempt to reuse the PCH from precompiled preamble, if it was built.
-  const PrecompiledPreamble *PreambleForCompletion = nullptr;
   if (Preamble) {
     auto Bounds =
         ComputePreambleBounds(*CI->getLangOpts(), ContentsBuffer.get(), 0);
-    if (Preamble->Preamble.CanReuse(*CI, ContentsBuffer.get(), Bounds,
-                                    VFS.get()))
-      PreambleForCompletion = &Preamble->Preamble;
+    if (!Preamble->CanReuse(*CI, ContentsBuffer.get(), Bounds, VFS.get()))
+      Preamble = nullptr;
   }
 
-  auto Clang = prepareCompilerInstance(std::move(CI), PreambleForCompletion,
+  auto Clang = prepareCompilerInstance(std::move(CI), Preamble,
                                        std::move(ContentsBuffer), PCHs, VFS,
                                        DummyDiagsConsumer);
   auto &DiagOpts = Clang->getDiagnosticOpts();
   DiagOpts.IgnoreWarnings = true;
 
   auto &FrontendOpts = Clang->getFrontendOpts();
   FrontendOpts.SkipFunctionBodies = true;
 
clangd::codeComplete(PathRef FileName, tooling::CompileCommand Command,
   if (!Action.Execute()) {
     // FIXME(ibiryukov): log errors
   }
   Action.EndSourceFile();
 
   return Items;
 }
 
-std::vector<DiagWithFixIts> ClangdUnit::getLocalDiagnostics() const {
-  if (!Unit)
-    return {}; // Parsing failed.
-
-  std::vector<DiagWithFixIts> Result;
-  auto PreambleDiagsSize = Preamble ? Preamble->Diags.size() : 0;
-  const auto &Diags = Unit->getDiagnostics();
-  Result.reserve(PreambleDiagsSize + Diags.size());
-
-  if (Preamble)
-    Result.insert(Result.end(), Preamble->Diags.begin(), Preamble->Diags.end());
-  Result.insert(Result.end(), Diags.begin(), Diags.end());
-  return Result;
-}
-
-void ClangdUnit::dumpAST(llvm::raw_ostream &OS) const {
-  if (!Unit) {
-    OS << "<no-ast-in-clang>";
-    return; // Parsing failed.
-  }
-  Unit->getASTContext().getTranslationUnitDecl()->dump(OS, true);
+void clangd::dumpAST(ParsedAST &AST, llvm::raw_ostream &OS) {
+  AST.getASTContext().getTranslationUnitDecl()->dump(OS, true);
 }
 
-llvm::Optional<ClangdUnit::ParsedAST>
-ClangdUnit::ParsedAST::Build(std::unique_ptr<clang::CompilerInvocation> CI,
-                             const PrecompiledPreamble *Preamble,
-                             ArrayRef<serialization::DeclID> PreambleDeclIDs,
-                             std::unique_ptr<llvm::MemoryBuffer> Buffer,
-                             std::shared_ptr<PCHContainerOperations> PCHs,
-                             IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
+llvm::Optional<ParsedAST>
+ParsedAST::Build(std::unique_ptr<clang::CompilerInvocation> CI,
 
   std::vector<DiagWithFixIts> ASTDiags;
   StoreDiagsConsumer UnitDiagsConsumer(/*ref*/ ASTDiags);
 
   auto Clang =
       prepareCompilerInstance(std::move(CI), Preamble, std::move(Buffer), PCHs,
                               VFS, /*ref*/ UnitDiagsConsumer);
 
std::vector<Location> takeLocations() {
     // This can happen when nodes in the AST are visited twice.
     std::sort(DeclarationLocations.begin(), DeclarationLocations.end());
     auto last =
         std::unique(DeclarationLocations.begin(), DeclarationLocations.end());
     DeclarationLocations.erase(last, DeclarationLocations.end());
     return std::move(DeclarationLocations);
   }
 
-  bool handleDeclOccurence(const Decl* D, index::SymbolRoleSet Roles,
-      ArrayRef<index::SymbolRelation> Relations, FileID FID, unsigned Offset,
-      index::IndexDataConsumer::ASTNodeInfo ASTNode) override
-      {
+  bool
+  handleDeclOccurence(const Decl *D, index::SymbolRoleSet Roles,
+                      ArrayRef<index::SymbolRelation> Relations, FileID FID,
+                      unsigned Offset,
+                      index::IndexDataConsumer::ASTNodeInfo ASTNode) override {
     if (isSearchedLocation(FID, Offset)) {
       addDeclarationLocation(D->getSourceRange());
     }
     return true;
   }
 
 private:
   bool isSearchedLocation(FileID FID, unsigned Offset) const {
if (!Lexer::getRawToken(SearchedLocation, Result, AST.getSourceManager(),
         SourceLocation BeforeSearchedLocation = getMacroArgExpandedLocation(
             AST.getSourceManager(),
             AST.getSourceManager().getFileEntryForID(DecLoc.first),
             DecLoc.second - 1);
         MacroDefinition MacroDef =
             PP.getMacroDefinitionAtLoc(IdentifierInfo, BeforeSearchedLocation);
         MacroInfo *MacroInf = MacroDef.getMacroInfo();
         if (MacroInf) {
-          addDeclarationLocation(
-              SourceRange(MacroInf->getDefinitionLoc(),
+          addDeclarationLocation(SourceRange(MacroInf->getDefinitionLoc(),
                                              MacroInf->getDefinitionEndLoc()));
         }
       }
     }
   }
 };
-} // namespace
-
-std::vector<Location> ClangdUnit::findDefinitions(Position Pos) {
-  if (!Unit)
-    return {}; // Parsing failed.
-
-  const SourceManager &SourceMgr = Unit->getASTContext().getSourceManager();
-  const FileEntry *FE = SourceMgr.getFileEntryForID(SourceMgr.getMainFileID());
-  if (!FE)
-    return {};
-
-  SourceLocation SourceLocationBeg = getBeginningOfIdentifier(Pos, FE);
-
-  auto DeclLocationsFinder = std::make_shared<DeclarationLocationsFinder>(
-      llvm::errs(), SourceLocationBeg, Unit->getASTContext(),
-      Unit->getPreprocessor());
-  index::IndexingOptions IndexOpts;
-  IndexOpts.SystemSymbolFilter =
-      index::IndexingOptions::SystemSymbolFilterKind::All;
-  IndexOpts.IndexFunctionLocals = true;
-
-  indexTopLevelDecls(Unit->getASTContext(), Unit->getTopLevelDecls(),
-                     DeclLocationsFinder, IndexOpts);
-
-  return DeclLocationsFinder->takeLocations();
-}
-
-SourceLocation ClangdUnit::getBeginningOfIdentifier(const Position &Pos,
-                                                    const FileEntry *FE) const {
 
+SourceLocation getBeginningOfIdentifier(ParsedAST &Unit, const Position &Pos,
+                                        const FileEntry *FE) {
   // The language server protocol uses zero-based line and column numbers.
   // Clang uses one-based numbers.
 
-  const ASTContext &AST = Unit->getASTContext();
+  const ASTContext &AST = Unit.getASTContext();
   const SourceManager &SourceMgr = AST.getSourceManager();
 
   SourceLocation InputLocation =
       getMacroArgExpandedLocation(SourceMgr, FE, Pos);
   if (Pos.character == 0) {
     return InputLocation;
   }
 
SourceLocation getBeginningOfIdentifier(ParsedAST &Unit, const Position &Pos,
   if (Lexer::getRawToken(PeekBeforeLocation, Result, SourceMgr,
                          AST.getLangOpts(), false)) {
     // getRawToken failed, just use InputLocation.
     return InputLocation;
   }
 
   if (Result.is(tok::raw_identifier)) {
     return Lexer::GetBeginningOfToken(PeekBeforeLocation, SourceMgr,
-                                      Unit->getASTContext().getLangOpts());
+                                      AST.getLangOpts());
   }
 
   return InputLocation;
 }
+} // namespace
+
+std::vector<Location> clangd::findDefinitions(ParsedAST &AST, Position Pos) {
+  const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
+  const FileEntry *FE = SourceMgr.getFileEntryForID(SourceMgr.getMainFileID());
+  if (!FE)
+    return {};
+
+  SourceLocation SourceLocationBeg = getBeginningOfIdentifier(AST, Pos, FE);
+
+  auto DeclLocationsFinder = std::make_shared<DeclarationLocationsFinder>(
+      llvm::errs(), SourceLocationBeg, AST.getASTContext(),
+      AST.getPreprocessor());
+  index::IndexingOptions IndexOpts;
+  IndexOpts.SystemSymbolFilter =
+      index::IndexingOptions::SystemSymbolFilterKind::All;
+  IndexOpts.IndexFunctionLocals = true;
+
+  indexTopLevelDecls(AST.getASTContext(), AST.getTopLevelDecls(),
+                     DeclLocationsFinder, IndexOpts);
+
+  return DeclLocationsFinder->takeLocations();
+}
 
-void ClangdUnit::ParsedAST::ensurePreambleDeclsDeserialized() {
+void ParsedAST::ensurePreambleDeclsDeserialized() {
   if (PendingTopLevelDecls.empty())
     return;
 
   std::vector<const Decl *> Resolved;
   Resolved.reserve(PendingTopLevelDecls.size());
 
   ExternalASTSource &Source = *getASTContext().getExternalSource();
   for (serialization::DeclID TopLevelDecl : PendingTopLevelDecls) {
     // Resolve the declaration ID to an actual declaration, possibly
     // deserializing the declaration in the process.
     if (Decl *D = Source.GetExternalDecl(TopLevelDecl))
       Resolved.push_back(D);
   }
 
   TopLevelDecls.reserve(TopLevelDecls.size() + PendingTopLevelDecls.size());
   TopLevelDecls.insert(TopLevelDecls.begin(), Resolved.begin(), Resolved.end());
 
   PendingTopLevelDecls.clear();
 }
 
-ClangdUnit::ParsedAST::ParsedAST(ParsedAST &&Other) = default;
+ParsedAST::ParsedAST(ParsedAST &&Other) = default;
 
-ClangdUnit::ParsedAST &ClangdUnit::ParsedAST::
-operator=(ParsedAST &&Other) = default;
+ParsedAST &ParsedAST::operator=(ParsedAST &&Other) = default;
 
-ClangdUnit::ParsedAST::~ParsedAST() {
+ParsedAST::~ParsedAST() {
   if (Action) {
     Action->EndSourceFile();
   }
 }
 
-ASTContext &ClangdUnit::ParsedAST::getASTContext() {
-  return Clang->getASTContext();
-}
+ASTContext &ParsedAST::getASTContext() { return Clang->getASTContext(); }
 
-const ASTContext &ClangdUnit::ParsedAST::getASTContext() const {
+const ASTContext &ParsedAST::getASTContext() const {
   return Clang->getASTContext();
 }
 
-Preprocessor &ClangdUnit::ParsedAST::getPreprocessor() {
-  return Clang->getPreprocessor();
-}
+Preprocessor &ParsedAST::getPreprocessor() { return Clang->getPreprocessor(); }
 
-const Preprocessor &ClangdUnit::ParsedAST::getPreprocessor() const {
+const Preprocessor &ParsedAST::getPreprocessor() const {
   return Clang->getPreprocessor();
 }
 
-ArrayRef<const Decl *> ClangdUnit::ParsedAST::getTopLevelDecls() {
+ArrayRef<const Decl *> ParsedAST::getTopLevelDecls() {
   ensurePreambleDeclsDeserialized();
   return TopLevelDecls;
 }
 
-const std::vector<DiagWithFixIts> &
-ClangdUnit::ParsedAST::getDiagnostics() const {
+const std::vector<DiagWithFixIts> &ParsedAST::getDiagnostics() const {
   return Diags;
 }
 
-ClangdUnit::ParsedAST::ParsedAST(
-    std::unique_ptr<CompilerInstance> Clang,
+ParsedAST::ParsedAST(std::unique_ptr<CompilerInstance> Clang,
                      std::unique_ptr<FrontendAction> Action,
                      std::vector<const Decl *> TopLevelDecls,
                      std::vector<serialization::DeclID> PendingTopLevelDecls,
                      std::vector<DiagWithFixIts> Diags)
     : Clang(std::move(Clang)), Action(std::move(Action)),
       Diags(std::move(Diags)), TopLevelDecls(std::move(TopLevelDecls)),
       PendingTopLevelDecls(std::move(PendingTopLevelDecls)) {
   assert(this->Clang);
   assert(this->Action);
 }
 
-ClangdUnit::PreambleData::PreambleData(
-    PrecompiledPreamble Preamble,
+ParsedASTWrapper::ParsedASTWrapper(ParsedASTWrapper &&Wrapper)
+    : AST(std::move(Wrapper.AST)) {}
+
+ParsedASTWrapper::ParsedASTWrapper(llvm::Optional<ParsedAST> AST)
+    : AST(std::move(AST)) {}
+
+PreambleData::PreambleData(PrecompiledPreamble Preamble,
                            std::vector<serialization::DeclID> TopLevelDeclIDs,
                            std::vector<DiagWithFixIts> Diags)
     : Preamble(std::move(Preamble)),
       TopLevelDeclIDs(std::move(TopLevelDeclIDs)), Diags(std::move(Diags)) {}
+
+std::shared_ptr<CppFile>
+CppFile::Create(PathRef FileName, tooling::CompileCommand Command,
+                std::shared_ptr<PCHContainerOperations> PCHs) {
+  return std::shared_ptr<CppFile>(
+      new CppFile(FileName, std::move(Command), std::move(PCHs)));
+}
+
+CppFile::CppFile(PathRef FileName, tooling::CompileCommand Command,
+                 std::shared_ptr<PCHContainerOperations> PCHs)
+    : FileName(FileName), Command(std::move(Command)), RebuildCounter(0),
+      RebuildInProgress(false), PCHs(std::move(PCHs)) {
+
+  std::lock_guard<std::mutex> Lock(Mutex);
+  LatestAvailablePreamble = nullptr;
+  PreamblePromise.set_value(nullptr);
+  PreambleFuture = PreamblePromise.get_future();
+
+  ASTPromise.set_value(ParsedASTWrapper(llvm::None));
+  ASTFuture = ASTPromise.get_future();
+}
+
+void CppFile::cancelRebuilds() {
+  std::unique_lock<std::mutex> Lock(Mutex);
+  // Cancel an ongoing rebuild, if any, and wait for it to finish.
+  ++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<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(ParsedASTWrapper(llvm::None));
+}
+
+llvm::Optional<std::vector<DiagWithFixIts>>
+CppFile::rebuild(StringRef NewContents,
+                 IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
+  return deferRebuild(NewContents, std::move(VFS)).get();
+}
+
+std::future<llvm::Optional<std::vector<DiagWithFixIts>>>
+CppFile::deferRebuild(StringRef NewContents,
+                      IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
+  std::shared_ptr<const PreambleData> OldPreamble;
+  std::shared_ptr<PCHContainerOperations> PCHs;
+  unsigned RequestRebuildCounter;
+  {
+    std::unique_lock<std::mutex> Lock(Mutex);
+    // Increase RebuildCounter to cancel all ongoing FinishRebuild operations.
+    // They will try to exit as early as possible and won't call set_value on
+    // our promises.
+    RequestRebuildCounter = ++this->RebuildCounter;
+    PCHs = this->PCHs;
+
+    // Remember the preamble to be used during rebuild.
+    OldPreamble = this->LatestAvailablePreamble;
+    // Setup std::promises and std::futures for Preamble and AST. Corresponding
+    // futures will wait until the rebuild process is finished.
+    if (futureIsReady(this->PreambleFuture)) {
+      this->PreamblePromise =
+          std::promise<std::shared_ptr<const PreambleData>>();
+      this->PreambleFuture = this->PreamblePromise.get_future();
+    }
+    if (futureIsReady(this->ASTFuture)) {
+      this->ASTPromise = std::promise<ParsedASTWrapper>();
+      this->ASTFuture = this->ASTPromise.get_future();
+    }
+  } // unlock Mutex.
+
+  // 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)
+      -> llvm::Optional<std::vector<DiagWithFixIts>> {
+    // Only one execution of this method is possible at a time.
+    // RebuildGuard will wait for any ongoing rebuilds to finish and will put us
+    // into a state for doing a rebuild.
+    RebuildGuard Rebuild(*That, RequestRebuildCounter);
+    if (Rebuild.wasCancelledBeforeConstruction())
+      return llvm::None;
+
+    std::vector<const char *> ArgStrs;
+    for (const auto &S : That->Command.CommandLine)
+      ArgStrs.push_back(S.c_str());
+
+    VFS->setCurrentWorkingDirectory(That->Command.Directory);
+
+    std::unique_ptr<CompilerInvocation> CI;
+    {
+      // FIXME(ibiryukov): store diagnostics from CommandLine when we start
+      // reporting them.
+      EmptyDiagsConsumer CommandLineDiagsConsumer;
+      IntrusiveRefCntPtr<DiagnosticsEngine> CommandLineDiagsEngine =
+          CompilerInstance::createDiagnostics(new DiagnosticOptions,
+                                              &CommandLineDiagsConsumer, false);
+      CI = createCompilerInvocation(ArgStrs, CommandLineDiagsEngine, VFS);
+    }
+    assert(CI && "Couldn't create CompilerInvocation");
+
+    std::unique_ptr<llvm::MemoryBuffer> ContentsBuffer =
+        llvm::MemoryBuffer::getMemBufferCopy(NewContents, That->FileName);
+
+    // A helper function to rebuild the preamble or reuse the existing one. Does
+    // not mutate any fields, only does the actual computation.
+    auto DoRebuildPreamble = [&]() -> std::shared_ptr<const PreambleData> {
+      auto Bounds =
+          ComputePreambleBounds(*CI->getLangOpts(), ContentsBuffer.get(), 0);
+      if (OldPreamble && OldPreamble->Preamble.CanReuse(
+                             *CI, ContentsBuffer.get(), Bounds, VFS.get())) {
+        return OldPreamble;
+      }
+
+      std::vector<DiagWithFixIts> PreambleDiags;
+      StoreDiagsConsumer PreambleDiagnosticsConsumer(/*ref*/ PreambleDiags);
+      IntrusiveRefCntPtr<DiagnosticsEngine> PreambleDiagsEngine =
+          CompilerInstance::createDiagnostics(
+              &CI->getDiagnosticOpts(), &PreambleDiagnosticsConsumer, false);
+      CppFilePreambleCallbacks SerializedDeclsCollector;
+      auto BuiltPreamble = PrecompiledPreamble::Build(
+          *CI, ContentsBuffer.get(), Bounds, *PreambleDiagsEngine, VFS, PCHs,
+          SerializedDeclsCollector);
+
+      if (BuiltPreamble) {
+        return std::make_shared<PreambleData>(
+            std::move(*BuiltPreamble),
+            SerializedDeclsCollector.takeTopLevelDeclIDs(),
+            std::move(PreambleDiags));
+      } else {
+        return nullptr;
+      }
+    };
+
+    // Compute updated Preamble.
+    std::shared_ptr<const PreambleData> NewPreamble = DoRebuildPreamble();
+    // Publish the new Preamble.
+    {
+      std::lock_guard<std::mutex> Lock(That->Mutex);
+      // We always set LatestAvailablePreamble to the new value, hoping that it
+      // will still be usable in the further requests.
+      That->LatestAvailablePreamble = NewPreamble;
+      if (RequestRebuildCounter != That->RebuildCounter)
+        return llvm::None; // Our rebuild request was cancelled, do nothing.
+      That->PreamblePromise.set_value(NewPreamble);
+    } // unlock Mutex
+
+    // Prepare the Preamble and supplementary data for rebuilding AST.
+    const PrecompiledPreamble *PreambleForAST = nullptr;
+    ArrayRef<serialization::DeclID> SerializedPreambleDecls = llvm::None;
+    std::vector<DiagWithFixIts> Diagnostics;
+    if (NewPreamble) {
+      PreambleForAST = &NewPreamble->Preamble;
+      SerializedPreambleDecls = NewPreamble->TopLevelDeclIDs;
+      Diagnostics.insert(Diagnostics.begin(), NewPreamble->Diags.begin(),
+                         NewPreamble->Diags.end());
+    }
+
+    // Compute updated AST.
+    llvm::Optional<ParsedAST> NewAST =
+        ParsedAST::Build(std::move(CI), PreambleForAST, SerializedPreambleDecls,
+                         std::move(ContentsBuffer), PCHs, VFS);
+
+    if (NewAST) {
+      Diagnostics.insert(Diagnostics.end(), NewAST->getDiagnostics().begin(),
+                         NewAST->getDiagnostics().end());
+    } else {
+      // Don't report even Preamble diagnostics if we coulnd't build AST.
+      Diagnostics.clear();
+    }
+
+    // Publish the new AST.
+    {
+      std::lock_guard<std::mutex> Lock(That->Mutex);
+      if (RequestRebuildCounter != That->RebuildCounter)
+        return Diagnostics; // Our rebuild request was cancelled, don't set
+                            // ASTPromise.
+
+      That->ASTPromise.set_value(ParsedASTWrapper(std::move(NewAST)));
+    } // unlock Mutex
+
+    return Diagnostics;
+  };
+
+  return std::async(std::launch::deferred, FinishRebuild, NewContents.str());
+}
+
+std::shared_future<std::shared_ptr<const PreambleData>>
+CppFile::getPreamble() const {
+  std::lock_guard<std::mutex> Lock(Mutex);
+  return PreambleFuture;
+}
+
+std::shared_ptr<const PreambleData> CppFile::getPossiblyStalePreamble() const {
+  std::lock_guard<std::mutex> Lock(Mutex);
+  return LatestAvailablePreamble;
+}
+
+std::shared_future<ParsedASTWrapper> CppFile::getAST() const {
+  std::lock_guard<std::mutex> Lock(Mutex);
+  return ASTFuture;
+}
+
+tooling::CompileCommand const &CppFile::getCompileCommand() const {
+  return Command;
+}
+
+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; });
+
+  WasCancelledBeforeConstruction = File.RebuildCounter != RequestRebuildCounter;
+  if (WasCancelledBeforeConstruction)
+    return;
+
+  File.RebuildInProgress = true;
+}
+
+bool CppFile::RebuildGuard::wasCancelledBeforeConstruction() const {
+  return WasCancelledBeforeConstruction;
+}
+
+CppFile::RebuildGuard::~RebuildGuard() {
+  if (WasCancelledBeforeConstruction)
+    return;
+
+  std::unique_lock<std::mutex> Lock(File.Mutex);
+  assert(File.RebuildInProgress);
+  File.RebuildInProgress = false;
+
+  if (File.RebuildCounter == RequestRebuildCounter) {
+    // Our rebuild request was successful.
+    assert(futureIsReady(File.ASTFuture));
+    assert(futureIsReady(File.PreambleFuture));
+  } else {
+    // Our rebuild request was cancelled, because further reparse was requested.
+    assert(!futureIsReady(File.ASTFuture));
+    assert(!futureIsReady(File.PreambleFuture));
+  }
+
+  Lock.unlock();
+  File.RebuildCond.notify_all();
+}

clangd/ClangdUnitStore.h

-//===--- ClangdUnitStore.h - A ClangdUnits container -------------*-C++-*-===//
+//===--- ClangdUnitStore.h - A container of CppFiles -------------*-C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===---------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDUNITSTORE_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDUNITSTORE_H
 
 #include <mutex>
 
 #include "ClangdUnit.h"
 #include "GlobalCompilationDatabase.h"
 #include "Path.h"
 #include "clang/Tooling/CompilationDatabase.h"
 
 namespace clang {
 namespace clangd {
 
-/// Thread-safe collection of ASTs built for specific files. Provides
-/// synchronized access to ASTs.
-class ClangdUnitStore {
+/// Thread-safe mapping from FileNames to CppFile.
+class CppFileCollection {
 public:
-  /// Run specified \p Action on the ClangdUnit for \p File.
-  /// If the file is not present in ClangdUnitStore, a new ClangdUnit will be
-  /// created from the \p FileContents. If the file is already present in the
-  /// store, ClangdUnit::reparse will be called with the new contents before
-  /// running \p Action.
-  template <class Func>
-  void runOnUnit(PathRef File, StringRef FileContents, StringRef ResourceDir,
+  std::shared_ptr<CppFile>
+  getOrCreateFile(PathRef File, PathRef ResourceDir,
                   GlobalCompilationDatabase &CDB,
                   std::shared_ptr<PCHContainerOperations> PCHs,
-                 IntrusiveRefCntPtr<vfs::FileSystem> VFS, Func Action) {
-    runOnUnitImpl(File, FileContents, ResourceDir, CDB, PCHs,
-                  /*ReparseBeforeAction=*/true, VFS,
-                  std::forward<Func>(Action));
-  }
-
-  /// Run specified \p Action on the ClangdUnit for \p File.
-  /// If the file is not present in ClangdUnitStore, a new ClangdUnit will be
-  /// created from the \p FileContents. If the file is already present in the
-  /// store, the \p Action will be run directly on it.
-  template <class Func>
-  void runOnUnitWithoutReparse(PathRef File, StringRef FileContents,
-                               StringRef ResourceDir,
-                               GlobalCompilationDatabase &CDB,
-                               std::shared_ptr<PCHContainerOperations> PCHs,
-                               IntrusiveRefCntPtr<vfs::FileSystem> VFS,
-                               Func Action) {
-    runOnUnitImpl(File, FileContents, ResourceDir, CDB, PCHs,
-                  /*ReparseBeforeAction=*/false, VFS,
-                  std::forward<Func>(Action));
-  }
-
-  /// Run the specified \p Action on the ClangdUnit for \p File.
-  /// Unit for \p File should exist in the store.
-  template <class Func> void runOnExistingUnit(PathRef File, Func Action) {
+                  IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
     std::lock_guard<std::mutex> Lock(Mutex);
 
     auto It = OpenedFiles.find(File);
-    assert(It != OpenedFiles.end() && "File is not in OpenedFiles");
+    if (It == OpenedFiles.end()) {
+      auto Command = getCompileCommand(CDB, File, ResourceDir);
 
-    Action(It->second);
+      It = OpenedFiles
+               .try_emplace(File, CppFile::Create(File, std::move(Command),
+                                                  std::move(PCHs)))
+               .first;
+    }
+    return It->second;
   }
 
-  /// Remove ClangdUnit for \p File, if any
-  void removeUnitIfPresent(PathRef File);
-
-private:
-  /// Run specified \p Action on the ClangdUnit for \p File.
-  template <class Func>
-  void runOnUnitImpl(PathRef File, StringRef FileContents,
-                     StringRef ResourceDir, GlobalCompilationDatabase &CDB,
-                     std::shared_ptr<PCHContainerOperations> PCHs,
-                     bool ReparseBeforeAction,
-                     IntrusiveRefCntPtr<vfs::FileSystem> VFS, Func Action) {
+  std::shared_ptr<CppFile> getFile(PathRef File) {
     std::lock_guard<std::mutex> Lock(Mutex);
 
     auto It = OpenedFiles.find(File);
-    if (It == OpenedFiles.end()) {
-      auto Commands = getCompileCommands(CDB, File);
-      assert(!Commands.empty() &&
-             "getCompileCommands should add default command");
-
-      It = OpenedFiles
-               .insert(std::make_pair(File, ClangdUnit(File, FileContents,
-                                                       ResourceDir, PCHs,
-                                                       Commands, VFS)))
-               .first;
-    } else if (ReparseBeforeAction) {
-      It->second.reparse(FileContents, VFS);
-    }
-    return Action(It->second);
+    if (It == OpenedFiles.end())
+      return nullptr;
+    return It->second;
   }
 
-  std::vector<tooling::CompileCommand>
-  getCompileCommands(GlobalCompilationDatabase &CDB, PathRef File);
+  /// Removes a CppFile, stored for \p File, if it's inside collection and
+  /// returns it.
+  std::shared_ptr<CppFile> removeIfPresent(PathRef File);
+
+private:
+  tooling::CompileCommand getCompileCommand(GlobalCompilationDatabase &CDB,
+                                            PathRef File, PathRef ResourceDir);
 
   std::mutex Mutex;
-  llvm::StringMap<ClangdUnit> OpenedFiles;
+  llvm::StringMap<std::shared_ptr<CppFile>> OpenedFiles;
 };
 } // namespace clangd
 } // namespace clang
 
 #endif

clangd/ClangdUnitStore.cpp

 //===--- ClangdUnitStore.cpp - A ClangdUnits container -----------*-C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "ClangdUnitStore.h"
 #include "llvm/Support/Path.h"
 
 using namespace clang::clangd;
 using namespace clang;
 
-void ClangdUnitStore::removeUnitIfPresent(PathRef File) {
+std::shared_ptr<CppFile> CppFileCollection::removeIfPresent(PathRef File) {
   std::lock_guard<std::mutex> Lock(Mutex);
 
   auto It = OpenedFiles.find(File);
   if (It == OpenedFiles.end())
-    return;
+    return nullptr;
+
+  std::shared_ptr<CppFile> Result = It->second;
   OpenedFiles.erase(It);
+  return Result;
 }
 
-std::vector<tooling::CompileCommand>
-ClangdUnitStore::getCompileCommands(GlobalCompilationDatabase &CDB,
-                                    PathRef File) {
+tooling::CompileCommand
+CppFileCollection::getCompileCommand(GlobalCompilationDatabase &CDB,
+                                     PathRef File, PathRef ResourceDir) {
   std::vector<tooling::CompileCommand> Commands = CDB.getCompileCommands(File);
   if (Commands.empty())
     // Add a fake command line if we know nothing.
     Commands.push_back(getDefaultCompileCommand(File));
-  return Commands;
+
+  // Inject the resource dir.
+  // FIXME: Don't overwrite it if it's already there.
+  Commands.front().CommandLine.push_back("-resource-dir=" +
+                                         std::string(ResourceDir));
+  return std::move(Commands.front());
 }

clangd/GlobalCompilationDatabase.h

 struct CompileCommand;
 } // namespace tooling
 
 namespace clangd {
 
 /// Returns a default compile command to use for \p File.
 tooling::CompileCommand getDefaultCompileCommand(PathRef File);
 
-/// Provides compilation arguments used for building ClangdUnit.
+/// Provides compilation arguments used for parsing C and C++ files.
 class GlobalCompilationDatabase {
 public:
   virtual ~GlobalCompilationDatabase() = default;
 
   virtual std::vector<tooling::CompileCommand>
   getCompileCommands(PathRef File) = 0;
 
   /// FIXME(ibiryukov): add facilities to track changes to compilation flags of

unittests/clangd/ClangdTests.cpp

 #include "clang/Config/config.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Regex.h"
 #include "gtest/gtest.h"
 #include <algorithm>
+#include <chrono>
 #include <iostream>
 #include <string>
 #include <vector>
 
 namespace clang {
 namespace vfs {
 
 /// An implementation of vfs::FileSystem that only allows access to
IntrusiveRefCntPtr<vfs::FileSystem> getTempOnlyFS() {
   return new vfs::FilteredFileSystem(std::move(TmpDirs),
                                      vfs::getRealFileSystem());
 }
 } // namespace vfs
 
 namespace clangd {
 namespace {
 
+// Don't wait for async ops in clangd test more than that to avoid blocking
+// indefinitely in case of bugs.
+static const std::chrono::seconds DefaultFutureTimeout =
+    std::chrono::seconds(10);
+
 class ErrorCheckingDiagConsumer : public DiagnosticsConsumer {
 public:
-  void onDiagnosticsReady(PathRef File,
-                          Tagged<std::vector<DiagWithFixIts>> Diagnostics) override {
+  void
+  onDiagnosticsReady(PathRef File,
     bool HadError = false;
     for (const auto &DiagAndFixIts : Diagnostics.Value) {
       // FIXME: severities returned by clangd should have a descriptive
       // diagnostic severity enum
       const int ErrorSeverity = 1;
       HadError = DiagAndFixIts.Diag.severity == ErrorSeverity;
     }
 
     std::lock_guard<std::mutex> Lock(Mutex);
     HadErrorInLastDiags = HadError;
     LastVFSTag = Diagnostics.Tag;
   }
 
   bool hadErrorInLastDiags() {
     std::lock_guard<std::mutex> Lock(Mutex);
     return HadErrorInLastDiags;
   }
 
-  VFSTag lastVFSTag() {
-    return LastVFSTag;
-  }
+  VFSTag lastVFSTag() { return LastVFSTag; }
 
 private:
   std::mutex Mutex;
   bool HadErrorInLastDiags = false;
   VFSTag LastVFSTag = VFSTag();
 };
 
 class MockCompilationDatabase : public GlobalCompilationDatabase {
ClangdServer Server(CDB, DiagConsumer, FS,
                         /*RunSynchronously=*/false);
     for (const auto &FileWithContents : ExtraFiles)
       FS.Files[getVirtualTestFilePath(FileWithContents.first)] =
           FileWithContents.second;
 
     auto SourceFilename = getVirtualTestFilePath(SourceFileRelPath);
 
     FS.ExpectedFile = SourceFilename;
-    Server.addDocument(SourceFilename, SourceContents);
+
+    // Have to sync reparses because RunSynchronously is false.
+    auto AddDocFuture = Server.addDocument(SourceFilename, SourceContents);
 
     auto Result = dumpASTWithoutMemoryLocs(Server, SourceFilename);
+
+    // Wait for reparse to finish before checking for errors.
+    EXPECT_EQ(AddDocFuture.wait_for(DefaultFutureTimeout),
+              std::future_status::ready);
     EXPECT_EQ(ExpectErrors, DiagConsumer.hadErrorInLastDiags());
     return Result;
   }
 };
 
 TEST_F(ClangdVFSTest, Parse) {
   // FIXME: figure out a stable format for AST dumps, so that we can check the
   // output of the dump itself is equal to the expected one, not just that it's
)cpp";
 
   auto FooCpp = getVirtualTestFilePath("foo.cpp");
   auto FooH = getVirtualTestFilePath("foo.h");
 
   FS.Files[FooH] = "int a;";
   FS.Files[FooCpp] = SourceContents;
   FS.ExpectedFile = FooCpp;
 
-  Server.addDocument(FooCpp, SourceContents);
+  // To sync reparses before checking for errors.
+  std::future<void> ParseFuture;
+
+  ParseFuture = Server.addDocument(FooCpp, SourceContents);
   auto DumpParse1 = dumpASTWithoutMemoryLocs(Server, FooCpp);
+  ASSERT_EQ(ParseFuture.wait_for(DefaultFutureTimeout),
+            std::future_status::ready);
   EXPECT_FALSE(DiagConsumer.hadErrorInLastDiags());
 
-  Server.addDocument(FooCpp, "");
+  ParseFuture = Server.addDocument(FooCpp, "");
   auto DumpParseEmpty = dumpASTWithoutMemoryLocs(Server, FooCpp);
+  ASSERT_EQ(ParseFuture.wait_for(DefaultFutureTimeout),
+            std::future_status::ready);
   EXPECT_FALSE(DiagConsumer.hadErrorInLastDiags());
 
-  Server.addDocument(FooCpp, SourceContents);
+  ParseFuture = Server.addDocument(FooCpp, SourceContents);
   auto DumpParse2 = dumpASTWithoutMemoryLocs(Server, FooCpp);
+  ASSERT_EQ(ParseFuture.wait_for(DefaultFutureTimeout),
+            std::future_status::ready);
   EXPECT_FALSE(DiagConsumer.hadErrorInLastDiags());
 
   EXPECT_EQ(DumpParse1, DumpParse2);
   EXPECT_NE(DumpParse1, DumpParseEmpty);
 }
 
 TEST_F(ClangdVFSTest, ReparseOnHeaderChange) {
   MockFSProvider FS;
)cpp";
 
   auto FooCpp = getVirtualTestFilePath("foo.cpp");
   auto FooH = getVirtualTestFilePath("foo.h");
 
   FS.Files[FooH] = "int a;";
   FS.Files[FooCpp] = SourceContents;
   FS.ExpectedFile = FooCpp;
 
-  Server.addDocument(FooCpp, SourceContents);
+  // To sync reparses before checking for errors.
+  std::future<void> ParseFuture;
+
+  ParseFuture = Server.addDocument(FooCpp, SourceContents);
   auto DumpParse1 = dumpASTWithoutMemoryLocs(Server, FooCpp);
+  ASSERT_EQ(ParseFuture.wait_for(DefaultFutureTimeout),
+            std::future_status::ready);
   EXPECT_FALSE(DiagConsumer.hadErrorInLastDiags());
 
   FS.Files[FooH] = "";
-  Server.forceReparse(FooCpp);
+  ParseFuture = Server.forceReparse(FooCpp);
   auto DumpParseDifferent = dumpASTWithoutMemoryLocs(Server, FooCpp);
+  ASSERT_EQ(ParseFuture.wait_for(DefaultFutureTimeout),
+            std::future_status::ready);
   EXPECT_TRUE(DiagConsumer.hadErrorInLastDiags());
 
   FS.Files[FooH] = "int a;";
-  Server.forceReparse(FooCpp);
+  ParseFuture = Server.forceReparse(FooCpp);
   auto DumpParse2 = dumpASTWithoutMemoryLocs(Server, FooCpp);
+  EXPECT_EQ(ParseFuture.wait_for(DefaultFutureTimeout),
+            std::future_status::ready);
   EXPECT_FALSE(DiagConsumer.hadErrorInLastDiags());
 
   EXPECT_EQ(DumpParse1, DumpParse2);
   EXPECT_NE(DumpParse1, DumpParseDifferent);
 }
 
 TEST_F(ClangdVFSTest, CheckVersions) {
   MockFSProvider FS;
   ErrorCheckingDiagConsumer DiagConsumer;
   MockCompilationDatabase CDB(/*AddFreestandingFlag=*/true);
   ClangdServer Server(CDB, DiagConsumer, FS,
                       /*RunSynchronously=*/true);
 
   auto FooCpp = getVirtualTestFilePath("foo.cpp");
   const auto SourceContents = "int a;";
   FS.Files[FooCpp] = SourceContents;
   FS.ExpectedFile = FooCpp;
 
+  // No need to sync reparses, because RunSynchronously is set
+  // to true.
   FS.Tag = "123";
   Server.addDocument(FooCpp, SourceContents);
-  EXPECT_EQ(DiagConsumer.lastVFSTag(), FS.Tag);
   EXPECT_EQ(Server.codeComplete(FooCpp, Position{0, 0}).Tag, FS.Tag);
+  EXPECT_EQ(DiagConsumer.lastVFSTag(), FS.Tag);
 
   FS.Tag = "321";
   Server.addDocument(FooCpp, SourceContents);
   EXPECT_EQ(DiagConsumer.lastVFSTag(), FS.Tag);
   EXPECT_EQ(Server.codeComplete(FooCpp, Position{0, 0}).Tag, FS.Tag);
 }
 
 // Only enable this test on Unix
TEST_F(ClangdVFSTest, SearchLibDir) {
 
   auto FooCpp = getVirtualTestFilePath("foo.cpp");
   const auto SourceContents = R"cpp(
 #include <string>
 mock_string x;
 )cpp";
   FS.Files[FooCpp] = SourceContents;
 
+  // No need to sync reparses, because RunSynchronously is set
+  // to true.
   Server.addDocument(FooCpp, SourceContents);
   EXPECT_FALSE(DiagConsumer.hadErrorInLastDiags());
 
   const auto SourceContentsWithError = R"cpp(
 #include <string>
 std::string x;
 )cpp";
   Server.addDocument(FooCpp, SourceContentsWithError);
)cpp";
   // size the same.
   // We complete on the 3rd line (2nd in zero-based numbering), because raw
   // string literal of the SourceContents starts with a newline(it's easy to
   // miss).
   Position CompletePos = {2, 8};
   FS.Files[FooCpp] = SourceContents;
   FS.ExpectedFile = FooCpp;
 
+  // No need to sync reparses here as there are no asserts on diagnostics (or
+  // other async operations).
   Server.addDocument(FooCpp, SourceContents);
 
   {
     auto CodeCompletionResults1 =
         Server.codeComplete(FooCpp, CompletePos, None).Value;
     EXPECT_TRUE(ContainsItem(CodeCompletionResults1, "aba"));
     EXPECT_FALSE(ContainsItem(CodeCompletionResults1, "cbc"));
   }