Method ComputeAsync

tensorflow/core/kernels/barrier_ops.cc:581–621

 protected:
  void ComputeAsync(OpKernelContext* ctx, Barrier* barrier,
                    DoneCallback callback) override {
    const Tensor* Tnum_elements;
    OP_REQUIRES_OK_ASYNC(ctx, ctx->input("num_elements", &Tnum_elements),
                         callback);
    OP_REQUIRES_ASYNC(ctx, TensorShapeUtils::IsScalar(Tnum_elements->shape()),
                      errors::InvalidArgument("num_elements must be a scalar."),
                      callback);
    const int32 num_elements = Tnum_elements->scalar<int32>()();

    DataTypeVector expected_inputs = {DT_STRING_REF, DT_INT32};
    // The first output is the insertion index, the second output is the key.
    DataTypeVector expected_outputs = {DT_INT64, DT_STRING};
    for (DataType dt : barrier->component_types()) {
      expected_outputs.push_back(dt);
    }
    OP_REQUIRES_OK_ASYNC(
        ctx, ctx->MatchSignature(expected_inputs, expected_outputs), callback);

    barrier->TryTakeMany(
        num_elements, allow_small_batch_, timeout_, ctx,
        [ctx, callback](const Tensor& indices, const Tensor& keys,
                        const Barrier::Tuple& values) {
          if (!ctx->status().ok()) {
            callback();
            return;
          }
          // At this point, indices, keys, and values
          // have all been written to successfully.
          OP_REQUIRES_OK_ASYNC(ctx, ctx->set_output("indices", indices),
                               callback);
          OP_REQUIRES_OK_ASYNC(ctx, ctx->set_output("keys", keys), callback);
          OpOutputList values_output;
          OP_REQUIRES_OK_ASYNC(ctx, ctx->output_list("values", &values_output),
                               callback);
          for (size_t i = 0; i < values.size(); ++i) {
            values_output.set(i, values[i]);
          }
          callback();
        });
  }

 private:
  int64 timeout_;