--- a/clang-tools-extra/clangd/TUScheduler.cpp	Tue Jun 15 19:13:43 2021 +0900
+++ b/clang-tools-extra/clangd/TUScheduler.cpp	Tue Jun 15 19:15:29 2021 +0900
@@ -56,6 +56,7 @@
 #include "support/Cancellation.h"
 #include "support/Context.h"
 #include "support/Logger.h"
+#include "support/MemoryTree.h"
 #include "support/Path.h"
 #include "support/Threading.h"
 #include "support/Trace.h"
@@ -69,13 +70,16 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Allocator.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 <atomic>
 #include <chrono>
+#include <condition_variable>
 #include <functional>
 #include <memory>
 #include <mutex>
@@ -177,7 +181,132 @@
   std::vector<KVPair> LRU; /* GUARDED_BY(Mut) */
 };
 
+/// A map from header files to an opened "proxy" file that includes them.
+/// If you open the header, the compile command from the proxy file is used.
+///
+/// This inclusion information could also naturally live in the index, but there
+/// are advantages to using open files instead:
+///  - it's easier to achieve a *stable* choice of proxy, which is important
+///    to avoid invalidating the preamble
+///  - context-sensitive flags for libraries with multiple configurations
+///    (e.g. C++ stdlib sensitivity to -std version)
+///  - predictable behavior, e.g. guarantees that go-to-def landing on a header
+///    will have a suitable command available
+///  - fewer scaling problems to solve (project include graphs are big!)
+///
+/// Implementation details:
+/// - We only record this for mainfiles where the command was trustworthy
+///   (i.e. not inferred). This avoids a bad inference "infecting" other files.
+/// - Once we've picked a proxy file for a header, we stick with it until the
+///   proxy file is invalidated *and* a new candidate proxy file is built.
+///   Switching proxies is expensive, as the compile flags will (probably)
+///   change and therefore we'll end up rebuilding the header's preamble.
+/// - We don't capture the actual compile command, but just the filename we
+///   should query to get it. This avoids getting out of sync with the CDB.
+///
+/// All methods are threadsafe. In practice, update() comes from preamble
+/// threads, remove()s mostly from the main thread, and get() from ASTWorker.
+/// Writes are rare and reads are cheap, so we don't expect much contention.
+class TUScheduler::HeaderIncluderCache {
+  // We should be be a little careful how we store the include graph of open
+  // files, as each can have a large number of transitive headers.
+  // This representation is O(unique transitive source files).
+  llvm::BumpPtrAllocator Arena;
+  struct Association {
+    llvm::StringRef MainFile;
+    // Circular-linked-list of associations with the same mainFile.
+    // Null indicates that the mainfile was removed.
+    Association *Next;
+  };
+  llvm::StringMap<Association, llvm::BumpPtrAllocator &> HeaderToMain;
+  llvm::StringMap<Association *, llvm::BumpPtrAllocator &> MainToFirst;
+  std::atomic<size_t> UsedBytes; // Updated after writes.
+  mutable std::mutex Mu;
+
+  void invalidate(Association *First) {
+    Association *Current = First;
+    do {
+      Association *Next = Current->Next;
+      Current->Next = nullptr;
+      Current = Next;
+    } while (Current != First);
+  }
+
+  // Create the circular list and return the head of it.
+  Association *associate(llvm::StringRef MainFile,
+                         llvm::ArrayRef<std::string> Headers) {
+    Association *First = nullptr, *Prev = nullptr;
+    for (const std::string &Header : Headers) {
+      auto &Assoc = HeaderToMain[Header];
+      if (Assoc.Next)
+        continue; // Already has a valid association.
+
+      Assoc.MainFile = MainFile;
+      Assoc.Next = Prev;
+      Prev = &Assoc;
+      if (!First)
+        First = &Assoc;
+    }
+    if (First)
+      First->Next = Prev;
+    return First;
+  }
+
+  void updateMemoryUsage() {
+    auto StringMapHeap = [](const auto &Map) {
+      // StringMap stores the hashtable on the heap.
+      // It contains pointers to the entries, and a hashcode for each.
+      return Map.getNumBuckets() * (sizeof(void *) + sizeof(unsigned));
+    };
+    size_t Usage = Arena.getTotalMemory() + StringMapHeap(MainToFirst) +
+                   StringMapHeap(HeaderToMain) + sizeof(*this);
+    UsedBytes.store(Usage, std::memory_order_release);
+  }
+
+public:
+  HeaderIncluderCache() : HeaderToMain(Arena), MainToFirst(Arena) {
+    updateMemoryUsage();
+  }
+
+  // Associate each header with MainFile (unless already associated).
+  // Headers not in the list will have their associations removed.
+  void update(PathRef MainFile, llvm::ArrayRef<std::string> Headers) {
+    std::lock_guard<std::mutex> Lock(Mu);
+    auto It = MainToFirst.try_emplace(MainFile, nullptr);
+    Association *&First = It.first->second;
+    if (First)
+      invalidate(First);
+    First = associate(It.first->first(), Headers);
+    updateMemoryUsage();
+  }
+
+  // Mark MainFile as gone.
+  // This will *not* disassociate headers with MainFile immediately, but they
+  // will be eligible for association with other files that get update()d.
+  void remove(PathRef MainFile) {
+    std::lock_guard<std::mutex> Lock(Mu);
+    Association *&First = MainToFirst[MainFile];
+    if (First)
+      invalidate(First);
+  }
+
+  /// Get the mainfile associated with Header, or the empty string if none.
+  std::string get(PathRef Header) const {
+    std::lock_guard<std::mutex> Lock(Mu);
+    return HeaderToMain.lookup(Header).MainFile.str();
+  }
+
+  size_t getUsedBytes() const {
+    return UsedBytes.load(std::memory_order_acquire);
+  }
+};
+
 namespace {
+
+bool isReliable(const tooling::CompileCommand &Cmd) {
+  return Cmd.Heuristic.empty();
+}
+
 /// Threadsafe manager for updating a TUStatus and emitting it after each
 /// update.
 class SynchronizedTUStatus {
@@ -219,10 +348,12 @@
 public:
   PreambleThread(llvm::StringRef FileName, ParsingCallbacks &Callbacks,
                  bool StorePreambleInMemory, bool RunSync,
-                 SynchronizedTUStatus &Status, ASTWorker &AW)
+                 SynchronizedTUStatus &Status,
+                 TUScheduler::HeaderIncluderCache &HeaderIncluders,
+                 ASTWorker &AW)
       : FileName(FileName), Callbacks(Callbacks),
         StoreInMemory(StorePreambleInMemory), RunSync(RunSync), Status(Status),
-        ASTPeer(AW) {}
+        ASTPeer(AW), HeaderIncluders(HeaderIncluders) {}
 
   /// 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
@@ -239,10 +370,14 @@
       return;
     }
     {
-      std::lock_guard<std::mutex> Lock(Mutex);
-      // If shutdown is issued, don't bother building.
-      if (Done)
-        return;
+      std::unique_lock<std::mutex> Lock(Mutex);
+      // If NextReq was requested with WantDiagnostics::Yes we cannot just drop
+      // that on the floor. Block until we start building it. This won't
+      // dead-lock as we are blocking the caller thread, while builds continue
+      // on preamble thread.
+      ReqCV.wait(Lock, [this] {
+        return !NextReq || NextReq->WantDiags != WantDiagnostics::Yes;
+      });
       NextReq = std::move(Req);
     }
     // Let the worker thread know there's a request, notify_one is safe as there
