| 70 | } |
| 71 | |
| 72 | void BufRendezvous::ProvideBuf(const string& key, Device* dev, |
| 73 | DeviceContext* dev_ctx, const Tensor* v, |
| 74 | const AllocatorAttributes& attr, |
| 75 | const ProducerCallback& done) { |
| 76 | Hook* h = nullptr; |
| 77 | Status providebuf_status; |
| 78 | do { |
| 79 | mutex_lock l(mu_); |
| 80 | if (!status_.ok()) { |
| 81 | providebuf_status = status_; |
| 82 | break; |
| 83 | } else { |
| 84 | auto it = hook_table_.find(key); |
| 85 | if (it == hook_table_.end()) { |
| 86 | h = new Hook; |
| 87 | it = hook_table_.insert(std::make_pair(key, h)).first; |
| 88 | } else { |
| 89 | if (it->second->prod_cb != nullptr) { |
| 90 | providebuf_status = errors::Internal( |
| 91 | "BufRendezvous::ProvideBuf already called for key ", key); |
| 92 | break; |
| 93 | } |
| 94 | h = it->second; |
| 95 | } |
| 96 | // Populate Hook with all of the prod values. |
| 97 | h->prod_dev = dev; |
| 98 | h->prod_ctx = dev_ctx; |
| 99 | h->prod_value = v; |
| 100 | h->prod_attr = attr; |
| 101 | h->prod_cb = done; |
| 102 | // If consumer is waiting, kick off right away, removing Hook from table. |
| 103 | if (h->cons_cb != nullptr) { |
| 104 | hook_table_.erase(it); |
| 105 | } else { |
| 106 | h = nullptr; |
| 107 | } |
| 108 | } |
| 109 | } while (false); |
| 110 | if (h) { |
| 111 | h->cons_cb(Status::OK(), h); |
| 112 | } |
| 113 | if (!providebuf_status.ok()) { |
| 114 | done(providebuf_status); |
| 115 | } |
| 116 | } |
| 117 | |
| 118 | void BufRendezvous::ConsumeBuf(const string& key, const string& device_name, |
| 119 | const uint64 device_incarnation, |