[clangd] Refactored threading in ClangdServer

[ilya-biryukov]

Summary

We now provide an abstraction of Scheduler that abstracts threading
and resource management in ClangdServer.
No changes to behavior are intended with an exception of changed error
messages.
This patch is preliminary work to allow a revamped threading
implementation that will move the threading code out of CppFile.

[ilya-biryukov]

Here's also a combined diff from both this and child revision: D42177

[bkramer]

clangd/ClangdServer.h

107

What's so simple about it? Why not clangd::ThreadPool?

Also there's llvm::ThreadPool, what's the difference between them?

[ilya-biryukov]

clangd/ClangdServer.h

107

Will rename it to ThreadPool.
Differences are:

• llvm::ThreadPool always process requests in FIFO order, we allow LIFO here (for code completion).
• llvm::ThreadPool will run tasks synchronously when LLVM_ENABLE_THREADS is set to 0. I'm not sure that makes sense for clangd, which has a runtime switch for that (-run-synchronously flag)
• llvm::ThreadPool will not process any tasks when ThreadsCount is set to 0, our implementation processes the tasks synchronously instead.

I'll also be adding per-unit queues in the latter commit (aka thread affinity) to our thead pool, so it'll have more differences. I suggest waiting a day or two before I send the patch for review.

Another minor difference is:

• llvm::ThreadPool creates a std::packaged_task and std::future for each task, our implementation simply runs the provided actions. The latter means less book-keeping and is more efficient, but I don't think it matters.

[sammccall]

So, I simultaneously think this is basically ready to land, and I want substantial changes :-)

This is much better already than what we have, and where I think we can further improve the design, this is a natural point on the way.

My preferred next steps would be:

• we discuss a little bit the directions I've outlined below, without the assumption that they belong in this patch.
• to the extent you agree, you make the less-invasive changes here (e.g. adding a VersionConstraint API, but only actually supporting the cases you've implemented)
• once we're on the same page for the other stuff, I'm happy to pick up any amount of it myself - whatever's not going to step on your toes

clangd/ClangdServer.cpp

37

Would be nice to have parallel names to the Scheduler methods, e.g. blockingASTRead() and blockingPreambleRead()

222

Having trouble with this one.
is this the same as "We currently block the calling thread until the AST is available, to avoid..."?

clangd/ClangdServer.h

164

This is a nice abstraction, so much better than dealing with Cppfile! A couple of observations:

1. all methods refer to a single file, neither the contracts nor the implementation have any interactions between files. An API that reflects this seems more natural. e.g. it naturally provides operations like "are we tracking this file", and makes it natural to be able to e.g. lock at the per-file level. e.g.
   
   class WorkingSet { shared_ptr<TranslationUnit> get(Path); shared_ptr<TranslationUnit> getOrCreate(Path) } class TranslationUnit { void update(ParseInputs); void read(Callback); }

This seems like it would make it easier to explain what the semantics of scheduleRemove are, too.

2. The callbacks from individual methods seem more powerful than needed, and encourage deeper coupling. Basically, the inputs (changes) and outputs (diagnostics, reads) don't seem to want to interact with the same code. This suggests decoupling them: the changes are a sequence of input events, diagnostics are a sequence of output events, reads look much as they do now.

One benefit here is that the versioning relationship between inputs and outputs no longer show up in the function signatures (by coupling an input to a matching output). Expressing them as data makes it easier to tweak them.

3. It's not spelled out how this interacts with drafts: whether text<->version is maintained here or externally, and what the contracts around versions are. There are no options offered, so I would guess that scheduleUpdate delivers new-version-or-nothing, and scheduleASTRead delivers... current-or-newer? current-or-nothing?

I think it would *probably* be clearer to have versions minted by the external DraftStore, that way we can decouple "we know about the contents of this file" from "we're building this file". e.g. we probably want wildly different policies for discarding resources of old versions, when "resources" = source code vs "resources" = ASTs and preambles.

4. Scheduler (or anything it decomposes into) is important and isolated enough that it deserves its own header.

169

this is a performance optimization, right?
I think the scheduler does enough of the compiling that giving it one of its own isn't too wasteful.

173

This seems like a smell - the command is sometimes an input from callsites that want builds, and sometimes an output to callsites that want builds. Also this class is all about threading, scheduling, and versions, and this method bypasses all of that.

(If this is a hack to keep code complete working while we refactor, that's fine - but the docs should say that and we should know what next steps look like!)

ISTM the issue here is coupling updates to the source with updates to the compile command.
Ultimately we should indeed be compiling a given version with a fixed CC, but that doesn't seem like the clearest interface for callers.

One option is:

• • Scheduler has a reference to the CDB (or a std::function wrapping it), and needs to come up with commands itself
  • it caches commands whenever it can, and has an "invalidateCommand(Path)" to drop its cache. with_file_lock { invalidateCommand(), scheduleUpdate() } is the equivalent to forceRebuild.
• it provides a scheduleSourceRead which is like scheduleASTRead but provides source, compile command, and the latest preamble available without blocking. This would be used for operations like codecomplete that can't use scheduleASTRead.

184

The callback here is a bit confusing and YAGNI (both clangd and our known embedders discard the returned future).
It seems enough to return synchronously and know that subsequent reads that ask for = or >= version are going to get nothing back.

200

can we pass in a consistency policy here? even if we only support a subset of {ast, preamble} x policy for now, a richer API would enable experiments/tradeoffs later. e.g.

enum class VersionConstraint {
  ANY,     // returns the first data that's available, typically immediately
  EQUAL, // returns the data for the previous update
  GREATER_EQUAL, // Returns the data for the previous update, or some subsequent update.
}

(I guess the semantics you've described for AST read are EQUAL, and ANY for preamble read)

206

these two names seem a bit confusing - might be easier as just CppFiles, Threads?

The names seem to be ~synonyms for the types, which I don't think is better than echoing the types. (Often there's a completely separate good name, but I don't think so here)

[ilya-biryukov]

• Renamed SimpleThreadPool to ThreadPool
• Removed PCHs from Scheduler's constructor
• Renamed waitFor(AST|Preamble)Action to blocking(AST|Preamble)Read
• Updates
• Remove getLastCommand from Scheduler's interface, pass Inputs to actions instead

[ilya-biryukov]

clangd/ClangdServer.cpp

222

Yes. Used your wording, thanks.

clangd/ClangdServer.h

169

It is probably redundant, yes. Removed it.

173

I've opted for passing action inputs to the read methods, as discussed offline. The implementation turned out to be ugly, we now store a separate StringMap<ParseInputs> in Scheduler to keep track of the latest inputs and not put that into CppFile.

184

I'm fine with removing the callback, but I'd like this change to focus on internals of ClangdServer without changing its interface, happy to remove the callback afterwards.
Added a FIXME for that.

200

I would rather not extend the API beyond what it can do at the moment, but happy to explore this approach after threading gets rid of the legacy stuff that we have now.

206

• I'd go with Files instead of CppFiles. WDYT?
• Executor sounds better to me than Threads. Could you elaborate why you find it confusing?

[ilya-biryukov]

clangd/ClangdServer.h

164

1. Given the nature of the LSP, the users of the interface will probably always call only a single function of TranslationUnit, so we won't win much in terms of the code clarity.

scheduleUpdate(File, Inputs) --> get(File).update(Inputs)

That adds some complexity to the interface, though. I'd opt for not doing that in the initial version.

2. One place where I find these callbacks useful are tests where we could wait for latest addDocument to complete. A more pressing concern is that the interface of ClangdServer does not allow to remove the callbacks (addDocument returns a future) and I would really love to keep the public interface of ClangdServer the same for the first iteration.

3. I would err on the side of saying that scheduleASTRead delivers current version. This gives correct results for the current callers of the interface. If we do current-or-newer I'd start with adding the versions to the interface, so that the callers of the API would have enough context to know whether the results correspond to the latest version or not. I'd really love to keep version out of the initial patch, happy to chat about them for follow-up patches.

4. Totally agree. I planned to move it into a separate header after the changes land. Happy to do that with a follow-up patch.

[sammccall]

Agreed we can defer lots of stuff in order to keep this patch compact.
Generally the things I think we can get right before landing:

• names and file organization for new things
• documentation including where we want to get to

clangd/ClangdServer.cpp

245

I think you want to take a reference here, and then capture by value

312

this fixme doesn't make sense if we're removing the callback

330

nit: you dropped a move here

376

nit: "Inp" in InpPreamble is pretty awkward. In this context, Read might work? or IP.

379

InpPreamble->Preamble->Preamble is pretty hard to understand. Can we find better names for these things? Or at least unpack it on one line?

clangd/ClangdServer.h

164

1. splitting the functionality up into per-file stuff can definitely be deferred until after this patch. Can you add a FIXME?

It should probably happen before the new implementation though, so it pushes that impl in the right direction.

2. The right abstraction for tests is probably "wait until the scheduler is idle", I think, and that can be implemented without each of these individual async methods being observable.

Agreed that we should keep ClangdServer's interface the same for this patch, so we can't change this yet. FIXME? (If you agree with the testing strategy, that'd be worth mentioning)

3. Always delivering exactly the sequenced version is fine, but please spell this out.

4. I'd like this newly added class to have the right name and file structure when it first lands, rather than trying to move it later.

165

I'm not sure about so many little structs, vs providing one with some fields being optional.
That would make a unified read API easier. WDYT? we can defer this.

177

This class has important responsibilities beyond threading itself, which "Scheduler" suggests.

I can't think of a perfectly coherent name, options that seem reasonable:

• TUManager - pretty bland/vague, but gets what this class is mostly about
• Workshop - kind of a silly metaphor, but unlikely to be confused with something else
• Clearinghouse - another silly metaphor, maybe more accurate but more obscure

177

Worth saying something abouth the threading properties here:

• Scheduler is not threadsafe, only the main thread should be providing updates and scheduling tasks.
  • callbacks are run on a large threadpool, and it's appropriate to do slow, blocking work in them

184

Sounds good. ScheduleUpdate -> update, though? Since this callback will go away, it's not really observable that anything is being scheduled.

195

similarly scheduleRemove -> remove

200

I'm happy for this not to land in this patch, but with some of the other API changes it should go in before/with the new implementation

202

I like "schedule" in the description, but it seems to in-the-weeds for a function name - ideally the verb would be what the caller is trying to do, not how we implement it.

runWithAST? or if you want to emphasize the async nature, runSoonWithAST?

206

executor often names an abstraction that's related to but not the same as a threadpool (c.f. java or google3). Not a big deal, but it doesn't seem to add any extra semantics over "threads" which would be less ambiguous.

216

duplicate private

[ilya-biryukov]

• Adressed review comments.
• Move threading code to separate files.

[ilya-biryukov]

clangd/ClangdServer.cpp

245

Makes sense, less copies. Thanks.

376

Used IP

379

It's now IP->PreambleData->Preamble. Not great, but should make more sense.

clangd/ClangdServer.h

164

1. I was trying to argue that per-file API might not be the best choice, since it certainly adds complexity both to the interface and to the implementation without changing the behavior too much. Happy to discuss this further.
2. Added the FIXMEs.
3. Done.
4. Done.

165

Small structs seem better to me. I'd argue they give more type safety, even if it means more typing.
Happy to discuss this in more detail.

177

Added comments.

177

It'd be nice to have some mention of the fact that the class handles threading responsibilities. None of the options seem to capture this.
I don't have good suggestions either, though.

184

update seems good.

202

runWithAST sounds good. Thanks

[ilya-biryukov]

clangd/ClangdServer.h

177

Rename of the Scheduler seems to be the only thing blocking this patch from landing.
I'm happy to go with either of the suggested alternatives or leave as is, I couldn't come up with anything better.

@sammccall, what option would you prefer?

[sammccall]

clangd/ClangdServer.h

177

For a descriptive name I think we need something TU-related in there. If we want it to cover threads too, I can't think of anything better than TUScheduler.

I'm also happy with Workshop or similar, which prods you to work out what the class actually does. Maybe other people will find this confusing. Will solicit opinions.

...

Chatted with @ioeric and @hokein - Workshop is too silly. Suggestions were along the lines of what we've been discussing: TUTaskManager or similar.
So I'd probably go with TUScheduler, it's not horrendously long and it's close enough to Scheduler that those of us who spent last week discussing it won't get confused :-)

[ilya-biryukov]

• Renamed Scheduler to TUScheduler
• Use make_scope_exit instead of onScopeExit

[ilya-biryukov]

• Consume error in dumpAST

[ilya-biryukov]

Should be ready now. Will land as soon as the review is accepted.

[sammccall]

Can you please remove the threading/ subdirectory?
It seems premature for these two files, and TUScheduler doesn't fit. It's unclear that there will be more.

I'd suggest renaming Threadpool.h -> Threading.h, CancellationFlag might fit in there, though up to you.

clangd/threading/TUScheduler.h

1 ↗

(On Diff #131926)

this class needs tests

clangd/threading/ThreadPool.h

1 ↗

(On Diff #131926)

this class needs tests

[ilya-biryukov]

• Remove threading/ dir, moved everything to the top-level
• Rename ThreadPool.h to Threading.h

[ilya-biryukov]

clangd/threading/TUScheduler.h

1 ↗

(On Diff #131926)

Will do :-(

clangd/threading/ThreadPool.h

1 ↗

(On Diff #131926)

Will do :-(

[sammccall]

As discussed offline, basically the only thing to do for testing ThreadPool is to bash it with a workload, and some sort of whole-program stress test seems ideal for this and will also give some coverage to other components (and we should run under tsan!).

TUScheduler on the other hand is a big important class with a clear interface and contract, and is about to get a new implementation - testcases verifying the contracts will be extremely valuable.
These tests don't really need to be concurrency-heavy I think.

clangd/TUScheduler.h

10 ↗

(On Diff #131941)

header guards are stale, sorry!

clangd/Threading.h

10 ↗

(On Diff #131941)

this one also stale

[ilya-biryukov]

• Properly ignore errors.
• Run all requests to completion when destroying ThreadPool.
• Added simple tests for TUScheduler.
• Fixed include guards.

[ilya-biryukov]

All comments should be addressed now. Let me know if I missed anything else.

clangd/threading/TUScheduler.h

1 ↗

(On Diff #131926)

Added a simple test for it. Please take a look and let me know if you have more ideas on how we should test it.

clangd/threading/ThreadPool.h

1 ↗

(On Diff #131926)

We have ClangdThreadingTest.StressTest and TUSchedulerTests that both run concurrent operations on ThreadPool.
As you pointed out, this should provide enough coverage for ThreadPool, so I didn't create any extra tests for it.

[sammccall]

Ship it!

clangd/ClangdServer.cpp

37

Nit: these names got out of sync again

clangd/TUScheduler.h

23 ↗

(On Diff #132100)

just use llvm::hardware_concurrency()?

clangd/Threading.h

30 ↗

(On Diff #132100)

add a comment to the destructor saying what it blocks on?

[ilya-biryukov]

Addressed last review comments:

• Rename blockingRead to blockingRun
• Added a comment to ThreadPool's destructor

[ilya-biryukov]

clangd/TUScheduler.h

23 ↗

(On Diff #132100)

It can return 0, which will cause clangd to run synchronously. This function is only called when someone wants to have at least one worker thread for async processing.

We can change it if you want, but I'd rather leave it as is in this patch.

clangd/ClangdServer.h

 public:
   /// \return getRealFileSystem() tagged with default tag, i.e. VFSTag()
   Tagged<IntrusiveRefCntPtr<vfs::FileSystem>>
   getTaggedFileSystem(PathRef File) override;
 };
 
 class ClangdServer;
 
-/// Returns a number of a default async threads to use for ClangdScheduler.
+/// Returns a number of a default async threads to use for Scheduler.
 /// Returned value is always >= 1 (i.e. will not cause requests to be processed
 /// synchronously).
 unsigned getDefaultAsyncThreadsCount();
 
-/// Handles running WorkerRequests of ClangdServer on a number of worker
-/// threads.
-class ClangdScheduler {
+/// A simple fixed-size thread pool implementation.
+class SimpleThreadPool {
 public:
   /// If \p AsyncThreadsCount is 0, requests added using addToFront and addToEnd
   /// will be processed synchronously on the calling thread.
   // Otherwise, \p AsyncThreadsCount threads will be created to schedule the
   // requests.
-  ClangdScheduler(unsigned AsyncThreadsCount);
-  ~ClangdScheduler();
+  SimpleThreadPool(unsigned AsyncThreadsCount);
+  ~SimpleThreadPool();
 
   /// 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;
if (RunSynchronously) {
       RequestQueue.push_back(
           BindWithForward(std::forward<Func>(F), std::forward<Args>(As)...));
     }
     RequestCV.notify_one();
   }
 
 private:
   bool RunSynchronously;
-  std::mutex Mutex;
+  mutable std::mutex Mutex;
   /// We run some tasks on separate threads(parsing, CppFile cleanup).
   /// These threads looks into RequestQueue to find requests to handle and
   /// terminate when Done is set to true.
   std::vector<std::thread> Workers;
   /// Setting Done to true will make the worker threads terminate.
   bool Done = false;
-  /// A queue of requests. Elements of this vector are async computations (i.e.
-  /// results of calling std::async(std::launch::deferred, ...)).
+  /// A queue of requests.
   std::deque<UniqueFunction<void()>> RequestQueue;
   /// Condition variable to wake up worker threads.
   std::condition_variable RequestCV;
 };
 
+/// Handles running tasks for ClangdServer and managing the resources (e.g.,
+/// preambles and ASTs) for opened files.
+class Scheduler {
+public:
+  Scheduler(unsigned AsyncThreadsCount, bool StorePreamblesInMemory,
+            std::shared_ptr<PCHContainerOperations> PCHs,
+            ASTParsedCallback ASTCallback);
+
+  /// Get the CompileCommand passed at the last scheduleUpdate call for \p File.
+  llvm::Optional<tooling::CompileCommand> getLatestCommand(PathRef File) const;
+
+  /// Schedule an update for \p File. Adds \p File to a list of tracked files if
+  /// \p File was not part of it before.
+  void scheduleUpdate(
+      Context Ctx, PathRef File, ParseInputs Inputs,
+      UniqueFunction<void(Context, llvm::Optional<std::vector<DiagWithFixIts>>)>
+          OnUpdated);
+
+  /// Remove \p File from the list of tracked files and schedule removal of its
+  /// resources. \p Action will be called when resources are freed.
+  /// If an error occurs during processing, it is forwarded to the \p Action
+  /// callback.
+  void scheduleRemove(PathRef File, UniqueFunction<void(llvm::Error)> Action);
+
+  /// Schedule an async read of the AST. \p Action will be called when AST is
+  /// ready.
+  /// If an error occurs during processing, it is forwarded to the \p Action
+  /// callback. If there was no error, AST is never null.
+  void
+  scheduleASTRead(PathRef File,
+                  UniqueFunction<void(llvm::Expected<ParsedAST &>)> Action);
+
+  /// Schedule an async read of the Preamble. Preamble passed to \p Action may
+  /// be stale.
+  /// If an error occurs during processing, it is forwarded to the \p Action
+  /// callback.
+  void schedulePreambleRead(
+      PathRef File,
+      UniqueFunction<void(llvm::Expected<const PreambleData *>)> Action);
+
+private:
+  CppFileCollection Units;
+  SimpleThreadPool Executor;
+};
+
 /// Provides API to manage ASTs for a collection of C++ files and request
 /// various language features.
 /// Currently supports async diagnostics, code completion, formatting and goto
 /// definition.
 class ClangdServer {
 public:
   /// Creates a new ClangdServer instance.
   /// To process parsing requests asynchronously, ClangdServer will spawn \p
   /// AsyncThreadsCount worker threads. However, if \p AsyncThreadsCount is 0,
   /// all requests will be processed on the calling thread.
   ///
   /// ClangdServer uses \p FSProvider to get an instance of vfs::FileSystem for
   /// each parsing request. Results of code completion and diagnostics also
   /// include a tag, that \p FSProvider returns along with the vfs::FileSystem.
   ///
   /// The value of \p ResourceDir will be used to search for internal headers
private:
   /// FIXME: This stats several files to find a .clang-format file. I/O can be
   /// slow. Think of a way to cache this.
   llvm::Expected<tooling::Replacements>
   formatCode(llvm::StringRef Code, PathRef File,
              ArrayRef<tooling::Range> Ranges);
 
   std::future<Context>
   scheduleReparseAndDiags(Context Ctx, PathRef File, VersionedDraft Contents,
-                          std::shared_ptr<CppFile> Resources,
                           Tagged<IntrusiveRefCntPtr<vfs::FileSystem>> TaggedFS,
                           bool UpdateCompileCommand);
 
-  std::future<Context>
-  scheduleCancelRebuild(Context Ctx, std::shared_ptr<CppFile> Resources);
-
   GlobalCompilationDatabase &CDB;
   DiagnosticsConsumer &DiagConsumer;
   FileSystemProvider &FSProvider;
   DraftStore DraftMgr;
   // The index used to look up symbols. This could be:
   //   - null (all index functionality is optional)
   //   - the dynamic index owned by ClangdServer (FileIdx)
   //   - the static index passed to the constructor
   //   - a merged view of a static and dynamic index (MergedIndex)
   SymbolIndex *Index;
   // If present, an up-to-date of symbols in open files. Read via Index.
   std::unique_ptr<FileIndex> FileIdx;
   // If present, a merged view of FileIdx and an external index. Read via Index.
   std::unique_ptr<SymbolIndex> MergedIndex;
-  CppFileCollection Units;
   std::string ResourceDir;
   // If set, this represents the workspace path.
   llvm::Optional<std::string> RootPath;
   std::shared_ptr<PCHContainerOperations> PCHs;
-  bool StorePreamblesInMemory;
   /// Used to serialize diagnostic callbacks.
   /// FIXME(ibiryukov): get rid of an extra map and put all version counters
   /// into CppFile.
   std::mutex DiagnosticsMutex;
   /// Maps from a filename to the latest version of reported diagnostics.
   llvm::StringMap<DocVersion> ReportedDiagnosticVersions;
   // WorkScheduler has to be the last member, because its destructor has to be
   // called before all other members to stop the worker thread that references
-  // ClangdServer
-  ClangdScheduler WorkScheduler;
+  // ClangdServer.
+  Scheduler 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 {
 
+// Issues an async read of AST and waits for results.
+template <class Ret, class Func>
+Ret waitForASTAction(Scheduler &S, PathRef File, Func &&F) {
+  std::packaged_task<Ret(llvm::Expected<ParsedAST &>)> Task(
+      std::forward<Func>(F));
+  auto Future = Task.get_future();
+  S.scheduleASTRead(File, std::move(Task));
+  return Future.get();
+}
+
+// Issues an async read of preamble and waits for results.
+template <class Ret, class Func>
+Ret waitForPreambleAction(Scheduler &S, PathRef File, Func &&F) {
+  std::packaged_task<Ret(llvm::Expected<const PreambleData *>)> Task(
+      std::forward<Func>(F));
+  auto Future = Task.get_future();
+  S.schedulePreambleRead(File, std::move(Task));
+  return Future.get();
+}
+
 tooling::CompileCommand getCompileCommand(GlobalCompilationDatabase &CDB,
                                           PathRef File, PathRef ResourceDir) {
   llvm::Optional<tooling::CompileCommand> C = CDB.getCompileCommand(File);
   if (!C) // FIXME: Suppress diagnostics? Let the user know?
     C = CDB.getFallbackCommand(File);
 
   // Inject the resource dir.
   // FIXME: Don't overwrite it if it's already there.
unsigned clangd::getDefaultAsyncThreadsCount() {
   unsigned HardwareConcurrency = std::thread::hardware_concurrency();
   // C++ standard says that hardware_concurrency()
   // may return 0, fallback to 1 worker thread in
   // that case.
   if (HardwareConcurrency == 0)
     return 1;
   return HardwareConcurrency;
 }
-
-ClangdScheduler::ClangdScheduler(unsigned AsyncThreadsCount)
+SimpleThreadPool::SimpleThreadPool(unsigned AsyncThreadsCount)
     : RunSynchronously(AsyncThreadsCount == 0) {
   if (RunSynchronously) {
     // Don't start the worker thread if we're running synchronously
     return;
   }
 
   Workers.reserve(AsyncThreadsCount);
   for (unsigned I = 0; I < AsyncThreadsCount; ++I) {
Workers.push_back(std::thread([this, I]() {
           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
+          // SimpleThreadPool 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();
       }
     }));
   }
 }
 
-ClangdScheduler::~ClangdScheduler() {
+SimpleThreadPool::~SimpleThreadPool() {
   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_all();
 
   for (auto &Worker : Workers)
     Worker.join();
 }
 
+Scheduler::Scheduler(unsigned AsyncThreadsCount, bool StorePreamblesInMemory,
+                     std::shared_ptr<PCHContainerOperations> PCHs,
+                     ASTParsedCallback ASTCallback)
+    : Units(StorePreamblesInMemory, std::move(PCHs), std::move(ASTCallback)),
+      Executor(AsyncThreadsCount) {}
+
+llvm::Optional<tooling::CompileCommand>
+Scheduler::getLatestCommand(PathRef File) const {
+  auto Resources = Units.getFile(File);
+  if (!Resources)
+    return llvm::None;
+  return Resources->getLatestCommand();
+}
+
+void Scheduler::scheduleUpdate(
+    Context Ctx, PathRef File, ParseInputs Inputs,
+    UniqueFunction<void(Context Ctx,
+                        llvm::Optional<std::vector<DiagWithFixIts>>)>
+        OnUpdated) {
+  auto Resources = Units.getOrCreateFile(File);
+  auto DeferredRebuild = Resources->deferRebuild(std::move(Inputs));
+
+  Executor.addToFront(
+      [](Context Ctx, decltype(OnUpdated) OnUpdated,
+         decltype(DeferredRebuild) DeferredRebuild) {
+        auto Diags = DeferredRebuild(Ctx);
+        OnUpdated(std::move(Ctx), Diags);
+      },
+      std::move(Ctx), std::move(OnUpdated), std::move(DeferredRebuild));
+}
+
+void Scheduler::scheduleRemove(PathRef File,
+                               UniqueFunction<void(llvm::Error)> Action) {
+  auto Resources = Units.removeIfPresent(File);
+  if (!Resources) {
+    Action(llvm::make_error<llvm::StringError>(
+        "trying to remove non-added document", llvm::errc::invalid_argument));
+    return;
+  }
+
+  auto DeferredCancel = Resources->deferCancelRebuild();
+  Executor.addToFront(
+      [](decltype(Action) Action, decltype(DeferredCancel) DeferredCancel) {
+        DeferredCancel();
+        Action(llvm::Error::success());
+      },
+      std::move(Action), std::move(DeferredCancel));
+}
+
+void Scheduler::scheduleASTRead(
+    PathRef File, UniqueFunction<void(llvm::Expected<ParsedAST &>)> Action) {
+  auto Resources = Units.getFile(File);
+  if (!Resources) {
+    Action(llvm::make_error<llvm::StringError>(
+        "trying to get AST for non-added document",
+        llvm::errc::invalid_argument));
+    return;
+  }
+  // We currently do all the reads of the AST on a running thread to avoid
+  // inconsistent states coming from subsequent updates.
+  Resources->getAST().get()->runUnderLock([&](ParsedAST *AST) {
+    if (AST)
+      Action(*AST);
+    else
+      Action(llvm::make_error<llvm::StringError>(
+          "Could not build AST for the latest file update",
+          llvm::errc::invalid_argument));
+  });
+}
+
+void Scheduler::schedulePreambleRead(
+    PathRef File,
+    UniqueFunction<void(llvm::Expected<const PreambleData *>)> Action) {
+  std::shared_ptr<CppFile> Resources = Units.getFile(File);
+  if (!Resources) {
+    Action(llvm::make_error<llvm::StringError>(
+        "trying to get preamble for non-added document",
+        llvm::errc::invalid_argument));
+    return;
+  }
+
+  std::shared_ptr<const PreambleData> Preamble =
+      Resources->getPossiblyStalePreamble();
+  Executor.addToFront(
+      [Resources, Preamble](decltype(Action) Action) mutable {
+        if (!Preamble)
+          Preamble = Resources->getPossiblyStalePreamble();
+
+        Action(Preamble.get());
+      },
+      std::move(Action));
+}
+
 ClangdServer::ClangdServer(GlobalCompilationDatabase &CDB,
                            DiagnosticsConsumer &DiagConsumer,
                            FileSystemProvider &FSProvider,
                            unsigned AsyncThreadsCount,
                            bool StorePreamblesInMemory,
                            bool BuildDynamicSymbolIndex, SymbolIndex *StaticIdx,
                            llvm::Optional<StringRef> ResourceDir)
     : CDB(CDB), DiagConsumer(DiagConsumer), FSProvider(FSProvider),
       FileIdx(BuildDynamicSymbolIndex ? new FileIndex() : nullptr),
-      // Pass a callback into `Units` to extract symbols from a newly parsed
-      // file and rebuild the file index synchronously each time an AST is
-      // parsed.
+      ResourceDir(ResourceDir ? ResourceDir->str() : getStandardResourceDir()),
+      PCHs(std::make_shared<PCHContainerOperations>()),
+      // Pass a callback into `WorkScheduler` to extract symbols from a newly
+      // parsed file and rebuild the file index synchronously each time an AST
+      // is parsed.
       // FIXME(ioeric): this can be slow and we may be able to index on less
       // critical paths.
-      Units(FileIdx
-                ? [this](const Context &Ctx, PathRef Path,
+      WorkScheduler(
+          AsyncThreadsCount, StorePreamblesInMemory, this->PCHs,
+          FileIdx ? [this](const Context &Ctx, PathRef Path,
                            ParsedAST *AST) { FileIdx->update(Ctx, Path, AST); }
-                : ASTParsedCallback()),
-      ResourceDir(ResourceDir ? ResourceDir->str() : getStandardResourceDir()),
-      PCHs(std::make_shared<PCHContainerOperations>()),
-      StorePreamblesInMemory(StorePreamblesInMemory),
-      WorkScheduler(AsyncThreadsCount) {
+                  : ASTParsedCallback()) {
   if (FileIdx && StaticIdx) {
     MergedIndex = mergeIndex(FileIdx.get(), StaticIdx);
     Index = MergedIndex.get();
   } else if (FileIdx)
     Index = FileIdx.get();
   else if (StaticIdx)
     Index = StaticIdx;
   else
     Index = nullptr;
 }
 
 void ClangdServer::setRootPath(PathRef RootPath) {
   std::string NewRootPath = llvm::sys::path::convert_to_slash(
       RootPath, llvm::sys::path::Style::posix);
   if (llvm::sys::fs::is_directory(NewRootPath))
     this->RootPath = NewRootPath;
 }
 
 std::future<Context> ClangdServer::addDocument(Context Ctx, PathRef File,
                                                StringRef Contents) {
   DocVersion Version = DraftMgr.updateDraft(File, Contents);
-
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
-  std::shared_ptr<CppFile> Resources =
-      Units.getOrCreateFile(File, ResourceDir, StorePreamblesInMemory, PCHs);
   return scheduleReparseAndDiags(std::move(Ctx), File,
                                  VersionedDraft{Version, Contents.str()},
-                                 std::move(Resources), std::move(TaggedFS),
-                                 /*UpdateCompileCommand=*/false);
+                                 TaggedFS, /*UpdateCompileCommand=*/false);
 }
 
 std::future<Context> ClangdServer::removeDocument(Context Ctx, PathRef File) {
   DraftMgr.removeDraft(File);
-  std::shared_ptr<CppFile> Resources = Units.removeIfPresent(File);
-  return scheduleCancelRebuild(std::move(Ctx), std::move(Resources));
+
+  std::promise<Context> DonePromise;
+  std::future<Context> DoneFuture = DonePromise.get_future();
+
+  auto Callback = BindWithForward(
+      [](Context Ctx, std::promise<Context> DonePromise, llvm::Error Err) {
+        // FIXME(ibiryukov): allow to report error to the caller.
+        // Discard error, if any.
+        (void)(bool) Err;
+        DonePromise.set_value(std::move(Ctx));
+      },
+      std::move(Ctx), std::move(DonePromise));
+
+  WorkScheduler.scheduleRemove(File, std::move(Callback));
+  return DoneFuture;
 }
 
 std::future<Context> ClangdServer::forceReparse(Context Ctx, PathRef File) {
   auto FileContents = DraftMgr.getDraft(File);
   assert(FileContents.Draft &&
          "forceReparse() was called for non-added document");
 
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
-  std::shared_ptr<CppFile> Resources =
-      Units.getOrCreateFile(File, ResourceDir, StorePreamblesInMemory, PCHs);
-  return scheduleReparseAndDiags(std::move(Ctx), File, FileContents,
-                                 std::move(Resources), std::move(TaggedFS),
-                                 /*UpdateCompileCommand=*/true);
+  return scheduleReparseAndDiags(std::move(Ctx), File, std::move(FileContents),
+                                 TaggedFS, /*UpdateCompileCommand=*/true);
 }
 
 std::future<std::pair<Context, Tagged<CompletionList>>>
 ClangdServer::codeComplete(Context Ctx, PathRef File, Position Pos,
                            const clangd::CodeCompleteOptions &Opts,
                            llvm::Optional<StringRef> OverridenContents,
                            IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS) {
   using ResultType = std::pair<Context, Tagged<CompletionList>>;

 void ClangdServer::codeComplete(
     Context Ctx, PathRef File, Position Pos,
     const clangd::CodeCompleteOptions &Opts,
     UniqueFunction<void(Context, Tagged<CompletionList>)> Callback,
     llvm::Optional<StringRef> OverridenContents,
     IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS) {
   using CallbackType = UniqueFunction<void(Context, Tagged<CompletionList>)>;
 
-  std::string Contents;
-  if (OverridenContents) {
-    Contents = *OverridenContents;
-  } else {
-    auto FileContents = DraftMgr.getDraft(File);
-    assert(FileContents.Draft &&
-           "codeComplete is called for non-added document");
-
-    Contents = std::move(*FileContents.Draft);
-  }
-
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
   if (UsedFS)
     *UsedFS = TaggedFS.Value;
 
-  std::shared_ptr<CppFile> Resources = Units.getFile(File);
-  assert(Resources && "Calling completion on non-added file");
-
-  // Remember the current Preamble and use it when async task starts executing.
-  // At the point when async task starts executing, we may have a different
-  // Preamble in Resources. However, we assume the Preamble that we obtain here
-  // is reusable in completion more often.
-  std::shared_ptr<const PreambleData> Preamble =
-      Resources->getPossiblyStalePreamble();
   // Copy completion options for passing them to async task handler.
   auto CodeCompleteOpts = Opts;
   if (!CodeCompleteOpts.Index) // Respect overridden index.
     CodeCompleteOpts.Index = Index;
 
-  // Copy File, as it is a PathRef that will go out of scope before Task is
-  // executed.
-  Path FileStr = File;
+  auto Command = WorkScheduler.getLatestCommand(File);
+  assert(Command && "codeComplete called for non-added document");
+
+  std::string Contents;
+  if (OverridenContents) {
+    Contents = OverridenContents->str();
+  } else {
+    VersionedDraft Latest = DraftMgr.getDraft(File);
+    assert(Latest.Draft && "codeComplete called for non-added document");
+    Contents = *Latest.Draft;
+  }
+
   // Copy PCHs to avoid accessing this->PCHs concurrently
   std::shared_ptr<PCHContainerOperations> PCHs = this->PCHs;
+  auto Task = [PCHs, Pos, TaggedFS, Command, CodeCompleteOpts](
+                  Context Ctx, std::string Contents, Path File,
+                  CallbackType Callback,
+                  llvm::Expected<const PreambleData *> PreambleErr) {
+    assert(PreambleErr &&
+           "error when trying to read preamble for codeComplete");
+    auto Preamble = *PreambleErr;
 
-  assert(Resources->getLatestCommand() &&
-         "CppFile is in inconsistent state, missing CompileCommand");
-  tooling::CompileCommand CompileCommand = *Resources->getLatestCommand();
-
-  // A task that will be run asynchronously.
-  auto Task =
-      // 'mutable' to reassign Preamble variable.
-      [FileStr, Preamble, Resources, Contents, Pos, CodeCompleteOpts, TaggedFS,
-       PCHs, CompileCommand](Context Ctx, CallbackType Callback) mutable {
-        if (!Preamble) {
-          // Maybe we built some preamble before processing this request.
-          Preamble = Resources->getPossiblyStalePreamble();
-        }
-        // FIXME(ibiryukov): even if Preamble is non-null, we may want to check
-        // both the old and the new version in case only one of them matches.
-        CompletionList Result = clangd::codeComplete(
-            Ctx, FileStr, CompileCommand,
-            Preamble ? &Preamble->Preamble : nullptr, Contents, Pos,
-            TaggedFS.Value, PCHs, CodeCompleteOpts);
+        Ctx, File, *Command, Preamble ? &Preamble->Preamble : nullptr, Contents,
+        Pos, TaggedFS.Value, PCHs, CodeCompleteOpts);
 
-        Callback(std::move(Ctx),
-                 make_tagged(std::move(Result), std::move(TaggedFS.Tag)));
-      };
 
-  WorkScheduler.addToFront(std::move(Task), std::move(Ctx),
-                           std::move(Callback));
+  WorkScheduler.schedulePreambleRead(
+      File, BindWithForward(Task, std::move(Ctx), std::move(Contents),
+                            File.str(), std::move(Callback)));
 }
 
 llvm::Expected<Tagged<SignatureHelp>>
 ClangdServer::signatureHelp(const Context &Ctx, PathRef File, Position Pos,
                             llvm::Optional<StringRef> OverridenContents,
                             IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS) {
-  std::string DraftStorage;
-  if (!OverridenContents) {
-    auto FileContents = DraftMgr.getDraft(File);
-    if (!FileContents.Draft)
+  auto TaggedFS = FSProvider.getTaggedFileSystem(File);
+  if (UsedFS)
+    *UsedFS = TaggedFS.Value;
+
+  std::string Contents;
+  if (OverridenContents) {
+    Contents = OverridenContents->str();
+  } else {
+    VersionedDraft Latest = DraftMgr.getDraft(File);
+    if (!Latest.Draft)
       return llvm::make_error<llvm::StringError>(
           "signatureHelp is called for non-added document",
           llvm::errc::invalid_argument);
-
-    DraftStorage = std::move(*FileContents.Draft);
-    OverridenContents = DraftStorage;
+    Contents = std::move(*Latest.Draft);
   }
 
-  auto TaggedFS = FSProvider.getTaggedFileSystem(File);
-  if (UsedFS)
-    *UsedFS = TaggedFS.Value;
-
-  std::shared_ptr<CppFile> Resources = Units.getFile(File);
-  if (!Resources)
+  auto Command = WorkScheduler.getLatestCommand(File);
+  if (!Command)
     return llvm::make_error<llvm::StringError>(
         "signatureHelp is called for non-added document",
         llvm::errc::invalid_argument);
 
-  assert(Resources->getLatestCommand() &&
-         "CppFile is in inconsistent state, missing CompileCommand");
-  tooling::CompileCommand CompileCommand = *Resources->getLatestCommand();
-
-  auto Preamble = Resources->getPossiblyStalePreamble();
-  auto Result = clangd::signatureHelp(
-      Ctx, File, CompileCommand, Preamble ? &Preamble->Preamble : nullptr,
-      *OverridenContents, Pos, TaggedFS.Value, PCHs);
-  return make_tagged(std::move(Result), TaggedFS.Tag);
+  auto Action = [=, &Ctx](llvm::Expected<const PreambleData *> PreambleErr)
+      -> Expected<Tagged<SignatureHelp>> {
+    if (!PreambleErr)
+      return PreambleErr.takeError();
+    auto Preamble = *PreambleErr;
+
+    return make_tagged(
+        clangd::signatureHelp(Ctx, File, *Command,
+                              Preamble ? &Preamble->Preamble : nullptr,
+                              Contents, Pos, TaggedFS.Value, PCHs),
+        TaggedFS.Tag);
+  };
+  return waitForPreambleAction<Expected<Tagged<SignatureHelp>>>(WorkScheduler,
+                                                                File, Action);
 }
 
 llvm::Expected<tooling::Replacements>
 ClangdServer::formatRange(StringRef Code, PathRef File, Range Rng) {
   size_t Begin = positionToOffset(Code, Rng.start);
   size_t Len = positionToOffset(Code, Rng.end) - Begin;
   return formatCode(Code, File, {tooling::Range(Begin, Len)});
 }
ClangdServer::formatOnType(StringRef Code, PathRef File, Position Pos) {
   size_t Len = CursorPos - PreviousLBracePos;
 
   return formatCode(Code, File, {tooling::Range(PreviousLBracePos, Len)});
 }
 
 Expected<std::vector<tooling::Replacement>>
 ClangdServer::rename(const Context &Ctx, PathRef File, Position Pos,
                      llvm::StringRef NewName) {
-  std::shared_ptr<CppFile> Resources = Units.getFile(File);
+  using RetType = Expected<std::vector<tooling::Replacement>>;
+  auto Action = [=](Expected<ParsedAST &> AST) -> RetType {
+    if (!AST)
+      return AST.takeError();
+
     RefactoringResultCollector ResultCollector;
-  Resources->getAST().get()->runUnderLock([&](ParsedAST *AST) {
     const SourceManager &SourceMgr = AST->getASTContext().getSourceManager();
     const FileEntry *FE =
         SourceMgr.getFileEntryForID(SourceMgr.getMainFileID());
     if (!FE)
-      return;
+      return llvm::make_error<llvm::StringError>(
+          "rename called for non-added document", llvm::errc::invalid_argument);
     SourceLocation SourceLocationBeg =
         clangd::getBeginningOfIdentifier(*AST, Pos, FE);
     tooling::RefactoringRuleContext Context(
         AST->getASTContext().getSourceManager());
     Context.setASTContext(AST->getASTContext());
     auto Rename = clang::tooling::RenameOccurrences::initiate(
         Context, SourceRange(SourceLocationBeg), NewName.str());
-    if (!Rename) {
-      ResultCollector.Result = Rename.takeError();
-      return;
-    }
+    if (!Rename)
+      return Rename.takeError();
+
     Rename->invoke(ResultCollector, Context);
-  });
+
     assert(ResultCollector.Result.hasValue());
     if (!ResultCollector.Result.getValue())
       return ResultCollector.Result->takeError();
 
     std::vector<tooling::Replacement> Replacements;
     for (const tooling::AtomicChange &Change : ResultCollector.Result->get()) {
       tooling::Replacements ChangeReps = Change.getReplacements();
       for (const auto &Rep : ChangeReps) {
-      // FIXME: Right now we only support renaming the main file, so we drop
-      // replacements not for the main file. In the future, we might consider to
-      // support:
+        // FIXME: Right now we only support renaming the main file, so we
+        // drop replacements not for the main file. In the future, we might
+        // consider to support:
         //   * rename in any included header
         //   * rename only in the "main" header
         //   * provide an error if there are symbols we won't rename (e.g.
         //     std::vector)
         //   * rename globally in project
         //   * rename in open files
         if (Rep.getFilePath() == File)
           Replacements.push_back(Rep);
       }
     }
     return Replacements;
+  };
+  return waitForASTAction<RetType>(WorkScheduler, File, std::move(Action));
 }
 
 llvm::Optional<std::string> ClangdServer::getDocument(PathRef File) {
   auto Latest = DraftMgr.getDraft(File);
   if (!Latest.Draft)
     return llvm::None;
   return std::move(*Latest.Draft);
 }
 
 std::string ClangdServer::dumpAST(PathRef File) {
-  std::shared_ptr<CppFile> Resources = Units.getFile(File);
-  if (!Resources)
-    return "<non-added file>";
+  auto Action = [](llvm::Expected<ParsedAST &> AST) -> std::string {
+    if (!AST)
+      return "<no-ast>";
 
     std::string Result;
-  Resources->getAST().get()->runUnderLock([&Result](ParsedAST *AST) {
+
     llvm::raw_string_ostream ResultOS(Result);
-    if (AST) {
-      clangd::dumpAST(*AST, ResultOS);
-    } else {
-      ResultOS << "<no-ast>";
-    }
     ResultOS.flush();
-  });
+
     return Result;
+  };
+  return waitForASTAction<std::string>(WorkScheduler, File, std::move(Action));
 }
 
 llvm::Expected<Tagged<std::vector<Location>>>
 ClangdServer::findDefinitions(const Context &Ctx, PathRef File, Position Pos) {
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
 
-  std::shared_ptr<CppFile> Resources = Units.getFile(File);
-  if (!Resources)
-    return llvm::make_error<llvm::StringError>(
-        "findDefinitions called on non-added file",
-        llvm::errc::invalid_argument);
-
-  std::vector<Location> Result;
-  Resources->getAST().get()->runUnderLock([Pos, &Result, &Ctx](ParsedAST *AST) {
+  using RetType = llvm::Expected<Tagged<std::vector<Location>>>;
+  auto Action = [=, &Ctx](llvm::Expected<ParsedAST &> AST) -> RetType {
     if (!AST)
-      return;
-    Result = clangd::findDefinitions(Ctx, *AST, Pos);
-  });
+      return AST.takeError();
+    auto Result = clangd::findDefinitions(Ctx, *AST, Pos);
     return make_tagged(std::move(Result), TaggedFS.Tag);
+  };
+  return waitForASTAction<RetType>(WorkScheduler, File, Action);
 }
 
 llvm::Optional<Path> ClangdServer::switchSourceHeader(PathRef Path) {
 
   StringRef SourceExtensions[] = {".cpp", ".c", ".cc", ".cxx",
                                   ".c++", ".m", ".mm"};
   StringRef HeaderExtensions[] = {".h", ".hh", ".hpp", ".hxx", ".inc"};
 
ClangdServer::findDocumentHighlights(const Context &Ctx, PathRef File,
   auto FileContents = DraftMgr.getDraft(File);
   if (!FileContents.Draft)
     return llvm::make_error<llvm::StringError>(
         "findDocumentHighlights called on non-added file",
         llvm::errc::invalid_argument);
 
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
 
-  std::shared_ptr<CppFile> Resources = Units.getFile(File);
-  if (!Resources)
-    return llvm::make_error<llvm::StringError>(
-        "findDocumentHighlights called on non-added file",
-        llvm::errc::invalid_argument);
-
-  std::vector<DocumentHighlight> Result;
-  llvm::Optional<llvm::Error> Err;
-  Resources->getAST().get()->runUnderLock([Pos, &Ctx, &Err,
-                                           &Result](ParsedAST *AST) {
-    if (!AST) {
-      Err = llvm::make_error<llvm::StringError>("Invalid AST",
-                                                llvm::errc::invalid_argument);
-      return;
-    }
-    Result = clangd::findDocumentHighlights(Ctx, *AST, Pos);
-  });
-
-  if (Err)
-    return std::move(*Err);
-  return make_tagged(Result, TaggedFS.Tag);
+  using RetType = llvm::Expected<Tagged<std::vector<DocumentHighlight>>>;
+  auto Action = [=, &Ctx](llvm::Expected<ParsedAST &> AST) -> RetType {
+    if (!AST)
+      return AST.takeError();
+    auto Result = clangd::findDocumentHighlights(Ctx, *AST, Pos);
+    return make_tagged(std::move(Result), TaggedFS.Tag);
+  };
+  return waitForASTAction<RetType>(WorkScheduler, File, Action);
 }
 
 std::future<Context> ClangdServer::scheduleReparseAndDiags(
     Context Ctx, PathRef File, VersionedDraft Contents,
-    std::shared_ptr<CppFile> Resources,
     Tagged<IntrusiveRefCntPtr<vfs::FileSystem>> TaggedFS,
     bool UpdateCompileCommand) {
+  DocVersion Version = Contents.Version;
 
-  llvm::Optional<tooling::CompileCommand> FileCmd =
-      Resources->getLatestCommand();
-  tooling::CompileCommand CompileCommand =
-      (UpdateCompileCommand || !FileCmd)
+  auto PrevCommand =
+      !UpdateCompileCommand ? WorkScheduler.getLatestCommand(File) : llvm::None;
+  tooling::CompileCommand Command =
+      UpdateCompileCommand || !PrevCommand
           ? getCompileCommand(CDB, File, ResourceDir)
-          : *FileCmd;
+          : std::move(*PrevCommand);
+
+  using OptDiags = llvm::Optional<std::vector<DiagWithFixIts>>;
+  Path FileStr = File.str();
+  VFSTag Tag = std::move(TaggedFS.Tag);
 
-  ParseInputs Inputs(std::move(CompileCommand), std::move(TaggedFS.Value),
-                     *std::move(Contents.Draft));
-  assert(Contents.Draft && "Draft must have contents");
-  UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>(const Context &)>
-      DeferredRebuild = Resources->deferRebuild(std::move(Inputs));
   std::promise<Context> DonePromise;
   std::future<Context> DoneFuture = DonePromise.get_future();
 
-  DocVersion Version = Contents.Version;
-  Path FileStr = File;
-  VFSTag Tag = TaggedFS.Tag;
-  auto ReparseAndPublishDiags =
-      [this, FileStr, Version,
-       Tag](UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>(
-                const Context &)>
-                DeferredRebuild,
-            std::promise<Context> DonePromise, Context Ctx) -> void {
+  auto Callback = [this, Version, FileStr,
+                   Tag](std::promise<Context> DonePromise, Context Ctx,
+                        OptDiags Diags) {
     auto Guard = onScopeExit([&]() { DonePromise.set_value(std::move(Ctx)); });
-
-    auto CurrentVersion = DraftMgr.getVersion(FileStr);
-    if (CurrentVersion != Version)
-      return; // This request is outdated
-
-    auto Diags = DeferredRebuild(Ctx);
     if (!Diags)
-      return; // A new reparse was requested before this one completed.
+      return;
 
-    // We need to serialize access to resulting diagnostics to avoid calling
-    // `onDiagnosticsReady` in the wrong order.
+    // We need to serialize access to resulting diagnostics to avoid
+    // calling `onDiagnosticsReady` in the wrong order.
     std::lock_guard<std::mutex> DiagsLock(DiagnosticsMutex);
     DocVersion &LastReportedDiagsVersion = ReportedDiagnosticVersions[FileStr];
     // FIXME(ibiryukov): get rid of '<' comparison here. In the current
-    // implementation diagnostics will not be reported after version counters'
-    // overflow. This should not happen in practice, since DocVersion is a
-    // 64-bit unsigned integer.
+    // implementation diagnostics will not be reported after version
+    // counters' overflow. This should not happen in practice, since
+    // DocVersion is a 64-bit unsigned integer.
     if (Version < LastReportedDiagsVersion)
       return;
     LastReportedDiagsVersion = Version;
 
-    DiagConsumer.onDiagnosticsReady(Ctx, FileStr,
-                                    make_tagged(std::move(*Diags), Tag));
+    DiagConsumer.onDiagnosticsReady(
+        Ctx, FileStr, make_tagged(std::move(*Diags), std::move(Tag)));
   };
 
-  WorkScheduler.addToFront(std::move(ReparseAndPublishDiags),
-                           std::move(DeferredRebuild), std::move(DonePromise),
-                           std::move(Ctx));
-  return DoneFuture;
-}
-
-std::future<Context>
-ClangdServer::scheduleCancelRebuild(Context Ctx,
-                                    std::shared_ptr<CppFile> Resources) {
-  std::promise<Context> DonePromise;
-  std::future<Context> DoneFuture = DonePromise.get_future();
-  if (!Resources) {
-    // No need to schedule any cleanup.
-    DonePromise.set_value(std::move(Ctx));
-    return DoneFuture;
-  }
-
-  UniqueFunction<void()> DeferredCancel = Resources->deferCancelRebuild();
-  auto CancelReparses = [Resources](std::promise<Context> DonePromise,
-                                    UniqueFunction<void()> DeferredCancel,
-                                    Context Ctx) {
-    DeferredCancel();
-    DonePromise.set_value(std::move(Ctx));
-  };
-  WorkScheduler.addToFront(std::move(CancelReparses), std::move(DonePromise),
-                           std::move(DeferredCancel), std::move(Ctx));
+  WorkScheduler.scheduleUpdate(
+      std::move(Ctx), File,
+      ParseInputs(Command, std::move(TaggedFS.Value),
+                  std::move(*Contents.Draft)),
+      BindWithForward(Callback, std::move(DonePromise)));
   return DoneFuture;
 }
 
 void ClangdServer::onFileEvent(const DidChangeWatchedFilesParams &Params) {
   // FIXME: Do nothing for now. This will be used for indexing and potentially
   // invalidating other caches.
 }

clangd/ClangdUnit.h

 };
 
 /// Information required to run clang (e.g., to parse AST or do code
 /// completion).
 struct ParseInputs {
   ParseInputs(tooling::CompileCommand CompileCommand,
               IntrusiveRefCntPtr<vfs::FileSystem> FS, std::string Contents);
 
-  /// Compilation arguments.
   tooling::CompileCommand CompileCommand;
-  /// FileSystem.
   IntrusiveRefCntPtr<vfs::FileSystem> FS;
-  /// Contents to parse.
   std::string Contents;
 };
 
 /// 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.

clangd/ClangdUnitStore.h

 
 class Logger;
 
 /// Thread-safe mapping from FileNames to CppFile.
 class CppFileCollection {
 public:
   /// \p ASTCallback is called when a file is parsed synchronously. This should
   /// not be expensive since it blocks diagnostics.
-  explicit CppFileCollection(ASTParsedCallback ASTCallback)
-      : ASTCallback(std::move(ASTCallback)) {}
+  explicit CppFileCollection(bool StorePreamblesInMemory,
+                             std::shared_ptr<PCHContainerOperations> PCHs,
+                             ASTParsedCallback ASTCallback)
+      : ASTCallback(std::move(ASTCallback)), PCHs(std::move(PCHs)),
+        StorePreamblesInMemory(StorePreamblesInMemory) {}
 
-  std::shared_ptr<CppFile>
-  getOrCreateFile(PathRef File, PathRef ResourceDir,
-                  bool StorePreamblesInMemory,
-                  std::shared_ptr<PCHContainerOperations> PCHs) {
+  std::shared_ptr<CppFile> getOrCreateFile(PathRef File) {
     std::lock_guard<std::mutex> Lock(Mutex);
     auto It = OpenedFiles.find(File);
     if (It == OpenedFiles.end()) {
       It = OpenedFiles
                .try_emplace(File, CppFile::Create(File, StorePreamblesInMemory,
-                                                  std::move(PCHs), ASTCallback))
+                                                  PCHs, ASTCallback))
                .first;
     }
     return It->second;
   }
 
-  std::shared_ptr<CppFile> getFile(PathRef File) {
+  std::shared_ptr<CppFile> getFile(PathRef File) const {
     std::lock_guard<std::mutex> Lock(Mutex);
-
     auto It = OpenedFiles.find(File);
     if (It == OpenedFiles.end())
       return nullptr;
     return It->second;
   }
 
   /// Removes a CppFile, stored for \p File, if it's inside collection and
   /// returns it.
   std::shared_ptr<CppFile> removeIfPresent(PathRef File);
 
 private:
-  std::mutex Mutex;
+  mutable std::mutex Mutex;
   llvm::StringMap<std::shared_ptr<CppFile>> OpenedFiles;
   ASTParsedCallback ASTCallback;
+  std::shared_ptr<PCHContainerOperations> PCHs;
+  bool StorePreamblesInMemory;
 };
 } // namespace clangd
 } // namespace clang
 
 #endif