@@ -265,6 +400,10 @@
 
       {
         WithContext Guard(std::move(CurrentReq->Ctx));
+        // Note that we don't make use of the ContextProvider here.
+        // Preamble tasks are always scheduled by ASTWorker tasks, and we
+        // reuse the context/config that was created at that level.
+
         // Build the preamble and let the waiters know about it.
         build(std::move(*CurrentReq));
       }
@@ -341,6 +480,7 @@
 
   SynchronizedTUStatus &Status;
   ASTWorker &ASTPeer;
+  TUScheduler::HeaderIncluderCache &HeaderIncluders;
 };
 
 class ASTWorkerHandle;
@@ -358,8 +498,9 @@
 class ASTWorker {
   friend class ASTWorkerHandle;
   ASTWorker(PathRef FileName, const GlobalCompilationDatabase &CDB,
-            TUScheduler::ASTCache &LRUCache, Semaphore &Barrier, bool RunSync,
-            DebouncePolicy UpdateDebounce, bool StorePreamblesInMemory,
+            TUScheduler::ASTCache &LRUCache,
+            TUScheduler::HeaderIncluderCache &HeaderIncluders,
+            Semaphore &Barrier, bool RunSync, const TUScheduler::Options &Opts,
             ParsingCallbacks &Callbacks);
 
 public:
@@ -370,19 +511,21 @@
   /// request, it is used to limit the number of actively running threads.
   static ASTWorkerHandle
   create(PathRef FileName, const GlobalCompilationDatabase &CDB,
-         TUScheduler::ASTCache &IdleASTs, AsyncTaskRunner *Tasks,
-         Semaphore &Barrier, DebouncePolicy UpdateDebounce,
-         bool StorePreamblesInMemory, ParsingCallbacks &Callbacks);
+         TUScheduler::ASTCache &IdleASTs,
+         TUScheduler::HeaderIncluderCache &HeaderIncluders,
+         AsyncTaskRunner *Tasks, Semaphore &Barrier,
+         const TUScheduler::Options &Opts, ParsingCallbacks &Callbacks);
   ~ASTWorker();
 
-  void update(ParseInputs Inputs, WantDiagnostics);
+  void update(ParseInputs Inputs, WantDiagnostics, bool ContentChanged);
   void
   runWithAST(llvm::StringRef Name,
              llvm::unique_function<void(llvm::Expected<InputsAndAST>)> Action,
              TUScheduler::ASTActionInvalidation);
   bool blockUntilIdle(Deadline Timeout) const;
 
-  std::shared_ptr<const PreambleData> getPossiblyStalePreamble() const;
+  std::shared_ptr<const PreambleData> getPossiblyStalePreamble(
+      std::shared_ptr<const ASTSignals> *ASTSignals = nullptr) const;
 
   /// Used to inform ASTWorker about a new preamble build by PreambleThread.
   /// Diagnostics are only published through this callback. This ensures they
@@ -409,21 +552,36 @@
   bool isASTCached() const;
 
 private:
+  // Details of an update request that are relevant to scheduling.
+  struct UpdateType {
+    // Do we want diagnostics from this version?
+    // If Yes, we must always build this version.
+    // If No, we only need to build this version if it's read.
+    // If Auto, we build if it's read or if the debounce expires.
+    WantDiagnostics Diagnostics;
+    // Did the main-file content of the document change?
+    // If so, we're allowed to cancel certain invalidated preceding reads.
+    bool ContentChanged;
+  };
+
   /// Publishes diagnostics for \p Inputs. It will build an AST or reuse the
   /// cached one if applicable. Assumes LatestPreamble is compatible for \p
   /// Inputs.
   void generateDiagnostics(std::unique_ptr<CompilerInvocation> Invocation,
                            ParseInputs Inputs, std::vector<Diag> CIDiags);
 
+  void updateASTSignals(ParsedAST &AST);
+
   // 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,
