[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

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.

11

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

12

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

19

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

23

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.

25

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)

29

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)

29

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.

33

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

33

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

33

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

34

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

36

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)

43

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

47

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

50

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

53

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.

235

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)

259

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.

338

this seems to never reuse?

368

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 :)

385

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)

623

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

623

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.

679

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.

693

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

623

ASTTask should be a member function, I think.

Moved it into a meber named publishDiagnostics

693

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

9–41

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

9–41

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

10

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

47–48

changes latest inputs -> replaces the current parser inputs

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

48

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

72

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

73

it it isn't -> until?

76

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.

84

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

90

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.

95

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.

250

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

252

nit: WantDiags is an enum, no move

272

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)

370

this sentence doesn't parse :-)

388

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.

448

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

448

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

450

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

generateDiagnostics?

517

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

576

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

589

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

626

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

640

incomplete comment

671

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

48

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

76

using compatible.

388

PW was for PeerWorker :P

589

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

626

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

671

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

250

comment is now obsolete, remove it and consider inlining Req

378

this is no longer locked or concurrently accessed right?

432

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...

454

Unused, I think?

983

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

1000

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

1052

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

250

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

378

right, deleted annottation and added a comment.

clang-tools-extra/clangd/Preamble.h

 #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 isPreambleUsable(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

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 isPreambleUsable(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.
+// manages an AST and a preamble for that file. All operations that modify or
+// read the AST are run on a separate dedicated thread asynchronously in FIFO
+// order. There is a second thread for preamble builds, ASTWorker thread issues
+// updates on that preamble thread as preamble becomes stale.
+//
+// As a consequence of that, read requests to the AST always get processed in
+// order and get the exact version of the file when they were issued. Those ASTs
+// are always built with a compatible preamble, built with latest update that
+// modified preamble section of the file.
+// FIXME: Get rid of the synchronization between ASTWorker::update and
+// PreambleThread builds by using patched ASTs instead of compatible preambles
+// for reads. Only keep the blocking behaviour for picky read requests, e.g.
+// Tweaks.
+//
+// Updates are processed in two phases, by issuing a preamble and an ast build.
+// The preamble build gets issued into a different worker and might be
+// overwritten by subsequent updates. An AST will be built for an update if
+// latest built preamble is compatible for it or a new preamble gets built for
+// that update.
+// AST builds are always executed on ASTWorker thread, either immediately
+// following an update request with a compatible preamble, or through a new
+// request, inserted at front of the queue, from PreambleThread after building a
+// new preamble, we'll call the ASTs built in latter case "golden AST"s.
+//
+// Diagnostics and index is updated as a result of building AST for write
+// requests. Progress on those updates are ensured with golden ASTs, as they'll
+// be built even when ASTWorker queue is full. Since ASTs aren't built with
+// incompatible preambles, we can partition diagnostic updates into epochs of
+// preambles. Each epoch starts with a golden AST(unless that update had
+// WantDiagnostics::No), and then moves forward with updates to the AST. Since
+// preambles are always built with newer versions of the file and updates in the
+// ASTWorker are in order, we ensure diagnostics are always built with newer
+// versions of 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.
 //
-// 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.
+// ASTWorker is also responsible for serving preambles. Since preambles can be
+// read concurrently, we run each of async preamble reads on its own thread.
 //
 // To limit the concurrent load that clangd produces we maintain a semaphore
 // that keeps more than a fixed number of threads from running concurrently.
 //
 // 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.
 
 #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/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>
 
 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:
+  using PreambleCallback =
+      std::function<void(std::unique_ptr<CompilerInvocation>, ParseInputs,
+                         std::shared_ptr<const PreambleData>)>;
   PreambleThread(llvm::StringRef FileName, ParsingCallbacks &Callbacks,
                  bool StorePreambleInMemory, bool RunSync,
                  SynchronizedTUStatus &Status)
       : FileName(FileName), Callbacks(Callbacks),
         StoreInMemory(StorePreambleInMemory), RunSync(RunSync), Status(Status) {
   }
 
-  size_t getUsedBytes() const {
-    auto Preamble = latest();
-    return Preamble ? Preamble->Preamble.getSize() : 0;
-  }
-
   /// 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(const CompilerInvocation &CI, ParseInputs PI,
+              PreambleCallback PC) {
     // Make possibly expensive copy while not holding the lock.
-    Request Req = {std::make_unique<CompilerInvocation>(*CI), std::move(PI)};
+    Request Req = {std::make_unique<CompilerInvocation>(CI), std::move(PI),
+                   std::move(PC), Context::current().clone()};
     if (RunSync) {
       build(std::move(Req));
       return;
     }
     {
       std::lock_guard<std::mutex> Lock(Mutex);
-      assert(!Done && "Build request to PreambleWorker after stop");
+      // ASTWorker might've updates left in it even after PreambleWorker is
+      // stopped. Ignore those as it should be able to serve remaining requests
+      // with stale preambles.
+      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);
   }
 
   /// Signals the run loop to exit.
   void stop() {
     dlog("Stopping preamble worker for {0}", 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;
+    PreambleCallback PC;
+    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();
+
+    if (LastInputVersion) {
+      vlog("Rebuilding invalidated preamble for {0} version {1} (previous was "
+           "version {2})",
+           FileName, Inputs.Version, *LastInputVersion);
+    } else {
+      vlog("Building first preamble for {0} version {1}", FileName,
+           Inputs.Version);
+    }
+    LastInputVersion = Inputs.Version;
 
     Status.update([&](TUStatus &Status) {
       Status.PreambleActivity = PreambleAction::Building;
     });
 
     auto Preamble = clang::clangd::buildPreamble(
-        FileName, std::move(*Req.CI), OldPreamble, Inputs, StoreInMemory,
+        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);
         });
