fix: Faster implementation of work queue (#2887)

This commit is contained in:
Alex Kremer
2026-01-13 12:21:20 +00:00
committed by GitHub
parent f33f15c02d
commit c6be761f33
4 changed files with 198 additions and 127 deletions

View File

@@ -25,9 +25,7 @@
#include "util/prometheus/Label.hpp"
#include "util/prometheus/Prometheus.hpp"
#include <boost/asio/post.hpp>
#include <boost/asio/spawn.hpp>
#include <boost/asio/strand.hpp>
#include <boost/json/object.hpp>
#include <chrono>
@@ -39,6 +37,27 @@
namespace rpc {
void
WorkQueue::OneTimeCallable::setCallable(std::function<void()> func)
{
func_ = std::move(func);
}
void
WorkQueue::OneTimeCallable::operator()()
{
if (not called_) {
func_();
called_ = true;
}
}
WorkQueue::OneTimeCallable::
operator bool() const
{
return func_.operator bool();
}
WorkQueue::WorkQueue(DontStartProcessingTag, std::uint32_t numWorkers, uint32_t maxSize)
: queued_{PrometheusService::counterInt(
"work_queue_queued_total_number",
@@ -56,8 +75,6 @@ WorkQueue::WorkQueue(DontStartProcessingTag, std::uint32_t numWorkers, uint32_t
"The current number of tasks in the queue"
)}
, ioc_{numWorkers}
, strand_{ioc_.get_executor()}
, waitTimer_(ioc_)
{
if (maxSize != 0)
maxSize_ = maxSize;
@@ -77,12 +94,14 @@ WorkQueue::~WorkQueue()
void
WorkQueue::startProcessing()
{
util::spawn(strand_, [this](auto yield) {
ASSERT(not hasDispatcher_, "Dispatcher already running");
ASSERT(not processingStarted_, "Attempt to start processing work queue more than once");
processingStarted_ = true;
hasDispatcher_ = true;
dispatcherLoop(yield);
});
// Spawn workers for all tasks that were queued before processing started
auto const numTasks = size();
for (auto i = 0uz; i < numTasks; ++i) {
util::spawn(ioc_, [this](auto yield) { executeTask(yield); });
}
}
bool
@@ -98,93 +117,28 @@ WorkQueue::postCoro(TaskType func, bool isWhiteListed, Priority priority)
return false;
}
++curSize_.get();
auto needsWakeup = false;
{
auto state = dispatcherState_.lock();
needsWakeup = std::exchange(state->isIdle, false);
auto state = queueState_.lock();
state->push(priority, std::move(func));
}
if (needsWakeup)
boost::asio::post(strand_, [this] { waitTimer_.cancel(); });
++curSize_.get();
if (not processingStarted_)
return true;
util::spawn(ioc_, [this](auto yield) { executeTask(yield); });
return true;
}
void
WorkQueue::dispatcherLoop(boost::asio::yield_context yield)
{
LOG(log_.info()) << "WorkQueue dispatcher starting";
// all ongoing tasks must be completed before stopping fully
while (not stopping_ or size() > 0) {
std::optional<TaskType> task;
{
auto state = dispatcherState_.lock();
if (state->empty()) {
state->isIdle = true;
} else {
task = state->popNext();
}
}
if (not stopping_ and not task.has_value()) {
waitTimer_.expires_at(std::chrono::steady_clock::time_point::max());
boost::system::error_code ec;
waitTimer_.async_wait(yield[ec]);
} else if (task.has_value()) {
util::spawn(
ioc_,
[this, spawnedAt = std::chrono::system_clock::now(), task = std::move(*task)](auto yield) mutable {
auto const takenAt = std::chrono::system_clock::now();
auto const waited =
std::chrono::duration_cast<std::chrono::microseconds>(takenAt - spawnedAt).count();
++queued_.get();
durationUs_.get() += waited;
LOG(log_.info()) << "WorkQueue wait time: " << waited << ", queue size: " << size();
task(yield);
--curSize_.get();
}
);
}
}
LOG(log_.info()) << "WorkQueue dispatcher shutdown requested - time to execute onTasksComplete";
{
auto onTasksComplete = onQueueEmpty_.lock();
ASSERT(onTasksComplete->operator bool(), "onTasksComplete must be set when stopping is true.");
onTasksComplete->operator()();
}
LOG(log_.info()) << "WorkQueue dispatcher finished";
}
void
WorkQueue::requestStop(std::function<void()> onQueueEmpty)
{
auto handler = onQueueEmpty_.lock();
*handler = std::move(onQueueEmpty);
handler->setCallable(std::move(onQueueEmpty));
stopping_ = true;
auto needsWakeup = false;
{
auto state = dispatcherState_.lock();
needsWakeup = std::exchange(state->isIdle, false);
}
if (needsWakeup)
boost::asio::post(strand_, [this] { waitTimer_.cancel(); });
}
void
@@ -194,6 +148,12 @@ WorkQueue::stop()
requestStop();
ioc_.join();
{
auto onTasksComplete = onQueueEmpty_.lock();
ASSERT(onTasksComplete->operator bool(), "onTasksComplete must be set when stopping is true.");
onTasksComplete->operator()();
}
}
WorkQueue
@@ -227,4 +187,29 @@ WorkQueue::size() const
return curSize_.get().value();
}
void
WorkQueue::executeTask(boost::asio::yield_context yield)
{
std::optional<TaskWithTimestamp> taskWithTimestamp;
{
auto state = queueState_.lock();
taskWithTimestamp = state->popNext();
}
ASSERT(
taskWithTimestamp.has_value(),
"Queue should not be empty as we spawn a coro with executeTask for each postCoro."
);
auto const takenAt = std::chrono::system_clock::now();
auto const waited =
std::chrono::duration_cast<std::chrono::microseconds>(takenAt - taskWithTimestamp->queuedAt).count();
++queued_.get();
durationUs_.get() += waited;
LOG(log_.info()) << "WorkQueue wait time: " << waited << ", queue size: " << size();
taskWithTimestamp->task(yield);
--curSize_.get();
}
} // namespace rpc

