[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

100

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

100

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

36

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

150

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

100

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.

105

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.

109

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.

120

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.

136

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)

142

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

150

Yes. Used your wording, thanks.

clangd/ClangdServer.h

105

It is probably redundant, yes. Removed it.

109

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.

120

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.

136

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.

142

• 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

100

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

150

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

171

nit: you dropped a move here

173

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

226

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

229

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

100

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.

101

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.

113

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

113

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

120

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

131

similarly scheduleRemove -> remove

136

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

138

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?

142

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.

152

duplicate private

[ilya-biryukov]

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

[ilya-biryukov]

clangd/ClangdServer.cpp

173

Makes sense, less copies. Thanks.

226

Used IP

229

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

clangd/ClangdServer.h

100

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.

101

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.

113

Added comments.

113

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.

120

update seems good.

138

runWithAST sounds good. Thanks

[ilya-biryukov]

clangd/ClangdServer.h

113

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

113

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

11

header guards are stale, sorry!

clangd/Threading.h

11

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

36

Nit: these names got out of sync again

clangd/TUScheduler.h

23

just use llvm::hardware_concurrency()?

clangd/Threading.h

30

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

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/CMakeLists.txt

add_clang_library(clangDaemon
   FuzzyMatch.cpp
   GlobalCompilationDatabase.cpp
   JSONExpr.cpp
   JSONRPCDispatcher.cpp
   Logger.cpp
   Protocol.cpp
   ProtocolHandlers.cpp
   SourceCode.cpp
+  Threading.cpp
   Trace.cpp
+  TUScheduler.cpp
   URI.cpp
   XRefs.cpp
   index/FileIndex.cpp
   index/Index.cpp
   index/MemIndex.cpp
   index/Merge.cpp
   index/SymbolCollector.cpp
   index/SymbolYAML.cpp

clangd/ClangdServer.h

 #include "ClangdUnit.h"
 #include "ClangdUnitStore.h"
 #include "CodeComplete.h"
 #include "CompileArgsCache.h"
 #include "DraftStore.h"
 #include "Function.h"
 #include "GlobalCompilationDatabase.h"
 #include "Protocol.h"
+#include "TUScheduler.h"
 #include "index/FileIndex.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringRef.h"
-#include <condition_variable>
 #include <functional>
-#include <mutex>
 #include <string>
-#include <thread>
 #include <type_traits>
 #include <utility>
 
 namespace clang {
 class PCHContainerOperations;
 
 namespace clangd {
 

 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.
-/// 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 {
-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();
-
-  /// 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(
-          BindWithForward(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(
-          BindWithForward(std::forward<Func>(F), std::forward<Args>(As)...));
-    }
-    RequestCV.notify_one();
-  }
-
-private:
-  bool RunSynchronously;
-  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, ...)).
-  std::deque<UniqueFunction<void()>> RequestQueue;
-  /// Condition variable to wake up worker threads.
-  std::condition_variable RequestCV;
-};
-
 /// 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
   /// (overriding defaults and -resource-dir compiler flag). If \p ResourceDir
   /// is None, ClangdServer will call CompilerInvocation::GetResourcePath() to
   /// obtain the standard resource directory.
   ///
   /// ClangdServer uses \p CDB to obtain compilation arguments for parsing. Note
   /// that ClangdServer only obtains compilation arguments once for each newly
   /// added file (i.e., when processing a first call to addDocument) and reuses
   /// those arguments for subsequent reparses. However, ClangdServer will check
   /// if compilation arguments changed on calls to forceReparse().
   ///
   /// After each parsing request finishes, ClangdServer reports diagnostics to
   /// \p DiagConsumer. Note that a callback to \p DiagConsumer happens on a
   /// worker thread. Therefore, instances of \p DiagConsumer must properly
   /// synchronize access to shared state.
   ///
   /// \p StorePreamblesInMemory defines whether the Preambles generated by
   /// clangd are stored in-memory or on disk.
   ///
   /// If \p BuildDynamicSymbolIndex is true, ClangdServer builds a dynamic
   /// in-memory index for symbols in all opened files and uses the index to
   /// augment code completion results.
   ///
   /// If \p StaticIdx is set, ClangdServer uses the index for global code
   /// completion.
   ClangdServer(GlobalCompilationDatabase &CDB,
                DiagnosticsConsumer &DiagConsumer,
                FileSystemProvider &FSProvider, unsigned AsyncThreadsCount,
                bool StorePreamblesInMemory,
                bool BuildDynamicSymbolIndex = false,
                SymbolIndex *StaticIdx = nullptr,
                llvm::Optional<StringRef> ResourceDir = llvm::None);
 
   /// Set the root path of the workspace.
   void setRootPath(PathRef RootPath);
 
   /// 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.
   /// \return A future that will become ready when the rebuild (including
   /// diagnostics) is finished.