-    {
-      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);
-    }
+    // FIXME: Invalidate NextReq, if this Preamble is compatible with that.
+    Req.PC(std::move(Req.CI), std::move(Req.Inputs), std::move(Preamble));
   }
 
   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) */
+  mutable std::condition_variable ReqCV; /* GUARDED_BY(Mutex) */
 
-  Notification BuiltFirst;
   const Path FileName;
   ParsingCallbacks &Callbacks;
   const bool StoreInMemory;
   const bool RunSync;
+  // Used only for logging.
+  llvm::Optional<std::string> LastInputVersion;
 
   SynchronizedTUStatus &Status;
 };
 
 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.
 /// The ASTWorker that manages the AST is shared by both the processing thread
public:
   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;
 
   /// 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:
+  /// 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) */
   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
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(); });
   }
 
   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)
#ifndef NDEBUG
   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);
 
     bool RanCallbackForPrevInputs = RanASTCallback;
     {
       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);
       // Report the diagnostics we collected when parsing the command line.
       Callbacks.onFailedAST(FileName, Inputs.Version,
                             std::move(CompilerInvocationDiags), RunPublish);
+      // Make sure anyone waiting for the preamble gets notified it could not be
+      // built.
+      BuiltFirstPreamble.notify();
       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;
-        });
-        return;
-      }
+    // Task that builds an AST and runs callbacks. Must be run with a preamble
+    // that's compatible with Inputs. Might be run immediately or in future,
+    // after a compatible preamble becomes ready. The built AST will only be
+    // cached if it is compatible with latest inputs of the ASTWorker, which
+    // might not be the case in case of deferred builds.
+    auto ASTTask = [this, WantDiags,
+                    CIDiags = std::move(CompilerInvocationDiags),
+                    TaskName = std::move(TaskName),
+                    RunPublish = std::move(RunPublish)](
+                       std::shared_ptr<const PreambleData> Preamble,
+                       std::unique_ptr<CompilerInvocation> Invocation,
+                       ParseInputs Inputs) mutable {
+      if (Preamble != LatestPreamble) {
+        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);
+        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.
+        // No need to rebuild the AST if we won't send the diagnostics.
         if (!CanPublishResults)
           return;
       }
