[llvm] [Support] [Debuginfod] waitQueueSize for ThreadPool

[noajshu]

Summary

The Debuginfod server (D114845 + 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.

[mehdi_amini]

That seems like a very ad-hoc API to me, can you handle this in your client code? Some atomic counter that you manage in your tasks should do the trick just as well?

[mysterymath]

In D126815#3554505, @mehdi_amini wrote:

That seems like a very ad-hoc API to me, can you handle this in your client code? Some atomic counter that you manage in your tasks should do the trick just as well?

Originally the client code did have such counters; I suggested the change to ThreadPool in code review.

If the producer of ThreadPool.async calls runs too much faster than the threads can consume the requests, then the size of the queue is unbounded. With a regular producer/consumer queue, you could check the size before producing, but there wasn't any way to do this in ThreadPool. Maintaining counters in the client is akin to keeping track of the queue size outside of the queue, which seems kludgey.

The suggestion was just for some kind of feedback mechanism from ThreadPool to callers; I don't have any strong opinions about the specifics of the API. Does a cleaner API come to mind?

[mysterymath]

Thinking about this a bit more, one alternative would be to have a bounded version of async that only submit to queue no larger than k, blocking otherwise. Another would be an optional MaxQueueSize that would cause async to block.

[noajshu]

I agree with @mysterymath there is more than one way to support this usage pattern and I also agree with @mehdi_amini that the user code could keep track on its own.
It seems like this ought to be a common scenario, that there should be an unbounded number of files ( / modules / sections / etc.) to be processed, but bounded memory. User code will be slightly simpler if ThreadPool provides a standard API to handle such cases.

I suppose another direction would be to abandon the producer-consumer model altogether. In D114845, this could be done (in part) by letting the threads share access to the recursive_directory_iterator, acquiring a lock then advancing it on their own to harvest jobs from the filesystem.

[mysterymath]

In D126815#3568630, @noajshu wrote:

I suppose another direction would be to abandon the producer-consumer model altogether. In D114845, this could be done (in part) by letting the threads share access to the recursive_directory_iterator, acquiring a lock then advancing it on their own to harvest jobs from the filesystem.

I like that; it should decouple those commits from this one, and it should be quite a bit cleaner than maintaining size counters.

[noajshu]

In D126815#3568645, @mysterymath wrote:

In D126815#3568630, @noajshu wrote:

I suppose another direction would be to abandon the producer-consumer model altogether. In D114845, this could be done (in part) by letting the threads share access to the recursive_directory_iterator, acquiring a lock then advancing it on their own to harvest jobs from the filesystem.

I like that; it should decouple those commits from this one, and it should be quite a bit cleaner than maintaining size counters.

Sounds good, thanks! I have updated D114845 with this change. I will wait until at least next week before merging to leave time for comments.

[noajshu]

I have just noticed the new Task Group feature of ThreadPool. Waiting on the queue size to decrease (as this diff allows) would appear an ad-hoc practice in light of the clean alternative of using a task group. I will therefore abandon this diff D126815 altogether. Thank you for the helpful comments!

llvm/include/llvm/Support/ThreadPool.h

public:
   auto async(Func &&F) -> std::shared_future<decltype(F())> {
     return asyncImpl(std::function<decltype(F())()>(std::forward<Func>(F)));
   }
 
   /// Blocking wait for all the threads to complete and the queue to be empty.
   /// 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!
   unsigned getThreadCount() const { return MaxThreadCount; }
 
   /// Returns true if the current thread is a worker thread of this thread pool.
   bool isWorkerThread() const;
 
#endif
 
   /// Tasks waiting for execution in the pool.
   std::queue<std::function<void()>> Tasks;
 
   /// Locking and signaling for accessing the Tasks queue.
   std::mutex QueueLock;
   std::condition_variable QueueCondition;
 
+  /// Signaling for queue size decreases
+  std::condition_variable QueueSizeDecreaseCondition;
+
   /// Signaling for job completion
   std::condition_variable CompletionCondition;
 
   /// Keep track of the number of thread actually busy
   unsigned ActiveThreads = 0;
 
 #if LLVM_ENABLE_THREADS // avoids warning for unused variable
   /// Signal for the destruction of the pool, asking thread to exit.

llvm/lib/Support/ThreadPool.cpp

Threads.emplace_back([this, ThreadID] {
           // Yeah, we have a task, grab it and release the lock on the queue
 
           // We first need to signal that we are active before popping the queue
           // in order for wait() to properly detect that even if the queue is
           // empty, there is still a task in flight.
           ++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();
 
         bool Notify;
         {
           // Adjust `ActiveThreads`, in case someone waits on ThreadPool::wait()
           std::lock_guard<std::mutex> LockGuard(QueueLock);

 }
 
 void ThreadPool::wait() {
   // Wait for all threads to complete and the queue to be empty
   std::unique_lock<std::mutex> LockGuard(QueueLock);
   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();
   for (const llvm::thread &Thread : Threads)
     if (CurrentThreadId == Thread.get_id())
       return true;
   return false;
 }