[clangd] Build ASTs only with fresh preambles or after building a new preamble

[kadircet]

Summary

This is another step for out-of-order preamble builds. To keep the
diagnostic behavior same, we only build ASTs either with "reusable" preambles,
the ones that are fully applicable to a given ParseInput, or after building a
new preamble. Which is the same behaviour as what we do today.

ASTs built through preamble callbacks are not cached as they are built on a
different thread and ASTWorker heavily relies on being the only thread updating
cached ASTs. This results in possibly building some ASTs twice (when there's an
immediate read after a preamble built without any write in between).

[sammccall]

So I think we need a careful description of the new model somewhere. Not so much on specific changes/constraints parts of the code operate under, but what it's trying to do.

My best understanding is:

• In general, read requests are processed in order by the ASTWorker, and get exactly the version of the file when the request was received. That is, the ASTworker queue interleaves updates and reads in the same order as LSP.
• Preambles are built on the PreambleWorker for a certain version of the file, and may be compatible or incompatible with subsequent versions. If a read is "picky" it may block the ASTWorker queue waiting for a compatible preamble, otherwise it will "patch up" an incompatible one.
• Diagnostics and indexes are updated opportunistically using onMainAST when an AST is built, but only if the used preamble is compatible. Publishing diagnostics from patched preambles is not allowed, but we also do not block waiting for an up-to-date preamble in order to generate publishable diagnostics. (Exception: if WantDiagnostics=Yes, we block).
• To ensure forward progress with diagnostics if the queue is full, when preamble is built we immediately build a "golden" AST from that version and publish diagnostics. To ensure diagnostics do not "go backwards", opportunistic diagnostics are suppressed if they don't use the latest preamble (the one with the last golden AST). This means the sequence of ASTs producing diagnostics is neatly partitioned into "epochs" of the preambles, the first in the epoch is always the golden version. (Exception: if WantDiagnostics=no, diagnostics are not emitted for the golden AST). Because these versions are built out-of-order (not interleaved with reads per LSP), the golden AST is not cached and reused for reads.

Is this about right?

Based on this, I do think the model would be much easier to reason about and verify if the golden AST was built on the ASTWorker thread in a queued task, rather than on the PreambleWorker thread in a callback. This means we understand allowed sequencing by looking at the queue rather than reasoning about mutexes/threads with multiple possible interleavings. By pushing the golden AST task to the front of the queue, it's more explicit that this is a priorities/scheduling decision. Following what happened in logs is probably easier due to less interleaving.
(The callback could enqueue the task, or the PreambleWorker could just enqueue the task directly at that point - not sure the callback indirection buys anything).

One thing that seems complicated in the model is that the AST build needs to decide whether to block on a compatible preamble, but it's the following read that determines whether it needs to. In the worst case, the picky read hasn't even been scheduled yet. I think picky reads probably need to block on the target preamble and build their AST themselves if the cached one isn't suitable. If they "miss" their preamble (i.e. preambleVersion >= readVersion but the preamble isn't compatible) then I think we should fail the request (i.e. picky requests may be invalidated by subsequent preamble edits).

[kadircet]

In D76725#1940282, @sammccall wrote:

So I think we need a careful description of the new model somewhere. Not so much on specific changes/constraints parts of the code operate under, but what it's trying to do.

My best understanding is:

• In general, read requests are processed in order by the ASTWorker, and get exactly the version of the file when the request was received. That is, the ASTworker queue interleaves updates and reads in the same order as LSP.
• Preambles are built on the PreambleWorker for a certain version of the file, and may be compatible or incompatible with subsequent versions. If a read is "picky" it may block the ASTWorker queue waiting for a compatible preamble, otherwise it will "patch up" an incompatible one.
• Diagnostics and indexes are updated opportunistically using onMainAST when an AST is built, but only if the used preamble is compatible. Publishing diagnostics from patched preambles is not allowed, but we also do not block waiting for an up-to-date preamble in order to generate publishable diagnostics. (Exception: if WantDiagnostics=Yes, we block).
• To ensure forward progress with diagnostics if the queue is full, when preamble is built we immediately build a "golden" AST from that version and publish diagnostics. To ensure diagnostics do not "go backwards", opportunistic diagnostics are suppressed if they don't use the latest preamble (the one with the last golden AST). This means the sequence of ASTs producing diagnostics is neatly partitioned into "epochs" of the preambles, the first in the epoch is always the golden version. (Exception: if WantDiagnostics=no, diagnostics are not emitted for the golden AST). Because these versions are built out-of-order (not interleaved with reads per LSP), the golden AST is not cached and reused for reads.

Is this about right?

Yes that's most of what this patch does. Tried to explain/sum up those in file comments.

Based on this, I do think the model would be much easier to reason about and verify if the golden AST was built on the ASTWorker thread in a queued task, rather than on the PreambleWorker thread in a callback. This means we understand allowed sequencing by looking at the queue rather than reasoning about mutexes/threads with multiple possible interleavings. By pushing the golden AST task to the front of the queue, it's more explicit that this is a priorities/scheduling decision. Following what happened in logs is probably easier due to less interleaving.
(The callback could enqueue the task, or the PreambleWorker could just enqueue the task directly at that point - not sure the callback indirection buys anything).

Done. It is still the callback that enqueues the task, but moved storage of latest build preamble from PreambleThread to ASTWorker to prevent a possible race between currently running update in astworker and build of preamble in preamblethread, as discussed offline.
This also enabled us to cache golden ASTs in case ASTWorker hasn't received any updates in between which gets rid of the additional AST build cost.

One thing that seems complicated in the model is that the AST build needs to decide whether to block on a compatible preamble, but it's the following read that determines whether it needs to. In the worst case, the picky read hasn't even been scheduled yet. I think picky reads probably need to block on the target preamble and build their AST themselves if the cached one isn't suitable. If they "miss" their preamble (i.e. preambleVersion >= readVersion but the preamble isn't compatible) then I think we should fail the request (i.e. picky requests may be invalidated by subsequent preamble edits).

Agreed, but this is not in the scope of this patch, will address that in the upcoming patches after dropping the synchronization between preamble and ast thread. In addition to that i am planning to ensure updates with WantDiagnostics::Yes gets build by making them non-overwritable in PreambleThread. That is any subsequent build requests will block (astworker thread) until update with WantDiags::Yes starts building.

[kadircet]

• Build golden asts in ASTWorker thread, rather than preamble thread.
• Explain new model in file comments.

[kadircet]

• Invalidate cached AST in case of preamble/input changes when building golden ASTs

[sammccall]

clang-tools-extra/clangd/TUScheduler.cpp

7–8

This is explaining a pretty complicated thing, so I think it's particularly important to clearly organize the explanation, use consistent terminology, and avoid including unneccesary details.

I'd suggest introducing with paragraphs in this order:

• TUScheduler and the idea of a worker per TU.
• ASTWorker thread, the queue of interleaved updates and reads, elision of dead updates.
• Compatibility of preambles and updates, picky reads.
• PreambleWorker thread, elision of preambles, golden ASTs.
• Published ASTs, epochs/monotonicity.

10

nit: name the other thread (there is a second PreambleWorker thread ...)

11

"issues updates" is vague. What about "... enqueues rebuilds on the PreambleWorker thread as preamble becomes stale. Currently, it then immediately blocks on that request."

11–12

I'm not sure why this paragraph is here rather than near the description of the ASTWorker queue.

16–18

I'm not sure what "serving preambles" means.

18

You haven't defined "compatible" anywhere, and it's not obvious what it means.

19–20

I think this no longer belongs here - merge with TUScheduler::Options::AsyncThreadsCount or just drop

22

The wording here (particularly the word "phases") implies sequencing: an update results in a preamble build followed by an ast build in sequence, when in fact it usually results in an ast build and preamble build in parallel (the former usually finishing first) with a second ast build sequenced after the preamble build.

24

This isn't true, an AST is built for an update if it is needed (a read is based on it, wantdiagnostics=yes, it changed the preamble, or wantdiagnostics=maybe and the debounce expired)

28

I'm not sure what compatible means here, but we will certainly build ASTs for incompatible preambles.
(Or is this describing the intermediate state as of this patch, with the plan to rewrite it all next patch? I think we should rather describe the new model, and then document the current deviations from it)

28

I'm not sure what "immediately following" is relative to: if it's after the update() call, it's not immediate - it has to wait in the queue.

28–29

This section is basically obsolete - elision of dead updates is (should be) mentioned above, and the details are stale (e.g. it refers to debounce as a possibility) and documented in-line.

32

"and *the* index *are* updated"... "building *the* AST"

32

you're using "write" and "update" interchangeably here, it would be clearer to be consistent

32

I think "as a result of" is just "after"?

33

Now we're using "updates" to mean something different...

35

It's not clear what this sentence is saying: some code is doing this partitioning?

(I know you're describing a property of the output of the process, but it's not clear from the text)

212

This never varies, why use a callback rather than just calling a function on the ASTWorker?
(I think I asked this in the previous version, but didn't see an answer)

236–238

Took me a while to understand what this comment is saying (why is it talking about ASTWorker queue state?)

Maybe:
// If we're shutting down, don't bother building preambles, reads will be cancelled.

313–314

this seems to never reuse?

313–314

That sounds racy to me - what happens if a new request arrives after build() invalidates but before build() returns?

Seems much more robust for PreambleWorker to hold the last built preamble, and ASTWorker to hold the last usable preamble. It's just a shared_ptr, what's the harm :)

332

storing this separately from the latest preamble (i.e. separately from where the decision is made to build/rebuild) seems a bit suspect (not actually wrong, but points towards storing the latest built preamble here too)

587

this function is getting too long/monolithic, we should find a way to split it up.
ASTTask should be a member function, I think.

587

I also lost track of the control flow halfway through the function, and can't work it out.
I don't know really what to concretely advise, but it needs to be a lot simpler, maybe by finding some better abstractions or just better understanding conceptually what this is doing.

635

You've called this a task, but it never gets scheduled on a queue, just called in various places. This seems like a helper function, and needs a better name.

649

Didn't we decide to make this available only after building the golden AST and publishing diagnostics?

[kadircet]

• Update description and address comments

[kadircet]

clang-tools-extra/clangd/TUScheduler.cpp

587

ASTTask should be a member function, I think.

Moved it into a meber named publishDiagnostics

649

as discussed offline, this is harmless since it is being executed on ASTWorker thread and makes new preamble available to code completion/signature help faster.

[kadircet]

• Make use of a separate queue for golden ASTs to prevent any races that might

occur due PreambleThread finishing multiple preambles before ASTWorker gets to
process others.