+      // Give up our ownership to old preamble before starting expensive AST
+      // build.
+      Preamble.reset();
+      // Used to check whether we can update AST cache and callback status.
+      bool InputsAreTheSame =
+          std::tie(FileInputs->CompileCommand, FileInputs->Contents) ==
+          std::tie(Inputs.CompileCommand, Inputs.Contents);
 
+      Status.update([&](TUStatus &Status) {
+        Status.ASTActivity.K = ASTAction::Building;
+        Status.ASTActivity.Name = std::move(TaskName);
+      });
       // Get the AST for diagnostics.
       llvm::Optional<std::unique_ptr<ParsedAST>> AST = IdleASTs.take(this);
-    auto RebuildStartTime = DebouncePolicy::clock::now();
+      llvm::Optional<DebouncePolicy::clock::duration> RebuildDuration;
       if (!AST) {
-      llvm::Optional<ParsedAST> NewAST =
-          buildAST(FileName, std::move(Invocation), CompilerInvocationDiags,
-                   Inputs, NewPreamble);
+        auto RebuildStartTime = DebouncePolicy::clock::now();
+        llvm::Optional<ParsedAST> NewAST = buildAST(
+            FileName, std::move(Invocation), CIDiags, Inputs, LatestPreamble);
+        RebuildDuration = DebouncePolicy::clock::now() - RebuildStartTime;
         // 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;
+        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.
+      // 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;
+          // 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 (RebuildDuration && 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);
+            RebuildTimes.push_back(*RebuildDuration);
           }
         }
         trace::Span Span("Running main AST callback");
-
         Callbacks.onMainAST(FileName, **AST, RunPublish);
+        // RanASTCallback might be referring to a different FileInputs in case
+        // of deferred builds, only update if that's not the case.
+        if (InputsAreTheSame)
           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.
+        // 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 iff it is compatible with
+      // current ASTWorker inputs.
+      if (InputsAreTheSame)
         IdleASTs.put(this, std::move(*AST));
     };
+
+    bool PreambleFresh =
+        !Inputs.ForceRebuild && LatestPreamble &&
+        isPreambleUsable(*LatestPreamble, Inputs, FileName, *Invocation);
+    // Do not build AST with stale preambles, as diagnostics might contain false
+    // positives and symbol information might be misleading. Instead defer the
+    // build until a compatible preamble is ready. It might be overwritten by a
+    // subsequent update.
+    if (!PreambleFresh) {
+      // Preamble is not compatible with current inputs.
+      IdleASTs.take(this);
+      PW.update(
+          *Invocation, Inputs,
+          [this, ASTTask = std::move(ASTTask)](
+              std::unique_ptr<CompilerInvocation> CI, ParseInputs PI,
+              std::shared_ptr<const PreambleData> Preamble) {
+            std::string TaskName = llvm::formatv("Build AST ({0})", PI.Version);
+            // Build golden AST with corresponding inputs.
+            auto DeferredTask = [Preamble = std::move(Preamble),
+                                 CI = std::move(CI), PI = std::move(PI),
+                                 ASTTask = std::move(ASTTask)]() mutable {
+              ASTTask(std::move(Preamble), std::move(CI), std::move(PI));
+            };
+            if (RunSync) {
+              DeferredTask();
+              return;
+            }
+            {
+              std::lock_guard<std::mutex> Lock(Mutex);
+              // We issue the update with a WantDiagnostics::Yes to make sure it
+              // doesn't get debounced and put it in front, as this update was
+              // issued before any requests in the queue.
+              Requests.push_front(
+                  {std::move(DeferredTask), std::move(TaskName),
+                   steady_clock::now(), Context::current().clone(),
+                   WantDiagnostics::Yes, TUScheduler::NoInvalidation, nullptr});
+            }
+            RequestsCV.notify_all();
+          });
+      // FIXME: We should build patched-up ASTs in case of stale preambles
+      // instead of blocking.
+      PW.blockUntilIdle(Deadline::infinity());
+      return;
+    }
+    vlog("Reusing preamble version {0} for version {1} of {2}",
+         LatestPreamble->Version, Inputs.Version, FileName);
+    if (!InputsAreTheSame) {
+      // Cannot use the cached AST if inputs have changed.
+      IdleASTs.take(this);
+    } else {
+      // 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.
+      // FIXME: should we refresh the cache in IdleASTs for the current file at
+      // this point?
+      if (RanCallbackForPrevInputs) {
+        RanASTCallback = true;
+        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;
+        });
+        return;
+      }
+    }
+    ASTTask(LatestPreamble, std::move(Invocation), std::move(Inputs));
+  };
   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 {
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);
 }
 
 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();