+  /// FIXME: don't return futures here, LSP does not require a response for this
+  /// request.
   std::future<Context> addDocument(Context Ctx, PathRef File,
                                    StringRef Contents);
   /// Remove \p File from list of tracked files, schedule a request to free
   /// resources associated with it.
   /// \return A future that will become ready when the file is removed and all
   /// associated resources are freed.
+  /// FIXME: don't return futures here, LSP does not require a response for this
+  /// request.
   std::future<Context> removeDocument(Context Ctx, PathRef File);
   /// Force \p File to be reparsed using the latest contents.
   /// Will also check if CompileCommand, provided by GlobalCompilationDatabase
   /// for \p File has changed. If it has, will remove currently stored Preamble
   /// and AST and rebuild them from scratch.
+  /// FIXME: don't return futures here, LSP does not require a response for this
+  /// request.
   std::future<Context> forceReparse(Context Ctx, PathRef File);
 
   /// DEPRECATED. Please use a callback-based version, this API is deprecated
   /// and will soon be removed.
   ///
   /// Run code completion for \p File at \p Pos.
   ///
   /// Request is processed asynchronously. You can use the returned future to
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);
 
-  std::future<Context>
-  scheduleCancelRebuild(Context Ctx, std::shared_ptr<CppFile> Resources);
-
   CompileArgsCache CompileArgs;
   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;
   // 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.
+  TUScheduler WorkScheduler;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clangd/ClangdServer.cpp

 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "clang/Tooling/Refactoring/RefactoringResultConsumer.h"
 #include "clang/Tooling/Refactoring/Rename/RenamingAction.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/FileSystem.h"
-#include "llvm/Support/FormatProviders.h"
-#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Path.h"
 #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 blockingASTRead(TUScheduler &S, PathRef File, Func &&F) {
+  std::packaged_task<Ret(llvm::Expected<InputsAndAST>)> Task(
+      std::forward<Func>(F));
+  auto Future = Task.get_future();
+  S.runWithAST(File, std::move(Task));
+  return Future.get();
+}
+
+// Issues an async read of preamble and waits for results.
+template <class Ret, class Func>
+Ret blockingPreambleRead(TUScheduler &S, PathRef File, Func &&F) {
+  std::packaged_task<Ret(llvm::Expected<InputsAndPreamble>)> Task(
+      std::forward<Func>(F));
+  auto Future = Task.get_future();
+  S.runWithPreamble(File, std::move(Task));
+  return Future.get();
+}
+
+void ignoreError(llvm::Error Err) {
+  handleAllErrors(std::move(Err), [](const llvm::ErrorInfoBase &) {});
+}
+
 std::string getStandardResourceDir() {
   static int Dummy; // Just an address in this process.
   return CompilerInvocation::GetResourcesPath("clangd", (void *)&Dummy);
 }
 
 class RefactoringResultCollector final
     : public tooling::RefactoringResultConsumer {
 public:

 
 } // namespace
 
 Tagged<IntrusiveRefCntPtr<vfs::FileSystem>>
 RealFileSystemProvider::getTaggedFileSystem(PathRef File) {
   return make_tagged(vfs::getRealFileSystem(), VFSTag());
 }
 
