Method ConsumeBuf

tensorflow/core/common_runtime/buf_rendezvous.cc:118–173

}

void BufRendezvous::ConsumeBuf(const string& key, const string& device_name,
                               const uint64 device_incarnation,
                               const ConsumerCallback& done) {
  // Check the incarnation in the request matches the current device
  // incarnation of the producer.
  Device* device;
  Status consumebuf_status = dev_mgr_->LookupDevice(device_name, &device);
  if (consumebuf_status.ok() &&
      device->attributes().incarnation() != device_incarnation) {
    consumebuf_status = errors::FailedPrecondition(
        "RecvBuf expects a different device incarnation: ", device_incarnation,
        " vs. ", device->attributes().incarnation(),
        ". Your worker job that contains the device (\"", device_name,
        "\") was probably restarted. Check your "
        "worker job for the reason why it was restarted.");
  }
  if (!consumebuf_status.ok()) {
    done(consumebuf_status, nullptr);
    return;
  }

  Hook* existing_hook = nullptr;
  do {
    mutex_lock l(mu_);
    if (!status_.ok()) {
      consumebuf_status = status_;
      break;
    }
    auto it = hook_table_.find(key);
    if (it != hook_table_.end()) {
      // Prepare to consume immediately.
      if (it->second->cons_cb) {
        consumebuf_status =
            errors::Internal("Second consumer arrived for key ", key);
        break;
      }
      existing_hook = it->second;
      hook_table_.erase(it);
      existing_hook->cons_cb = done;
    } else {
      // Hang consumer callback on the Hook.
      Hook* h = new Hook;
      hook_table_[key] = h;
      h->cons_cb = done;
      return;
    }
  } while (false);
  if (existing_hook) {
    existing_hook->cons_cb(Status::OK(), existing_hook);
    return;
  }
  if (!consumebuf_status.ok()) {
    done(consumebuf_status, nullptr);
    return;
  }
}

/*static*/