| 125 | } |
| 126 | |
| 127 | Status TaskSchedulerImpl::StartTaskGroup(size_t thread_id, int group_id, |
| 128 | int64_t total_num_tasks) { |
| 129 | ARROW_DCHECK(group_id >= 0 && group_id < static_cast<int>(task_groups_.size())); |
| 130 | TaskGroup& task_group = task_groups_[group_id]; |
| 131 | |
| 132 | bool aborted = aborted_.value.load(); |
| 133 | bool all_tasks_finished = false; |
| 134 | { |
| 135 | std::lock_guard<std::mutex> lock(mutex_); |
| 136 | |
| 137 | if (task_group.state_ == TaskGroupState::NOT_READY) { |
| 138 | task_group.num_tasks_present_ = total_num_tasks; |
| 139 | if (total_num_tasks == 0) { |
| 140 | task_group.state_ = TaskGroupState::ALL_TASKS_FINISHED; |
| 141 | all_tasks_finished = true; |
| 142 | } |
| 143 | task_group.state_ = TaskGroupState::READY; |
| 144 | } |
| 145 | } |
| 146 | |
| 147 | if (!aborted && all_tasks_finished) { |
| 148 | bool all_task_groups_finished = false; |
| 149 | RETURN_NOT_OK(OnTaskGroupFinished(thread_id, group_id, &all_task_groups_finished)); |
| 150 | if (all_task_groups_finished) { |
| 151 | return Status::OK(); |
| 152 | } |
| 153 | } |
| 154 | |
| 155 | if (!aborted) { |
| 156 | tasks_added_recently_.value.store(true); |
| 157 | return ScheduleMore(thread_id); |
| 158 | } else { |
| 159 | return Status::Cancelled("Scheduler cancelled"); |
| 160 | } |
| 161 | } |
| 162 | |
| 163 | std::vector<std::pair<int, int64_t>> TaskSchedulerImpl::PickTasks(int num_tasks, |
| 164 | int start_task_group) { |