Method TryTakeMany

tensorflow/core/kernels/barrier_ops.cc:201–251

  }

  void TryTakeMany(int num_elements, bool allow_small_batch, int64 timeout,
                   OpKernelContext* ctx,
                   const IndicesKeysValuesCallback& callback) {
    int num_elements_to_deliver = num_elements;
    {
      mutex_lock lock(mu_);
      if (closed_) {
        int available_elements = ready_size();
        if (allow_small_batch) {
          // We want to deliver a maximum of num_elements, if there are less
          // elements available, we deliver at most the available_elements. If
          // there are no
          // elements available, a call to TryTakeMany should fail with
          // OutOfRange. We trigger this error by setting the request here to 1.
          num_elements_to_deliver = std::min(num_elements, available_elements);
        } else {
          // We're happy to wait for additional elements to be completed.
          available_elements += incomplete_.size();
        }
        // If there are 0 available elements or less elements than the
        // number we can deliver, then we are done.
        if (available_elements < std::max(num_elements_to_deliver, 1)) {
          ctx->SetStatus(errors::OutOfRange(
              "Barrier '", name_, "' is closed and has ",
              "insufficient elements (requested ", num_elements_to_deliver,
              ", total size ", available_elements, ")"));
          callback(Tensor(DT_INT64), Tensor(DT_STRING), Tuple());
          return;
        }
      }
    }

    ready_queue_->TryDequeueMany(
        num_elements_to_deliver, ctx, allow_small_batch,
        [this, ctx, callback](const Tuple& t) {
          Tensor indices(DT_INT64);
          Tensor keys(DT_STRING);
          Tuple values;

          if (!ctx->status().ok()) {
            callback(indices, keys, values);
            return;
          }

          CHECK_EQ(t.size(), 2 + num_components());
          indices = t[0];
          keys = t[1];
          values.insert(values.begin(), t.begin() + 2, t.end());
          callback(indices, keys, values);
        });
  }

  void Close(OpKernelContext* ctx, bool cancel_pending_enqueues,
             const DoneCallback& callback) {