1 files changed, 52 insertions, 34 deletions

diff --git a/clang-tools-extra/clangd/index/Background.cpp b/clang-tools-extra/clangd/index/Background.cpp
index 897c513830f..c168038f775 100644
--- a/clang-tools-extra/clangd/index/Background.cpp
+++ b/clang-tools-extra/clangd/index/Background.cpp
@@ -22,6 +22,7 @@
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/SHA1.h"
@@ -135,14 +136,8 @@ BackgroundIndex::BackgroundIndex(
           })) {
   assert(ThreadPoolSize > 0 && "Thread pool size can't be zero.");
   assert(this->IndexStorageFactory && "Storage factory can not be null!");
-  while (ThreadPoolSize--) {
+  while (ThreadPoolSize--)
     ThreadPool.emplace_back([this] { run(); });
-    // Set priority to low, since background indexing is a long running task we
-    // do not want to eat up cpu when there are any other high priority threads.
-    // FIXME: In the future we might want a more general way of handling this to
-    // support tasks with various priorities.
-    setThreadPriority(ThreadPool.back(), ThreadPriority::Low);
-  }
 }
 
 BackgroundIndex::~BackgroundIndex() {
@@ -163,6 +158,7 @@ void BackgroundIndex::run() {
   WithContext Background(BackgroundContext.clone());
   while (true) {
     Optional<Task> Task;
+    ThreadPriority Priority;
     {
       std::unique_lock<std::mutex> Lock(QueueMu);
       QueueCV.wait(Lock, [&] { return ShouldStop || !Queue.empty(); });
@@ -172,10 +168,16 @@ void BackgroundIndex::run() {
         return;
       }
       ++NumActiveTasks;
-      Task = std::move(Queue.front());
+      std::tie(Task, Priority) = std::move(Queue.front());
       Queue.pop_front();
     }
+
+    if (Priority != ThreadPriority::Normal)
+      setCurrentThreadPriority(Priority);
     (*Task)();
+    if (Priority != ThreadPriority::Normal)
+      setCurrentThreadPriority(ThreadPriority::Normal);
+
     {
       std::unique_lock<std::mutex> Lock(QueueMu);
       assert(NumActiveTasks > 0 && "before decrementing");
@@ -193,44 +195,60 @@ bool BackgroundIndex::blockUntilIdleForTest(
 }
 
 void BackgroundIndex::enqueue(const std::vector<std::string> &ChangedFiles) {
-  enqueueTask([this, ChangedFiles] {
-    trace::Span Tracer("BackgroundIndexEnqueue");
-    // We're doing this asynchronously, because we'll read shards here too.
-    // FIXME: read shards here too.
-
-    log("Enqueueing {0} commands for indexing", ChangedFiles.size());
-    SPAN_ATTACH(Tracer, "files", int64_t(ChangedFiles.size()));
-
-    // We shuffle the files because processing them in a random order should
-    // quickly give us good coverage of headers in the project.
-    std::vector<unsigned> Permutation(ChangedFiles.size());
-    std::iota(Permutation.begin(), Permutation.end(), 0);
-    std::mt19937 Generator(std::random_device{}());
-    std::shuffle(Permutation.begin(), Permutation.end(), Generator);
-
-    for (const unsigned I : Permutation)
-      enqueue(ChangedFiles[I]);
-  });
+  enqueueTask(
+      [this, ChangedFiles] {
+        trace::Span Tracer("BackgroundIndexEnqueue");
+        // We're doing this asynchronously, because we'll read shards here too.
+        // FIXME: read shards here too.
+
+        log("Enqueueing {0} commands for indexing", ChangedFiles.size());
+        SPAN_ATTACH(Tracer, "files", int64_t(ChangedFiles.size()));
+
+        // We shuffle the files because processing them in a random order should
+        // quickly give us good coverage of headers in the project.
+        std::vector<unsigned> Permutation(ChangedFiles.size());
+        std::iota(Permutation.begin(), Permutation.end(), 0);
+        std::mt19937 Generator(std::random_device{}());
+        std::shuffle(Permutation.begin(), Permutation.end(), Generator);
+
+        for (const unsigned I : Permutation)
+          enqueue(ChangedFiles[I]);
+      },
+      ThreadPriority::Normal);
 }
 
 void BackgroundIndex::enqueue(const std::string &File) {
   ProjectInfo Project;
   if (auto Cmd = CDB.getCompileCommand(File, &Project)) {
     auto *Storage = IndexStorageFactory(Project.SourceRoot);
+    // Set priority to low, since background indexing is a long running
+    // task we do not want to eat up cpu when there are any other high
+    // priority threads.
     enqueueTask(Bind(
-        [this, File, Storage](tooling::CompileCommand Cmd) {
-          Cmd.CommandLine.push_back("-resource-dir=" + ResourceDir);
-          if (auto Error = index(std::move(Cmd), Storage))
-            log("Indexing {0} failed: {1}", File, std::move(Error));
-        },
-        std::move(*Cmd)));
+                    [this, File, Storage](tooling::CompileCommand Cmd) {
+                      Cmd.CommandLine.push_back("-resource-dir=" + ResourceDir);
+                      if (auto Error = index(std::move(Cmd), Storage))
+                        log("Indexing {0} failed: {1}", File, std::move(Error));
+                    },
+                    std::move(*Cmd)),
+                ThreadPriority::Low);
   }
 }
 
-void BackgroundIndex::enqueueTask(Task T) {
+void BackgroundIndex::enqueueTask(Task T, ThreadPriority Priority) {
   {
     std::lock_guard<std::mutex> Lock(QueueMu);
-    Queue.push_back(std::move(T));
+    auto I = Queue.end();
+    // We first store the tasks with Normal priority in the front of the queue.
+    // Then we store low priority tasks. Normal priority tasks are pretty rare,
+    // they should not grow beyond single-digit numbers, so it is OK to do
+    // linear search and insert after that.
+    if (Priority == ThreadPriority::Normal) {
+      I = llvm::find_if(Queue, [](const std::pair<Task, ThreadPriority> &Elem) {
+        return Elem.second == ThreadPriority::Low;
+      });
+    }
+    Queue.insert(I, {std::move(T), Priority});
   }
   QueueCV.notify_all();
 }