TaskQueue.h 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138
  1. //===-- llvm/Support/TaskQueue.h - A TaskQueue implementation ---*- C++ -*-===//
  2. //
  3. // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
  4. // See https://llvm.org/LICENSE.txt for license information.
  5. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
  6. //
  7. //===----------------------------------------------------------------------===//
  8. //
  9. // This file defines a crude C++11 based task queue.
  10. //
  11. //===----------------------------------------------------------------------===//
  12. #ifndef LLVM_SUPPORT_TASKQUEUE_H
  13. #define LLVM_SUPPORT_TASKQUEUE_H
  14. #include "llvm/Config/llvm-config.h"
  15. #include "llvm/Support/ThreadPool.h"
  16. #include "llvm/Support/thread.h"
  17. #include <atomic>
  18. #include <cassert>
  19. #include <condition_variable>
  20. #include <deque>
  21. #include <functional>
  22. #include <future>
  23. #include <memory>
  24. #include <mutex>
  25. #include <utility>
  26. namespace llvm {
  27. /// TaskQueue executes serialized work on a user-defined Thread Pool. It
  28. /// guarantees that if task B is enqueued after task A, task B begins after
  29. /// task A completes and there is no overlap between the two.
  30. class TaskQueue {
  31. // Because we don't have init capture to use move-only local variables that
  32. // are captured into a lambda, we create the promise inside an explicit
  33. // callable struct. We want to do as much of the wrapping in the
  34. // type-specialized domain (before type erasure) and then erase this into a
  35. // std::function.
  36. template <typename Callable> struct Task {
  37. using ResultTy = std::result_of_t<Callable()>;
  38. explicit Task(Callable C, TaskQueue &Parent)
  39. : C(std::move(C)), P(std::make_shared<std::promise<ResultTy>>()),
  40. Parent(&Parent) {}
  41. template<typename T>
  42. void invokeCallbackAndSetPromise(T*) {
  43. P->set_value(C());
  44. }
  45. void invokeCallbackAndSetPromise(void*) {
  46. C();
  47. P->set_value();
  48. }
  49. void operator()() noexcept {
  50. ResultTy *Dummy = nullptr;
  51. invokeCallbackAndSetPromise(Dummy);
  52. Parent->completeTask();
  53. }
  54. Callable C;
  55. std::shared_ptr<std::promise<ResultTy>> P;
  56. TaskQueue *Parent;
  57. };
  58. public:
  59. /// Construct a task queue with no work.
  60. TaskQueue(ThreadPool &Scheduler) : Scheduler(Scheduler) { (void)Scheduler; }
  61. /// Blocking destructor: the queue will wait for all work to complete.
  62. ~TaskQueue() {
  63. Scheduler.wait();
  64. assert(Tasks.empty());
  65. }
  66. /// Asynchronous submission of a task to the queue. The returned future can be
  67. /// used to wait for the task (and all previous tasks that have not yet
  68. /// completed) to finish.
  69. template <typename Callable>
  70. std::future<std::result_of_t<Callable()>> async(Callable &&C) {
  71. #if !LLVM_ENABLE_THREADS
  72. static_assert(false,
  73. "TaskQueue requires building with LLVM_ENABLE_THREADS!");
  74. #endif
  75. Task<Callable> T{std::move(C), *this};
  76. using ResultTy = std::result_of_t<Callable()>;
  77. std::future<ResultTy> F = T.P->get_future();
  78. {
  79. std::lock_guard<std::mutex> Lock(QueueLock);
  80. // If there's already a task in flight, just queue this one up. If
  81. // there is not a task in flight, bypass the queue and schedule this
  82. // task immediately.
  83. if (IsTaskInFlight)
  84. Tasks.push_back(std::move(T));
  85. else {
  86. Scheduler.async(std::move(T));
  87. IsTaskInFlight = true;
  88. }
  89. }
  90. return F;
  91. }
  92. private:
  93. void completeTask() {
  94. // We just completed a task. If there are no more tasks in the queue,
  95. // update IsTaskInFlight to false and stop doing work. Otherwise
  96. // schedule the next task (while not holding the lock).
  97. std::function<void()> Continuation;
  98. {
  99. std::lock_guard<std::mutex> Lock(QueueLock);
  100. if (Tasks.empty()) {
  101. IsTaskInFlight = false;
  102. return;
  103. }
  104. Continuation = std::move(Tasks.front());
  105. Tasks.pop_front();
  106. }
  107. Scheduler.async(std::move(Continuation));
  108. }
  109. /// The thread pool on which to run the work.
  110. ThreadPool &Scheduler;
  111. /// State which indicates whether the queue currently is currently processing
  112. /// any work.
  113. bool IsTaskInFlight = false;
  114. /// Mutex for synchronizing access to the Tasks array.
  115. std::mutex QueueLock;
  116. /// Tasks waiting for execution in the queue.
  117. std::deque<std::function<void()>> Tasks;
  118. };
  119. } // namespace llvm
  120. #endif // LLVM_SUPPORT_TASKQUEUE_H