| 32 | namespace bytedance::bolt::exec { |
| 33 | |
| 34 | void ExchangeClient::addRemoteTaskId(const std::string& taskId) { |
| 35 | RequestSpec requestSpec; |
| 36 | std::shared_ptr<ExchangeSource> toClose; |
| 37 | { |
| 38 | std::lock_guard<std::mutex> l(queue_->mutex()); |
| 39 | |
| 40 | bool duplicate = !remoteTaskIds_.insert(taskId).second; |
| 41 | if (duplicate) { |
| 42 | // Do not add sources twice. Presto protocol may add duplicate sources |
| 43 | // and the task updates have no guarantees of arriving in order. |
| 44 | return; |
| 45 | } |
| 46 | |
| 47 | std::shared_ptr<ExchangeSource> source; |
| 48 | try { |
| 49 | source = ExchangeSource::create(taskId, destination_, queue_, pool_); |
| 50 | } catch (const BoltException&) { |
| 51 | throw; |
| 52 | } catch (const std::exception& e) { |
| 53 | // Task ID can be very long. Truncate to 128 characters. |
| 54 | BOLT_FAIL( |
| 55 | "Failed to create ExchangeSource: {}. Task ID: {}.", |
| 56 | e.what(), |
| 57 | taskId.substr(0, 128)); |
| 58 | } |
| 59 | |
| 60 | if (closed_) { |
| 61 | toClose = std::move(source); |
| 62 | } else { |
| 63 | sources_.push_back(source); |
| 64 | queue_->addSourceLocked(); |
| 65 | // Put new source into 'producingSources_' queue to prioritise fetching |
| 66 | // from these to find out whether these are productive or not. |
| 67 | producingSources_.push(source); |
| 68 | |
| 69 | requestSpec = pickSourcesToRequestLocked(); |
| 70 | } |
| 71 | } |
| 72 | |
| 73 | // Outside of lock. |
| 74 | if (toClose) { |
| 75 | toClose->close(); |
| 76 | } else { |
| 77 | request(requestSpec); |
| 78 | } |
| 79 | } |
| 80 | |
| 81 | void ExchangeClient::noMoreRemoteTasks() { |
| 82 | queue_->noMoreSources(); |