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

[sammccall]

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.

[sammccall]

Add test for boosting in queue.

[kadircet]

clang-tools-extra/clangd/TUScheduler.h

145 ↗

(On Diff #209273)

Is the Adds \p File to a list of tracked files if ... part got deleted intentionally?

clang-tools-extra/clangd/index/Background.h

106 ↗

(On Diff #209273)

maybe store a QueuePriority instead of unsigned? Also it seems like the enum belongs to BackgroundQueue, rather then BackgroundIndex

[sammccall]

clang-tools-extra/clangd/TUScheduler.h

145 ↗

(On Diff #209273)

It's now covered by the last line.

clang-tools-extra/clangd/index/Background.h

106 ↗

(On Diff #209273)

This is a layering thing: the queue knows about tasks and tags, but doesn't know about TUs, filenames, or active files. I don't plan on actually splitting this header out, but it helps me reason about the code. Do you think it's too abstract?

[kadircet]

clang-tools-extra/clangd/index/Background.h

106 ↗

(On Diff #209273)

Nah that's OK.

Queue has become more priority aware after this patch, previously it only cared about the priority within the task struct. Therefore it become confusing to see the enum in BackgroundIndex itself.

clang-tools-extra/trunk/clangd/ClangdServer.cpp

void ClangdServer::addDocument(PathRef File, llvm::StringRef Contents,
   Opts.SuggestMissingIncludes = SuggestMissingIncludes;
 
   // Compile command is set asynchronously during update, as it can be slow.
   ParseInputs Inputs;
   Inputs.FS = FS;
   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); }
 
 llvm::StringRef ClangdServer::getDocument(PathRef File) const {
   return WorkScheduler.getContents(File);
 }
 

clang-tools-extra/trunk/clangd/TUScheduler.h

public:
   /// The order of results is unspecified.
   std::vector<std::pair<Path, std::size_t>> getUsedBytesPerFile() const;
 
   /// Returns a list of files with ASTs currently stored in memory. This method
   /// is not very reliable and is only used for test. E.g., the results will not
   /// 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
   /// if requested with WantDiags::Auto or WantDiags::Yes.
   void remove(PathRef File);
 
   /// Returns the current contents of the buffer for File, per last update().
   /// The returned StringRef may be invalidated by any write to TUScheduler.

clang-tools-extra/trunk/clangd/TUScheduler.cpp

for (auto &File : Files)
     if (!File.getValue()->Worker->blockUntilIdle(D))
       return false;
   if (PreambleTasks)
     if (!PreambleTasks->wait(D))
       return false;
   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(
         File, CDB, *IdleASTs,
         WorkerThreads ? WorkerThreads.getPointer() : nullptr, Barrier,
         UpdateDebounce, StorePreamblesInMemory, *Callbacks);
     FD = std::unique_ptr<FileData>(
         new FileData{Inputs.Contents, std::move(Worker)});
   } else {
     FD->Contents = Inputs.Contents;
   }
   FD->Worker->update(std::move(Inputs), WantDiags);
+  return NewFile;
 }
 
 void TUScheduler::remove(PathRef File) {
   bool Removed = Files.erase(File);
   if (!Removed)
     elog("Trying to remove file from TUScheduler that is not tracked: {0}",
          File);
 }

clang-tools-extra/trunk/clangd/index/Background.h

 class BackgroundQueue {
 public:
   /// A work item on the thread pool's queue.
   struct Task {
     explicit Task(std::function<void()> Run) : Run(std::move(Run)) {}
 
     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; }
   };
 
   // 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).
   void work(std::function<void()> OnIdle = nullptr);
 
   // Stop processing new tasks, allowing all work() calls to return soon.
   void stop();
 
   // Disables thread priority lowering to ensure progress on loaded systems.
   // Only affects tasks that run after the call.
   static void preventThreadStarvationInTests();
   LLVM_NODISCARD bool
   blockUntilIdleForTest(llvm::Optional<double> TimeoutSeconds);
 
 private:
   std::mutex Mu;
   unsigned NumActiveTasks = 0; // Only idle when queue is empty *and* no tasks.
   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
 // all commands in a compilation database. Indexing happens in the background.
 // FIXME: it should also persist its state on disk for fast start.
 // FIXME: it should watch for changes to files on disk.
 class BackgroundIndex : public SwapIndex {
 public:
public:
 
   // Enqueue translation units for indexing.
   // The indexing happens in a background thread, so the symbols will be
   // available sometime later.
   void enqueue(const std::vector<std::string> &ChangedFiles) {
     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() {
     Rebuilder.shutdown();
     Queue.stop();
   }
 
   // Wait until the queue is empty, to allow deterministic testing.
private:
   BackgroundQueue::Task
   changedFilesTask(const std::vector<std::string> &ChangedFiles);
   BackgroundQueue::Task indexFileTask(tooling::CompileCommand Cmd,
                                       BackgroundIndexStorage *Storage);
 
   // from lowest to highest priority
   enum QueuePriority {
     IndexFile,
+    IndexBoostedFile,
     LoadShards,
   };
   BackgroundQueue Queue;
   AsyncTaskRunner ThreadPool;
   GlobalCompilationDatabase::CommandChanged::Subscription CommandsChanged;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clang-tools-extra/trunk/clangd/index/Background.cpp

 #include "index/IndexAction.h"
 #include "index/MemIndex.h"
 #include "index/Ref.h"
 #include "index/Relation.h"
 #include "index/Serialization.h"
 #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"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/Path.h"
BackgroundQueue::Task T([this, ChangedFiles] {
     Queue.append(std::move(Tasks));
   });
 
   T.QueuePri = LoadShards;
   T.ThreadPri = llvm::ThreadPriority::Default;
   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) {
   BackgroundQueue::Task T([this, Storage, Cmd] {
     // We can't use llvm::StringRef here since we are going to
     // move from Cmd during the call below.
     const std::string FileName = Cmd.Filename;
     if (auto Error = index(std::move(Cmd), Storage))
       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.
 void BackgroundIndex::update(
     llvm::StringRef MainFile, IndexFileIn Index,
     const llvm::StringMap<ShardVersion> &ShardVersionsSnapshot,
     BackgroundIndexStorage *IndexStorage, bool HadErrors) {
   // Partition symbols/references into files.

clang-tools-extra/trunk/clangd/index/BackgroundQueue.cpp

void BackgroundQueue::stop() {
     ShouldStop = true;
   }
   CV.notify_all();
 }
 
 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());
   }
   CV.notify_all();
 }
 
 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);
   return wait(Lock, CV, timeoutSeconds(TimeoutSeconds),
               [&] { return Queue.empty() && NumActiveTasks == 0; });
 }
 
 } // namespace clangd
 } // namespace clang

clang-tools-extra/trunk/clangd/unittests/BackgroundIndexTests.cpp

TEST(BackgroundQueueTest, Priority) {
   Q.append(std::vector<BackgroundQueue::Task>(2, Hi));
 
   // After finishing, check the tasks that ran.
   ThreadPool.wait();
   EXPECT_GE(HiRan, 10u);
   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