[clangd] Decouple preambleworker from astworker, NFCI

[kadircet]

Summary

First step to enable deferred preamble builds. Not intending to land it
alone, will have follow-ups that will implement full deferred build
functionality and will land after all of them are ready.

[sammccall]

clang-tools-extra/clangd/TUScheduler.cpp

155

This doc comment doesn't mention "thread" once :-)

156

move documentation of requestBuild() to that function?

157

Move documentation of stop() semantics to stop()?

159

we may want to call this PreambleThread? Current name emphasizes peer relationship with ASTWorker, but current code has parent/child relationship.

171

or just update()?

201

I think all the members in this class with "preamble" in the name can drop that word without any loss, e.g. PreambleWorker::waitForFirst()

205

latest()?

218

Why do we rebuild the preamble if stop is requested?

265

as discussed a bit in chat, I think this part needs some more thought. Currently the ASTWorker and PreambleWorker emit data for the same file with some interleaving that's hard to reason about. The state needs an owner.

(e.g. consider a class that holds a TUStatus and exposes an Event<TUStatus>, with threadsafe setPreambleAction and setWorkerAction methods, or so)

291

just build(), to avoid confusion clangd::buildPreamble?

Also consider merging with updateLatestBuiltPreamble, both are small and private and they always go together.

295

(nullptr should work here?)

312

Yikes, lots of state :-)
I'd consider merging the two CVs - I find keeping the events separate and reasoning when you need notify_one vs _all doesn't seem to pay off, may just be the way I think about queues. Up to you.

322

nit: last vs latest, be consistent

325

please make this owning for simplicity, this isn't a lightweight object anyway

596

this doesn't seem correct (maybe ok in this patch because of the blocking, but not in general). You're assuming the last available preamble is the one that the last AST was built with.

I suppose you can't check the preamble of the current ParsedAST because it might not be cached, and you nevertheless want to skip rebuild if the diagnostics are going to be the same. I can't think of anything better than continuing to hold the shared_ptr for PreambleForLastBuiltAST or something like that.

[kadircet]

clang-tools-extra/clangd/TUScheduler.cpp

218

it was to be consistent with astworker, but i suppose it isn't necessary as ASTWorker does that to make sure each LSP request gets a reply, and it can be done with stale preambles.

Are there any other reasons for ASTWorker to empty the queue after stop?

295

right, it was copied over :/

312

i like them to be seperate as it makes naming easier :D but i agree that having less state also makes reasoning easier.

596

right, this is just a "hack" to keep this change NFC.

in the follow-up patches i am planning to signal whether latest built preamble is reusable for a given ParseInputs, and also signal what the AST should be patched with.

diagnostics(ast) will only be built if preamble is re-usable.

[kadircet]

• Address comments

[sammccall]

