Method TryEnqueueMany

tensorflow/core/kernels/priority_queue.cc:130–195

}

void PriorityQueue::TryEnqueueMany(const Tuple& tuple, OpKernelContext* ctx,
                                   DoneCallback callback) {
  const int64 batch_size = tuple[0].dim_size(0);
  if (batch_size == 0) {
    callback();
    return;
  }

  CancellationManager* cm = ctx->cancellation_manager();
  CancellationToken token = cm->get_cancellation_token();
  bool already_cancelled;
  {
    mutex_lock l(mu_);
    already_cancelled = !cm->RegisterCallback(
        token, [this, cm, token]() { Cancel(kEnqueue, cm, token); });
    if (!already_cancelled) {
      enqueue_attempts_.emplace_back(
          batch_size, callback, ctx, cm, token,
          [tuple, this, ctx](Attempt* attempt) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
            if (closed_) {
              attempt->context->SetStatus(
                  errors::Cancelled("PriorityQueue '", name_, "' is closed."));
              return kComplete;
            }
            RunResult result = kNoProgress;
            while (queues_[0].size() < static_cast<size_t>(capacity_)) {
              result = kProgress;
              const int index =
                  tuple[0].dim_size(0) - attempt->elements_requested;

              PersistentTensor priority_element;
              attempt->context->SetStatus(GetElementComponentFromBatch(
                  tuple, index, 0, attempt->context, &priority_element));
              if (!attempt->context->status().ok()) return kComplete;
              Tensor* priority_tensor = priority_element.AccessTensor(ctx);
              if (!TensorShapeUtils::IsScalar(priority_tensor->shape())) {
                attempt->context->SetStatus(errors::InvalidArgument(
                    "Expected the priority element to be a scalar, but "
                    "received shape: ",
                    priority_tensor->shape().DebugString()));
                return kComplete;
              }
              const int64 priority = priority_tensor->scalar<int64>()();
              for (int i = 0; i < num_components(); ++i) {
                PersistentTensor element;
                attempt->context->SetStatus(GetElementComponentFromBatch(
                    tuple, index, i, attempt->context, &element));
                if (!attempt->context->status().ok()) return kComplete;
                queues_[i].emplace(priority, element);
              }
              --attempt->elements_requested;
              if (attempt->elements_requested == 0) {
                return kComplete;
              }
            }
            return result;
          });
    }