[kadircet]

• Add assertion to explicitly spell out scheduling for golden ASTs.

[sammccall]

Had some discussion offline.

• having ASTWorker not worry about preamble validity before dispatching to preambleworker seems like a win
• for this, preambleworker needs to call preambleReady whether it's new or not, so ASTWorker can produce diagnostics
• AST reuse from diagnostics->request seems much more useful than the other way around (e.g. it reduces request latency), so don't bother with the latter. (And we can drop diagnostic computation in some cases)

This yields pseudocode like:

ASTWorker::update(): enqueue({
  currentInputs = inputs
  preambleworker::update(inputs)
})
ASTWorker::runWithAST(): enqueue({
  ast = cache.get()
  if (!ast) patch preamble and build ast
  action(ast)
  cache.put(ast)
})
PreambleWorker::update(): enqueue({
  if (!preamble.compatible(inputs))
    build preamble
  ASTWorker::preambleReady(preamble)
})
ASTWorker::preambleReady(): enqueue({
  preamble = p
  build ast
  publish ast.diagnostics
  if (inputs == currentInputs) cache.put(ast)
  else if (preamble != oldPreamble) cache.get() // force next read to use this preamble
})

(I'm not sure how simple the actual code can be. I do think defining the methods in that order may help readability)

clang-tools-extra/clangd/TUScheduler.cpp

8–10

nit: just "This ASTWorker manages updates..."

8–10

uber-nit: I'd say the scheduler *manages* workers and the worker *processes* updates.
(Just to avoid mixing metaphors)

9

nit: space before open parens (not a big deal, but occurs several times)

16–18

changes latest inputs -> replaces the current parser inputs

("changes" suggests some tricky mutation, and inputs isn't defined here)

17

You need to mention "building an AST" here, as you reference it below.

29–47

nit: we use "main file" or "main-file" much more often than "mainfile".

30

it it isn't -> until?

33

I don't think "usable" is a good name for this, because we *use* preambles that are not "usable".

I think "compatible" or "up-to-date" or "fresh" or so would have enough semantic distance to make this less confusing.

41

Nit: first and third sentences in this paragraph say the same thing, drop the third?

47

This gets a bit far into implementation details (uses of queues etc) and is hard to follow. Maybe:

When a new preamble is built, a "golden" AST is immediately built from that version of the file.
This ensures diagnostics get updated even if the queue is full.

52

This is talking about the subtleties of ordering/version guarantees - not really interesting from the outside, and not enough detail here to avoid reading the code. I'd suggest moving these to preambleReady or so.

224

Some params get copied explicitly, others get passed by value. Pick one strategy?

226

nit: WantDiags is an enum, no move

246–247

The prefix "PreambleThread" implies that there's something in the code called PreambleThread, and that there's one such thing, neither are true.

(I also think there's a risk in adding ad-hoc prefixes to things that it adds noise and becomes less readable to people with little context, though at dlog that's less of an issue)

314–316

this sentence doesn't parse :-)

333

Please don't use initialisms like this for members, they're hard to follow out of context.
I'd suggest ASTPeer (and PreamblePeer) or just Peer, to indicate these threads are tightly-coupled partners.

402

nit: move somewhere more appropriate (near getCurrentPreamble or update if modelling as public, else into the private section?)

402

nit: use a verb for mutating operations (updatePreamble)?

404

This name is confusing, almost all of its implementation is building AST.

generateDiagnostics?

477

PreambleRequests or GoldenASTRequests or so? The queue doens't actually contain preambles.

537

Placing this between ASTWorker and its implementation seems confusing - move it down?

550

This is going to trigger a rebuild of the fileindex - is it really necessary?

590

"Golden AST from preamble"?
Current text is not terribly descriptive...

604

incomplete comment

635

this isn't a safe assert - it does IO and could transiently become true/false

[kadircet]

• Update comments and re-order functions

[kadircet]

preambleReady part is a little bit different than you've described, it looks more like:

ASTWorker::preambleReady(): enqueue({
  if(preamble!=lastPreamble) {
     lastPreamble = preamble
     cache.get(); // force next read to use this preamble
  }
  build ast
  publish ast.diagnostics
  if (inputs == currentInputs) cache.put(ast)
})

also the last 3 steps are on a different function called generateDiagnostics since the code is a little bit long.
it immediately follows the preambleReady(now named updatePreamble).

clang-tools-extra/clangd/TUScheduler.cpp

17

just saying "it will issue a build for new inputs", which is explained in details in the next paragraph.

33

using compatible.

333

PW was for PeerWorker :P

550

astworker won't trigger that if inputs are the same.

590

this is not necessarily a golden ast now. just changing it to "Build AST"

635

right, thanks for catching it!

the one above is also not true (assert(LatestPreamble)) as preamble build might've failed. I suppose we would still want to move forward so that we can surface diagnostics saying why it failed.

[kadircet]

• Rebase

[sammccall]

Thanks for your patience!
This is very nice, only +100 lines on TUScheduler.cpp in the end!

clang-tools-extra/clangd/TUScheduler.cpp

224

comment is now obsolete, remove it and consider inlining Req

325

this is no longer locked or concurrently accessed right?

378

nit: last sentence is a fragment, I think you want to swap the dot and comma:

...through this callback. This ensures... of the file, as the callback...

408

Unused, I think?

793

I'd suggest InputsAreLatest or something that hints at *why* they might differ

810

This is reusing the AST built for a read if the file (and preamble) hasn't changed, right?

I think it's helpful to explain where teh cache data might be coming from.

IIRC this is the case we thought wasn't so important. It competes with a different optimization: skipping clang-tidy, warning analysis, etc when building ASTs because we know their diagnostics will never be used.

This might be a good place to leave a comment like "FIXME: maybe we're better off never reusing this AST, so queued AST builds aren't required to produce diagnostics" or so

862

This explains the what but not the why.
(We may be building an older version of the source, as the queue raced ahead while we were waiting on the preamble. In that case the queue can't reuse the AST we built.)
Comment could go here or on InputsAreTheSame

[kadircet]

• Address comments and rebase

[kadircet]

clang-tools-extra/clangd/TUScheduler.cpp

224

i would rather keep it out-of-line since it is used in two places.

325

right, deleted annottation and added a comment.

clang-tools-extra/clangd/Preamble.h

 //===--- Preamble.h - Reusing expensive parts of the AST ---------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_PREAMBLE_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_PREAMBLE_H
 
 #include "CollectMacros.h"
 #include "Compiler.h"
 #include "Diagnostics.h"
 #include "FS.h"
 #include "Headers.h"
+#include "Path.h"
 #include "index/CanonicalIncludes.h"
+#include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/PrecompiledPreamble.h"
 #include "clang/Tooling/CompilationDatabase.h"
 
 #include <memory>
 #include <string>
 #include <vector>
 
 namespace clang {
struct PreambleData {
   CanonicalIncludes CanonIncludes;
 };
 
 using PreambleParsedCallback =
     std::function<void(ASTContext &, std::shared_ptr<clang::Preprocessor>,
                        const CanonicalIncludes &)>;
 
 /// Build a preamble for the new inputs unless an old one can be reused.
-/// If \p OldPreamble can be reused, it is returned unchanged.
-/// If \p OldPreamble is null, always builds the preamble.
 /// If \p PreambleCallback is set, it will be run on top of the AST while
-/// building the preamble. Note that if the old preamble was reused, no AST is
-/// built and, therefore, the callback will not be executed.
+/// building the preamble.
 std::shared_ptr<const PreambleData>
 buildPreamble(PathRef FileName, CompilerInvocation CI,
-              std::shared_ptr<const PreambleData> OldPreamble,
               const ParseInputs &Inputs, bool StoreInMemory,
               PreambleParsedCallback PreambleCallback);
 
+/// Returns true if \p Preamble is reusable for \p Inputs. Note that it will
+/// return true when some missing headers are now available.
+/// FIXME: Should return more information about the delta between \p Preamble
+/// and \p Inputs, e.g. new headers.
+bool isPreambleCompatible(const PreambleData &Preamble,
+                          const ParseInputs &Inputs, PathRef FileName,
+                          const CompilerInvocation &CI);
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_PREAMBLE_H

clang-tools-extra/clangd/Preamble.cpp

 //===--- Preamble.cpp - Reusing expensive parts of the AST ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
PreambleData::PreambleData(const ParseInputs &Inputs,
     : Version(Inputs.Version), CompileCommand(Inputs.CompileCommand),
       Preamble(std::move(Preamble)), Diags(std::move(Diags)),
       Includes(std::move(Includes)), Macros(std::move(Macros)),
       StatCache(std::move(StatCache)), CanonIncludes(std::move(CanonIncludes)) {
 }
 
 std::shared_ptr<const PreambleData>
 buildPreamble(PathRef FileName, CompilerInvocation CI,
-              std::shared_ptr<const PreambleData> OldPreamble,
               const ParseInputs &Inputs, bool StoreInMemory,
               PreambleParsedCallback PreambleCallback) {
   // Note that we don't need to copy the input contents, preamble can live
   // without those.
   auto ContentsBuffer =
       llvm::MemoryBuffer::getMemBuffer(Inputs.Contents, FileName);
   auto Bounds =
       ComputePreambleBounds(*CI.getLangOpts(), ContentsBuffer.get(), 0);
 
-  if (OldPreamble &&
-      compileCommandsAreEqual(Inputs.CompileCommand,
-                              OldPreamble->CompileCommand) &&
-      OldPreamble->Preamble.CanReuse(CI, ContentsBuffer.get(), Bounds,
-                                     Inputs.FS.get())) {
-    vlog("Reusing preamble version {0} for version {1} of {2}",
-         OldPreamble->Version, Inputs.Version, FileName);
-    return OldPreamble;
-  }
-  if (OldPreamble)
-    vlog("Rebuilding invalidated preamble for {0} version {1} "
-         "(previous was version {2})",
-         FileName, Inputs.Version, OldPreamble->Version);
-  else
-    vlog("Building first preamble for {0} verson {1}", FileName,
-         Inputs.Version);
-
   trace::Span Tracer("BuildPreamble");
   SPAN_ATTACH(Tracer, "File", FileName);
   StoreDiags PreambleDiagnostics;
   llvm::IntrusiveRefCntPtr<DiagnosticsEngine> PreambleDiagsEngine =
       CompilerInstance::createDiagnostics(&CI.getDiagnosticOpts(),
                                           &PreambleDiagnostics, false);
 
   // Skip function bodies when building the preamble to speed up building
return std::make_shared<PreambleData>(
         SerializedDeclsCollector.takeCanonicalIncludes());
   } else {
     elog("Could not build a preamble for file {0} version {1}", FileName,
          Inputs.Version);
     return nullptr;
   }
 }
 
+bool isPreambleCompatible(const PreambleData &Preamble,
+                          const ParseInputs &Inputs, PathRef FileName,
+                          const CompilerInvocation &CI) {
+  auto ContentsBuffer =
+      llvm::MemoryBuffer::getMemBuffer(Inputs.Contents, FileName);
+  auto Bounds =
+      ComputePreambleBounds(*CI.getLangOpts(), ContentsBuffer.get(), 0);
+  return compileCommandsAreEqual(Inputs.CompileCommand,
+                                 Preamble.CompileCommand) &&
+         Preamble.Preamble.CanReuse(CI, ContentsBuffer.get(), Bounds,
+                                    Inputs.FS.get());
+}
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/TUScheduler.cpp

 //===--- TUScheduler.cpp -----------------------------------------*-C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
-// For each file, managed by TUScheduler, we create a single ASTWorker that
-// manages an AST for that file. All operations that modify or read the AST are
-// run on a separate dedicated thread asynchronously in FIFO order.
+// TUScheduler manages a worker per active file. This ASTWorker processes
+// updates (modifications to file contents) and reads (actions performed on
+// preamble/AST) to the file.
 //
-// We start processing each update immediately after we receive it. If two or
-// more updates come subsequently without reads in-between, we attempt to drop
-// an older one to not waste time building the ASTs we don't need.
+// Each ASTWorker owns a dedicated thread to process updates and reads to the
+// relevant file. Any request gets queued in FIFO order to be processed by that
+// thread.
 //
-// The processing thread of the ASTWorker is also responsible for building the
-// preamble. However, unlike AST, the same preamble can be read concurrently, so
-// we run each of async preamble reads on its own thread.
+// An update request replaces current praser inputs to ensure any subsequent
+// read sees the version of the file they were requested. It will also issue a
+// build for new inputs.
 //
-// To limit the concurrent load that clangd produces we maintain a semaphore
-// that keeps more than a fixed number of threads from running concurrently.
+// ASTWorker processes the file in two parts, a preamble and a main-file
+// section. A preamble can be reused between multiple versions of the file until
+// invalidated by a modification to a header, compile commands or modification
+// to relevant part of the current file. Such a preamble is called compatible.
+// An update is considered dead if no read was issued for that version and
+// diagnostics weren't requested by client or could be generated for a later
+// version of the file. ASTWorker eliminates such requests as they are
+// redundant.
 //
-// Rationale for cancelling updates.
-// LSP clients can send updates to clangd on each keystroke. Some files take
-// significant time to parse (e.g. a few seconds) and clangd can get starved by
-// the updates to those files. Therefore we try to process only the last update,
-// if possible.
-// Our current strategy to do that is the following:
-// - For each update we immediately schedule rebuild of the AST.
-// - Rebuild of the AST checks if it was cancelled before doing any actual work.
-//   If it was, it does not do an actual rebuild, only reports llvm::None to the
-//   callback
-// - When adding an update, we cancel the last update in the queue if it didn't
-//   have any reads.
-// There is probably a optimal ways to do that. One approach we might take is
-// the following:
-// - For each update we remember the pending inputs, but delay rebuild of the
-//   AST for some timeout.
-// - If subsequent updates come before rebuild was started, we replace the
-//   pending inputs and reset the timer.
-// - If any reads of the AST are scheduled, we start building the AST
-//   immediately.
+// In the presence of stale (non-compatible) preambles, ASTWorker won't publish
+// diagnostics for update requests. Read requests will be served with ASTs build
+// with stale preambles, unless the read is picky and requires a compatible
+// preamble. In such cases it will block until new preamble is built.
+//
+// ASTWorker owns a PreambleThread for building preambles. If the preamble gets
+// invalidated by an update request, a new build will be requested on
+// PreambleThread. Since PreambleThread only receives requests for newer
+// versions of the file, in case of multiple requests it will only build the
+// last one and skip requests in between. Unless client force requested
+// diagnostics(WantDiagnostics::Yes).
+//
+// When a new preamble is built, a "golden" AST is immediately built from that
+// version of the file. This ensures diagnostics get updated even if the queue
+// is full.
+//
+// Some read requests might just need preamble. Since preambles can be read
+// concurrently, ASTWorker runs these requests on their own thread. These
+// requests will receive latest build preamble, which might possibly be stale.
 
 #include "TUScheduler.h"
 #include "Cancellation.h"
 #include "Compiler.h"
 #include "Context.h"
 #include "Diagnostics.h"
 #include "GlobalCompilationDatabase.h"
 #include "Logger.h"
 #include "ParsedAST.h"
 #include "Path.h"
 #include "Preamble.h"
 #include "Threading.h"
 #include "Trace.h"
 #include "index/CanonicalIncludes.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Tooling/CompilationDatabase.h"
+#include "llvm/ADT/FunctionExtras.h"
+#include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Threading.h"
 #include <algorithm>
+#include <chrono>
+#include <functional>
 #include <memory>
 #include <mutex>
 #include <queue>
 #include <string>
 #include <thread>
+#include <type_traits>
+#include <utility>
+#include <vector>
 
 namespace clang {
 namespace clangd {
 using std::chrono::steady_clock;
 
 namespace {
 class ASTWorker;
 } // namespace
private:
   const Path FileName;
 
   std::mutex StatusMu;
   TUStatus Status;
   bool CanPublish = true;
   ParsingCallbacks &Callbacks;
 };
 
-/// Responsible for building and providing access to the preamble of a TU.
-/// Whenever the thread is idle and the preamble is outdated, it starts to build
-/// a fresh preamble from the latest inputs. If RunSync is true, preambles are
-/// built synchronously in update() instead.
+/// Responsible for building preambles. Whenever the thread is idle and the
+/// preamble is outdated, it starts to build a fresh preamble from the latest
+/// inputs. If RunSync is true, preambles are built synchronously in update()
+/// instead.
 class PreambleThread {
 public:
   PreambleThread(llvm::StringRef FileName, ParsingCallbacks &Callbacks,
                  bool StorePreambleInMemory, bool RunSync,
-                 SynchronizedTUStatus &Status)
+                 SynchronizedTUStatus &Status, ASTWorker &AW)
       : FileName(FileName), Callbacks(Callbacks),
-        StoreInMemory(StorePreambleInMemory), RunSync(RunSync), Status(Status) {
-  }
-
-  size_t getUsedBytes() const {
-    auto Preamble = latest();
-    return Preamble ? Preamble->Preamble.getSize() : 0;
-  }
+        StoreInMemory(StorePreambleInMemory), RunSync(RunSync), Status(Status),
+        ASTPeer(AW) {}
 
   /// It isn't guaranteed that each requested version will be built. If there
   /// are multiple update requests while building a preamble, only the last one
   /// will be built.
-  void update(CompilerInvocation *CI, ParseInputs PI) {
-    // If compiler invocation was broken, just fail out early.
-    if (!CI) {
-      // Make sure anyone waiting for the preamble gets notified it could not be
-      // built.
-      BuiltFirst.notify();
-      return;
-    }
+  void update(std::unique_ptr<CompilerInvocation> CI, ParseInputs PI,
+              std::vector<Diag> CIDiags, WantDiagnostics WantDiags) {
     // Make possibly expensive copy while not holding the lock.
-    Request Req = {std::make_unique<CompilerInvocation>(*CI), std::move(PI)};
+    Request Req = {std::move(CI), std::move(PI), std::move(CIDiags), WantDiags,
+                   Context::current().clone()};
     if (RunSync) {
       build(std::move(Req));
+      Status.update([&](TUStatus &Status) {
+        Status.PreambleActivity = PreambleAction::Idle;
+      });
       return;
     }
     {
       std::lock_guard<std::mutex> Lock(Mutex);
-      assert(!Done && "Build request to PreambleWorker after stop");
+      // If shutdown is issued, don't bother building.
+      if (Done)
+        return;
       NextReq = std::move(Req);
     }
     // Let the worker thread know there's a request, notify_one is safe as there
     // should be a single worker thread waiting on it.
     ReqCV.notify_all();
   }
 
-  /// Blocks until at least a single request has been processed. Note that it
-  /// will unblock even after an unsuccessful build.
-  void waitForFirst() const { BuiltFirst.wait(); }
-
-  /// Returns the latest built preamble, might be null if no preamble has been
-  /// built or latest attempt resulted in a failure.
-  std::shared_ptr<const PreambleData> latest() const {
-    std::lock_guard<std::mutex> Lock(Mutex);
-    return LatestBuild;
-  }
-
   void run() {
-    dlog("Starting preamble worker for {0}", FileName);
     while (true) {
       {
         std::unique_lock<std::mutex> Lock(Mutex);
         assert(!CurrentReq && "Already processing a request?");
         // Wait until stop is called or there is a request.
         ReqCV.wait(Lock, [this] { return NextReq || Done; });
         if (Done)
           break;
         CurrentReq = std::move(*NextReq);
         NextReq.reset();
       }
+
+      WithContext Guard(std::move(CurrentReq->Ctx));
       // Build the preamble and let the waiters know about it.
       build(std::move(*CurrentReq));
       bool IsEmpty = false;
       {
         std::lock_guard<std::mutex> Lock(Mutex);
         CurrentReq.reset();
         IsEmpty = !NextReq.hasValue();
       }
       if (IsEmpty) {
         // We don't perform this above, before waiting for a request to make
         // tests more deterministic. As there can be a race between this thread
         // and client thread(clangdserver).
         Status.update([](TUStatus &Status) {
           Status.PreambleActivity = PreambleAction::Idle;
         });
       }
-      BuiltFirst.notify();
       ReqCV.notify_all();
     }
-    // We are no longer going to build any preambles, let the waiters know that.
-    BuiltFirst.notify();
-    dlog("Preamble worker for {0} finished", FileName);
+    dlog("Preamble worker for {0} stopped", FileName);
   }
 
   /// Signals the run loop to exit.
   void stop() {
-    dlog("Stopping preamble worker for {0}", FileName);
+    dlog("Preamble worker for {0} received stop", FileName);
     {
       std::lock_guard<std::mutex> Lock(Mutex);
       Done = true;
+      NextReq.reset();
     }
     // Let the worker thread know that it should stop.
     ReqCV.notify_all();
   }
 
   bool blockUntilIdle(Deadline Timeout) const {
     std::unique_lock<std::mutex> Lock(Mutex);
     return wait(Lock, ReqCV, Timeout, [&] { return !NextReq && !CurrentReq; });
   }
 
 private:
   /// Holds inputs required for building a preamble. CI is guaranteed to be
   /// non-null.
   struct Request {
     std::unique_ptr<CompilerInvocation> CI;
     ParseInputs Inputs;
+    std::vector<Diag> CIDiags;
+    WantDiagnostics WantDiags;
+    Context Ctx;
   };
 
   bool isDone() {
     std::lock_guard<std::mutex> Lock(Mutex);
     return Done;
   }
 
-  /// Builds a preamble for Req and caches it. Might re-use the latest built
-  /// preamble if it is valid for Req. Also signals waiters about the build.
-  /// FIXME: We shouldn't cache failed preambles, if we've got a successful
-  /// build before.
-  void build(Request Req) {
-    assert(Req.CI && "Got preamble request with null compiler invocation");
-    const ParseInputs &Inputs = Req.Inputs;
-    std::shared_ptr<const PreambleData> OldPreamble =
-        Inputs.ForceRebuild ? nullptr : latest();
-
-    Status.update([&](TUStatus &Status) {
-      Status.PreambleActivity = PreambleAction::Building;
-    });
-
-    auto Preamble = clang::clangd::buildPreamble(
-        FileName, std::move(*Req.CI), OldPreamble, Inputs, StoreInMemory,
-        [this, Version(Inputs.Version)](
-            ASTContext &Ctx, std::shared_ptr<clang::Preprocessor> PP,
-            const CanonicalIncludes &CanonIncludes) {
-          Callbacks.onPreambleAST(FileName, Version, Ctx, std::move(PP),
-                                  CanonIncludes);
-        });
-    {
-      std::lock_guard<std::mutex> Lock(Mutex);
-      // LatestBuild might be the last reference to old preamble, do not trigger
-      // destructor while holding the lock.
-      std::swap(LatestBuild, Preamble);
-    }
-  }
+  /// Builds a preamble for \p Req, might reuse LatestBuild if possible.
+  /// Notifies ASTWorker after build finishes.
+  void build(Request Req);
 
   mutable std::mutex Mutex;
   bool Done = false;                  /* GUARDED_BY(Mutex) */
   llvm::Optional<Request> NextReq;    /* GUARDED_BY(Mutex) */
   llvm::Optional<Request> CurrentReq; /* GUARDED_BY(Mutex) */
   // Signaled whenever a thread populates NextReq or worker thread builds a
   // Preamble.
   mutable std::condition_variable ReqCV;           /* GUARDED_BY(Mutex) */
   std::shared_ptr<const PreambleData> LatestBuild; /* GUARDED_BY(Mutex) */
 
-  Notification BuiltFirst;
   const Path FileName;
   ParsingCallbacks &Callbacks;
   const bool StoreInMemory;
   const bool RunSync;
 
   SynchronizedTUStatus &Status;
+  ASTWorker &ASTPeer;
 };
 
 class ASTWorkerHandle;
 
 /// Owns one instance of the AST, schedules updates and reads of it.
 /// Also responsible for building and providing access to the preamble.
 /// Each ASTWorker processes the async requests sent to it on a separate
 /// dedicated thread.
public:
   void
   runWithAST(llvm::StringRef Name,
              llvm::unique_function<void(llvm::Expected<InputsAndAST>)> Action,
              TUScheduler::ASTActionInvalidation);
   bool blockUntilIdle(Deadline Timeout) const;
 
   std::shared_ptr<const PreambleData> getPossiblyStalePreamble() const;
 
+  /// Used to inform ASTWorker about a new preamble build by PreambleThread.
+  /// Diagnostics are only published through this callback, which ensures they
+  /// are always for newer versions of the file. As this callback gets called in
+  /// the same order as update requests.
+  void updatePreamble(std::unique_ptr<CompilerInvocation> CI, ParseInputs PI,
+                      std::shared_ptr<const PreambleData> Preamble,
+                      std::vector<Diag> CIDiags, WantDiagnostics WantDiags);
+
   /// Obtain a preamble reflecting all updates so far. Threadsafe.
   /// It may be delivered immediately, or later on the worker thread.
   void getCurrentPreamble(
       llvm::unique_function<void(std::shared_ptr<const PreambleData>)>);
   /// Returns compile command from the current file inputs.
   tooling::CompileCommand getCurrentCompileCommand() const;
 
   /// Wait for the first build of preamble to finish. Preamble itself can be
   /// accessed via getPossiblyStalePreamble(). Note that this function will
   /// return after an unsuccessful build of the preamble too, i.e. result of
   /// getPossiblyStalePreamble() can be null even after this function returns.
   void waitForFirstPreamble() const;
 
   std::size_t getUsedBytes() const;
   bool isASTCached() const;
 
 private:
+  /// Publishes diagnostics for \p Inputs. It will build an AST or reuse the
+  /// cached one if applicable. Assumes LatestPreamble is compatible for \p
+  /// Inputs.
+  void generateDiagnostics(std::unique_ptr<CompilerInvocation> Invocation,
+                           ParseInputs Inputs, std::vector<Diag> CIDiags);
+
+  /// Returns the latest built preamble, might be null if no preamble has been
+  /// built or latest attempt resulted in a failure.
+  std::shared_ptr<const PreambleData> latestPreamble() const {
+    std::lock_guard<std::mutex> Lock(Mutex);
+    return LatestPreamble;
+  }
+
   // Must be called exactly once on processing thread. Will return after
   // stop() is called on a separate thread and all pending requests are
   // processed.
   void run();
   /// Signal that run() should finish processing pending requests and exit.
   void stop();
   /// Adds a new task to the end of the request queue.
   void startTask(llvm::StringRef Name, llvm::unique_function<void()> Task,
private:
   const bool RunSync;
   /// Time to wait after an update to see whether another update obsoletes it.
   const DebouncePolicy UpdateDebounce;
   /// File that ASTWorker is responsible for.
   const Path FileName;
   const GlobalCompilationDatabase &CDB;
   /// Callback invoked when preamble or main file AST is built.
   ParsingCallbacks &Callbacks;
+  /// Latest build preamble for current TU.
+  std::shared_ptr<const PreambleData> LatestPreamble;
+  Notification BuiltFirstPreamble;
 
   Semaphore &Barrier;
   /// Whether the 'onMainAST' callback ran for the current FileInputs.
   bool RanASTCallback = false;
   /// Guards members used by both TUScheduler and the worker thread.
   mutable std::mutex Mutex;
   /// File inputs, currently being used by the worker.
   /// Inputs are written and read by the worker thread, compile command can also
   /// be consumed by clients of ASTWorker.
   std::shared_ptr<const ParseInputs> FileInputs; /* GUARDED_BY(Mutex) */
   /// Times of recent AST rebuilds, used for UpdateDebounce computation.
   llvm::SmallVector<DebouncePolicy::clock::duration, 8>
       RebuildTimes; /* GUARDED_BY(Mutex) */
   /// Set to true to signal run() to finish processing.
   bool Done;                              /* GUARDED_BY(Mutex) */
   std::deque<Request> Requests;           /* GUARDED_BY(Mutex) */
+  std::queue<Request> PreambleRequests;   /* GUARDED_BY(Mutex) */
   llvm::Optional<Request> CurrentRequest; /* GUARDED_BY(Mutex) */
   mutable std::condition_variable RequestsCV;
   /// Guards the callback that publishes results of AST-related computations
   /// (diagnostics, highlightings) and file statuses.
   std::mutex PublishMu;
   // Used to prevent remove document + add document races that lead to
   // out-of-order callbacks for publishing results of onMainAST callback.
   //
   // The lifetime of the old/new ASTWorkers will overlap, but their handles
   // don't. When the old handle is destroyed, the old worker will stop reporting
   // any results to the user.
   bool CanPublishResults = true; /* GUARDED_BY(PublishMu) */
 
   SynchronizedTUStatus Status;
-  PreambleThread PW;
+  PreambleThread PreamblePeer;
 };
 
 /// A smart-pointer-like class that points to an active ASTWorker.
 /// In destructor, signals to the underlying ASTWorker that no new requests will
 /// be sent and the processing loop may exit (after running all pending
 /// requests).
 class ASTWorkerHandle {
   friend class ASTWorker;
public:
   /// be schedule via the returned reference, i.e. only reads of the preamble
   /// are possible.
   std::shared_ptr<const ASTWorker> lock() { return Worker; }
 
 private:
   std::shared_ptr<ASTWorker> Worker;
 };
 
 ASTWorkerHandle
 ASTWorker::create(PathRef FileName, const GlobalCompilationDatabase &CDB,
                   TUScheduler::ASTCache &IdleASTs, AsyncTaskRunner *Tasks,
                   Semaphore &Barrier, DebouncePolicy UpdateDebounce,
                   bool StorePreamblesInMemory, ParsingCallbacks &Callbacks) {
   std::shared_ptr<ASTWorker> Worker(
       new ASTWorker(FileName, CDB, IdleASTs, Barrier, /*RunSync=*/!Tasks,
                     UpdateDebounce, StorePreamblesInMemory, Callbacks));
   if (Tasks) {
     Tasks->runAsync("ASTWorker:" + llvm::sys::path::filename(FileName),
                     [Worker]() { Worker->run(); });
     Tasks->runAsync("PreambleWorker:" + llvm::sys::path::filename(FileName),
-                    [Worker]() { Worker->PW.run(); });
+                    [Worker]() { Worker->PreamblePeer.run(); });
   }
 
   return ASTWorkerHandle(std::move(Worker));
 }
 
 ASTWorker::ASTWorker(PathRef FileName, const GlobalCompilationDatabase &CDB,
                      TUScheduler::ASTCache &LRUCache, Semaphore &Barrier,
                      bool RunSync, DebouncePolicy UpdateDebounce,
                      bool StorePreamblesInMemory, ParsingCallbacks &Callbacks)
     : IdleASTs(LRUCache), RunSync(RunSync), UpdateDebounce(UpdateDebounce),
       FileName(FileName), CDB(CDB), Callbacks(Callbacks), Barrier(Barrier),
       Done(false), Status(FileName, Callbacks),
-      PW(FileName, Callbacks, StorePreamblesInMemory, RunSync, Status) {
+      PreamblePeer(FileName, Callbacks, StorePreamblesInMemory,
+                   // FIXME: Run PreamblePeer asynchronously once ast patching
+                   // is available.
+                   /*RunSync=*/true, Status, *this) {
   auto Inputs = std::make_shared<ParseInputs>();
   // Set a fallback command because compile command can be accessed before
   // `Inputs` is initialized. Other fields are only used after initialization
   // from client inputs.
   Inputs->CompileCommand = CDB.getFallbackCommand(FileName);
   FileInputs = std::move(Inputs);
 }
 
 ASTWorker::~ASTWorker() {
   // Make sure we remove the cached AST, if any.
   IdleASTs.take(this);
 #ifndef NDEBUG
   std::lock_guard<std::mutex> Lock(Mutex);
   assert(Done && "handle was not destroyed");
   assert(Requests.empty() && !CurrentRequest &&
          "unprocessed requests when destroying ASTWorker");
 #endif
 }
 
 void ASTWorker::update(ParseInputs Inputs, WantDiagnostics WantDiags) {
   std::string TaskName = llvm::formatv("Update ({0})", Inputs.Version);
   auto Task = [=]() mutable {
-    auto RunPublish = [&](llvm::function_ref<void()> Publish) {
-      // Ensure we only publish results from the worker if the file was not
-      // removed, making sure there are not race conditions.
-      std::lock_guard<std::mutex> Lock(PublishMu);
-      if (CanPublishResults)
-        Publish();
-    };
-
     // Get the actual command as `Inputs` does not have a command.
     // FIXME: some build systems like Bazel will take time to preparing
     // environment to build the file, it would be nice if we could emit a
     // "PreparingBuild" status to inform users, it is non-trivial given the
     // current implementation.
     if (auto Cmd = CDB.getCompileCommand(FileName))
       Inputs.CompileCommand = *Cmd;
     else
       // FIXME: consider using old command if it's not a fallback one.
       Inputs.CompileCommand = CDB.getFallbackCommand(FileName);
-    auto PrevInputs = getCurrentFileInputs();
-    // Will be used to check if we can avoid rebuilding the AST.
+
     bool InputsAreTheSame =
-        std::tie(PrevInputs->CompileCommand, PrevInputs->Contents) ==
+        std::tie(FileInputs->CompileCommand, FileInputs->Contents) ==
         std::tie(Inputs.CompileCommand, Inputs.Contents);
+    // Cached AST is invalidated.
+    if (!InputsAreTheSame) {
+      IdleASTs.take(this);
+      RanASTCallback = false;
+    }
 
-    bool RanCallbackForPrevInputs = RanASTCallback;
+    // Update current inputs so that subsequent reads can see them.
     {
       std::lock_guard<std::mutex> Lock(Mutex);
       FileInputs = std::make_shared<ParseInputs>(Inputs);
     }
-    RanASTCallback = false;
+
     log("ASTWorker building file {0} version {1} with command {2}\n[{3}]\n{4}",
         FileName, Inputs.Version, Inputs.CompileCommand.Heuristic,
         Inputs.CompileCommand.Directory,
         llvm::join(Inputs.CompileCommand.CommandLine, " "));
-    // Rebuild the preamble and the AST.
+
     StoreDiags CompilerInvocationDiagConsumer;
     std::vector<std::string> CC1Args;
     std::unique_ptr<CompilerInvocation> Invocation = buildCompilerInvocation(
         Inputs, CompilerInvocationDiagConsumer, &CC1Args);
-    // This is true for now, as we always block until new preamble is build.
-    // Once we start to block preambles out-of-order we need to make sure
-    // OldPreamble refers to the preamble that was used to build last AST.
-    auto OldPreamble = PW.latest();
-    PW.update(Invocation.get(), Inputs);
-    // FIXME make use of the stale preambles instead.
-    PW.blockUntilIdle(Deadline::infinity());
-    auto NewPreamble = PW.latest();
     // Log cc1 args even (especially!) if creating invocation failed.
     if (!CC1Args.empty())
       vlog("Driver produced command: cc1 {0}", llvm::join(CC1Args, " "));
     std::vector<Diag> CompilerInvocationDiags =
         CompilerInvocationDiagConsumer.take();
     if (!Invocation) {
       elog("Could not build CompilerInvocation for file {0}", FileName);
       // Remove the old AST if it's still in cache.
       IdleASTs.take(this);
+      RanASTCallback = false;
       // Report the diagnostics we collected when parsing the command line.
       Callbacks.onFailedAST(FileName, Inputs.Version,
-                            std::move(CompilerInvocationDiags), RunPublish);
-      return;
-    }
-
-    bool CanReuseAST = InputsAreTheSame && (OldPreamble == NewPreamble);
-    // Before doing the expensive AST reparse, we want to release our reference
-    // to the old preamble, so it can be freed if there are no other references
-    // to it.
-    OldPreamble.reset();
-    Status.update([&](TUStatus &Status) {
-      Status.ASTActivity.K = ASTAction::Building;
-      Status.ASTActivity.Name = std::move(TaskName);
-    });
-    if (!CanReuseAST) {
-      IdleASTs.take(this); // Remove the old AST if it's still in cache.
-    } else {
-      // We don't need to rebuild the AST, check if we need to run the callback.
-      if (RanCallbackForPrevInputs) {
-        RanASTCallback = true;
-        // Take a shortcut and don't report the diagnostics, since they should
-        // not changed. All the clients should handle the lack of OnUpdated()
-        // call anyway to handle empty result from buildAST.
-        // FIXME(ibiryukov): the AST could actually change if non-preamble
-        // includes changed, but we choose to ignore it.
-        // FIXME(ibiryukov): should we refresh the cache in IdleASTs for the
-        // current file at this point?
-        log("Skipping rebuild of the AST for {0}, inputs are the same.",
-            FileName);
-
-        Status.update([](TUStatus &Status) {
-          Status.Details.ReuseAST = true;
-          Status.Details.BuildFailed = false;
+                            std::move(CompilerInvocationDiags),
+                            [&](llvm::function_ref<void()> Publish) {
+                              // Ensure we only publish results from the worker
+                              // if the file was not removed, making sure there
+                              // are not race conditions.
+                              std::lock_guard<std::mutex> Lock(PublishMu);
+                              if (CanPublishResults)
+                                Publish();
                             });
-        return;
-      }
-    }
+      // Make sure anyone waiting for the preamble gets notified it could not be
+      // built.
+      BuiltFirstPreamble.notify();
 
-    // We only need to build the AST if diagnostics were requested.
-    if (WantDiags == WantDiagnostics::No)
-      return;
-
-    {
-      std::lock_guard<std::mutex> Lock(PublishMu);
-      // No need to rebuild the AST if we won't send the diagnostics. However,
-      // note that we don't prevent preamble rebuilds.
-      if (!CanPublishResults)
-        return;
-    }
-
-    // Get the AST for diagnostics.
-    llvm::Optional<std::unique_ptr<ParsedAST>> AST = IdleASTs.take(this);
-    auto RebuildStartTime = DebouncePolicy::clock::now();
-    if (!AST) {
-      llvm::Optional<ParsedAST> NewAST =
-          buildAST(FileName, std::move(Invocation), CompilerInvocationDiags,
-                   Inputs, NewPreamble);
-      // buildAST fails.
-      Status.update([&](TUStatus &Status) {
-        Status.Details.ReuseAST = false;
-        Status.Details.BuildFailed = !NewAST.hasValue();
-      });
-      AST = NewAST ? std::make_unique<ParsedAST>(std::move(*NewAST)) : nullptr;
-    } else {
-      // We are reusing the AST.
-      Status.update([](TUStatus &Status) {
-        Status.Details.ReuseAST = true;
-        Status.Details.BuildFailed = false;
-      });
-    }
-
-    // We want to report the diagnostics even if this update was cancelled.
-    // It seems more useful than making the clients wait indefinitely if they
-    // spam us with updates.
-    // Note *AST can still be null if buildAST fails.
-    if (*AST) {
-      {
-        // Try to record the AST-build time, to inform future update debouncing.
-        // This is best-effort only: if the lock is held, don't bother.
-        auto RebuildDuration = DebouncePolicy::clock::now() - RebuildStartTime;
-        std::unique_lock<std::mutex> Lock(Mutex, std::try_to_lock);
-        if (Lock.owns_lock()) {
-          // Do not let RebuildTimes grow beyond its small-size (i.e. capacity).
-          if (RebuildTimes.size() == RebuildTimes.capacity())
-            RebuildTimes.erase(RebuildTimes.begin());
-          RebuildTimes.push_back(RebuildDuration);
-        }
-      }
-      trace::Span Span("Running main AST callback");
-
-      Callbacks.onMainAST(FileName, **AST, RunPublish);
-      RanASTCallback = true;
-    } else {
-      // Failed to build the AST, at least report diagnostics from the command
-      // line if there were any.
-      // FIXME: we might have got more errors while trying to build the AST,
-      //        surface them too.
-      Callbacks.onFailedAST(FileName, Inputs.Version, CompilerInvocationDiags,
-                            RunPublish);
-    }
-    // Stash the AST in the cache for further use.
-    IdleASTs.put(this, std::move(*AST));
+    PreamblePeer.update(std::move(Invocation), std::move(Inputs),
+                        std::move(CompilerInvocationDiags), WantDiags);
   };
   startTask(TaskName, std::move(Task), WantDiags, TUScheduler::NoInvalidation);
 }
 
 void ASTWorker::runWithAST(
     llvm::StringRef Name,
     llvm::unique_function<void(llvm::Expected<InputsAndAST>)> Action,
     TUScheduler::ASTActionInvalidation Invalidation) {
   auto Task = [=, Action = std::move(Action)]() mutable {
     if (isCancelled())
       return Action(llvm::make_error<CancelledError>());
     llvm::Optional<std::unique_ptr<ParsedAST>> AST = IdleASTs.take(this);
     auto CurrentInputs = getCurrentFileInputs();
     if (!AST) {
       StoreDiags CompilerInvocationDiagConsumer;
       std::unique_ptr<CompilerInvocation> Invocation = buildCompilerInvocation(
           *CurrentInputs, CompilerInvocationDiagConsumer);
       // Try rebuilding the AST.
       vlog("ASTWorker rebuilding evicted AST to run {0}: {1} version {2}", Name,
            FileName, CurrentInputs->Version);
+      // FIXME: We might need to build a patched ast once preamble thread starts
+      // running async. Currently getPossiblyStalePreamble below will always
+      // return a compatible preamble as ASTWorker::update blocks.
       llvm::Optional<ParsedAST> NewAST =
           Invocation
               ? buildAST(FileName,
                          std::make_unique<CompilerInvocation>(*Invocation),
                          CompilerInvocationDiagConsumer.take(), *CurrentInputs,
                          getPossiblyStalePreamble())
               : None;
       AST = NewAST ? std::make_unique<ParsedAST>(std::move(*NewAST)) : nullptr;
     }
     // Make sure we put the AST back into the LRU cache.
     auto _ = llvm::make_scope_exit(
         [&AST, this]() { IdleASTs.put(this, std::move(*AST)); });
     // Run the user-provided action.
     if (!*AST)
       return Action(llvm::make_error<llvm::StringError>(
           "invalid AST", llvm::errc::invalid_argument));
     vlog("ASTWorker running {0} on version {2} of {1}", Name, FileName,
          CurrentInputs->Version);
     Action(InputsAndAST{*CurrentInputs, **AST});
   };
   startTask(Name, std::move(Task), /*UpdateType=*/None, Invalidation);
 }
 
+void PreambleThread::build(Request Req) {
+  assert(Req.CI && "Got preamble request with null compiler invocation");
+  const ParseInputs &Inputs = Req.Inputs;
+
+  Status.update([&](TUStatus &Status) {
+    Status.PreambleActivity = PreambleAction::Building;
+  });
+  auto _ = llvm::make_scope_exit([this, &Req] {
+    ASTPeer.updatePreamble(std::move(Req.CI), std::move(Req.Inputs),
+                           LatestBuild, std::move(Req.CIDiags),
+                           std::move(Req.WantDiags));
+  });
+
+  if (!LatestBuild || Inputs.ForceRebuild) {
+    vlog("Building first preamble for {0} version {1}", FileName,
+         Inputs.Version);
+  } else if (isPreambleCompatible(*LatestBuild, Inputs, FileName, *Req.CI)) {
+    vlog("Reusing preamble version {0} for version {1} of {2}",
+         LatestBuild->Version, Inputs.Version, FileName);
+    return;
+  } else {
+    vlog("Rebuilding invalidated preamble for {0} version {1} (previous was "
+         "version {2})",
+         FileName, Inputs.Version, LatestBuild->Version);
+  }
+
+  LatestBuild = clang::clangd::buildPreamble(
+      FileName, *Req.CI, Inputs, StoreInMemory,
+      [this, Version(Inputs.Version)](ASTContext &Ctx,
+                                      std::shared_ptr<clang::Preprocessor> PP,
+                                      const CanonicalIncludes &CanonIncludes) {
+        Callbacks.onPreambleAST(FileName, Version, Ctx, std::move(PP),
+                                CanonIncludes);
+      });
+}
+
+void ASTWorker::updatePreamble(std::unique_ptr<CompilerInvocation> CI,
+                               ParseInputs PI,
+                               std::shared_ptr<const PreambleData> Preamble,
+                               std::vector<Diag> CIDiags,
+                               WantDiagnostics WantDiags) {
+  std::string TaskName = llvm::formatv("Build AST for ({0})", PI.Version);
+  // Store preamble and build diagnostics with new preamble if requested.
+  auto Task = [this, Preamble = std::move(Preamble), CI = std::move(CI),
+               PI = std::move(PI), CIDiags = std::move(CIDiags),
+               WantDiags = std::move(WantDiags)]() mutable {
+    // Update the preamble inside ASTWorker queue to ensure atomicity. As a task
+    // running inside ASTWorker assumes internals won't change until it
+    // finishes.
+    if (Preamble != LatestPreamble) {
+      // Cached AST is no longer valid.
+      IdleASTs.take(this);
+      RanASTCallback = false;
+      std::lock_guard<std::mutex> Lock(Mutex);
+      // LatestPreamble might be the last reference to old preamble, do not
+      // trigger destructor while holding the lock.
+      std::swap(LatestPreamble, Preamble);
+    }
+    // Give up our ownership to old preamble before starting expensive AST
+    // build.
+    Preamble.reset();
+    BuiltFirstPreamble.notify();
+    // We only need to build the AST if diagnostics were requested.
+    if (WantDiags == WantDiagnostics::No)
+      return;
+    // Report diagnostics with the new preamble to ensure progress. Otherwise
+    // diagnostics might get stale indefinitely if user keeps invalidating the
+    // preamble.
+    generateDiagnostics(std::move(CI), std::move(PI), std::move(CIDiags));
+  };
+  if (RunSync) {
+    Task();
+    return;
+  }
+  {
+    std::lock_guard<std::mutex> Lock(Mutex);
+    PreambleRequests.push({std::move(Task), std::move(TaskName),
+                           steady_clock::now(), Context::current().clone(),
+                           llvm::None, TUScheduler::NoInvalidation, nullptr});
+  }
+  RequestsCV.notify_all();
+}
+
+void ASTWorker::generateDiagnostics(
+    std::unique_ptr<CompilerInvocation> Invocation, ParseInputs Inputs,
+    std::vector<Diag> CIDiags) {
+  assert(Invocation);
+  // No need to rebuild the AST if we won't send the diagnostics.
+  {
+    std::lock_guard<std::mutex> Lock(PublishMu);
+    if (!CanPublishResults)
+      return;
+  }
+  // Used to check whether we can update AST cache.
+  bool InputsAreTheSame =
+      std::tie(FileInputs->CompileCommand, FileInputs->Contents) ==
+      std::tie(Inputs.CompileCommand, Inputs.Contents);
+  // Take a shortcut and don't report the diagnostics, since they should be the
+  // same. All the clients should handle the lack of OnUpdated() call anyway to
+  // handle empty result from buildAST.
+  // FIXME: the AST could actually change if non-preamble includes changed,
+  // but we choose to ignore it.
+  if (InputsAreTheSame && RanASTCallback)
+    return;
+
+  // Get the AST for diagnostics, either build it or use the cached one.
+  std::string TaskName = llvm::formatv("Build AST ({0})", Inputs.Version);
+  Status.update([&](TUStatus &Status) {
+    Status.ASTActivity.K = ASTAction::Building;
+    Status.ASTActivity.Name = std::move(TaskName);
+  });
+  llvm::Optional<std::unique_ptr<ParsedAST>> AST = IdleASTs.take(this);
+  if (!AST) {
+    auto RebuildStartTime = DebouncePolicy::clock::now();
+    llvm::Optional<ParsedAST> NewAST = clang::clangd::buildAST(
+        FileName, std::move(Invocation), CIDiags, Inputs, LatestPreamble);
+    auto RebuildDuration = DebouncePolicy::clock::now() - RebuildStartTime;
+    // Try to record the AST-build time, to inform future update debouncing.
+    // This is best-effort only: if the lock is held, don't bother.
+    std::unique_lock<std::mutex> Lock(Mutex, std::try_to_lock);
+    if (Lock.owns_lock()) {
+      // Do not let RebuildTimes grow beyond its small-size (i.e.
+      // capacity).
+      if (RebuildTimes.size() == RebuildTimes.capacity())
+        RebuildTimes.erase(RebuildTimes.begin());
+      RebuildTimes.push_back(RebuildDuration);
+      Lock.unlock();
+    }
+    Status.update([&](TUStatus &Status) {
+      Status.Details.ReuseAST = false;
+      Status.Details.BuildFailed = !NewAST.hasValue();
+    });
+    AST = NewAST ? std::make_unique<ParsedAST>(std::move(*NewAST)) : nullptr;
+  } else {
+    assert(InputsAreTheSame && !RanASTCallback &&
+           "forgot to invalidate cached ast?");
+    log("Skipping rebuild of the AST for {0}, inputs are the same.", FileName);
+    Status.update([](TUStatus &Status) {
+      Status.Details.ReuseAST = true;
+      Status.Details.BuildFailed = false;
+    });
+  }
+
+  // Publish diagnostics.
+  auto RunPublish = [&](llvm::function_ref<void()> Publish) {
+    // Ensure we only publish results from the worker if the file was not
+    // removed, making sure there are not race conditions.
+    std::lock_guard<std::mutex> Lock(PublishMu);
+    if (CanPublishResults)
+      Publish();
+  };
+  if (*AST) {
+    trace::Span Span("Running main AST callback");
+    Callbacks.onMainAST(FileName, **AST, RunPublish);
+  } else {
+    // Failed to build the AST, at least report diagnostics from the
+    // command line if there were any.
+    // FIXME: we might have got more errors while trying to build the
+    // AST, surface them too.
+    Callbacks.onFailedAST(FileName, Inputs.Version, CIDiags, RunPublish);
+  }
+
+  // Stash the AST in the cache for further use if it was build for current
+  // inputs.
+  if (InputsAreTheSame) {
+    RanASTCallback = *AST != nullptr;
+    IdleASTs.put(this, std::move(*AST));
+  }
+}
+
 std::shared_ptr<const PreambleData>
 ASTWorker::getPossiblyStalePreamble() const {
-  return PW.latest();
+  std::lock_guard<std::mutex> Lock(Mutex);
+  return LatestPreamble;
 }
 
 void ASTWorker::getCurrentPreamble(
     llvm::unique_function<void(std::shared_ptr<const PreambleData>)> Callback) {
   // We could just call startTask() to throw the read on the queue, knowing
   // it will run after any updates. But we know this task is cheap, so to
   // improve latency we cheat: insert it on the queue after the last update.
   std::unique_lock<std::mutex> Lock(Mutex);
Requests.insert(LastUpdate.base(),
                           Context::current().clone(),
                           /*UpdateType=*/None,
                           /*InvalidationPolicy=*/TUScheduler::NoInvalidation,
                           /*Invalidate=*/nullptr});
   Lock.unlock();
   RequestsCV.notify_all();
 }
 
-void ASTWorker::waitForFirstPreamble() const { PW.waitForFirst(); }
+void ASTWorker::waitForFirstPreamble() const { BuiltFirstPreamble.wait(); }
 
 std::shared_ptr<const ParseInputs> ASTWorker::getCurrentFileInputs() const {
   std::unique_lock<std::mutex> Lock(Mutex);
   return FileInputs;
 }
 
 tooling::CompileCommand ASTWorker::getCurrentCompileCommand() const {
   std::unique_lock<std::mutex> Lock(Mutex);
void ASTWorker::stop() {
     std::lock_guard<std::mutex> Lock(PublishMu);
     CanPublishResults = false;
   }
   {
     std::lock_guard<std::mutex> Lock(Mutex);
     assert(!Done && "stop() called twice");
     Done = true;
   }
+  // We are no longer going to build any preambles, let the waiters know that.
+  BuiltFirstPreamble.notify();
+  PreamblePeer.stop();
   Status.stop();
   RequestsCV.notify_all();
 }
 
 void ASTWorker::startTask(llvm::StringRef Name,
                           llvm::unique_function<void()> Task,
                           llvm::Optional<WantDiagnostics> UpdateType,
                           TUScheduler::ASTActionInvalidation Invalidation) {
if (RunSync) {
     Requests.push_back({std::move(Task), std::string(Name), steady_clock::now(),
                         std::move(Ctx), UpdateType, Invalidation,
                         std::move(Invalidate)});
   }
   RequestsCV.notify_all();
 }
 
 void ASTWorker::run() {
-  auto _ = llvm::make_scope_exit([this] { PW.stop(); });
   while (true) {
     {
       std::unique_lock<std::mutex> Lock(Mutex);
       assert(!CurrentRequest && "A task is already running, multiple workers?");
       for (auto Wait = scheduleLocked(); !Wait.expired();
            Wait = scheduleLocked()) {
+        assert(PreambleRequests.empty() &&
+               "Preamble updates should be scheduled immediately");
         if (Done) {
           if (Requests.empty())
             return;
           else     // Even though Done is set, finish pending requests.
             break; // However, skip delays to shutdown fast.
         }
 
         // Tracing: we have a next request, attribute this sleep to it.
while (true) {
                         std::chrono::duration_cast<std::chrono::milliseconds>(
                             Wait.time() - steady_clock::now())
                             .count());
           }
         }
 
         wait(Lock, RequestsCV, Wait);
       }
+      // Any request in ReceivedPreambles is at least as old as the
+      // Requests.front(), so prefer them first to preserve LSP order.
+      if (!PreambleRequests.empty()) {
+        CurrentRequest = std::move(PreambleRequests.front());
+        PreambleRequests.pop();
+      } else {
         CurrentRequest = std::move(Requests.front());
         Requests.pop_front();
+      }
     } // unlock Mutex
 
     // It is safe to perform reads to CurrentRequest without holding the lock as
     // only writer is also this thread.
     {
       std::unique_lock<Semaphore> Lock(Barrier, std::try_to_lock);
       if (!Lock.owns_lock()) {
         Status.update([&](TUStatus &Status) {
// only writer is also this thread.
       });
       CurrentRequest->Action();
     }
 
     bool IsEmpty = false;
     {
       std::lock_guard<std::mutex> Lock(Mutex);
       CurrentRequest.reset();
-      IsEmpty = Requests.empty();
+      IsEmpty = Requests.empty() && PreambleRequests.empty();
     }
     if (IsEmpty) {
       Status.update([&](TUStatus &Status) {
         Status.ASTActivity.K = ASTAction::Idle;
         Status.ASTActivity.Name = "";
       });
     }
     RequestsCV.notify_all();
   }
 }
 
 Deadline ASTWorker::scheduleLocked() {
+  // Process new preambles immediately.
+  if (!PreambleRequests.empty())
+    return Deadline::zero();
   if (Requests.empty())
     return Deadline::infinity(); // Wait for new requests.
   // Handle cancelled requests first so the rest of the scheduler doesn't.
   for (auto I = Requests.begin(), E = Requests.end(); I != E; ++I) {
     if (!isCancelled(I->Ctx)) {
       // Cancellations after the first read don't affect current scheduling.
       if (I->UpdateType == None)
         break;
for (; Next != Requests.end(); ++Next)
         return true; // Prefer later diagnostics.
     return false;
   }
   llvm_unreachable("Unknown WantDiagnostics");
 }
 
 bool ASTWorker::blockUntilIdle(Deadline Timeout) const {
   std::unique_lock<std::mutex> Lock(Mutex);
-  return wait(Lock, RequestsCV, Timeout,
-              [&] { return Requests.empty() && !CurrentRequest; });
+  return wait(Lock, RequestsCV, Timeout, [&] {
+    return PreambleRequests.empty() && Requests.empty() && !CurrentRequest;
+  });
 }
 
 // Render a TUAction to a user-facing string representation.
 // TUAction represents clangd-internal states, we don't intend to expose them
 // to users (say C++ programmers) directly to avoid confusion, we use terms that
 // are familiar by C++ programmers.
 std::string renderTUAction(const PreambleAction PA, const ASTAction &AA) {
   llvm::SmallVector<std::string, 2> Result;
if (!PreambleTasks) {
     Action(InputsAndPreamble{It->second->Contents,
                              It->second->Worker->getCurrentCompileCommand(),
                              Preamble.get()});
     return;
   }
 
   // Future is populated if the task needs a specific preamble.
   std::future<std::shared_ptr<const PreambleData>> ConsistentPreamble;
+  // FIXME: Currently this only holds because ASTWorker blocks after issuing a
+  // preamble build. Get rid of consistent reads or make them build on the
+  // calling thread instead.
   if (Consistency == Consistent) {
     std::promise<std::shared_ptr<const PreambleData>> Promise;
     ConsistentPreamble = Promise.get_future();
     It->second->Worker->getCurrentPreamble(
         [Promise = std::move(Promise)](
             std::shared_ptr<const PreambleData> Preamble) mutable {
           Promise.set_value(std::move(Preamble));
         });

clang-tools-extra/clangd/unittests/CodeCompleteTests.cpp

 //===-- CodeCompleteTests.cpp -----------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
CodeCompleteResult completions(const TestTU &TU, Position Point,
 
   auto Inputs = TU.inputs();
   IgnoreDiagnostics Diags;
   auto CI = buildCompilerInvocation(Inputs, Diags);
   if (!CI) {
     ADD_FAILURE() << "Couldn't build CompilerInvocation";
     return {};
   }
-  auto Preamble =
-      buildPreamble(testPath(TU.Filename), *CI, /*OldPreamble=*/nullptr, Inputs,
+  auto Preamble = buildPreamble(testPath(TU.Filename), *CI, Inputs,
                                 /*InMemory=*/true, /*Callback=*/nullptr);
   return codeComplete(testPath(TU.Filename), Inputs.CompileCommand,
                       Preamble.get(), TU.Code, Point, Inputs.FS, Opts);
 }
 
 // Runs code completion.
 CodeCompleteResult completions(llvm::StringRef Text,
                                std::vector<Symbol> IndexSymbols = {},
                                clangd::CodeCompleteOptions Opts = {},
EXPECT_THAT(
           AllOf(Named("complete_class"), Kind(CompletionItemKind::Class)),
           AllOf(Named("complete_function"), Kind(CompletionItemKind::Function)),
           AllOf(Named("complete_type_alias"),
                 Kind(CompletionItemKind::Interface)),
           AllOf(Named("complete_variable"), Kind(CompletionItemKind::Variable)),
           AllOf(Named("complete_static_member"),
                 Kind(CompletionItemKind::Property))));
 
-   Results = completions(
       R"cpp(
         enum Color {
           Red
         };
         Color u = ^
       )cpp");
-   EXPECT_THAT(Results.Completions,
-               Contains(
-                   AllOf(Named("Red"), Kind(CompletionItemKind::EnumMember))));
+  EXPECT_THAT(
+      Results.Completions,
+      Contains(AllOf(Named("Red"), Kind(CompletionItemKind::EnumMember))));
 }
 
 TEST(CompletionTest, NoDuplicates) {
   auto Results = completions(
       R"cpp(
           class Adapter {
           };
 
SignatureHelp signatures(llvm::StringRef Text, Position Point,
   auto TU = TestTU::withCode(Text);
   auto Inputs = TU.inputs();
   IgnoreDiagnostics Diags;
   auto CI = buildCompilerInvocation(Inputs, Diags);
   if (!CI) {
     ADD_FAILURE() << "Couldn't build CompilerInvocation";
     return {};
   }
-  auto Preamble =
-      buildPreamble(testPath(TU.Filename), *CI, /*OldPreamble=*/nullptr, Inputs,
+  auto Preamble = buildPreamble(testPath(TU.Filename), *CI, Inputs,
                                 /*InMemory=*/true, /*Callback=*/nullptr);
   if (!Preamble) {
     ADD_FAILURE() << "Couldn't build Preamble";
     return {};
   }
   return signatureHelp(testPath(TU.Filename), Inputs.CompileCommand, *Preamble,
                        Text, Point, Inputs.FS, Index.get());
 }
 

 }
 
 TEST(CompletionTest, FixItForArrowToDot) {
   MockFSProvider FS;
   MockCompilationDatabase CDB;
 
   CodeCompleteOptions Opts;
   Opts.IncludeFixIts = true;
-  const char* Code =
+  const char *Code =
       R"cpp(
         class Auxilary {
          public:
           void AuxFunction();
         };
         class ClassWithPtr {
          public:
           void MemberFunction();
if (!C.FixIts.empty()) {
       EXPECT_THAT(C.FixIts, ElementsAre(ReplacementEdit));
     }
   }
 }
 
 TEST(CompletionTest, FixItForDotToArrow) {
   CodeCompleteOptions Opts;
   Opts.IncludeFixIts = true;
-  const char* Code =
+  const char *Code =
       R"cpp(
         class Auxilary {
          public:
           void AuxFunction();
         };
         class ClassWithPtr {
          public:
           void MemberFunction();
constexpr const char *TestCodes[] = {
         }
       )cpp",
       R"cpp(
         #include "foo/[[a^/]]foo.h"
       )cpp",
       R"cpp(
         #include "foo/abc/[[fo^o.h"]]
       )cpp",
-      };
   for (const auto &Text : TestCodes) {
     Annotations TestCode(Text);
     TU.Code = TestCode.code().str();
     auto Results = completions(TU, TestCode.point());
     if (Results.Completions.size() != 1) {
       ADD_FAILURE() << "Results.Completions.size() != 1" << Text;
       continue;
     }
TEST(CompletionTest, NoInsertIncludeIfOnePresent) {
   TU.AdditionalFiles["foo.h"] = "";
 
   std::string DeclFile = URI::create(testPath("foo")).toString();
   Symbol Sym = func("Func");
   Sym.CanonicalDeclaration.FileURI = DeclFile.c_str();
   Sym.IncludeHeaders.emplace_back("\"foo.h\"", 2);
   Sym.IncludeHeaders.emplace_back("\"bar.h\"", 1000);
 
-  EXPECT_THAT(
-      completions(TU, Test.point(), {Sym}).Completions,
-      UnorderedElementsAre(
-          AllOf(Named("Func"), HasInclude("\"foo.h\""), Not(InsertInclude()))));
+  EXPECT_THAT(completions(TU, Test.point(), {Sym}).Completions,
+              UnorderedElementsAre(AllOf(Named("Func"), HasInclude("\"foo.h\""),
+                                         Not(InsertInclude()))));
 }
 
 TEST(CompletionTest, MergeMacrosFromIndexAndSema) {
   Symbol Sym;
   Sym.Name = "Clangd_Macro_Test";
   Sym.ID = SymbolID("c:foo.cpp@8@macro@Clangd_Macro_Test");
   Sym.SymInfo.Kind = index::SymbolKind::Macro;
   Sym.Flags |= Symbol::IndexedForCodeCompletion;

clang-tools-extra/clangd/unittests/FileIndexTests.cpp

 //===-- FileIndexTests.cpp  ---------------------------*- C++ -*-----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
TEST(FileIndexTest, RebuildWithPreamble) {
   )cpp";
 
   // Rebuild the file.
   IgnoreDiagnostics IgnoreDiags;
   auto CI = buildCompilerInvocation(PI, IgnoreDiags);
 
   FileIndex Index;
   bool IndexUpdated = false;
-  buildPreamble(FooCpp, *CI, /*OldPreamble=*/nullptr, PI,
+  buildPreamble(FooCpp, *CI, PI,
                 /*StoreInMemory=*/true,
                 [&](ASTContext &Ctx, std::shared_ptr<Preprocessor> PP,
                     const CanonicalIncludes &CanonIncludes) {
                   EXPECT_FALSE(IndexUpdated)
                       << "Expected only a single index update";
                   IndexUpdated = true;
                   Index.updatePreamble(FooCpp, /*Version=*/"null", Ctx,
                                        std::move(PP), CanonIncludes);
TEST(FileIndexTest, ReferencesInMainFileWithPreamble) {
 
   // Expect to see references in main file, references in headers are excluded
   // because we only index main AST.
   EXPECT_THAT(getRefs(Index, findSymbol(TU.headerSymbols(), "Foo").ID),
               RefsAre({RefRange(Main.range())}));
 }
 
 TEST(FileIndexTest, MergeMainFileSymbols) {
-  const char* CommonHeader = "void foo();";
+  const char *CommonHeader = "void foo();";
   TestTU Header = TestTU::withCode(CommonHeader);
   TestTU Cpp = TestTU::withCode("void foo() {}");
   Cpp.Filename = "foo.cpp";
   Cpp.HeaderFilename = "foo.h";
   Cpp.HeaderCode = CommonHeader;
 
   FileIndex Index;
   auto HeaderAST = Header.build();

clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp

 //===-- TUSchedulerTests.cpp ------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 

 #include "Threading.h"
 #include "clang/Basic/DiagnosticDriver.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <algorithm>
+#include <atomic>
 #include <bits/stdint-uintn.h>
 #include <chrono>
 #include <utility>
 
 namespace clang {
 namespace clangd {
 namespace {
 
std::atomic<int> CallbackCount(0);
     });
     Inputs.Contents = "#include <B>";
     Inputs.Version = "B";
     S.update(Path, Inputs, WantDiagnostics::Yes);
 
     S.runWithPreamble("StaleRead", Path, TUScheduler::Stale,
                       [&](Expected<InputsAndPreamble> Pre) {
                         ASSERT_TRUE(bool(Pre));
+                        ASSERT_TRUE(Pre->Preamble);
                         EXPECT_EQ(Pre->Preamble->Version, "A");
                         EXPECT_THAT(includes(Pre->Preamble),
                                     ElementsAre("<A>"));
                         InconsistentReadDone.notify();
                         ++CallbackCount;
                       });
     S.runWithPreamble("ConsistentRead", Path, TUScheduler::Consistent,
                       [&](Expected<InputsAndPreamble> Pre) {
                         ASSERT_TRUE(bool(Pre));
+                        ASSERT_TRUE(Pre->Preamble);
                         EXPECT_EQ(Pre->Preamble->Version, "B");
                         EXPECT_THAT(includes(Pre->Preamble),
                                     ElementsAre("<B>"));
                         ++CallbackCount;
                       });
+    S.blockUntilIdle(timeoutSeconds(10));
   }
   EXPECT_EQ(2, CallbackCount);
 }
 
 TEST_F(TUSchedulerTests, Cancellation) {
   // We have the following update/read sequence
   //   U0
   //   U1(WantDiags=Yes) <-- cancelled
TEST_F(TUSchedulerTests, ForceRebuild) {
   auto Header = testPath("foo.h");
 
   auto SourceContents = R"cpp(
       #include "foo.h"
       int b = a;
     )cpp";
 
   ParseInputs Inputs = getInputs(Source, SourceContents);
+  std::atomic<size_t> DiagCount(0);
 
   // Update the source contents, which should trigger an initial build with
   // the header file missing.
   updateWithDiags(
-      S, Source, Inputs, WantDiagnostics::Yes, [](std::vector<Diag> Diags) {
+      S, Source, Inputs, WantDiagnostics::Yes,
+      [&DiagCount](std::vector<Diag> Diags) {
+        ++DiagCount;
         EXPECT_THAT(Diags,
                     ElementsAre(Field(&Diag::Message, "'foo.h' file not found"),
                                 Field(&Diag::Message,
                                       "use of undeclared identifier 'a'")));
       });
 
   // Add the header file. We need to recreate the inputs since we changed a
   // file from underneath the test FS.
   Files[Header] = "int a;";
   Timestamps[Header] = time_t(1);
   Inputs = getInputs(Source, SourceContents);
 
   // The addition of the missing header file shouldn't trigger a rebuild since
   // we don't track missing files.
   updateWithDiags(
-      S, Source, Inputs, WantDiagnostics::Yes, [](std::vector<Diag> Diags) {
+      S, Source, Inputs, WantDiagnostics::Yes,
+      [&DiagCount](std::vector<Diag> Diags) {
+        ++DiagCount;
         ADD_FAILURE() << "Did not expect diagnostics for missing header update";
       });
 
   // Forcing the reload should should cause a rebuild which no longer has any
   // errors.
   Inputs.ForceRebuild = true;
-  updateWithDiags(
-      S, Source, Inputs, WantDiagnostics::Yes,
-      [](std::vector<Diag> Diags) { EXPECT_THAT(Diags, IsEmpty()); });
+  updateWithDiags(S, Source, Inputs, WantDiagnostics::Yes,
+                  [&DiagCount](std::vector<Diag> Diags) {
+                    ++DiagCount;
+                    EXPECT_THAT(Diags, IsEmpty());
+                  });
 
   ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
+  EXPECT_EQ(DiagCount, 2U);
 }
 TEST_F(TUSchedulerTests, NoChangeDiags) {
   TUScheduler S(CDB, optsForTest(), captureDiags());
 
   auto FooCpp = testPath("foo.cpp");
   auto Contents = "int a; int b;";
 
   updateWithDiags(
TEST_F(TUSchedulerTests, TUStatus) {
 
   EXPECT_THAT(CaptureTUStatus.allStatus(),
               ElementsAre(
                   // Everything starts with ASTWorker starting to execute an
                   // update
                   TUState(PreambleAction::Idle, ASTAction::RunningAction),
                   // We build the preamble
                   TUState(PreambleAction::Building, ASTAction::RunningAction),
-                  // Preamble worker goes idle
+                  // We built the preamble, and issued ast built on ASTWorker
+                  // thread. Preambleworker goes idle afterwards.
+                  TUState(PreambleAction::Idle, ASTAction::RunningAction),
+                  // Start task for building the ast, as a result of building
+                  // preamble, on astworker thread.
                   TUState(PreambleAction::Idle, ASTAction::RunningAction),
-                  // We start building the ast
+                  // AST build starts.
                   TUState(PreambleAction::Idle, ASTAction::Building),
-                  // Built finished succesffully
+                  // AST built finished successfully
                   TUState(PreambleAction::Idle, ASTAction::Building),
-                  // Rnning go to def
+                  // Running go to def
                   TUState(PreambleAction::Idle, ASTAction::RunningAction),
-                  // both workers go idle
+                  // ASTWorker goes idle.
                   TUState(PreambleAction::Idle, ASTAction::Idle)));
 }
 
 TEST_F(TUSchedulerTests, CommandLineErrors) {
   // We should see errors from command-line parsing inside the main file.
   CDB.ExtraClangFlags = {"-fsome-unknown-flag"};
 
   // (!) 'Ready' must live longer than TUScheduler.

clang-tools-extra/clangd/unittests/TestTU.cpp

 //===--- TestTU.cpp - Scratch source files for testing --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
ParseInputs TestTU::inputs() const {
   std::string ThunkContents = "#import \"" + FullHeaderName + "\"\n";
 
   llvm::StringMap<std::string> Files(AdditionalFiles);
   Files[FullFilename] = Code;
   Files[FullHeaderName] = HeaderCode;
   Files[ImportThunk] = ThunkContents;
 
   ParseInputs Inputs;
-  auto& Argv = Inputs.CompileCommand.CommandLine;
+  auto &Argv = Inputs.CompileCommand.CommandLine;
   Argv = {"clang"};
   // FIXME: this shouldn't need to be conditional, but it breaks a
   // GoToDefinition test for some reason (getMacroArgExpandedLocation fails).
   if (!HeaderCode.empty()) {
     Argv.push_back("-include");
     Argv.push_back(ImplicitHeaderGuard ? ImportThunk : FullHeaderName);
     // ms-compatibility changes the meaning of #import.
     // The default is OS-dependent (on on windows), ensure it's off.

 }
 
 ParsedAST TestTU::build() const {
   auto Inputs = inputs();
   StoreDiags Diags;
   auto CI = buildCompilerInvocation(Inputs, Diags);
   assert(CI && "Failed to build compilation invocation.");
   auto Preamble =
-      buildPreamble(testPath(Filename), *CI,
-                    /*OldPreamble=*/nullptr, Inputs,
+      buildPreamble(testPath(Filename), *CI, Inputs,
                     /*StoreInMemory=*/true, /*PreambleCallback=*/nullptr);
   auto AST = buildAST(testPath(Filename), std::move(CI), Diags.take(), Inputs,
                       Preamble);
   if (!AST.hasValue()) {
     ADD_FAILURE() << "Failed to build code:\n" << Code;
     llvm_unreachable("Failed to build TestTU!");
   }
   // Check for error diagnostics and report gtest failures (unless expected).
   // This guards against accidental syntax errors silently subverting tests.
   // error-ok is awfully primitive - using clang -verify would be nicer.
   // Ownership and layering makes it pretty hard.
   bool ErrorOk = [&, this] {
     llvm::StringLiteral Marker = "error-ok";
     if (llvm::StringRef(Code).contains(Marker) ||
         llvm::StringRef(HeaderCode).contains(Marker))
       return true;
-    for (const auto& KV : this->AdditionalFiles)
+    for (const auto &KV : this->AdditionalFiles)
       if (llvm::StringRef(KV.second).contains(Marker))
         return true;
     return false;
   }();
   if (!ErrorOk) {
     for (const auto &D : AST->getDiagnostics())
       if (D.Severity >= DiagnosticsEngine::Error) {
         ADD_FAILURE()