-                 llvm::Optional<WantDiagnostics> UpdateType,
+                 llvm::Optional<UpdateType> Update,
                  TUScheduler::ASTActionInvalidation);
 
   /// Determines the next action to perform.
@@ -439,24 +597,25 @@
     std::string Name;
     steady_clock::time_point AddTime;
     Context Ctx;
-    llvm::Optional<WantDiagnostics> UpdateType;
+    llvm::Optional<Context> QueueCtx;
+    llvm::Optional<UpdateType> Update;
     TUScheduler::ASTActionInvalidation InvalidationPolicy;
     Canceler Invalidate;
   };
 
   /// Handles retention of ASTs.
   TUScheduler::ASTCache &IdleASTs;
+  TUScheduler::HeaderIncluderCache &HeaderIncluders;
   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;
+  /// Callback to create processing contexts for tasks.
+  const std::function<Context(llvm::StringRef)> ContextProvider;
   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.
@@ -468,16 +627,26 @@
   /// thread are not locked, as it's the only writer.
   ParseInputs FileInputs; /* GUARDED_BY(Mutex) */
   /// Times of recent AST rebuilds, used for UpdateDebounce computation.
-  llvm::SmallVector<DebouncePolicy::clock::duration, 8>
+  llvm::SmallVector<DebouncePolicy::clock::duration>
       RebuildTimes; /* GUARDED_BY(Mutex) */
   /// Set to true to signal run() to finish processing.
   bool Done;                              /* GUARDED_BY(Mutex) */
   std::deque<Request> Requests;           /* GUARDED_BY(Mutex) */
-  std::queue<Request> PreambleRequests;   /* GUARDED_BY(Mutex) */
   llvm::Optional<Request> CurrentRequest; /* GUARDED_BY(Mutex) */
+  /// Signalled whenever a new request has been scheduled or processing of a
+  /// request has completed.
   mutable std::condition_variable RequestsCV;
+  std::shared_ptr<const ASTSignals> LatestASTSignals; /* GUARDED_BY(Mutex) */
+  /// Latest build preamble for current TU.
+  /// None means no builds yet, null means there was an error while building.
+  /// Only written by ASTWorker's thread.
+  llvm::Optional<std::shared_ptr<const PreambleData>> LatestPreamble;
+  std::deque<Request> PreambleRequests; /* GUARDED_BY(Mutex) */
+  /// Signaled whenever LatestPreamble changes state or there's a new
+  /// PreambleRequest.
+  mutable std::condition_variable PreambleCV;
   /// Guards the callback that publishes results of AST-related computations
-  /// (diagnostics, highlightings) and file statuses.
+  /// (diagnostics) 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.
@@ -537,12 +706,14 @@
 
 ASTWorkerHandle
 ASTWorker::create(PathRef FileName, const GlobalCompilationDatabase &CDB,
-                  TUScheduler::ASTCache &IdleASTs, AsyncTaskRunner *Tasks,
-                  Semaphore &Barrier, DebouncePolicy UpdateDebounce,
-                  bool StorePreamblesInMemory, ParsingCallbacks &Callbacks) {
+                  TUScheduler::ASTCache &IdleASTs,
+                  TUScheduler::HeaderIncluderCache &HeaderIncluders,
+                  AsyncTaskRunner *Tasks, Semaphore &Barrier,
+                  const TUScheduler::Options &Opts,
+                  ParsingCallbacks &Callbacks) {
   std::shared_ptr<ASTWorker> Worker(
-      new ASTWorker(FileName, CDB, IdleASTs, Barrier, /*RunSync=*/!Tasks,
-                    UpdateDebounce, StorePreamblesInMemory, Callbacks));
+      new ASTWorker(FileName, CDB, IdleASTs, HeaderIncluders, Barrier,
+                    /*RunSync=*/!Tasks, Opts, Callbacks));
   if (Tasks) {
     Tasks->runAsync("ASTWorker:" + llvm::sys::path::filename(FileName),
                     [Worker]() { Worker->run(); });
@@ -554,16 +725,17 @@
 }
 
 ASTWorker::ASTWorker(PathRef FileName, const GlobalCompilationDatabase &CDB,
-                     TUScheduler::ASTCache &LRUCache, Semaphore &Barrier,
-                     bool RunSync, DebouncePolicy UpdateDebounce,
-                     bool StorePreamblesInMemory, ParsingCallbacks &Callbacks)
-    : IdleASTs(LRUCache), RunSync(RunSync), UpdateDebounce(UpdateDebounce),
-      FileName(FileName), CDB(CDB), Callbacks(Callbacks), Barrier(Barrier),
-      Done(false), Status(FileName, Callbacks),
-      PreamblePeer(FileName, Callbacks, StorePreamblesInMemory,
-                   // FIXME: Run PreamblePeer asynchronously once ast patching
-                   // is available.
-                   /*RunSync=*/true, Status, *this) {
+                     TUScheduler::ASTCache &LRUCache,
+                     TUScheduler::HeaderIncluderCache &HeaderIncluders,
+                     Semaphore &Barrier, bool RunSync,
+                     const TUScheduler::Options &Opts,
+                     ParsingCallbacks &Callbacks)
+    : IdleASTs(LRUCache), HeaderIncluders(HeaderIncluders), RunSync(RunSync),
+      UpdateDebounce(Opts.UpdateDebounce), FileName(FileName),
+      ContextProvider(Opts.ContextProvider), CDB(CDB), Callbacks(Callbacks),
+      Barrier(Barrier), Done(false), Status(FileName, Callbacks),
+      PreamblePeer(FileName, Callbacks, Opts.StorePreamblesInMemory, RunSync,
+                   Status, HeaderIncluders, *this) {
   // Set a fallback command because compile command can be accessed before
   // `Inputs` is initialized. Other fields are only used after initialization
   // from client inputs.
@@ -581,7 +753,8 @@
 #endif
 }
 
