| 18 | namespace tensorflow { |
| 19 | |
| 20 | Status AsyncIoRendezvous::Send(uint64 key_hash, const TensorPayload& val) { |
| 21 | VLOG(2) << "Send " << this << " " << key_hash; |
| 22 | |
| 23 | MutexedItemQueue* mq = nullptr; |
| 24 | { |
| 25 | mutex_lock l(mu_); |
| 26 | mq = table_.at(key_hash); |
| 27 | } |
| 28 | CHECK(mq); |
| 29 | ItemQueue* queue = &mq->queue; |
| 30 | mutex& mu = mq->mu; |
| 31 | mu.lock(); |
| 32 | if (queue->empty() || queue->front()->IsSendValue()) { |
| 33 | // There is no waiter for this message. Append the message |
| 34 | // into the queue. The waiter will pick it up when arrives. |
| 35 | // Only send-related fields need to be filled. |
| 36 | VLOG(2) << "Enqueue Send Item (key hash:" << key_hash << ", queue:" << queue |
| 37 | << "). "; |
| 38 | Item* item = new Item; |
| 39 | item->value = val; |
| 40 | queue->push_back(item); |
| 41 | mu.unlock(); |
| 42 | return Status::OK(); |
| 43 | } |
| 44 | |
| 45 | VLOG(2) << "Consume Recv Item (key hash:" << key_hash << "). "; |
| 46 | // There is an earliest waiter to consume this message. |
| 47 | Item* item = queue->front(); |
| 48 | queue->pop_front(); |
| 49 | mu.unlock(); |
| 50 | |
| 51 | // Notify the waiter by invoking its done closure, outside the |
| 52 | // lock. |
| 53 | CHECK(!item->IsSendValue()); |
| 54 | item->waiter(Status::OK(), val); |
| 55 | delete item; |
| 56 | return Status::OK(); |
| 57 | } |
| 58 | |
| 59 | void AsyncIoRendezvous::RecvAsync(uint64 key_hash, DoneCallback done) { |
| 60 | VLOG(2) << "Recv " << this << " " << key_hash; |