View File

@@ -25,15 +25,12 @@
#include "util/prometheus/Counter.hpp"
#include "util/prometheus/Gauge.hpp"
#include <boost/asio.hpp>
#include <boost/asio/spawn.hpp>
#include <boost/asio/steady_timer.hpp>
#include <boost/asio/strand.hpp>
#include <boost/asio/thread_pool.hpp>
#include <boost/json.hpp>
#include <boost/json/object.hpp>
#include <atomic>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <functional>
@@ -64,7 +61,13 @@ struct Reportable {
*/
class WorkQueue : public Reportable {
using TaskType = std::function<void(boost::asio::yield_context)>;
using QueueType = std::queue<TaskType>;
struct TaskWithTimestamp {
TaskType task;
std::chrono::system_clock::time_point queuedAt;
};
using QueueType = std::queue<TaskWithTimestamp>;
public:
/**
@@ -76,22 +79,21 @@ public:
};
private:
struct DispatcherState {
struct QueueState {
QueueType high;
QueueType normal;
bool isIdle = false;
size_t highPriorityCounter = 0;
void
push(Priority priority, auto&& task)
push(Priority priority, TaskType&& task)
{
auto& queue = [this, priority] -> QueueType& {
if (priority == Priority::High)
return high;
return normal;
}();
queue.push(std::forward<decltype(task)>(task));
queue.push(TaskWithTimestamp{.task = std::move(task), .queuedAt = std::chrono::system_clock::now()});
}
[[nodiscard]] bool
@@ -100,21 +102,21 @@ private:
return high.empty() and normal.empty();
}
[[nodiscard]] std::optional<TaskType>
[[nodiscard]] std::optional<TaskWithTimestamp>
popNext()
{
if (not high.empty() and (highPriorityCounter < kTAKE_HIGH_PRIO or normal.empty())) {
auto task = std::move(high.front());
auto taskWithTimestamp = std::move(high.front());
high.pop();
++highPriorityCounter;
return task;
return taskWithTimestamp;
}
if (not normal.empty()) {
auto task = std::move(normal.front());
auto taskWithTimestamp = std::move(normal.front());
normal.pop();
highPriorityCounter = 0;
return task;
return taskWithTimestamp;
}
return std::nullopt;
@@ -133,14 +135,26 @@ private:
util::Logger log_{"RPC"};
boost::asio::thread_pool ioc_;
boost::asio::strand<boost::asio::thread_pool::executor_type> strand_;
bool hasDispatcher_ = false;
std::atomic_bool stopping_;
std::atomic_bool processingStarted_{false};
util::Mutex<std::function<void()>> onQueueEmpty_;
util::Mutex<DispatcherState> dispatcherState_;
boost::asio::steady_timer waitTimer_;
class OneTimeCallable {
std::function<void()> func_;
bool called_{false};
public:
void
setCallable(std::function<void()> func);
void
operator()();
explicit
operator bool() const;
};
util::Mutex<OneTimeCallable> onQueueEmpty_;
util::Mutex<QueueState> queueState_;
public:
struct DontStartProcessingTag {};
@@ -234,7 +248,7 @@ public:
private:
void
dispatcherLoop(boost::asio::yield_context yield);
executeTask(boost::asio::yield_context yield);
};
} // namespace rpc