-void ASTWorker::update(ParseInputs Inputs, WantDiagnostics WantDiags) {
+void ASTWorker::update(ParseInputs Inputs, WantDiagnostics WantDiags,
+                       bool ContentChanged) {
   std::string TaskName = llvm::formatv("Update ({0})", Inputs.Version);
   auto Task = [=]() mutable {
     // Get the actual command as `Inputs` does not have a command.
@@ -589,10 +762,25 @@
     // 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;
+    auto Cmd = CDB.getCompileCommand(FileName);
+    // If we don't have a reliable command for this file, it may be a header.
+    // Try to find a file that includes it, to borrow its command.
+    if (!Cmd || !isReliable(*Cmd)) {
+      std::string ProxyFile = HeaderIncluders.get(FileName);
+      if (!ProxyFile.empty()) {
+        auto ProxyCmd = CDB.getCompileCommand(ProxyFile);
+        if (!ProxyCmd || !isReliable(*ProxyCmd)) {
+          // This command is supposed to be reliable! It's probably gone.
+          HeaderIncluders.remove(ProxyFile);
+        } else {
+          // We have a reliable command for an including file, use it.
+          Cmd = tooling::transferCompileCommand(std::move(*ProxyCmd), FileName);
+        }
+      }
+    }
+    if (Cmd)
+      Inputs.CompileCommand = std::move(*Cmd);
     else
-      // FIXME: consider using old command if it's not a fallback one.
       Inputs.CompileCommand = CDB.getFallbackCommand(FileName);
 
     bool InputsAreTheSame =
@@ -613,7 +801,7 @@
     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, " "));
+        printArgv(Inputs.CompileCommand.CommandLine));
 
     StoreDiags CompilerInvocationDiagConsumer;
     std::vector<std::string> CC1Args;
@@ -621,7 +809,7 @@
         Inputs, CompilerInvocationDiagConsumer, &CC1Args);
     // Log cc1 args even (especially!) if creating invocation failed.
     if (!CC1Args.empty())
-      vlog("Driver produced command: cc1 {0}", llvm::join(CC1Args, " "));
+      vlog("Driver produced command: cc1 {0}", printArgv(CC1Args));
     std::vector<Diag> CompilerInvocationDiags =
         CompilerInvocationDiagConsumer.take();
     if (!Invocation) {
@@ -640,17 +828,30 @@
                               if (CanPublishResults)
                                 Publish();
                             });
+      // Note that this might throw away a stale preamble that might still be
+      // useful, but this is how we communicate a build error.
+      LatestPreamble.emplace();
       // Make sure anyone waiting for the preamble gets notified it could not be
       // built.
-      BuiltFirstPreamble.notify();
+      PreambleCV.notify_all();
       return;
     }
 
     PreamblePeer.update(std::move(Invocation), std::move(Inputs),
                         std::move(CompilerInvocationDiags), WantDiags);
+    std::unique_lock<std::mutex> Lock(Mutex);
+    PreambleCV.wait(Lock, [this] {
+      // Block until we reiceve a preamble request, unless a preamble already
+      // exists, as patching an empty preamble would imply rebuilding it from
+      // scratch.
+      // We block here instead of the consumer to prevent any deadlocks. Since
+      // LatestPreamble is only populated by ASTWorker thread.
+      return LatestPreamble || !PreambleRequests.empty() || Done;
+    });
     return;
   };
-  startTask(TaskName, std::move(Task), WantDiags, TUScheduler::NoInvalidation);
+  startTask(TaskName, std::move(Task), UpdateType{WantDiags, ContentChanged},
+            TUScheduler::NoInvalidation);
 }
 
 void ASTWorker::runWithAST(
@@ -689,13 +890,12 @@
         [&AST, this]() { IdleASTs.put(this, std::move(*AST)); });
     // Run the user-provided action.
     if (!*AST)
-      return Action(llvm::make_error<llvm::StringError>(
-          "invalid AST", llvm::errc::invalid_argument));
+      return Action(error(llvm::errc::invalid_argument, "invalid AST"));
     vlog("ASTWorker running {0} on version {2} of {1}", Name, FileName,
          FileInputs.Version);
     Action(InputsAndAST{FileInputs, **AST});
   };
-  startTask(Name, std::move(Task), /*UpdateType=*/None, Invalidation);
+  startTask(Name, std::move(Task), /*Update=*/None, Invalidation);
 }
 
 void PreambleThread::build(Request Req) {
@@ -709,6 +909,7 @@
     ASTPeer.updatePreamble(std::move(Req.CI), std::move(Req.Inputs),
                            LatestBuild, std::move(Req.CIDiags),
                            std::move(Req.WantDiags));
+    Callbacks.onPreamblePublished(FileName);
   });
 
   if (!LatestBuild || Inputs.ForceRebuild) {
@@ -732,6 +933,8 @@
         Callbacks.onPreambleAST(FileName, Version, Ctx, std::move(PP),
                                 CanonIncludes);
       });
