Function TEST

cpp/src/arrow/acero/task_util_test.cc:104–177

This test simulates one of the current use patterns of the task scheduler. There are a number of groups. The groups are allocated to stages. All groups in a stage execute concurrently. When all groups in that stage finish the next stage is started.

// concurrently.  When all groups in that stage finish the next
// stage is started.
TEST(TaskScheduler, Stress) {
#ifndef ARROW_ENABLE_THREADING
  GTEST_SKIP() << "Test requires threading support";
#endif
  constexpr int kNumThreads = 8;
  constexpr int kNumGroups = 8;
  constexpr int kGroupsPerStage = 3;
  constexpr int kTasksPerGroup = 32;
  constexpr int kNumStages = (kNumGroups % kGroupsPerStage == 0)
                                 ? (kNumGroups / kGroupsPerStage)
                                 : (kNumGroups / kGroupsPerStage) + 1;
  constexpr int kTrailingGroups = (kNumGroups % kGroupsPerStage == 0)
                                      ? kGroupsPerStage
                                      : kNumGroups % kGroupsPerStage;

  ThreadIndexer thread_indexer;
  int num_threads = std::min(static_cast<int>(thread_indexer.Capacity()), kNumThreads);
  ASSERT_OK_AND_ASSIGN(std::shared_ptr<ThreadPool> thread_pool,
                       ThreadPool::Make(num_threads));

  std::array<std::atomic<int>, kNumStages - 1> stage_counters;
  for (int i = 0; i < kNumStages - 1; i++) {
    stage_counters[i].store(kGroupsPerStage);
  }
  std::atomic<int> final_counter(kTrailingGroups);
  std::mutex mutex;
  std::condition_variable finish_cv;

  std::vector<int> group_ids;
  auto scheduler = TaskScheduler::Make();

  std::function<void(std::size_t, int)> start_next_stage = [&](std::size_t thread_id,
                                                               int stage_index) {
    int start = stage_index * kGroupsPerStage;
    int end = std::min(kNumGroups, start + kGroupsPerStage);
    for (int i = start; i < end; i++) {
      ASSERT_OK(thread_pool->Spawn([&, i] {
        std::size_t my_thread_id = thread_indexer();
        SleepABit();
        ASSERT_OK(scheduler->StartTaskGroup(my_thread_id, group_ids[i], kTasksPerGroup));
      }));
    }
  };

  for (auto i = 0; i < kNumGroups; i++) {
    int next_stage = (i / kGroupsPerStage) + 1;
    TaskScheduler::TaskGroupContinuationImpl finish =
        MakeFinalContinuation(&final_counter, &mutex, &finish_cv);
    if (next_stage < kNumStages) {
      finish =
          MakeContinuation(&stage_counters[next_stage - 1], start_next_stage, next_stage);
    }
    group_ids.push_back(scheduler->RegisterTaskGroup(SlowTaskImpl, finish));
  }
  scheduler->RegisterEnd();

  TaskScheduler::AbortContinuationImpl abort = [] { FAIL() << "Unexpected abort"; };
  TaskScheduler::ScheduleImpl schedule =