[clangd] Add priorities to background index queue, extract to separate class

[kadircet]

LGTM, thanks

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

160

maybe move these below, next to Queue?

clang-tools-extra/clangd/CMakeLists.txt

add_clang_library(clangDaemon
   QueryDriverDatabase.cpp
   Threading.cpp
   Trace.cpp
   TUScheduler.cpp
   URI.cpp
   XRefs.cpp
 
   index/Background.cpp
-  index/BackgroundRebuild.cpp
   index/BackgroundIndexStorage.cpp
+  index/BackgroundQueue.cpp
+  index/BackgroundRebuild.cpp
   index/CanonicalIncludes.cpp
   index/FileIndex.cpp
   index/Index.cpp
   index/IndexAction.cpp
   index/MemIndex.cpp
   index/Merge.cpp
   index/Ref.cpp
   index/Relation.cpp

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

 #include "index/Serialization.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/Threading.h"
 #include <atomic>
 #include <condition_variable>
 #include <deque>
 #include <mutex>
+#include <queue>
 #include <string>
 #include <thread>
 #include <vector>
 
 namespace clang {
 namespace clangd {
 
 // Handles storage and retrieval of index shards. Both store and load
public:
   using Factory =
       llvm::unique_function<BackgroundIndexStorage *(llvm::StringRef)>;
 
   // Creates an Index Storage that saves shards into disk. Index storage uses
   // CDBDirectory + ".clangd/index/" as the folder to save shards.
   static Factory createDiskBackedStorageFactory();
 };
 
+// A priority queue of tasks which can be run on (external) worker threads.
+class BackgroundQueue {
+public:
+  /// A work item on the thread pool's queue.
+  struct Task {
+    template <typename Func>
+    explicit Task(Func &&F) : Run(std::forward<Func>(F)){};
+
+    std::function<void()> Run;
+    llvm::ThreadPriority ThreadPri = llvm::ThreadPriority::Background;
+    // Higher-priority tasks will run first.
+    unsigned QueuePri = 0;
+
+    bool operator<(const Task &O) const { return QueuePri < O.QueuePri; }
+  };
+
+  // Add tasks to the queue.
+  void push(Task);
+  void append(std::vector<Task>);
+
+  // 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
+};
+
 // 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:
   /// If BuildIndexPeriodMs is greater than 0, the symbol index will only be
   /// rebuilt periodically (one per \p BuildIndexPeriodMs); otherwise, index is
   /// rebuilt for each indexed file.
   BackgroundIndex(
       Context BackgroundContext, const FileSystemProvider &,
       const GlobalCompilationDatabase &CDB,
       BackgroundIndexStorage::Factory IndexStorageFactory,
       size_t ThreadPoolSize = llvm::heavyweight_hardware_concurrency());
   ~BackgroundIndex(); // Blocks while the current task finishes.
 
   // 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);
+  void enqueue(const std::vector<std::string> &ChangedFiles) {
+    Queue.push(changedFilesTask(ChangedFiles));
+  }
 
   // Cause background threads to stop after ther current task, any remaining
   // tasks will be discarded.
-  void stop();
+  void stop() {
+    Rebuilder.shutdown();
+    Queue.stop();
+  }
 
   // Wait until the queue is empty, to allow deterministic testing.
   LLVM_NODISCARD bool
-  blockUntilIdleForTest(llvm::Optional<double> TimeoutSeconds = 10);
-
-  // Disables thread priority lowering in background index to make sure it can
-  // progress on loaded systems. Only affects tasks that run after the call.
-  static void preventThreadStarvationInTests();
+  blockUntilIdleForTest(llvm::Optional<double> TimeoutSeconds = 10) {
+    return Queue.blockUntilIdleForTest(TimeoutSeconds);
+  }
 
 private:
   /// Represents the state of a single file when indexing was performed.
   struct ShardVersion {
     FileDigest Digest{{0}};
     bool HadErrors = false;
   };
 
   /// Given index results from a TU, only update symbols coming from files with
   /// different digests than \p ShardVersionsSnapshot. Also stores new index
   /// information on IndexStorage.
   void update(llvm::StringRef MainFile, IndexFileIn Index,
               const llvm::StringMap<ShardVersion> &ShardVersionsSnapshot,
               BackgroundIndexStorage *IndexStorage, bool HadErrors);
 
   // configuration
   const FileSystemProvider &FSProvider;
   const GlobalCompilationDatabase &CDB;
   Context BackgroundContext;
 
   // index state
+  BackgroundQueue::Task
+  changedFilesTask(const std::vector<std::string> &ChangedFiles);
+  BackgroundQueue::Task indexingTask(tooling::CompileCommand Cmd,
+                                     BackgroundIndexStorage *Storage);
   llvm::Error index(tooling::CompileCommand,
                     BackgroundIndexStorage *IndexStorage);
-  std::mutex IndexMu;
-  std::condition_variable IndexCV;
 
   FileSymbols IndexedSymbols;
   BackgroundIndexRebuilder Rebuilder;
   llvm::StringMap<ShardVersion> ShardVersions; // Key is absolute file path.
   std::mutex ShardVersionsMu;
 
   BackgroundIndexStorage::Factory IndexStorageFactory;
   struct Source {
     std::string Path;
     bool NeedsReIndexing;
     Source(llvm::StringRef Path, bool NeedsReIndexing)
         : Path(Path), NeedsReIndexing(NeedsReIndexing) {}
   };
   // Loads the shards for a single TU and all of its dependencies. Returns the
   // list of sources and whether they need to be re-indexed.
   std::vector<Source> loadShard(const tooling::CompileCommand &Cmd,
                                 BackgroundIndexStorage *IndexStorage,
                                 llvm::StringSet<> &LoadedShards);
   // Tries to load shards for the ChangedFiles.
   std::vector<std::pair<tooling::CompileCommand, BackgroundIndexStorage *>>
   loadShards(std::vector<std::string> ChangedFiles);
-  void enqueue(tooling::CompileCommand Cmd, BackgroundIndexStorage *Storage);
 
-  // queue management
-  using Task = std::function<void()>;
-  void run(); // Main loop executed by Thread. Runs tasks from Queue.
-  void enqueueTask(Task T, llvm::ThreadPriority Prioirty);
-  void enqueueLocked(tooling::CompileCommand Cmd,
-                     BackgroundIndexStorage *IndexStorage);
-  std::mutex QueueMu;
-  unsigned NumActiveTasks = 0; // Only idle when queue is empty *and* no tasks.
-  std::condition_variable QueueCV;
-  bool ShouldStop = false;
-  std::deque<std::pair<Task, llvm::ThreadPriority>> Queue;
+  // from lowest to highest priority
+  enum QueuePriority {
+    IndexFile,
+    LoadShards,
+  };
+  BackgroundQueue Queue;
   AsyncTaskRunner ThreadPool;
   GlobalCompilationDatabase::CommandChanged::Subscription CommandsChanged;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

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

 //===-- Background.cpp - Build an index in a background thread ------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "index/Background.h"
 #include "ClangdUnit.h"
 #include "Compiler.h"
+#include "Context.h"
 #include "Headers.h"
 #include "Logger.h"
 #include "Path.h"
 #include "SourceCode.h"
 #include "Symbol.h"
 #include "Threading.h"
 #include "Trace.h"
 #include "URI.h"
 #include "index/FileIndex.h"
 #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 "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"
 #include "llvm/Support/Threading.h"
 
+#include <algorithm>
 #include <atomic>
 #include <chrono>
 #include <condition_variable>
 #include <memory>
 #include <mutex>
 #include <numeric>
 #include <queue>
 #include <random>
 #include <string>
 #include <thread>
 
 namespace clang {
 namespace clangd {
 namespace {
 
-static std::atomic<bool> PreventStarvation = {false};
-
 // Resolves URI to file paths with cache.
 class URIToFileCache {
 public:
   URIToFileCache(llvm::StringRef HintPath) : HintPath(HintPath) {}
 
   llvm::StringRef resolve(llvm::StringRef FileURI) {
     auto I = URIToPathCache.try_emplace(FileURI);
     if (I.second) {
: SwapIndex(llvm::make_unique<MemIndex>()), FSProvider(FSProvider),
       IndexStorageFactory(std::move(IndexStorageFactory)),
       CommandsChanged(
           CDB.watch([&](const std::vector<std::string> &ChangedFiles) {
             enqueue(ChangedFiles);
           })) {
   assert(ThreadPoolSize > 0 && "Thread pool size can't be zero.");
   assert(this->IndexStorageFactory && "Storage factory can not be null!");
   for (unsigned I = 0; I < ThreadPoolSize; ++I) {
-    ThreadPool.runAsync("background-worker-" + llvm::Twine(I + 1),
-                        [this] { run(); });
+    ThreadPool.runAsync("background-worker-" + llvm::Twine(I + 1), [this] {
+      WithContext Ctx(this->BackgroundContext.clone());
+      Queue.work([&] { Rebuilder.idle(); });
+    });
   }
 }
 
 BackgroundIndex::~BackgroundIndex() {
   stop();
   ThreadPool.wait();
 }
 
-void BackgroundIndex::stop() {
-  Rebuilder.shutdown();
-  {
-    std::lock_guard<std::mutex> QueueLock(QueueMu);
-    std::lock_guard<std::mutex> IndexLock(IndexMu);
-    ShouldStop = true;
-  }
-  QueueCV.notify_all();
-  IndexCV.notify_all();
-}
-
-void BackgroundIndex::run() {
-  WithContext Background(BackgroundContext.clone());
-  while (true) {
-    llvm::Optional<Task> Task;
-    llvm::ThreadPriority Priority;
-    {
-      std::unique_lock<std::mutex> Lock(QueueMu);
-      QueueCV.wait(Lock, [&] { return ShouldStop || !Queue.empty(); });
-      if (ShouldStop) {
-        Queue.clear();
-        QueueCV.notify_all();
-        return;
-      }
-      ++NumActiveTasks;
-      std::tie(Task, Priority) = std::move(Queue.front());
-      Queue.pop_front();
-    }
-
-    if (Priority != llvm::ThreadPriority::Default && !PreventStarvation.load())
-      llvm::set_thread_priority(Priority);
-    (*Task)();
-    if (Priority != llvm::ThreadPriority::Default)
-      llvm::set_thread_priority(llvm::ThreadPriority::Default);
-
-    {
-      std::unique_lock<std::mutex> Lock(QueueMu);
-      if (NumActiveTasks == 1 && Queue.empty()) {
-        // We just finished the last item, the queue is going idle.
-        Lock.unlock();
-        Rebuilder.idle();
-        Lock.lock();
-      }
-      assert(NumActiveTasks > 0 && "before decrementing");
-      --NumActiveTasks;
-    }
-    QueueCV.notify_all();
-  }
-}
-
-bool BackgroundIndex::blockUntilIdleForTest(
-    llvm::Optional<double> TimeoutSeconds) {
-  std::unique_lock<std::mutex> Lock(QueueMu);
-  return wait(Lock, QueueCV, timeoutSeconds(TimeoutSeconds),
-              [&] { return Queue.empty() && NumActiveTasks == 0; });
-}
-
-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.
-        log("Enqueueing {0} commands for indexing", ChangedFiles.size());
-        SPAN_ATTACH(Tracer, "files", int64_t(ChangedFiles.size()));
+BackgroundQueue::Task BackgroundIndex::changedFilesTask(
+    const std::vector<std::string> &ChangedFiles) {
+  BackgroundQueue::Task T([this, ChangedFiles] {
 
-        auto NeedsReIndexing = loadShards(std::move(ChangedFiles));
-        // Run indexing for files that need to be updated.
-        std::shuffle(NeedsReIndexing.begin(), NeedsReIndexing.end(),
-                     std::mt19937(std::random_device{}()));
-        for (auto &Elem : NeedsReIndexing)
-          enqueue(std::move(Elem.first), Elem.second);
-      },
-      llvm::ThreadPriority::Default);
+    std::vector<BackgroundQueue::Task> Tasks;
+    Tasks.reserve(NeedsReIndexing.size());
+      Tasks.push_back(indexingTask(std::move(Elem.first), Elem.second));
+    Queue.append(std::move(Tasks));
+  });
+
+  T.QueuePri = LoadShards;
+  T.ThreadPri = llvm::ThreadPriority::Default;
+  return T;
 }
 
-void BackgroundIndex::enqueue(tooling::CompileCommand Cmd,
+BackgroundQueue::Task
+BackgroundIndex::indexingTask(tooling::CompileCommand Cmd,
                               BackgroundIndexStorage *Storage) {
-  enqueueTask(Bind(
-                  [this, Storage](tooling::CompileCommand 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));
-                  },
-                  std::move(Cmd)),
-              llvm::ThreadPriority::Background);
-}
-
-void BackgroundIndex::enqueueTask(Task T, llvm::ThreadPriority Priority) {
-  {
-    std::lock_guard<std::mutex> Lock(QueueMu);
-    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 == llvm::ThreadPriority::Default) {
-      I = llvm::find_if(
-          Queue, [](const std::pair<Task, llvm::ThreadPriority> &Elem) {
-            return Elem.second == llvm::ThreadPriority::Background;
-          });
-    }
-    Queue.insert(I, {std::move(T), Priority});
-  }
-  QueueCV.notify_all();
+  BackgroundQueue::Task T([this, Storage, Cmd] {
+      elog("Indexing {0} failed: {1}", FileName, std::move(Error));
+  T.QueuePri = IndexFile;
+  return T;
 }
 
 /// 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,
for (const auto &Dependency : Dependencies) {
         FilesToIndex.insert(Dependency.Path);
       break;
     }
   }
   Rebuilder.doneLoading();
   return NeedsReIndexing;
 }
 
-void BackgroundIndex::preventThreadStarvationInTests() {
-  PreventStarvation.store(true);
-}
-
 } // namespace clangd
 } // namespace clang

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

+//===-- BackgroundQueue.cpp - Task queue for background index -------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "index/Background.h"
+
+namespace clang {
+namespace clangd {
+
+static std::atomic<bool> PreventStarvation = {false};
+
+void BackgroundQueue::preventThreadStarvationInTests() {
+  PreventStarvation.store(true);
+}
+
+void BackgroundQueue::work(std::function<void()> OnIdle) {
+  while (true) {
+    llvm::Optional<Task> Task;
+    {
+      std::unique_lock<std::mutex> Lock(Mu);
+      CV.wait(Lock, [&] { return ShouldStop || !Queue.empty(); });
+      if (ShouldStop) {
+        Queue.clear();
+        CV.notify_all();
+        return;
+      }
+      ++NumActiveTasks;
+      std::pop_heap(Queue.begin(), Queue.end());
+      Task = std::move(Queue.back());
+      Queue.pop_back();
+    }
+
+    if (Task->ThreadPri != llvm::ThreadPriority::Default &&
+        !PreventStarvation.load())
+      llvm::set_thread_priority(Task->ThreadPri);
+    Task->Run();
+    if (Task->ThreadPri != llvm::ThreadPriority::Default)
+      llvm::set_thread_priority(llvm::ThreadPriority::Default);
+
+    {
+      std::unique_lock<std::mutex> Lock(Mu);
+      if (NumActiveTasks == 1 && Queue.empty() && OnIdle) {
+        // We just finished the last item, the queue is going idle.
+        Lock.unlock();
+        OnIdle();
+        Lock.lock();
+      }
+      assert(NumActiveTasks > 0 && "before decrementing");
+      --NumActiveTasks;
+    }
+    CV.notify_all();
+  }
+}
+
+void BackgroundQueue::stop() {
+  {
+    std::lock_guard<std::mutex> QueueLock(Mu);
+    ShouldStop = true;
+  }
+  CV.notify_all();
+}
+
+void BackgroundQueue::push(Task T) {
+  {
+    std::lock_guard<std::mutex> Lock(Mu);
+    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);
+    std::move(Tasks.begin(), Tasks.end(), std::back_inserter(Queue));
+    std::make_heap(Queue.begin(), Queue.end());
+  }
+  CV.notify_all();
+}
+
+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/clangd/tool/ClangdMain.cpp

llvm::cl::ParseCommandLineOptions(
       "\n\thttps://clang.llvm.org/extra/clangd.html"
       "\n\thttps://microsoft.github.io/language-server-protocol/");
   if (Test) {
     Sync = true;
     InputStyle = JSONStreamStyle::Delimited;
     LogLevel = Logger::Verbose;
     PrettyPrint = true;
     // Ensure background index makes progress.
-    BackgroundIndex::preventThreadStarvationInTests();
+    BackgroundQueue::preventThreadStarvationInTests();
   }
   if (Test || EnableTestScheme) {
     static URISchemeRegistry::Add<TestScheme> X(
         "test", "Test scheme for clangd lit tests.");
   }
 
   if (!Sync && WorkerThreadsCount == 0) {
     llvm::errs() << "A number of worker threads cannot be 0. Did you mean to "

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

loadShard(llvm::StringRef ShardIdentifier) const override {
     return llvm::make_unique<IndexFileIn>(std::move(*IndexFile));
   }
 
   mutable llvm::StringSet<> AccessedPaths;
 };
 
 class BackgroundIndexTest : public ::testing::Test {
 protected:
-  BackgroundIndexTest() { BackgroundIndex::preventThreadStarvationInTests(); }
+  BackgroundIndexTest() { BackgroundQueue::preventThreadStarvationInTests(); }
 };
 
 TEST_F(BackgroundIndexTest, NoCrashOnErrorFile) {
   MockFSProvider FS;
   FS.Files[testPath("root/A.cc")] = "error file";
   llvm::StringMap<std::string> Storage;
   size_t CacheHits = 0;
   MemoryShardStorage MSS(Storage, CacheHits);
TEST_F(BackgroundIndexRebuilderTest, LoadingShards) {
   // No rebuilding for indexed files while loading.
   Rebuilder.startLoading();
   for (unsigned I = 0; I < 3 * BackgroundIndexRebuilder::TUsBeforeRebuild; ++I)
     EXPECT_FALSE(checkRebuild([&] { Rebuilder.indexedTU(); }));
   // But they get indexed when we're done, even if no shards were loaded.
   EXPECT_TRUE(checkRebuild([&] { Rebuilder.doneLoading(); }));
 }
 
+TEST(BackgroundQueueTest, Priority) {
+  // Create high and low priority tasks.
+  // Once a bunch of high priority tasks have run, the queue is stopped.
+  // So the low priority tasks should never run.
+  BackgroundQueue Q;
+  std::atomic<unsigned> HiRan(0), LoRan(0);
+  BackgroundQueue::Task Lo([&] { ++LoRan; });
+  BackgroundQueue::Task Hi([&] {
+    if (++HiRan >= 10)
+      Q.stop();
+  });
+  Hi.QueuePri = 100;
+
+  // Enqueuing the low-priority ones first shouldn't make them run first.
+  Q.append(std::vector<BackgroundQueue::Task>(30, Lo));
+  for (unsigned I = 0; I < 30; ++I)
+    Q.push(Hi);
+
+  AsyncTaskRunner ThreadPool;
+  for (unsigned I = 0; I < 5; ++I)
+    ThreadPool.runAsync("worker", [&] { Q.work(); });
+  // We should test enqueue with active workers, but it's hard to avoid races.
+  // Just make sure we don't crash.
+  Q.push(Lo);
+  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);
+}
+
 } // namespace clangd
 } // namespace clang