+  PW.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) {
     assert(!Done && "running a task after stop()");
     trace::Span Tracer(Name + ":" + llvm::sys::path::filename(FileName));
     Task();
     return;
   }
 
   {
     std::lock_guard<std::mutex> Lock(Mutex);
     assert(!Done && "running a task after stop()");
     // Cancel any requests invalidated by this request.
     if (UpdateType) {
       for (auto &R : llvm::reverse(Requests)) {
         if (R.InvalidationPolicy == TUScheduler::InvalidateOnUpdate)
           R.Invalidate();
         if (R.UpdateType)
           break; // Older requests were already invalidated by the older update.
       }
     }
 
     // Allow this request to be cancelled if invalidated.
     Context Ctx = Context::current().derive(kFileBeingProcessed, FileName);
     Canceler Invalidate = nullptr;
     if (Invalidation) {
       WithContext WC(std::move(Ctx));
       std::tie(Ctx, Invalidate) = cancelableTask();
     }
     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()) {
         if (Done) {
           if (Requests.empty())
while (true) {
     // only writer is also this thread.
     {
       std::unique_lock<Semaphore> Lock(Barrier, std::try_to_lock);
       if (!Lock.owns_lock()) {
         Status.update([&](TUStatus &Status) {
           Status.ASTActivity.K = ASTAction::Queued;
           Status.ASTActivity.Name = CurrentRequest->Name;
         });
         Lock.lock();
       }
       WithContext Guard(std::move(CurrentRequest->Ctx));
       trace::Span Tracer(CurrentRequest->Name);
       Status.update([&](TUStatus &Status) {
         Status.ASTActivity.K = ASTAction::RunningAction;
         Status.ASTActivity.Name = CurrentRequest->Name;
       });
       CurrentRequest->Action();

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

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

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, 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),
-                  // We start building the ast
+                  // Start task for building the ast, as a result of building
+                  // preamble, on astworker thread.
+                  TUState(PreambleAction::Idle, ASTAction::RunningAction),
+                  // 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

ParsedAST TestTU::build() const {
   Inputs.Opts.ClangTidyOpts.WarningsAsErrors = ClangTidyWarningsAsErrors;
   Inputs.Index = ExternalIndex;
   if (Inputs.Index)
     Inputs.Opts.SuggestMissingIncludes = true;
   StoreDiags Diags;
   auto CI = buildCompilerInvocation(Inputs, Diags);
   assert(CI && "Failed to build compilation invocation.");
   auto Preamble =
-      buildPreamble(FullFilename, *CI,
-                    /*OldPreamble=*/nullptr, Inputs,
+      buildPreamble(FullFilename, *CI, Inputs,
                     /*StoreInMemory=*/true, /*PreambleCallback=*/nullptr);
   auto AST =
       buildAST(FullFilename, 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).