+  if (LatestBuild && isReliable(LatestBuild->CompileCommand))
+    HeaderIncluders.update(FileName, LatestBuild->Includes.allHeaders());
 }
 
 void ASTWorker::updatePreamble(std::unique_ptr<CompilerInvocation> CI,
@@ -747,7 +950,7 @@
     // Update the preamble inside ASTWorker queue to ensure atomicity. As a task
     // running inside ASTWorker assumes internals won't change until it
     // finishes.
-    if (Preamble != LatestPreamble) {
+    if (!LatestPreamble || Preamble != *LatestPreamble) {
       ++PreambleBuildCount;
       // Cached AST is no longer valid.
       IdleASTs.take(this);
@@ -755,12 +958,16 @@
       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);
+      if (LatestPreamble)
+        std::swap(*LatestPreamble, Preamble);
+      else
+        LatestPreamble = std::move(Preamble);
     }
+    // Notify anyone waiting for a preamble.
+    PreambleCV.notify_all();
     // Give up our ownership to old preamble before starting expensive AST
     // build.
     Preamble.reset();
-    BuiltFirstPreamble.notify();
     // We only need to build the AST if diagnostics were requested.
     if (WantDiags == WantDiagnostics::No)
       return;
@@ -775,13 +982,25 @@
   }
   {
     std::lock_guard<std::mutex> Lock(Mutex);
-    PreambleRequests.push({std::move(Task), std::move(TaskName),
-                           steady_clock::now(), Context::current().clone(),
-                           llvm::None, TUScheduler::NoInvalidation, nullptr});
+    PreambleRequests.push_back({std::move(Task), std::move(TaskName),
+                                steady_clock::now(), Context::current().clone(),
+                                llvm::None, llvm::None,
+                                TUScheduler::NoInvalidation, nullptr});
   }
+  PreambleCV.notify_all();
   RequestsCV.notify_all();
 }
 
+void ASTWorker::updateASTSignals(ParsedAST &AST) {
+  auto Signals = std::make_shared<const ASTSignals>(ASTSignals::derive(AST));
+  // Existing readers of ASTSignals will have their copy preserved until the
+  // read is completed. The last reader deletes the old ASTSignals.
+  {
+    std::lock_guard<std::mutex> Lock(Mutex);
+    std::swap(LatestASTSignals, Signals);
+  }
+}
+
 void ASTWorker::generateDiagnostics(
     std::unique_ptr<CompilerInvocation> Invocation, ParseInputs Inputs,
     std::vector<Diag> CIDiags) {
@@ -789,6 +1008,7 @@
   static constexpr trace::Metric ASTAccessForDiag(
       "ast_access_diag", trace::Metric::Counter, "result");
   assert(Invocation);
+  assert(LatestPreamble);
   // No need to rebuild the AST if we won't send the diagnostics.
   {
     std::lock_guard<std::mutex> Lock(PublishMu);
@@ -821,7 +1041,7 @@
   if (!AST || !InputsAreLatest) {
     auto RebuildStartTime = DebouncePolicy::clock::now();
     llvm::Optional<ParsedAST> NewAST = ParsedAST::build(
-        FileName, Inputs, std::move(Invocation), CIDiags, LatestPreamble);
+        FileName, Inputs, std::move(Invocation), CIDiags, *LatestPreamble);
     auto RebuildDuration = DebouncePolicy::clock::now() - RebuildStartTime;
     ++ASTBuildCount;
     // Try to record the AST-build time, to inform future update debouncing.
@@ -859,6 +1079,7 @@
   if (*AST) {
     trace::Span Span("Running main AST callback");
     Callbacks.onMainAST(FileName, **AST, RunPublish);
+    updateASTSignals(**AST);
   } else {
     // Failed to build the AST, at least report diagnostics from the
     // command line if there were any.
@@ -876,13 +1097,18 @@
   }
 }
 
-std::shared_ptr<const PreambleData>
-ASTWorker::getPossiblyStalePreamble() const {
+std::shared_ptr<const PreambleData> ASTWorker::getPossiblyStalePreamble(
+    std::shared_ptr<const ASTSignals> *ASTSignals) const {
   std::lock_guard<std::mutex> Lock(Mutex);
-  return LatestPreamble;
+  if (ASTSignals)
+    *ASTSignals = LatestASTSignals;
+  return LatestPreamble ? *LatestPreamble : nullptr;
 }
 
-void ASTWorker::waitForFirstPreamble() const { BuiltFirstPreamble.wait(); }
+void ASTWorker::waitForFirstPreamble() const {
+  std::unique_lock<std::mutex> Lock(Mutex);
+  PreambleCV.wait(Lock, [this] { return LatestPreamble.hasValue() || Done; });
+}
 
 tooling::CompileCommand ASTWorker::getCurrentCompileCommand() const {
   std::unique_lock<std::mutex> Lock(Mutex);
@@ -896,9 +1122,9 @@
   // Note that we don't report the size of ASTs currently used for processing
   // the in-flight requests. We used this information for debugging purposes
   // only, so this should be fine.
-  Result.UsedBytes = IdleASTs.getUsedBytes(this);
+  Result.UsedBytesAST = IdleASTs.getUsedBytes(this);
   if (auto Preamble = getPossiblyStalePreamble())
-    Result.UsedBytes += Preamble->Preamble.getSize();
+    Result.UsedBytesPreamble = Preamble->Preamble.getSize();
   return Result;
 }
 
@@ -914,20 +1140,21 @@
     assert(!Done && "stop() called twice");
     Done = true;
   }
+  PreamblePeer.stop();
   // We are no longer going to build any preambles, let the waiters know that.
-  BuiltFirstPreamble.notify();
-  PreamblePeer.stop();
+  PreambleCV.notify_all();
   Status.stop();
   RequestsCV.notify_all();
 }
 
 void ASTWorker::startTask(llvm::StringRef Name,
                           llvm::unique_function<void()> Task,
-                          llvm::Optional<WantDiagnostics> UpdateType,
+                          llvm::Optional<UpdateType> Update,
                           TUScheduler::ASTActionInvalidation Invalidation) {
   if (RunSync) {
     assert(!Done && "running a task after stop()");
     trace::Span Tracer(Name + ":" + llvm::sys::path::filename(FileName));
+    WithContext WithProvidedContext(ContextProvider(FileName));
     Task();
     return;
   }
@@ -936,11 +1163,11 @@
     std::lock_guard<std::mutex> Lock(Mutex);
     assert(!Done && "running a task after stop()");
     // Cancel any requests invalidated by this request.
-    if (UpdateType) {
+    if (Update && Update->ContentChanged) {
       for (auto &R : llvm::reverse(Requests)) {
         if (R.InvalidationPolicy == TUScheduler::InvalidateOnUpdate)
           R.Invalidate();
-        if (R.UpdateType)
+        if (R.Update && R.Update->ContentChanged)
           break; // Older requests were already invalidated by the older update.
       }
     }
@@ -953,9 +1180,34 @@
       std::tie(Ctx, Invalidate) = cancelableTask(
           /*Reason=*/static_cast<int>(ErrorCode::ContentModified));
     }
+    // Trace the time the request spends in the queue, and the requests that
+    // it's going to wait for.
+    llvm::Optional<Context> QueueCtx;
+    if (trace::enabled()) {
+      // Tracers that follow threads and need strict nesting will see a tiny
+      // instantaneous event "we're enqueueing", and sometime later it runs.
+      WithContext WC(Ctx.clone());
+      trace::Span Tracer("Queued:" + Name);
+      if (Tracer.Args) {
+        if (CurrentRequest)
+          SPAN_ATTACH(Tracer, "CurrentRequest", CurrentRequest->Name);
+        llvm::json::Array PreambleRequestsNames;
+        for (const auto &Req : PreambleRequests)
+          PreambleRequestsNames.push_back(Req.Name);
+        SPAN_ATTACH(Tracer, "PreambleRequestsNames",
+                    std::move(PreambleRequestsNames));
+        llvm::json::Array RequestsNames;
+        for (const auto &Req : Requests)
+          RequestsNames.push_back(Req.Name);
+        SPAN_ATTACH(Tracer, "RequestsNames", std::move(RequestsNames));
+      }
+      // For tracers that follow contexts, keep the trace span's context alive
+      // until we dequeue the request, so they see the full duration.
+      QueueCtx = Context::current().clone();
+    }
     Requests.push_back({std::move(Task), std::string(Name), steady_clock::now(),
-                        std::move(Ctx), UpdateType, Invalidation,
-                        std::move(Invalidate)});
+                        std::move(Ctx), std::move(QueueCtx), Update,
+                        Invalidation, std::move(Invalidate)});
   }
   RequestsCV.notify_all();
 }
