Method RegisterCall

tensorflow/core/distributed_runtime/base_rendezvous_mgr.cc:869–898

}

void BaseRemoteRendezvous::RegisterCall(BaseRecvTensorCall* call,
                                        const Rendezvous::Args& args) {
  CancellationManager* cm = args.cancellation_manager;
  {
    mutex_lock l(mu_);
    if (!status_.ok()) {
      call->StartAbort(status_);
      return;
    }
    bool already_cancelled = false;
    InactiveCallback callback = [] {};
    if (cm != nullptr) {
      auto token = cm->get_cancellation_token();
      already_cancelled = !cm->RegisterCallback(token, [this, call] {
        {
          mutex_lock l(mu_);
          if (active_.find(call) == active_.end()) return;
          call->StartAbort(
              errors::Cancelled("RecvFromRemoteAsync is cancelled."));
        }
      });
      callback = [cm, token] { cm->TryDeregisterCallback(token); };
    }
    if (already_cancelled) {
      call->StartAbort(errors::Cancelled("RecvFromRemoteAsync is cancelled."));
    } else {
      CHECK(active_.emplace(call, callback).second);
    }
  }
}

void BaseRemoteRendezvous::DeregisterCall(BaseRecvTensorCall* call) {
  mutex_lock l(mu_);