-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)
-    : 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]() {
-      llvm::set_thread_name(llvm::formatv("scheduler/{0}", I));
-      while (true) {
-        UniqueFunction<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();
-      }
-    }));
-  }
-}
-
-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_all();
-
-  for (auto &Worker : Workers)
-    Worker.join();
-}
-
 ClangdServer::ClangdServer(GlobalCompilationDatabase &CDB,
                            DiagnosticsConsumer &DiagConsumer,
                            FileSystemProvider &FSProvider,
                            unsigned AsyncThreadsCount,
                            bool StorePreamblesInMemory,
                            bool BuildDynamicSymbolIndex, SymbolIndex *StaticIdx,
                            llvm::Optional<StringRef> ResourceDir)
     : CompileArgs(CDB,
                   ResourceDir ? ResourceDir->str() : getStandardResourceDir()),
       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.
+      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,
+          FileIdx ? [this](const Context &Ctx, PathRef Path,
                            ParsedAST *AST) { FileIdx->update(Ctx, Path, AST); }
-                : ASTParsedCallback()),
-      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, StorePreamblesInMemory, PCHs);
   return scheduleReparseAndDiags(std::move(Ctx), File,
                                  VersionedDraft{Version, Contents.str()},
-                                 std::move(Resources), std::move(TaggedFS));
+                                 std::move(TaggedFS));
 }
 
 std::future<Context> ClangdServer::removeDocument(Context Ctx, PathRef File) {
   DraftMgr.removeDraft(File);
   CompileArgs.invalidate(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) {
+        if (Err)
+          ignoreError(std::move(Err));
+        DonePromise.set_value(std::move(Ctx));
+      },
+      std::move(Ctx), std::move(DonePromise));
+
+  WorkScheduler.remove(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");
 
   // forceReparse promises to request new compilation flags from CDB, so we
   // remove any cahced flags.
   CompileArgs.invalidate(File);
 
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
-  std::shared_ptr<CppFile> Resources =
-      Units.getOrCreateFile(File, StorePreamblesInMemory, PCHs);
-  return scheduleReparseAndDiags(std::move(Ctx), File, FileContents,
-                                 std::move(Resources), std::move(TaggedFS));
+  return scheduleReparseAndDiags(std::move(Ctx), File, std::move(FileContents),
+                                 std::move(TaggedFS));
 }
 
 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>>;
 
   std::promise<ResultType> ResultPromise;

 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;
+  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;
-  tooling::CompileCommand CompileCommand = CompileArgs.getCompileCommand(File);
-  // 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);
+  auto Task = [PCHs, Pos, TaggedFS, CodeCompleteOpts](
+                  Context Ctx, std::string Contents, Path File,
+                  CallbackType Callback, llvm::Expected<InputsAndPreamble> IP) {
+    assert(IP && "error when trying to read preamble for codeComplete");
+    auto PreambleData = IP->Preamble;
+    auto &Command = IP->Inputs.CompileCommand;
+
+        Ctx, File, Command, PreambleData ? &PreambleData->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.runWithPreamble(
+      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)
-      return llvm::make_error<llvm::StringError>(
-          "signatureHelp is called for non-added document",
-          llvm::errc::invalid_argument);
-
-    DraftStorage = std::move(*FileContents.Draft);
-    OverridenContents = DraftStorage;
-  }
-
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
   if (UsedFS)
     *UsedFS = TaggedFS.Value;
 
-  std::shared_ptr<CppFile> Resources = Units.getFile(File);
-  if (!Resources)
+  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);
+    Contents = std::move(*Latest.Draft);
+  }
 
