3 files changed, 50 insertions, 2 deletions

diff --git a/clang-tools-extra/clangd/index/Background.cpp b/clang-tools-extra/clangd/index/Background.cpp
index c5200e861b9..458e6fc355f 100644
--- a/clang-tools-extra/clangd/index/Background.cpp
+++ b/clang-tools-extra/clangd/index/Background.cpp
@@ -27,6 +27,7 @@
 #include "index/SymbolCollector.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
+#include "clang/Driver/Types.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
@@ -171,6 +172,11 @@ BackgroundQueue::Task BackgroundIndex::changedFilesTask(
   return T;
 }
 
+static llvm::StringRef filenameWithoutExtension(llvm::StringRef Path) {
+  Path = llvm::sys::path::filename(Path);
+  return Path.drop_back(llvm::sys::path::extension(Path).size());
+}
+
 BackgroundQueue::Task
 BackgroundIndex::indexFileTask(tooling::CompileCommand Cmd,
                                BackgroundIndexStorage *Storage) {
@@ -182,9 +188,19 @@ BackgroundIndex::indexFileTask(tooling::CompileCommand Cmd,
       elog("Indexing {0} failed: {1}", FileName, std::move(Error));
   });
   T.QueuePri = IndexFile;
+  T.Tag = filenameWithoutExtension(Cmd.Filename);
   return T;
 }
 
+void BackgroundIndex::boostRelated(llvm::StringRef Path) {
+  namespace types = clang::driver::types;
+  auto Type =
+      types::lookupTypeForExtension(llvm::sys::path::extension(Path).substr(1));
+  // is this a header?
+  if (Type != types::TY_INVALID && types::onlyPrecompileType(Type))
+    Queue.boost(filenameWithoutExtension(Path), IndexBoostedFile);
+}
+
 /// Given index results from a TU, only update symbols coming from files that
 /// are different or missing from than \p ShardVersionsSnapshot. Also stores new
 /// index information on IndexStorage.
diff --git a/clang-tools-extra/clangd/index/Background.h b/clang-tools-extra/clangd/index/Background.h
index 99785f31a12..39fe8daeba8 100644
--- a/clang-tools-extra/clangd/index/Background.h
+++ b/clang-tools-extra/clangd/index/Background.h
@@ -69,8 +69,8 @@ public:
 
     std::function<void()> Run;
     llvm::ThreadPriority ThreadPri = llvm::ThreadPriority::Background;
-    // Higher-priority tasks will run first.
-    unsigned QueuePri = 0;
+    unsigned QueuePri = 0; // Higher-priority tasks will run first.
+    std::string Tag;       // Allows priority to be boosted later.
 
     bool operator<(const Task &O) const { return QueuePri < O.QueuePri; }
   };
@@ -78,6 +78,10 @@ public:
   // Add tasks to the queue.
   void push(Task);
   void append(std::vector<Task>);
+  // Boost priority of current and new tasks with matching Tag, if they are
+  // lower priority.
+  // Reducing the boost of a tag affects future tasks but not current ones.
+  void boost(llvm::StringRef Tag, unsigned NewPriority);
 
   // Process items on the queue until the queue is stopped.
   // If the queue becomes empty, OnIdle will be called (on one worker).
@@ -98,6 +102,7 @@ private:
   std::condition_variable CV;
   bool ShouldStop = false;
   std::vector<Task> Queue; // max-heap
+  llvm::StringMap<unsigned> Boosts;
 };
 
 // Builds an in-memory index by by running the static indexer action over
@@ -123,6 +128,10 @@ public:
     Queue.push(changedFilesTask(ChangedFiles));
   }
 
+  /// Boosts priority of indexing related to Path.
+  /// Typically used to index TUs when headers are opened.
+  void boostRelated(llvm::StringRef Path);
+
   // Cause background threads to stop after ther current task, any remaining
   // tasks will be discarded.
   void stop() {
@@ -187,6 +196,7 @@ private:
   // from lowest to highest priority
   enum QueuePriority {
     IndexFile,
+    IndexBoostedFile,
     LoadShards,
   };
   BackgroundQueue Queue;
diff --git a/clang-tools-extra/clangd/index/BackgroundQueue.cpp b/clang-tools-extra/clangd/index/BackgroundQueue.cpp
index 487dcb8f146..914097225b9 100644
--- a/clang-tools-extra/clangd/index/BackgroundQueue.cpp
+++ b/clang-tools-extra/clangd/index/BackgroundQueue.cpp
@@ -67,6 +67,7 @@ void BackgroundQueue::stop() {
 void BackgroundQueue::push(Task T) {
   {
     std::lock_guard<std::mutex> Lock(Mu);
+    T.QueuePri = std::max(T.QueuePri, Boosts.lookup(T.Tag));
     Queue.push_back(std::move(T));
     std::push_heap(Queue.begin(), Queue.end());
   }
@@ -76,12 +77,33 @@ void BackgroundQueue::push(Task T) {
 void BackgroundQueue::append(std::vector<Task> Tasks) {
   {
     std::lock_guard<std::mutex> Lock(Mu);
+    for (Task &T : Tasks)
+      T.QueuePri = std::max(T.QueuePri, Boosts.lookup(T.Tag));
     std::move(Tasks.begin(), Tasks.end(), std::back_inserter(Queue));
     std::make_heap(Queue.begin(), Queue.end());
   }
   CV.notify_all();
 }
 
+void BackgroundQueue::boost(llvm::StringRef Tag, unsigned NewPriority) {
+  std::lock_guard<std::mutex> Lock(Mu);
+  unsigned &Boost = Boosts[Tag];
+  bool Increase = NewPriority > Boost;
+  Boost = NewPriority;
+  if (!Increase)
+    return; // existing tasks unaffected
+
+  unsigned Changes = 0;
+  for (Task &T : Queue)
+    if (Tag == T.Tag && NewPriority > T.QueuePri) {
+      T.QueuePri = NewPriority;
+      ++Changes;
+    }
+  if (Changes)
+    std::make_heap(Queue.begin(), Queue.end());
+  // No need to signal, only rearranged items in the queue.
+}
+
 bool BackgroundQueue::blockUntilIdleForTest(
     llvm::Optional<double> TimeoutSeconds) {
   std::unique_lock<std::mutex> Lock(Mu);