@@ -1001,13 +1253,16 @@
       // Requests.front(), so prefer them first to preserve LSP order.
       if (!PreambleRequests.empty()) {
         CurrentRequest = std::move(PreambleRequests.front());
-        PreambleRequests.pop();
+        PreambleRequests.pop_front();
       } else {
         CurrentRequest = std::move(Requests.front());
         Requests.pop_front();
       }
     } // unlock Mutex
 
+    // Inform tracing that the request was dequeued.
+    CurrentRequest->QueueCtx.reset();
+
     // It is safe to perform reads to CurrentRequest without holding the lock as
     // only writer is also this thread.
     {
@@ -1025,6 +1280,7 @@
         Status.ASTActivity.K = ASTAction::RunningAction;
         Status.ASTActivity.Name = CurrentRequest->Name;
       });
+      WithContext WithProvidedContext(ContextProvider(FileName));
       CurrentRequest->Action();
     }
 
@@ -1054,20 +1310,20 @@
   for (auto I = Requests.begin(), E = Requests.end(); I != E; ++I) {
     if (!isCancelled(I->Ctx)) {
       // Cancellations after the first read don't affect current scheduling.
-      if (I->UpdateType == None)
+      if (I->Update == None)
         break;
       continue;
     }
     // Cancelled reads are moved to the front of the queue and run immediately.
-    if (I->UpdateType == None) {
+    if (I->Update == None) {
       Request R = std::move(*I);
       Requests.erase(I);
       Requests.push_front(std::move(R));
       return Deadline::zero();
     }
     // Cancelled updates are downgraded to auto-diagnostics, and may be elided.
-    if (I->UpdateType == WantDiagnostics::Yes)
-      I->UpdateType = WantDiagnostics::Auto;
+    if (I->Update->Diagnostics == WantDiagnostics::Yes)
+      I->Update->Diagnostics = WantDiagnostics::Auto;
   }
 
   while (shouldSkipHeadLocked()) {
@@ -1080,7 +1336,7 @@
   // We debounce "maybe-unused" writes, sleeping in case they become dead.
   // But don't delay reads (including updates where diagnostics are needed).
   for (const auto &R : Requests)
-    if (R.UpdateType == None || R.UpdateType == WantDiagnostics::Yes)
+    if (R.Update == None || R.Update->Diagnostics == WantDiagnostics::Yes)
       return Deadline::zero();
   // Front request needs to be debounced, so determine when we're ready.
   Deadline D(Requests.front().AddTime + UpdateDebounce.compute(RebuildTimes));
@@ -1091,16 +1347,16 @@
 bool ASTWorker::shouldSkipHeadLocked() const {
   assert(!Requests.empty());
   auto Next = Requests.begin();
-  auto UpdateType = Next->UpdateType;
-  if (!UpdateType) // Only skip updates.
+  auto Update = Next->Update;
+  if (!Update) // Only skip updates.
     return false;
   ++Next;
   // An update is live if its AST might still be read.
   // That is, if it's not immediately followed by another update.
-  if (Next == Requests.end() || !Next->UpdateType)
+  if (Next == Requests.end() || !Next->Update)
     return false;
   // The other way an update can be live is if its diagnostics might be used.
-  switch (*UpdateType) {
+  switch (Update->Diagnostics) {
   case WantDiagnostics::Yes:
     return false; // Always used.
   case WantDiagnostics::No:
@@ -1108,8 +1364,7 @@
   case WantDiagnostics::Auto:
     // Used unless followed by an update that generates diagnostics.
     for (; Next != Requests.end(); ++Next)
-      if (Next->UpdateType == WantDiagnostics::Yes ||
-          Next->UpdateType == WantDiagnostics::Auto)
+      if (Next->Update && Next->Update->Diagnostics != WantDiagnostics::No)
         return true; // Prefer later diagnostics.
     return false;
   }
@@ -1117,10 +1372,24 @@
 }
 
 bool ASTWorker::blockUntilIdle(Deadline Timeout) const {
-  std::unique_lock<std::mutex> Lock(Mutex);
-  return wait(Lock, RequestsCV, Timeout, [&] {
-    return PreambleRequests.empty() && Requests.empty() && !CurrentRequest;
-  });
+  auto WaitUntilASTWorkerIsIdle = [&] {
+    std::unique_lock<std::mutex> Lock(Mutex);
+    return wait(Lock, RequestsCV, Timeout, [&] {
+      return PreambleRequests.empty() && Requests.empty() && !CurrentRequest;
+    });
+  };
+  // Make sure ASTWorker has processed all requests, which might issue new
+  // updates to PreamblePeer.
+  WaitUntilASTWorkerIsIdle();
+  // Now that ASTWorker processed all requests, ensure PreamblePeer has served
+  // all update requests. This might create new PreambleRequests for the
+  // ASTWorker.
+  PreamblePeer.blockUntilIdle(Timeout);
+  assert(Requests.empty() &&
+         "No new normal tasks can be scheduled concurrently with "
+         "blockUntilIdle(): ASTWorker isn't threadsafe");
+  // Finally make sure ASTWorker has processed all of the preamble updates.
+  return WaitUntilASTWorkerIsIdle();
 }
 
 // Render a TUAction to a user-facing string representation.
@@ -1153,7 +1422,7 @@
   }
   if (Result.empty())
     return "idle";
-  return llvm::join(Result, ",");
+  return llvm::join(Result, ", ");
 }
 
 } // namespace
