| 67 | } |
| 68 | |
| 69 | void PaddingFIFOQueue::TryDequeueMany(int num_elements, OpKernelContext* ctx, |
| 70 | bool allow_small_batch, |
| 71 | CallbackWithTuple callback) { |
| 72 | if (num_elements == 0) { |
| 73 | Tuple tuple; |
| 74 | tuple.reserve(num_components()); |
| 75 | for (int i = 0; i < num_components(); ++i) { |
| 76 | // TODO(josh11b,misard): Switch to allocate_output(). |
| 77 | // See similar comment in fifo_queue.cc |
| 78 | Tensor element; |
| 79 | // Here, ManyOutShape returns zeros for undetermined shapes, |
| 80 | // which is exactly what we want to use. |
| 81 | OP_REQUIRES_OK(ctx, ctx->allocate_temp(component_dtypes_[i], |
| 82 | ManyOutShape(i, 0), &element)); |
| 83 | tuple.emplace_back(element); |
| 84 | } |
| 85 | callback(tuple); |
| 86 | return; |
| 87 | } |
| 88 | |
| 89 | CancellationManager* cm = ctx->cancellation_manager(); |
| 90 | CancellationToken token = cm->get_cancellation_token(); |
| 91 | bool already_cancelled; |
| 92 | { |
| 93 | mutex_lock l(mu_); |
| 94 | already_cancelled = !cm->RegisterCallback( |
| 95 | token, [this, cm, token]() { Cancel(kDequeue, cm, token); }); |
| 96 | if (!already_cancelled) { |
| 97 | // TODO(josh11b): This makes two copies of callback, avoid this if possible. |
| 98 | dequeue_attempts_.emplace_back( |
| 99 | num_elements, [callback]() { callback(Tuple()); }, ctx, cm, token, |
| 100 | [callback, allow_small_batch, |
| 101 | this](Attempt* attempt) EXCLUSIVE_LOCKS_REQUIRED(mu_) { |
| 102 | int32 queue_size = queues_[0].size(); |
| 103 | if (closed_ && queue_size < attempt->elements_requested) { |
| 104 | // If we don't have enough for a full dequeue, we have |
| 105 | // to reset the attempt tuple. |
| 106 | if (!attempt->tuples.empty()) { |
| 107 | // Restore already-dequeued elements to the front of the queue. |
| 108 | for (int64 i = attempt->tuples.size() - 1; i >= 0; --i) { |
| 109 | for (int j = 0; j < num_components(); ++j) { |
| 110 | PersistentTensor element; |
| 111 | Status s = GetElementComponent(attempt->tuples[i], j, |
| 112 | attempt->context, &element); |
| 113 | if (!s.ok()) { |
| 114 | attempt->context->SetStatus( |
| 115 | errors::DataLoss("Failed to restore element from " |
| 116 | "partially-dequeued batch " |
| 117 | "to PaddingFIFOQueue: ", |
| 118 | s.error_message())); |
| 119 | } |
| 120 | queues_[j].push_front(element); |
| 121 | } |
| 122 | } |
| 123 | } |
| 124 | if (allow_small_batch && !queues_[0].empty()) { |
| 125 | // Request all remaining elements in the queue. |
| 126 | queue_size = queues_[0].size(); |
nothing calls this directly
no test coverage detected