-  auto Preamble = Resources->getPossiblyStalePreamble();
-  auto Result =
-      clangd::signatureHelp(Ctx, File, CompileArgs.getCompileCommand(File),
-                            Preamble ? &Preamble->Preamble : nullptr,
-                            *OverridenContents, Pos, TaggedFS.Value, PCHs);
-  return make_tagged(std::move(Result), TaggedFS.Tag);
+  auto Action = [=, &Ctx](llvm::Expected<InputsAndPreamble> IP)
+      -> Expected<Tagged<SignatureHelp>> {
+    if (!IP)
+      return IP.takeError();
+    auto PreambleData = IP->Preamble;
+    auto &Command = IP->Inputs.CompileCommand;
+
+    return make_tagged(
+        clangd::signatureHelp(Ctx, File, Command,
+                              PreambleData ? &PreambleData->Preamble : nullptr,
+                              Contents, Pos, TaggedFS.Value, PCHs),
+        TaggedFS.Tag);
+  };
+  return blockingPreambleRead<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<InputsAndAST> InpAST) -> RetType {
+    if (!InpAST)
+      return InpAST.takeError();
+    auto &AST = InpAST->AST;
+
     RefactoringResultCollector ResultCollector;
-  Resources->getAST().get()->runUnderLock([&](ParsedAST *AST) {
-    const SourceManager &SourceMgr = AST->getASTContext().getSourceManager();
+    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);
+        clangd::getBeginningOfIdentifier(AST, Pos, FE);
     tooling::RefactoringRuleContext Context(
-        AST->getASTContext().getSourceManager());
-    Context.setASTContext(AST->getASTContext());
+        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 blockingASTRead<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<InputsAndAST> InpAST) -> std::string {
+    if (!InpAST) {
+      ignoreError(InpAST.takeError());
+      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>";
-    }
+    clangd::dumpAST(InpAST->AST, ResultOS);
     ResultOS.flush();
-  });
+
     return Result;