@@ -1178,13 +1447,19 @@
 TUScheduler::TUScheduler(const GlobalCompilationDatabase &CDB,
                          const Options &Opts,
                          std::unique_ptr<ParsingCallbacks> Callbacks)
-    : CDB(CDB), StorePreamblesInMemory(Opts.StorePreamblesInMemory),
+    : CDB(CDB), Opts(Opts),
       Callbacks(Callbacks ? move(Callbacks)
                           : std::make_unique<ParsingCallbacks>()),
-      Barrier(Opts.AsyncThreadsCount),
+      Barrier(Opts.AsyncThreadsCount), QuickRunBarrier(Opts.AsyncThreadsCount),
       IdleASTs(
           std::make_unique<ASTCache>(Opts.RetentionPolicy.MaxRetainedASTs)),
-      UpdateDebounce(Opts.UpdateDebounce) {
+      HeaderIncluders(std::make_unique<HeaderIncluderCache>()) {
+  // Avoid null checks everywhere.
+  if (!Opts.ContextProvider) {
+    this->Opts.ContextProvider = [](llvm::StringRef) {
+      return Context::current().clone();
+    };
+  }
   if (0 < Opts.AsyncThreadsCount) {
     PreambleTasks.emplace();
     WorkerThreads.emplace();
@@ -1216,18 +1491,25 @@
                          WantDiagnostics WantDiags) {
   std::unique_ptr<FileData> &FD = Files[File];
   bool NewFile = FD == nullptr;
+  bool ContentChanged = false;
   if (!FD) {
     // Create a new worker to process the AST-related tasks.
-    ASTWorkerHandle Worker = ASTWorker::create(
-        File, CDB, *IdleASTs,
-        WorkerThreads ? WorkerThreads.getPointer() : nullptr, Barrier,
-        UpdateDebounce, StorePreamblesInMemory, *Callbacks);
+    ASTWorkerHandle Worker =
+        ASTWorker::create(File, CDB, *IdleASTs, *HeaderIncluders,
+                          WorkerThreads ? WorkerThreads.getPointer() : nullptr,
+                          Barrier, Opts, *Callbacks);
     FD = std::unique_ptr<FileData>(
         new FileData{Inputs.Contents, std::move(Worker)});
-  } else {
+    ContentChanged = true;
+  } else if (FD->Contents != Inputs.Contents) {
+    ContentChanged = true;
     FD->Contents = Inputs.Contents;
   }
-  FD->Worker->update(std::move(Inputs), WantDiags);
+  FD->Worker->update(std::move(Inputs), WantDiags, ContentChanged);
+  // There might be synthetic update requests, don't change the LastActiveFile
+  // in such cases.
+  if (ContentChanged)
+    LastActiveFile = File.str();
   return NewFile;
 }
 
@@ -1236,23 +1518,41 @@
   if (!Removed)
     elog("Trying to remove file from TUScheduler that is not tracked: {0}",
          File);
+  // We don't call HeaderIncluders.remove(File) here.
+  // If we did, we'd avoid potentially stale header/mainfile associations.
+  // However, it would mean that closing a mainfile could invalidate the
+  // preamble of several open headers.
 }
 
-llvm::StringMap<std::string> TUScheduler::getAllFileContents() const {
-  llvm::StringMap<std::string> Results;
-  for (auto &It : Files)
-    Results.try_emplace(It.getKey(), It.getValue()->Contents);
-  return Results;
+void TUScheduler::run(llvm::StringRef Name, llvm::StringRef Path,
+                      llvm::unique_function<void()> Action) {
+  runWithSemaphore(Name, Path, std::move(Action), Barrier);
+}
+
+void TUScheduler::runQuick(llvm::StringRef Name, llvm::StringRef Path,
+                           llvm::unique_function<void()> Action) {
+  // Use QuickRunBarrier to serialize quick tasks: we are ignoring
+  // the parallelism level set by the user, don't abuse it
+  runWithSemaphore(Name, Path, std::move(Action), QuickRunBarrier);
 }
 
-void TUScheduler::run(llvm::StringRef Name,
-                      llvm::unique_function<void()> Action) {
-  if (!PreambleTasks)
+void TUScheduler::runWithSemaphore(llvm::StringRef Name, llvm::StringRef Path,
+                                   llvm::unique_function<void()> Action,
+                                   Semaphore &Sem) {
+  if (Path.empty())
+    Path = LastActiveFile;
+  else
+    LastActiveFile = Path.str();
+  if (!PreambleTasks) {
+    WithContext WithProvidedContext(Opts.ContextProvider(Path));
     return Action();
-  PreambleTasks->runAsync(Name, [this, Ctx = Context::current().clone(),
+  }
+  PreambleTasks->runAsync(Name, [this, &Sem, Ctx = Context::current().clone(),
+                                 Path(Path.str()),
                                  Action = std::move(Action)]() mutable {
-    std::lock_guard<Semaphore> BarrierLock(Barrier);
+    std::lock_guard<Semaphore> BarrierLock(Sem);
     WithContext WC(std::move(Ctx));
+    WithContext WithProvidedContext(Opts.ContextProvider(Path));
     Action();
   });
 }
@@ -1267,6 +1567,7 @@
         "trying to get AST for non-added document", ErrorCode::InvalidParams));
     return;
   }
+  LastActiveFile = File.str();
 
   It->second->Worker->runWithAST(Name, std::move(Action), Invalidation);
 }
@@ -1281,40 +1582,45 @@
         ErrorCode::InvalidParams));
     return;
   }
