…lines 1–71 unchanged (71 lines)
72 }
73 delete h;
74 }
75 table->clear();
76}
77
78std::string BufRendezvous::Hook::DebugString() const {
79 return absl::StrCat(
80 "[dev:", (prod_dev ? prod_dev->name() : "none"),
81 ", ctx:", reinterpret_cast<uint64_t>(prod_ctx),
82 ", val:", reinterpret_cast<uint64_t>(prod_value),
83 ", pcb:", prod_cb ? reinterpret_cast<uint64_t>(&prod_cb) : 0,
84 ", ccb:", cons_cb ? reinterpret_cast<uint64_t>(&cons_cb) : 0, "]");
85}
86
87void BufRendezvous::ProvideBuf(const std::string& key, Device* dev,
88 DeviceContext* dev_ctx, const Tensor* v,
89 const AllocatorAttributes& attr,
90 const ProducerCallback& done,
91 CancellationManager* cancellation_manager) {
92 DVLOG(4) << "ProvideBuf: key = " << key;
93#ifndef NDEBUG
…lines 94–151 unchanged (58 lines)
152 }
153 if (!providebuf_status.ok()) {
154 done(providebuf_status);
155 }
156}
157
158void BufRendezvous::ConsumeBuf(const std::string& key,
159 const std::string& device_name,
160 const uint64_t device_incarnation,
161 const ConsumerCallback& done,
162 CancellationManager* cancellation_manager) {
163 DVLOG(4) << "ConsumeBuf: key = " << key << " device_name = " << device_name;
164#ifndef NDEBUG
165 if (VLOG_IS_ON(4)) {
166 LogContents();
…lines 167–230 unchanged (64 lines)
231 if (!consumebuf_status.ok()) {
232 done(consumebuf_status, nullptr);
233 return;
234 }
235}
236
237void BufRendezvous::CancelHook(const std::string& key) {
238 Hook* h = nullptr;
239 {
240 mutex_lock l(mu_);
241 auto it = hook_table_.find(key);
242 if (it == hook_table_.end()) return;
243 h = it->second;
244 hook_table_.erase(it);
245 }
246 if (h != nullptr) {
247 auto s = errors::Cancelled("Operation was cancelled for BufRendezvous key ",
248 key);
249 if (h->prod_cb != nullptr) {
250 h->prod_cb(s);
251 }
252 if (h->cons_cb != nullptr) {
253 h->cons_cb(s, /*Hook=*/nullptr);
254 }
255 delete h;
256 }
257}
258
259/*static*/
260void BufRendezvous::DoneWithHook(Hook* h) {
261 h->prod_cb(absl::OkStatus());
262 delete h;
263}
264
265void BufRendezvous::LogContents() {
266 mutex_lock l(mu_);
267 LOG(INFO) << strings::StrCat("BufRendezvous ",
268 strings::Hex(reinterpret_cast<uint64_t>(this)),
269 " step_id=", step_id_, " current contents:");
270 for (const auto& it : hook_table_) {
271 LOG(INFO) << it.first << ":" << it.second->DebugString();
272 }
273}
274
275} // namespace tensorflow
276