This basically looks good, just nits (mostly about comments, sorry I can't help myself).

I think the only real outstanding thing is TUStatus, which we're still discussing offline. Since this patch is intended to be NFC, I see two paths consistent with the spirit of that:

• work out how to make this emit the same sequence of events (probably in a fairly brittle way), document that, and sort it out in the next patch
• make TUStatus reflect the new model, without changing its essential design or scope. Currently TUStatus is currently (stateful bits about TU, worker thread activity), so the extension would be (stateful bits about TU, worker thread activity, preamble thread activity).

clang-tools-extra/clangd/TUScheduler.cpp

156

This sentence is hard to parse (processes, build, requests, and run are all ambiguous between nouns and verbs, and the pronouns, articles etc are missing).

Suggestion: "It processes updates asynchronously, building preambles on a dedicated thread. (This thread should be spawned externally and call the run() method)."

156

Actually, is "processing requests" really the clearest way to describe the model here?

What about "Whenever the thread is idle and the preamble is outdated, it starts to build a fresh preamble from the latest inputs."

157

nit: "which" appears to refer to RunSync here, and "provided" doesn't exactly mean "true".

Suggestion: If RunSync is true, preambles are built synchronously in update() instead.

171

request -> requested version

218

Yeah, I think that's the only reason and the asymmetry is worthwhile here (esp because one preamble is way more expensive than one AST).

If that's not clear in the ASTworker piece, please do add a comment!

224

I think the comment adds more confusion than it resolves: you seem to be describing why an alternative (reacquiring the lock after building, and then calling reset?) doesn't work, but never state what that alternative is.

Moreover, clearing NextReq after moving the value into CurrentReq seems natural, I'm not sure a comment is necessary (except maybe to complain that std::optional's move constructor doesn't do this!)

243

we've been stopped -> it should stop?

300

If LatestBuild is the last shared_ptr to the old preamble, we're destroying it here with the lock held, which probably does IO.

Instead, std::swap(Preamble, LatestBuild) so that the destructor runs when the function exits.

310

(not sure the "written by" comment is needed since this is just a regular guarded-by-mutex?)

596

right, this is just a "hack" to keep this change NFC.

can you add a comment about this? (In particular, why it's correct?)

in the follow-up patches i am planning to signal whether latest built preamble is reusable for a given ParseInputs

Does "reusable" mean "completely valid" (current semantics), or usable with some tweaks (e.g. added headers)? It would be good to define some precise terminology around this.

diagnostics(ast) will only be built if preamble is re-usable.

SG

[sammccall]

In D76125#1926773, @sammccall wrote:

I think the only real outstanding thing is TUStatus, which we're still discussing offline. Since this patch is intended to be NFC, I see two paths consistent with the spirit of that

I'd really like to find any way to unblock this that isn't "redesign the API" because I don't think we should do that in a hurry. So also open to other ideas here, or happy to take a stab at this part if it's helpful.

[sammccall]

clang-tools-extra/clangd/TUScheduler.cpp

412

why are we moving this?

433

(and this)

[kadircet]

clang-tools-extra/clangd/TUScheduler.cpp

265

sent out https://reviews.llvm.org/D76304

412

it has been moved into PreambleThread instead, same as the one below handling the failure.

596

reusable means completely valid.

[kadircet]

• Address comments

[kadircet]

• Reset CurrentReq inside run loop, rather than inside build.

[kadircet]

• Start preambleworker in sync mode instead of blocking

clang-tools-extra/clangd/Preamble.h

 
 /// 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.
 std::shared_ptr<const PreambleData>
-buildPreamble(PathRef FileName, CompilerInvocation &CI,
+buildPreamble(PathRef FileName, CompilerInvocation CI,
               std::shared_ptr<const PreambleData> OldPreamble,
               const ParseInputs &Inputs, bool StoreInMemory,
               PreambleParsedCallback PreambleCallback);
 
-
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_PREAMBLE_H

clang-tools-extra/clangd/Preamble.cpp

PreambleData::PreambleData(const ParseInputs &Inputs,
                            CanonicalIncludes CanonIncludes)
     : 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,
+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 =

clang-tools-extra/clangd/TUScheduler.cpp

 #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/Optional.h"
 #include "llvm/ADT/ScopeExit.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Threading.h"
 #include <algorithm>
 #include <memory>
 #include <mutex>
 #include <queue>
 #include <thread>
private:
   std::mutex Mut;
   unsigned MaxRetainedASTs;
   /// Items sorted in LRU order, i.e. first item is the most recently accessed
   /// one.
   std::vector<KVPair> LRU; /* GUARDED_BY(Mut) */
 };
 
 namespace {
+/// 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.
+class PreambleThread {
+public:
+  PreambleThread(llvm::StringRef FileName, ParsingCallbacks &Callbacks,
+                 bool StorePreambleInMemory, bool RunSync)
+      : FileName(FileName), Callbacks(Callbacks),
+        StoreInMemory(StorePreambleInMemory), RunSync(RunSync) {}
+
+  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) {
+      TUStatus::BuildDetails Details;
+      Details.BuildFailed = true;
+      std::string TaskName = llvm::formatv("Update ({0})", PI.Version);
+      emitTUStatus({TUAction::BuildingPreamble, std::move(TaskName)}, &Details);
+      // Make sure anyone waiting for the preamble gets notified it could not be
+      // built.
+      BuiltFirst.notify();
+      return;
+    }
+    // Make possibly expensive copy while not holding the lock.
+    Request Req = {std::make_unique<CompilerInvocation>(*CI), std::move(PI)};
+    if (RunSync) {
+      build(std::move(Req));
+      return;
+    }
+    {
+      std::lock_guard<std::mutex> Lock(Mutex);
+      assert(!Done && "Build request to PreambleWorker after stop");
+      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();
+      }
+      // Build the preamble and let the waiters know about it.
+      build(std::move(*CurrentReq));
+      {
+        std::lock_guard<std::mutex> Lock(Mutex);
+        CurrentReq.reset();
+      }
+      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;
+    }
+    // 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;
+  };
+
+  bool isDone() {
+    std::lock_guard<std::mutex> Lock(Mutex);
+    return Done;
+  }
+
+  /// Updates the TUStatus and emits it. Only called in the worker thread.
+  void emitTUStatus(TUAction Action,
+                    const TUStatus::BuildDetails *Details = nullptr) {
+    // Do not emit TU statuses when the worker is shutting down.
+    if (isDone())
+      return;
+    TUStatus Status({std::move(Action), {}});
+    if (Details)
+      Status.Details = *Details;
+    Callbacks.onFileUpdated(FileName, Status);
+  }
+
+  /// 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();
+
+    std::string TaskName = llvm::formatv("Update ({0})", Inputs.Version);
+    emitTUStatus({TUAction::BuildingPreamble, std::move(TaskName)});
+
+    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);
+    }
+    BuiltFirst.notify();
+  }
+
+  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;
+};
+
 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
 /// and the TUScheduler. The TUScheduler should discard an ASTWorker when
