| 291 | } |
| 292 | |
| 293 | void RecvAsync(const ParsedKey& key, const Args& recv_args, |
| 294 | DoneCallback done) override { |
| 295 | uint64 key_hash = KeyHash(key.FullKey()); |
| 296 | VLOG(2) << "Recv " << this << " " << key_hash << " " << key.FullKey(); |
| 297 | |
| 298 | mu_.lock(); |
| 299 | if (!status_.ok()) { |
| 300 | // Rendezvous has been aborted. |
| 301 | Status s = status_; |
| 302 | mu_.unlock(); |
| 303 | done(s, Args(), recv_args, Tensor(), false); |
| 304 | return; |
| 305 | } |
| 306 | |
| 307 | ItemQueue* queue = &table_[key_hash]; |
| 308 | if (queue->empty() || !queue->front()->IsSendValue()) { |
| 309 | // There is no message to pick up. |
| 310 | // Only recv-related fields need to be filled. |
| 311 | CancellationManager* cm = recv_args.cancellation_manager; |
| 312 | CancellationToken token = CancellationManager::kInvalidToken; |
| 313 | bool already_cancelled = false; |
| 314 | if (cm != nullptr) { |
| 315 | token = cm->get_cancellation_token(); |
| 316 | already_cancelled = !cm->RegisterCallback(token, [this, token, |
| 317 | key_hash] { |
| 318 | Item* item = nullptr; |
| 319 | { |
| 320 | mutex_lock l(mu_); |
| 321 | ItemQueue* queue = &table_[key_hash]; |
| 322 | if (!queue->empty() && !queue->front()->IsSendValue()) { |
| 323 | for (auto it = queue->begin(); it != queue->end(); it++) { |
| 324 | if ((*it)->cancellation_token == token) { |
| 325 | item = *it; |
| 326 | if (queue->size() == 1) { |
| 327 | table_.erase(key_hash); |
| 328 | } else { |
| 329 | queue->erase(it); |
| 330 | } |
| 331 | break; |
| 332 | } |
| 333 | } |
| 334 | } |
| 335 | } |
| 336 | |
| 337 | if (item != nullptr) { |
| 338 | item->waiter(StatusGroup::MakeDerived( |
| 339 | errors::Cancelled("RecvAsync is cancelled.")), |
| 340 | Args(), item->recv_args, Tensor(), /*is_dead=*/false); |
| 341 | delete item; |
| 342 | } |
| 343 | }); |
| 344 | } |
| 345 | if (already_cancelled) { |
| 346 | mu_.unlock(); |
| 347 | done(StatusGroup::MakeDerived( |
| 348 | errors::Cancelled("RecvAsync is cancelled.")), |
| 349 | Args(), recv_args, Tensor(), /*is_dead=*/false); |
| 350 | return; |