+  };
+  return blockingASTRead<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) {
-    if (!AST)
-      return;
-    Result = clangd::findDefinitions(Ctx, *AST, Pos);
-  });
+  using RetType = llvm::Expected<Tagged<std::vector<Location>>>;
+  auto Action = [=, &Ctx](llvm::Expected<InputsAndAST> InpAST) -> RetType {
+    if (!InpAST)
+      return InpAST.takeError();
+    auto Result = clangd::findDefinitions(Ctx, InpAST->AST, Pos);
     return make_tagged(std::move(Result), TaggedFS.Tag);
+  };
+  return blockingASTRead<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<InputsAndAST> InpAST) -> RetType {
+    if (!InpAST)
+      return InpAST.takeError();
+    auto Result = clangd::findDocumentHighlights(Ctx, InpAST->AST, Pos);
+    return make_tagged(std::move(Result), TaggedFS.Tag);
+  };
+  return blockingASTRead<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) {
-  assert(Contents.Draft && "Draft must have contents");
-  ParseInputs Inputs = {CompileArgs.getCompileCommand(File),
-                        std::move(TaggedFS.Value), *std::move(Contents.Draft)};
+  tooling::CompileCommand Command = CompileArgs.getCompileCommand(File);
+
+  using OptDiags = llvm::Optional<std::vector<DiagWithFixIts>>;
+
+  DocVersion Version = Contents.Version;
+  Path FileStr = File.str();
+  VFSTag Tag = std::move(TaggedFS.Tag);
 
-  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 =
         llvm::make_scope_exit([&]() { 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.
 
     // 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.
     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.update(std::move(Ctx), File,
+                       ParseInputs{std::move(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.
 }
 
 std::vector<std::pair<Path, std::size_t>>
 ClangdServer::getUsedBytesPerFile() const {
-  return Units.getUsedBytesPerFile();
+  return WorkScheduler.getUsedBytesPerFile();
 }

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, 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) 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);
 
   /// Gets used memory for each of the stored files.
   std::vector<std::pair<Path, std::size_t>> getUsedBytesPerFile() const;
 
 private:
   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

clangd/TUScheduler.h

+//===--- TUScheduler.h -------------------------------------------*-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_TUSCHEDULER_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_TUSCHEDULER_H
+
+#include "ClangdUnit.h"
+#include "ClangdUnitStore.h"
+#include "Function.h"
+#include "Threading.h"
+
+namespace clang {
+namespace clangd {
+/// Returns a number of a default async threads to use for TUScheduler.
+/// Returned value is always >= 1 (i.e. will not cause requests to be processed
+/// synchronously).
+unsigned getDefaultAsyncThreadsCount();
+
+struct InputsAndAST {
+  const ParseInputs &Inputs;
+  ParsedAST *
+};
+
+struct InputsAndPreamble {
+  const ParseInputs &Inputs;
+  const PreambleData *Preamble;
+};
+
+/// Handles running tasks for ClangdServer and managing the resources (e.g.,
+/// preambles and ASTs) for opened files.
+/// TUScheduler is not thread-safe, only one thread should be providing updates
+/// and scheduling tasks.
+/// Callbacks are run on a threadpool and it's appropriate to do slow work in
+/// them.
+class TUScheduler {
+public:
+  TUScheduler(unsigned AsyncThreadsCount, bool StorePreamblesInMemory,
+              ASTParsedCallback ASTCallback);
+
+  /// Returns estimated memory usage for each of the currently open files.
+  /// The order of results is unspecified.
+  std::vector<std::pair<Path, std::size_t>> getUsedBytesPerFile() 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.
+  /// FIXME(ibiryukov): remove the callback from this function.
+  void update(
+      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.
+  /// FIXME(ibiryukov): the callback passed to this function is not used, we
+  /// should remove it.
+  void remove(PathRef File, UniqueFunction<void(llvm::Error)> Action);
+
+  /// Schedule an async read of the AST. \p Action will be called when AST is
+  /// ready. The AST passed to \p Action refers to the version of \p File
+  /// tracked at the time of the call, even if new updates are received before
+  /// \p Action is executed.
+  /// If an error occurs during processing, it is forwarded to the \p Action
+  /// callback.
+  void runWithAST(PathRef File,
+                  UniqueFunction<void(llvm::Expected<InputsAndAST>)> Action);
+
+  /// Schedule an async read of the Preamble. Preamble passed to \p Action may
+  /// be built for any version of the file, callers must not rely on it being
+  /// consistent with the current version of the file.
+  /// If an error occurs during processing, it is forwarded to the \p Action
+  /// callback.
+  void runWithPreamble(
+      PathRef File,
+      UniqueFunction<void(llvm::Expected<InputsAndPreamble>)> Action);
+
+private:
+  const ParseInputs &getInputs(PathRef File);
+
+  llvm::StringMap<ParseInputs> CachedInputs;
+  CppFileCollection Files;
+  ThreadPool Threads;
+};
+} // namespace clangd
+} // namespace clang
+
+#endif

clangd/TUScheduler.cpp

+#include "TUScheduler.h"
+#include "clang/Frontend/PCHContainerOperations.h"
+#include "llvm/Support/Errc.h"
+
+namespace clang {
+namespace clangd {
+unsigned 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;
+}
+
+TUScheduler::TUScheduler(unsigned AsyncThreadsCount,
+                         bool StorePreamblesInMemory,
+                         ASTParsedCallback ASTCallback)
+
+    : Files(StorePreamblesInMemory, std::make_shared<PCHContainerOperations>(),
+            std::move(ASTCallback)),
+      Threads(AsyncThreadsCount) {}
+
+void TUScheduler::update(
+    Context Ctx, PathRef File, ParseInputs Inputs,
+    UniqueFunction<void(Context Ctx,
+                        llvm::Optional<std::vector<DiagWithFixIts>>)>
+        OnUpdated) {
+  CachedInputs[File] = Inputs;
+
+  auto Resources = Files.getOrCreateFile(File);
+  auto DeferredRebuild = Resources->deferRebuild(std::move(Inputs));
+
+  Threads.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 TUScheduler::remove(PathRef File,
+                         UniqueFunction<void(llvm::Error)> Action) {
+  CachedInputs.erase(File);
+
+  auto Resources = Files.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();
+  Threads.addToFront(
+      [](decltype(Action) Action, decltype(DeferredCancel) DeferredCancel) {
+        DeferredCancel();
+        Action(llvm::Error::success());
+      },
+      std::move(Action), std::move(DeferredCancel));
+}
+
+void TUScheduler::runWithAST(
+    PathRef File, UniqueFunction<void(llvm::Expected<InputsAndAST>)> Action) {
+  auto Resources = Files.getFile(File);
+  if (!Resources) {
+    Action(llvm::make_error<llvm::StringError>(
+        "trying to get AST for non-added document",
+        llvm::errc::invalid_argument));
+    return;
+  }
+
+  const ParseInputs &Inputs = getInputs(File);
+  // We currently block the calling thread until AST is available and run the
+  // action on the calling thread to avoid inconsistent states coming from
+  // subsequent updates.
+  // FIXME(ibiryukov): this should be moved to the worker threads.
+  Resources->getAST().get()->runUnderLock([&](ParsedAST *AST) {
+    if (AST)
+      Action(InputsAndAST{Inputs, *AST});
+    else
+      Action(llvm::make_error<llvm::StringError>(
+          "Could not build AST for the latest file update",
+          llvm::errc::invalid_argument));
+  });
+}
+
+void TUScheduler::runWithPreamble(
+    PathRef File,
+    UniqueFunction<void(llvm::Expected<InputsAndPreamble>)> Action) {
+  std::shared_ptr<CppFile> Resources = Files.getFile(File);
+  if (!Resources) {
+    Action(llvm::make_error<llvm::StringError>(
+        "trying to get preamble for non-added document",
+        llvm::errc::invalid_argument));
+    return;
+  }
+
+  const ParseInputs &Inputs = getInputs(File);
+  std::shared_ptr<const PreambleData> Preamble =
+      Resources->getPossiblyStalePreamble();
+  Threads.addToFront(
+      [Resources, Preamble, Inputs](decltype(Action) Action) mutable {
+        if (!Preamble)
+          Preamble = Resources->getPossiblyStalePreamble();
+
+        Action(InputsAndPreamble{Inputs, Preamble.get()});
+      },
+      std::move(Action));
+}
+
+const ParseInputs &TUScheduler::getInputs(PathRef File) {
+  auto It = CachedInputs.find(File);
+  assert(It != CachedInputs.end());
+  return It->second;
+}
+
+std::vector<std::pair<Path, std::size_t>>
+TUScheduler::getUsedBytesPerFile() const {
+  return Files.getUsedBytesPerFile();
+}
+} // namespace clangd
+} // namespace clang

clangd/Threading.h

+//===--- ThreadPool.h --------------------------------------------*- 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_THREADING_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_THREADING_H
+
+#include "Function.h"
+#include <condition_variable>
+#include <deque>
+#include <mutex>
+#include <thread>
+#include <vector>
+
+namespace clang {
+namespace clangd {
+/// A simple fixed-size thread pool implementation.
+class ThreadPool {
+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.
+  ThreadPool(unsigned AsyncThreadsCount);
+  ~ThreadPool();
+
+  /// 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(
+          BindWithForward(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(
+          BindWithForward(std::forward<Func>(F), std::forward<Args>(As)...));
+    }
+    RequestCV.notify_one();
+  }
+
+private:
+  bool RunSynchronously;
+  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.
+  std::deque<UniqueFunction<void()>> RequestQueue;
+  /// Condition variable to wake up worker threads.
+  std::condition_variable RequestCV;
+};
+} // namespace clangd
+} // namespace clang
+#endif

clangd/Threading.cpp

+#include "Threading.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/Threading.h"
+
+namespace clang {
+namespace clangd {
+ThreadPool::ThreadPool(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]() {
+      llvm::set_thread_name(llvm::formatv("scheduler/{0}", I));
+      while (true) {
+        UniqueFunction<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 (RequestQueue.empty()) {
+            assert(Done);
+            return;
+          }
+
+          // We process requests starting from the front of the queue. Users of
+          // ThreadPool 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();
+      }
+    }));
+  }
+}
+
+ThreadPool::~ThreadPool() {
+  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();
+}
+} // namespace clangd
+} // namespace clang

