| 390 | } |
| 391 | |
| 392 | void TaskSchedulerImpl::Abort(AbortContinuationImpl impl) { |
| 393 | bool all_finished = true; |
| 394 | DCHECK_EQ(aborted_.value.load(), false); |
| 395 | aborted_.value.store(true); |
| 396 | { |
| 397 | std::lock_guard<std::mutex> lock(mutex_); |
| 398 | abort_cont_impl_ = std::move(impl); |
| 399 | if (register_finished_) { |
| 400 | for (size_t i = 0; i < task_groups_.size(); ++i) { |
| 401 | TaskGroup& task_group = task_groups_[i]; |
| 402 | switch (task_group.state_) { |
| 403 | case TaskGroupState::NOT_READY: { |
| 404 | task_group.state_ = TaskGroupState::ALL_TASKS_FINISHED; |
| 405 | break; |
| 406 | } |
| 407 | case TaskGroupState::READY: { |
| 408 | int64_t expected = task_group.num_tasks_started_.value.load(); |
| 409 | for (;;) { |
| 410 | if (task_group.num_tasks_started_.value.compare_exchange_strong( |
| 411 | expected, task_group.num_tasks_present_)) { |
| 412 | break; |
| 413 | } |
| 414 | } |
| 415 | int64_t before_add = task_group.num_tasks_finished_.value.fetch_add( |
| 416 | task_group.num_tasks_present_ - expected); |
| 417 | if (before_add >= expected) { |
| 418 | task_group.state_ = TaskGroupState::ALL_TASKS_FINISHED; |
| 419 | } else { |
| 420 | all_finished = false; |
| 421 | task_group.state_ = TaskGroupState::ALL_TASKS_STARTED; |
| 422 | } |
| 423 | break; |
| 424 | } |
| 425 | case TaskGroupState::ALL_TASKS_STARTED: { |
| 426 | all_finished = false; |
| 427 | break; |
| 428 | } |
| 429 | default: |
| 430 | break; |
| 431 | } |
| 432 | } |
| 433 | } |
| 434 | } |
| 435 | if (all_finished) { |
| 436 | abort_cont_impl_(); |
| 437 | } |
| 438 | } |
| 439 | |
| 440 | std::unique_ptr<TaskScheduler> TaskScheduler::Make() { |
| 441 | std::unique_ptr<TaskSchedulerImpl> impl{new TaskSchedulerImpl()}; |