private:
   /// Handles retention of ASTs.
   TUScheduler::ASTCache &IdleASTs;
   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;
-  /// Whether to keep the built preambles in memory or on disk.
-  const bool StorePreambleInMemory;
   /// Callback invoked when preamble or main file AST is built.
   ParsingCallbacks &Callbacks;
   /// Only accessed by the worker thread.
   TUStatus Status;
 
   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) */
-  std::shared_ptr<const PreambleData> LastBuiltPreamble; /* GUARDED_BY(Mutex) */
+  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) */
-  /// Becomes ready when the first preamble build finishes.
-  Notification PreambleWasBuilt;
   /// 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
   // 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) */
+
+  PreambleThread PW;
 };
 
 /// 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;

 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("worker:" + llvm::sys::path::filename(FileName),
+  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)
     : IdleASTs(LRUCache), RunSync(RunSync), UpdateDebounce(UpdateDebounce),
       FileName(FileName), CDB(CDB),
-      StorePreambleInMemory(StorePreamblesInMemory),
       Callbacks(Callbacks), Status{TUAction(TUAction::Idle, ""),
                                    TUStatus::BuildDetails()},
-      Barrier(Barrier), Done(false) {
+      Barrier(Barrier), Done(false),
+      // FIXME: Run preambleworker async.
+      PW(FileName, Callbacks, StorePreamblesInMemory, /*RunSync=*/true) {
   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);
 }
 
bool InputsAreTheSame =
         std::tie(Inputs.CompileCommand, Inputs.Contents);
 
     bool RanCallbackForPrevInputs = RanASTCallback;
     {
       std::lock_guard<std::mutex> Lock(Mutex);
       FileInputs = std::make_shared<ParseInputs>(Inputs);
     }
     RanASTCallback = false;
-    emitTUStatus({TUAction::BuildingPreamble, TaskName});
     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);
+    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);
-      TUStatus::BuildDetails Details;
-      Details.BuildFailed = true;
-      emitTUStatus({TUAction::BuildingPreamble, TaskName}, &Details);
       // 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.
-      PreambleWasBuilt.notify();
       return;
     }
 
-    std::shared_ptr<const PreambleData> OldPreamble =
-        Inputs.ForceRebuild ? std::shared_ptr<const PreambleData>()
-                            : getPossiblyStalePreamble();
-    std::shared_ptr<const PreambleData> NewPreamble = buildPreamble(
-        FileName, *Invocation, OldPreamble, Inputs, StorePreambleInMemory,
-        [this, Version(Inputs.Version)](
-            ASTContext &Ctx, std::shared_ptr<clang::Preprocessor> PP,
-            const CanonicalIncludes &CanonIncludes) {
-          Callbacks.onPreambleAST(FileName, Version, Ctx, std::move(PP),
-                                  CanonIncludes);
-        });
-
     bool CanReuseAST = InputsAreTheSame && (OldPreamble == NewPreamble);
-    {
-      std::lock_guard<std::mutex> Lock(Mutex);
-      LastBuiltPreamble = 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();
-    PreambleWasBuilt.notify();
     emitTUStatus({TUAction::BuildingFile, 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
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 {
-  std::lock_guard<std::mutex> Lock(Mutex);
-  return LastBuiltPreamble;
+  return PW.latest();
 }
 
 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 { PreambleWasBuilt.wait(); }
+void ASTWorker::waitForFirstPreamble() const { PW.waitForFirst(); }
 
 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::emitTUStatus(TUAction Action,
   std::lock_guard<std::mutex> Lock(PublishMu);
   // Do not emit TU statuses when the ASTWorker is shutting down.
   if (CanPublishResults) {
     Callbacks.onFileUpdated(FileName, Status);
   }
 }
 
 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())