Method ScheduleMore

cpp/src/arrow/acero/task_util.cc:325–390

}

Status TaskSchedulerImpl::ScheduleMore(size_t thread_id, int num_tasks_finished) {
  if (aborted_.value.load()) {
    return Status::Cancelled("Scheduler cancelled");
  }

  ARROW_DCHECK(register_finished_);

  if (use_sync_execution_) {
    return ExecuteMore(thread_id, 1, true);
  }

  int num_new_tasks = num_tasks_finished;
  for (;;) {
    int expected = num_tasks_to_schedule_.value.load();
    if (num_tasks_to_schedule_.value.compare_exchange_strong(expected, 0)) {
      num_new_tasks += expected;
      break;
    }
  }
  if (num_new_tasks == 0) {
    return Status::OK();
  }

  const auto& tasks = PickTasks(num_new_tasks);
  if (static_cast<int>(tasks.size()) < num_new_tasks) {
    num_tasks_to_schedule_.value += num_new_tasks - static_cast<int>(tasks.size());
  }

  bool expected_might_have_missed_tasks = true;
  if (tasks_added_recently_.value.compare_exchange_strong(
          expected_might_have_missed_tasks, false)) {
    if (tasks.empty()) {
      // num_tasks_finished has already been added to num_tasks_to_schedule so
      // pass 0 here.
      return ScheduleMore(thread_id);
    }
  }

  for (size_t i = 0; i < tasks.size(); ++i) {
    int group_id = tasks[i].first;
    int64_t task_id = tasks[i].second;
    RETURN_NOT_OK(schedule_impl_([this, group_id, task_id](size_t thread_id) -> Status {
      bool task_group_finished = false;
      // PostExecuteTask must be called later if any error ocurres during task execution
      // (including ScheduleMore), so we preserve the status.
      auto status = [&]() {
        RETURN_NOT_OK(ScheduleMore(thread_id, 1));
        return ExecuteTask(thread_id, group_id, task_id, &task_group_finished);
      }();

      if (!status.ok()) {
        task_group_finished = PostExecuteTask(thread_id, group_id);
      }

      if (task_group_finished) {
        bool all_task_groups_finished = false;
        RETURN_NOT_OK(
            OnTaskGroupFinished(thread_id, group_id, &all_task_groups_finished));