unittests/clangd/CMakeLists.txt

add_extra_unittest(ClangdTests
   ContextTests.cpp
   FileIndexTests.cpp
   FuzzyMatchTests.cpp
   IndexTests.cpp
   JSONExprTests.cpp
   URITests.cpp
   TestFS.cpp
   TraceTests.cpp
+  TUSchedulerTests.cpp
   SourceCodeTests.cpp
   SymbolCollectorTests.cpp
   XRefsTests.cpp
   )
 
 target_link_libraries(ClangdTests
   PRIVATE
   clangBasic

unittests/clangd/TUSchedulerTests.cpp

+//===-- TUSchedulerTests.cpp ------------------------------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "TUScheduler.h"
+#include "TestFS.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include <algorithm>
+#include <utility>
+
+namespace clang {
+namespace clangd {
+
+using ::testing::Pair;
+using ::testing::Pointee;
+
+void ignoreUpdate(Context, llvm::Optional<std::vector<DiagWithFixIts>>) {}
+void ignoreError(llvm::Error Err) {
+  handleAllErrors(std::move(Err), [](const llvm::ErrorInfoBase &) {});
+}
+
+class TUSchedulerTests : public ::testing::Test {
+protected:
+  ParseInputs getInputs(PathRef File, std::string Contents) {
+    return ParseInputs{*CDB.getCompileCommand(File), buildTestFS(Files),
+                       std::move(Contents)};
+  }
+
+  void changeFile(PathRef File, std::string Contents) {
+    Files[File] = Contents;
+  }
+
+private:
+  llvm::StringMap<std::string> Files;
+  MockCompilationDatabase CDB;
+};
+
+TEST_F(TUSchedulerTests, MissingFiles) {
+  TUScheduler S(getDefaultAsyncThreadsCount(),
+                /*StorePreamblesInMemory=*/true,
+                /*ASTParsedCallback=*/nullptr);
+
+  auto Added = getVirtualTestFilePath("added.cpp");
+  changeFile(Added, "");
+
+  auto Missing = getVirtualTestFilePath("missing.cpp");
+  changeFile(Missing, "");
+
+  S.update(Context::empty(), Added, getInputs(Added, ""), ignoreUpdate);
+
+  // Assert each operation for missing file is an error (even if it's available
+  // in VFS).
+  S.runWithAST(Missing, [&](llvm::Expected<InputsAndAST> AST) {
+    ASSERT_FALSE(bool(AST));
+    ignoreError(AST.takeError());
+  });
+  S.runWithPreamble(Missing, [&](llvm::Expected<InputsAndPreamble> Preamble) {
+    ASSERT_FALSE(bool(Preamble));
+    ignoreError(Preamble.takeError());
+  });
+  S.remove(Missing, [&](llvm::Error Err) {
+    EXPECT_TRUE(bool(Err));
+    ignoreError(std::move(Err));
+  });
+
+  // Assert there aren't any errors for added file.
+  S.runWithAST(
+      Added, [&](llvm::Expected<InputsAndAST> AST) { EXPECT_TRUE(bool(AST)); });
+  S.runWithPreamble(Added, [&](llvm::Expected<InputsAndPreamble> Preamble) {
+    EXPECT_TRUE(bool(Preamble));
+  });
+  S.remove(Added, [&](llvm::Error Err) { EXPECT_FALSE(bool(Err)); });
+
+  // Assert that all operations fail after removing the file.
+  S.runWithAST(Added, [&](llvm::Expected<InputsAndAST> AST) {
+    ASSERT_FALSE(bool(AST));
+    ignoreError(AST.takeError());
+  });
+  S.runWithPreamble(Added, [&](llvm::Expected<InputsAndPreamble> Preamble) {
+    ASSERT_FALSE(bool(Preamble));
+    ignoreError(Preamble.takeError());
+  });
+  S.remove(Added, [&](llvm::Error Err) {
+    EXPECT_TRUE(bool(Err));
+    ignoreError(std::move(Err));
+  });
+}
+
+TEST_F(TUSchedulerTests, ManyUpdates) {
+  const int FilesCount = 3;
+  const int UpdatesPerFile = 10;
+
+  std::mutex Mut;
+  int TotalASTReads = 0;
+  int TotalPreambleReads = 0;
+  int TotalUpdates = 0;
+
+  // Run TUScheduler and collect some stats.
+  {
+    TUScheduler S(getDefaultAsyncThreadsCount(),
+                  /*StorePreamblesInMemory=*/true,
+                  /*ASTParsedCallback=*/nullptr);
+
+    std::vector<std::string> Files;
+    for (int I = 0; I < FilesCount; ++I) {
+      Files.push_back(
+          getVirtualTestFilePath("foo" + std::to_string(I) + ".cpp").str());
+      changeFile(Files.back(), "");
+    }
+
+    llvm::StringRef Contents1 = R"cpp(int a;)cpp";
+    llvm::StringRef Contents2 = R"cpp(int main() { return 1; })cpp";
+    llvm::StringRef Contents3 =
+        R"cpp(int a; int b; int sum() { return a + b; })cpp";
+
+    llvm::StringRef AllContents[] = {Contents1, Contents2, Contents3};
+    const int AllContentsSize = 3;
+
+    for (int FileI = 0; FileI < FilesCount; ++FileI) {
+      for (int UpdateI = 0; UpdateI < UpdatesPerFile; ++UpdateI) {
+        auto Contents = AllContents[(FileI + UpdateI) % AllContentsSize];
+
+        auto File = Files[FileI];
+        auto Inputs = getInputs(File, Contents.str());
+        static Key<std::pair<int, int>> FileAndUpdateKey;
+        auto Ctx = Context::empty().derive(FileAndUpdateKey,
+                                           std::make_pair(FileI, UpdateI));
+        S.update(std::move(Ctx), File, Inputs,
+                 [FileI, UpdateI, &Mut, &TotalUpdates](
+                     Context Ctx,
+                     llvm::Optional<std::vector<DiagWithFixIts>> Diags) {
+                   EXPECT_THAT(Ctx.get(FileAndUpdateKey),
+                               Pointee(Pair(FileI, UpdateI)));
+
+                   std::lock_guard<std::mutex> Lock(Mut);
+                   ++TotalUpdates;
+                 });
+
+        S.runWithAST(File, [Inputs, &Mut,
+                            &TotalASTReads](llvm::Expected<InputsAndAST> AST) {
+          ASSERT_TRUE((bool)AST);
+          EXPECT_EQ(AST->Inputs.FS, Inputs.FS);
+          EXPECT_EQ(AST->Inputs.Contents, Inputs.Contents);
+
+          std::lock_guard<std::mutex> Lock(Mut);
+          ++TotalASTReads;
+        });
+
+        S.runWithPreamble(
+            File, [Inputs, &Mut, &TotalPreambleReads](
+                      llvm::Expected<InputsAndPreamble> Preamble) {
+              ASSERT_TRUE((bool)Preamble);
+              EXPECT_EQ(Preamble->Inputs.FS, Inputs.FS);
+              EXPECT_EQ(Preamble->Inputs.Contents, Inputs.Contents);
+
+              std::lock_guard<std::mutex> Lock(Mut);
+              ++TotalPreambleReads;
+            });
+      }
+    }
+  } // TUScheduler destructor waits for all operations to finish.
+
+  std::lock_guard<std::mutex> Lock(Mut);
+  EXPECT_EQ(TotalUpdates, FilesCount * UpdatesPerFile);
+  EXPECT_EQ(TotalASTReads, FilesCount * UpdatesPerFile);
+  EXPECT_EQ(TotalPreambleReads, FilesCount * UpdatesPerFile);
+}
+
+} // namespace clangd
+} // namespace clang