| 87 | |
| 88 | template <typename T> |
| 89 | void TryInsertMany(const Tensor& keys, int component_index, |
| 90 | const Tensor& values, OpKernelContext* ctx, |
| 91 | const DoneCallback& callback) { |
| 92 | TensorShape element_shape = values.shape(); |
| 93 | OP_REQUIRES_ASYNC( |
| 94 | ctx, keys.NumElements() == 0 || element_shape.num_elements() > 0, |
| 95 | errors::InvalidArgument("Tensors with no elements are not supported ", |
| 96 | name_, ": received shape ", |
| 97 | element_shape.DebugString()), |
| 98 | callback); |
| 99 | if (element_shape.dims() > 0) element_shape.RemoveDim(0); |
| 100 | const std::size_t num_inserted = keys.NumElements(); |
| 101 | |
| 102 | // For each key, update the corresponding incomplete tuple with the |
| 103 | // the corresponding given value at component_index. |
| 104 | // This will be passed to the final callback at the very end. |
| 105 | bool new_elements = false; |
| 106 | |
| 107 | // Will be used for the final insert into the queue. |
| 108 | Tuple insert_tuple; |
| 109 | |
| 110 | { |
| 111 | mutex_lock lock(mu_); |
| 112 | if (closed_) { |
| 113 | OP_REQUIRES_ASYNC( |
| 114 | ctx, |
| 115 | !cancel_pending_enqueues_ && |
| 116 | (num_inserted == 0 || !incomplete_.empty()), |
| 117 | errors::Cancelled( |
| 118 | "Barrier ", name_, " is closed. Pending enqueues cancelled: ", |
| 119 | cancel_pending_enqueues_, |
| 120 | ". Number of new insertions: ", num_inserted, |
| 121 | ". Number of incomplete keys: ", incomplete_.size(), "."), |
| 122 | callback); |
| 123 | } |
| 124 | |
| 125 | // Step 1: insert into the incomplete map and identify which |
| 126 | // entries are, in fact, complete and ready for enqueueing. Store |
| 127 | // them in a vector |
| 128 | std::vector<Tuple> ready_tuples; |
| 129 | |
| 130 | for (int i = 0; i < num_inserted; ++i) { |
| 131 | OP_REQUIRES_OK_ASYNC( |
| 132 | ctx, |
| 133 | InsertOneLocked<T>(ctx, keys, values, element_shape, |
| 134 | component_index, i, &ready_tuples, |
| 135 | &new_elements), |
| 136 | callback); |
| 137 | } |
| 138 | |
| 139 | if (new_elements) ++input_index_; |
| 140 | |
| 141 | // This probably won't happen before the heat death of the |
| 142 | // universe, but who knows? Moore's law FTW. |
| 143 | OP_REQUIRES_ASYNC( |
| 144 | ctx, input_index_ != std::numeric_limits<int64>::max(), |
| 145 | errors::Internal( |
| 146 | "Barrier has had ", input_index_, |
nothing calls this directly
no test coverage detected