[PATCH] D126815: [llvm] [Support] [Debuginfod] waitQueueSize for ThreadPool

[Noah Shutty via Phabricator via llvm-commits]

noajshu created this revision.
noajshu added reviewers: mysterymath, phosek, MaskRay.
Herald added subscribers: StephenFan, hiraditya.
Herald added a project: All.
noajshu requested review of this revision.
Herald added a project: LLVM.
Herald added a subscriber: llvm-commits.

The Debuginfod server (D114845 <https://reviews.llvm.org/D114845> + D114846 <https://reviews.llvm.org/D114846>) scans the filesystem for valid binaries, using many threads via `ThreadPool`. Since the filesystem can hypothetically have very many files, simply submitting all jobs to the `ThreadPool` without waiting for any to finish processing can cause the queue to use unbounded memory.

This diff adds a small `waitQueueSize(size_t Size=0)`  to the `ThreadPool`, which blocks until the queue size is at most `Size`. This allows to keep the total queue size below a constant by waiting whenever it has grown too big.


https://reviews.llvm.org/D126815

Files:
  llvm/include/llvm/Support/ThreadPool.h
  llvm/lib/Support/ThreadPool.cpp


Index: llvm/lib/Support/ThreadPool.cpp
===================================================================
--- llvm/lib/Support/ThreadPool.cpp
+++ llvm/lib/Support/ThreadPool.cpp
@@ -54,6 +54,9 @@
           ++ActiveThreads;
           Task = std::move(Tasks.front());
           Tasks.pop();
+
+          // Notify the condition variable that the queue size has decreased.
+          QueueSizeDecreaseCondition.notify_one();
         }
         // Run the task we just grabbed
         Task();
@@ -80,6 +83,13 @@
   CompletionCondition.wait(LockGuard, [&] { return workCompletedUnlocked(); });
 }
 
+void ThreadPool::waitQueueSize(size_t Size) {
+  // Wait for the queue to have at most Size elements
+  std::unique_lock<std::mutex> LockGuard(QueueLock);
+  QueueSizeDecreaseCondition.wait(LockGuard,
+                                  [&] { return Tasks.size() <= Size; });
+}
+
 bool ThreadPool::isWorkerThread() const {
   std::unique_lock<std::mutex> LockGuard(ThreadsLock);
   llvm::thread::id CurrentThreadId = llvm::this_thread::get_id();
Index: llvm/include/llvm/Support/ThreadPool.h
===================================================================
--- llvm/include/llvm/Support/ThreadPool.h
+++ llvm/include/llvm/Support/ThreadPool.h
@@ -64,6 +64,9 @@
   /// It is an error to try to add new tasks while blocking on this call.
   void wait();
 
+  /// Blocking wait for the queue to have size at most Size.
+  void waitQueueSize(size_t Size = 0);
+
   // TODO: misleading legacy name warning!
   // Returns the maximum number of worker threads in the pool, not the current
   // number of threads!
@@ -153,6 +156,9 @@
   std::mutex QueueLock;
   std::condition_variable QueueCondition;
 
+  /// Signaling for queue size decreases
+  std::condition_variable QueueSizeDecreaseCondition;
+
   /// Signaling for job completion
   std::condition_variable CompletionCondition;
 


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D126815.433509.patch
Type: text/x-patch
Size: 1886 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20220601/bdfac7a6/attachment.bin>