[clangd] Prioritize indexing of files that share a basename with the open file.

Summary:
In practice, opening Foo.h will still often result in Foo.cpp making the
second index build instead of the first, as the rebuild policy doesn't
know to wait.

Reviewers: kadircet

Subscribers: ilya-biryukov, javed.absar, MaskRay, jkorous, arphaman, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D64575

git-svn-id: https://llvm.org/svn/llvm-project/clang-tools-extra/trunk@365888 91177308-0d34-0410-b5e6-96231b3b80d8

7 files changed, 97 insertions, 8 deletions

diff --git a/clangd/ClangdServer.cpp b/clangd/ClangdServer.cpp
index 142c62c9..10949ef0 100644
--- a/clangd/ClangdServer.cpp
+++ b/clangd/ClangdServer.cpp
@@ -151,7 +151,10 @@ void ClangdServer::addDocument(PathRef File, llvm::StringRef Contents,
   Inputs.Contents = Contents;
   Inputs.Opts = std::move(Opts);
   Inputs.Index = Index;
-  WorkScheduler.update(File, Inputs, WantDiags);
+  bool NewFile = WorkScheduler.update(File, Inputs, WantDiags);
+  // If we loaded Foo.h, we want to make sure Foo.cpp is indexed.
+  if (NewFile && BackgroundIdx)
+    BackgroundIdx->boostRelated(File);
 }
 
 void ClangdServer::removeDocument(PathRef File) { WorkScheduler.remove(File); }
diff --git a/clangd/TUScheduler.cpp b/clangd/TUScheduler.cpp
index 5e761547..a77d2269 100644
--- a/clangd/TUScheduler.cpp
+++ b/clangd/TUScheduler.cpp
@@ -860,9 +860,10 @@ bool TUScheduler::blockUntilIdle(Deadline D) const {
   return true;
 }
 
-void TUScheduler::update(PathRef File, ParseInputs Inputs,
+bool TUScheduler::update(PathRef File, ParseInputs Inputs,
                          WantDiagnostics WantDiags) {
   std::unique_ptr<FileData> &FD = Files[File];
+  bool NewFile = FD == nullptr;
   if (!FD) {
     // Create a new worker to process the AST-related tasks.
     ASTWorkerHandle Worker = ASTWorker::create(
@@ -875,6 +876,7 @@ void TUScheduler::update(PathRef File, ParseInputs Inputs,
     FD->Contents = Inputs.Contents;
   }
   FD->Worker->update(std::move(Inputs), WantDiags);
+  return NewFile;
 }
 
 void TUScheduler::remove(PathRef File) {
diff --git a/clangd/TUScheduler.h b/clangd/TUScheduler.h
index 4943d197..74973d0a 100644
--- a/clangd/TUScheduler.h
+++ b/clangd/TUScheduler.h
@@ -142,13 +142,14 @@ public:
   /// contain files that currently run something over their AST.
   std::vector<Path> getFilesWithCachedAST() const;
 
-  /// Schedule an update for \p File. Adds \p File to a list of tracked files if
-  /// \p File was not part of it before. The compile command in \p Inputs is
-  /// ignored; worker queries CDB to get the actual compile command.
+  /// Schedule an update for \p File.
+  /// The compile command in \p Inputs is ignored; worker queries CDB to get
+  /// the actual compile command.
   /// If diagnostics are requested (Yes), and the context is cancelled
   /// before they are prepared, they may be skipped if eventual-consistency
   /// permits it (i.e. WantDiagnostics is downgraded to Auto).
-  void update(PathRef File, ParseInputs Inputs, WantDiagnostics WD);
+  /// Returns true if the file was not previously tracked.
+  bool update(PathRef File, ParseInputs Inputs, WantDiagnostics WD);
 
   /// Remove \p File from the list of tracked files and schedule removal of its
   /// resources. Pending diagnostics for closed files may not be delivered, even
diff --git a/clangd/index/Background.cpp b/clangd/index/Background.cpp
index c5200e86..458e6fc3 100644
--- a/clangd/index/Background.cpp
+++ b/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/clangd/index/Background.h b/clangd/index/Background.h
index 99785f31..39fe8dae 100644
--- a/clangd/index/Background.h
+++ b/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/clangd/index/BackgroundQueue.cpp b/clangd/index/BackgroundQueue.cpp
index 487dcb8f..91409722 100644
--- a/clangd/index/BackgroundQueue.cpp
+++ b/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);
diff --git a/clangd/unittests/BackgroundIndexTests.cpp b/clangd/unittests/BackgroundIndexTests.cpp
index 55eb0bf8..15d064a9 100644
--- a/clangd/unittests/BackgroundIndexTests.cpp
+++ b/clangd/unittests/BackgroundIndexTests.cpp
@@ -678,5 +678,40 @@ TEST(BackgroundQueueTest, Priority) {
   EXPECT_EQ(LoRan, 0u);
 }
 
+TEST(BackgroundQueueTest, Boost) {
+  std::string Sequence;
+
+  BackgroundQueue::Task A([&] { Sequence.push_back('A'); });
+  A.Tag = "A";
+  A.QueuePri = 1;
+
+  BackgroundQueue::Task B([&] { Sequence.push_back('B'); });
+  B.QueuePri = 2;
+  B.Tag = "B";
+
+  {
+    BackgroundQueue Q;
+    Q.append({A, B});
+    Q.work([&] { Q.stop(); });
+    EXPECT_EQ("BA", Sequence) << "priority order";
+  }
+  Sequence.clear();
+  {
+    BackgroundQueue Q;
+    Q.boost("A", 3);
+    Q.append({A, B});
+    Q.work([&] { Q.stop(); });
+    EXPECT_EQ("AB", Sequence) << "A was boosted before enqueueing";
+  }
+  Sequence.clear();
+  {
+    BackgroundQueue Q;
+    Q.append({A, B});
+    Q.boost("A", 3);
+    Q.work([&] { Q.stop(); });
+    EXPECT_EQ("AB", Sequence) << "A was boosted after enqueueing";
+  }
+}
+
 } // namespace clangd
 } // namespace clang