infra: raise MaxTasksReachedError in TaskGroup (#2866)

Author: canepat

silkworm/infra/concurrency/task_group.cpp

@@ -25,6 +25,9 @@ void TaskGroup::spawn(const any_io_executor& executor, Task<void> task) {
     if (is_closed_) {
         throw SpawnAfterCloseError();
     }
+    if (tasks_.size() == max_tasks_) {
+        throw MaxTasksReachedError(max_tasks_);
+    }
 
     auto task_id = ++last_task_id_;
 
@@ -102,7 +105,7 @@ void TaskGroup::on_complete(size_t task_id, const std::exception_ptr& ex_ptr) {
         // if a task threw during cancellation - rethrow from wait()
         auto result_ex_ptr = (ex_ptr && !is_cancelled) ? ex_ptr : std::exception_ptr{};
 
-        bool ok = completions_.try_send({task_id, result_ex_ptr});
+        const bool ok = completions_.try_send({task_id, result_ex_ptr});
         if (!ok) {
             throw std::runtime_error("TaskGroup::on_complete: completions queue is full, unexpected max_tasks limit breach");
         }

silkworm/infra/concurrency/task_group.hpp

@@ -54,7 +54,8 @@ namespace silkworm::concurrency {
 class TaskGroup {
   public:
     TaskGroup(const boost::asio::any_io_executor& executor, size_t max_tasks)
-        : completions_(executor, max_tasks),
+        : max_tasks_(max_tasks),
+          completions_(executor, max_tasks),
           exceptions_(executor, 1) {}
 
     TaskGroup(const TaskGroup&) = delete;
@@ -64,6 +65,11 @@ class TaskGroup {
       public:
         SpawnAfterCloseError() : std::runtime_error("TaskGroup can't spawn after it was closed") {}
     };
+    class MaxTasksReachedError : public std::runtime_error {
+      public:
+        explicit MaxTasksReachedError(size_t max_tasks)
+            : std::runtime_error("TaskGroup can't spawn more than " + std::to_string(max_tasks) + " tasks") {}
+    };
 
     //! Similar to co_spawn, but also adds the task to this group until it completes.
     void spawn(const boost::asio::any_io_executor& executor, Task<void> task);
@@ -80,6 +86,7 @@ class TaskGroup {
     bool is_closed_{false};
     size_t last_task_id_{0};
     std::map<size_t, boost::asio::cancellation_signal> tasks_;
+    size_t max_tasks_{0};
     concurrency::Channel<std::pair<size_t, std::exception_ptr>> completions_;
     concurrency::Channel<std::exception_ptr> exceptions_;
 };

silkworm/infra/concurrency/task_group_test.cpp

@@ -38,14 +38,14 @@ static Task<void> async_throw() {
     throw TestException();
 }
 
-static Task<void> wait_until_cancelled(bool* is_cancelled) {
+static Task<void> wait_until_cancelled(bool& is_cancelled) {
     try {
         auto executor = co_await this_coro::executor;
         steady_timer timer(executor);
         timer.expires_after(1h);
         co_await timer.async_wait(use_awaitable);
     } catch (const boost::system::system_error&) {
-        *is_cancelled = true;
+        is_cancelled = true;
     }
 }
 
@@ -76,7 +76,7 @@ TEST_CASE("TaskGroup.1.wait_until_cancelled") {
     auto executor = runner.executor();
     TaskGroup group{executor, 1};
     bool is_cancelled = false;
-    group.spawn(executor, wait_until_cancelled(&is_cancelled));
+    group.spawn(executor, wait_until_cancelled(is_cancelled));
     CHECK_THROWS_AS(runner.run(group.wait() && async_throw()), TestException);
     CHECK(is_cancelled);
 }
@@ -86,9 +86,9 @@ TEST_CASE("TaskGroup.some.wait_until_cancelled") {
     auto executor = runner.executor();
     TaskGroup group{executor, 3};
     std::array<bool, 3> is_cancelled{};
-    group.spawn(executor, wait_until_cancelled(&is_cancelled[0]));
-    group.spawn(executor, wait_until_cancelled(&is_cancelled[1]));
-    group.spawn(executor, wait_until_cancelled(&is_cancelled[2]));
+    group.spawn(executor, wait_until_cancelled(is_cancelled[0]));
+    group.spawn(executor, wait_until_cancelled(is_cancelled[1]));
+    group.spawn(executor, wait_until_cancelled(is_cancelled[2]));
     CHECK_THROWS_AS(runner.run(group.wait() && async_throw()), TestException);
     CHECK(is_cancelled[0]);
     CHECK(is_cancelled[1]);
@@ -101,11 +101,11 @@ TEST_CASE("TaskGroup.some.mix") {
     TaskGroup group{executor, 6};
     std::array<bool, 3> is_cancelled{};
     group.spawn(executor, async_ok());
-    group.spawn(executor, wait_until_cancelled(&is_cancelled[0]));
+    group.spawn(executor, wait_until_cancelled(is_cancelled[0]));
     group.spawn(executor, async_ok());
-    group.spawn(executor, wait_until_cancelled(&is_cancelled[1]));
+    group.spawn(executor, wait_until_cancelled(is_cancelled[1]));
     group.spawn(executor, async_ok());
-    group.spawn(executor, wait_until_cancelled(&is_cancelled[2]));
+    group.spawn(executor, wait_until_cancelled(is_cancelled[2]));
     CHECK_THROWS_AS(runner.run(group.wait() && async_throw()), TestException);
     CHECK(is_cancelled[0]);
     CHECK(is_cancelled[1]);
@@ -156,11 +156,11 @@ TEST_CASE("TaskGroup.task_cancelled_exception_is_ignored") {
 TEST_CASE("TaskGroup.task_exception_during_cancellation_is_rethrown") {
     test_util::TaskRunner runner;
     auto executor = runner.executor();
-    TaskGroup group{executor, 1};
+    TaskGroup group{executor, 2};
 
     auto task = [&]() -> Task<void> {
         bool is_cancelled = false;
-        co_await wait_until_cancelled(&is_cancelled);
+        co_await wait_until_cancelled(is_cancelled);
         if (is_cancelled) {
             throw std::runtime_error("exception_during_cancellation");
         }
@@ -181,4 +181,31 @@ TEST_CASE("TaskGroup.task_exception_during_cancellation_is_rethrown") {
     CHECK(runner.run(test()));
 }
 
+TEST_CASE("TaskGroup.avoid_max_tasks_limit_breach") {
+    log::init(log::Settings{
+        .log_std_out = true,
+        .log_nocolor = true,
+        .log_threads = true,
+        .log_verbosity = log::Level::kInfo,
+    });
+    test_util::TaskRunner runner;
+    auto executor = runner.executor();
+    const size_t max_tasks = 10;
+    TaskGroup group{executor, max_tasks};
+
+    auto task = [&]() -> Task<void> {
+        bool is_cancelled = false;
+        co_await wait_until_cancelled(is_cancelled);
+        if (is_cancelled) {
+            throw std::runtime_error("exception_during_cancellation");
+        }
+    };
+    // Spawn max_tasks tasks
+    for (size_t i = 0; i < max_tasks; ++i) {
+        REQUIRE_NOTHROW(group.spawn(executor, task()));
+    }
+    // Trying to spawn one more *must* trigger MaxTasksReachedError exception
+    CHECK_THROWS_AS(group.spawn(executor, task()), TaskGroup::MaxTasksReachedError);
+}
+
 }  // namespace silkworm::concurrency