+  LastActiveFile = File.str();
 
   if (!PreambleTasks) {
     trace::Span Tracer(Name);
     SPAN_ATTACH(Tracer, "file", File);
+    std::shared_ptr<const ASTSignals> Signals;
     std::shared_ptr<const PreambleData> Preamble =
-        It->second->Worker->getPossiblyStalePreamble();
+        It->second->Worker->getPossiblyStalePreamble(&Signals);
+    WithContext WithProvidedContext(Opts.ContextProvider(File));
     Action(InputsAndPreamble{It->second->Contents,
                              It->second->Worker->getCurrentCompileCommand(),
-                             Preamble.get()});
+                             Preamble.get(), Signals.get()});
     return;
   }
 
   std::shared_ptr<const ASTWorker> Worker = It->second->Worker.lock();
-  auto Task =
-      [Worker, Consistency, Name = Name.str(), File = File.str(),
-       Contents = It->second->Contents,
-       Command = Worker->getCurrentCompileCommand(),
-       Ctx = Context::current().derive(kFileBeingProcessed, std::string(File)),
-       Action = std::move(Action), this]() mutable {
-        std::shared_ptr<const PreambleData> Preamble;
-        if (Consistency == PreambleConsistency::Stale) {
-          // Wait until the preamble is built for the first time, if preamble
-          // is required. This avoids extra work of processing the preamble
-          // headers in parallel multiple times.
-          Worker->waitForFirstPreamble();
-        }
-        Preamble = Worker->getPossiblyStalePreamble();
+  auto Task = [Worker, Consistency, Name = Name.str(), File = File.str(),
+               Contents = It->second->Contents,
+               Command = Worker->getCurrentCompileCommand(),
+               Ctx = Context::current().derive(kFileBeingProcessed,
+                                               std::string(File)),
+               Action = std::move(Action), this]() mutable {
+    std::shared_ptr<const PreambleData> Preamble;
+    if (Consistency == PreambleConsistency::Stale) {
+      // Wait until the preamble is built for the first time, if preamble
+      // is required. This avoids extra work of processing the preamble
+      // headers in parallel multiple times.
+      Worker->waitForFirstPreamble();
+    }
+    std::shared_ptr<const ASTSignals> Signals;
+    Preamble = Worker->getPossiblyStalePreamble(&Signals);
 
-        std::lock_guard<Semaphore> BarrierLock(Barrier);
-        WithContext Guard(std::move(Ctx));
-        trace::Span Tracer(Name);
-        SPAN_ATTACH(Tracer, "file", File);
-        Action(InputsAndPreamble{Contents, Command, Preamble.get()});
-      };
+    std::lock_guard<Semaphore> BarrierLock(Barrier);
+    WithContext Guard(std::move(Ctx));
+    trace::Span Tracer(Name);
+    SPAN_ATTACH(Tracer, "file", File);
+    WithContext WithProvidedContext(Opts.ContextProvider(File));
+    Action(InputsAndPreamble{Contents, Command, Preamble.get(), Signals.get()});
+  };
 
   PreambleTasks->runAsync("task:" + llvm::sys::path::filename(File),
                           std::move(Task));
@@ -1366,5 +1672,15 @@
   return P;
 }
 
+void TUScheduler::profile(MemoryTree &MT) const {
+  for (const auto &Elem : fileStats()) {
+    MT.detail(Elem.first())
+        .child("preamble")
+        .addUsage(Opts.StorePreamblesInMemory ? Elem.second.UsedBytesPreamble
+                                              : 0);
+    MT.detail(Elem.first()).child("ast").addUsage(Elem.second.UsedBytesAST);
+    MT.child("header_includer_cache").addUsage(HeaderIncluders->getUsedBytes());
+  }
+}
 } // namespace clangd
 } // namespace clang