| 136 | } |
| 137 | |
| 138 | Status InputReceived(ExecNode* input, ExecBatch batch) override { |
| 139 | std::vector<QueuedBatch> to_deliver; |
| 140 | bool should_finish = false; |
| 141 | { |
| 142 | std::lock_guard lk(mutex_); |
| 143 | int current_count = counter_.count(); |
| 144 | int adjusted_count = current_count + jitter_dist_(rng_); |
| 145 | QueuedBatch queued{adjusted_count, std::move(batch)}; |
| 146 | queue_.push(std::move(queued)); |
| 147 | while (!queue_.empty() && queue_.top().adjusted_order <= current_count) { |
| 148 | to_deliver.push_back(std::move(queue_.top())); |
| 149 | queue_.pop(); |
| 150 | } |
| 151 | if (counter_.Increment()) { |
| 152 | should_finish = true; |
| 153 | } |
| 154 | } |
| 155 | Dispatch(std::move(to_deliver)); |
| 156 | if (should_finish) { |
| 157 | Finish(); |
| 158 | } |
| 159 | return Status::OK(); |
| 160 | } |
| 161 | |
| 162 | Status InputFinished(ExecNode* input, int total_batches) override { |
| 163 | if (counter_.SetTotal(total_batches)) { |