| 69 | } |
| 70 | |
| 71 | void PriorityQueue::TryEnqueue(const Tuple& tuple, OpKernelContext* ctx, |
| 72 | DoneCallback callback) { |
| 73 | CancellationManager* cm = ctx->cancellation_manager(); |
| 74 | CancellationToken token = cm->get_cancellation_token(); |
| 75 | bool already_cancelled; |
| 76 | { |
| 77 | mutex_lock l(mu_); |
| 78 | already_cancelled = !cm->RegisterCallback( |
| 79 | token, [this, cm, token]() { Cancel(kEnqueue, cm, token); }); |
| 80 | if (!already_cancelled) { |
| 81 | enqueue_attempts_.emplace_back( |
| 82 | 1, callback, ctx, cm, token, |
| 83 | [tuple, this](Attempt* attempt) EXCLUSIVE_LOCKS_REQUIRED(mu_) { |
| 84 | if (closed_) { |
| 85 | attempt->context->SetStatus( |
| 86 | errors::Cancelled("PriorityQueue '", name_, "' is closed.")); |
| 87 | return kComplete; |
| 88 | } |
| 89 | if (queues_[0].size() < static_cast<size_t>(capacity_)) { |
| 90 | if (!TensorShapeUtils::IsScalar(tuple[0].shape())) { |
| 91 | attempt->context->SetStatus(errors::InvalidArgument( |
| 92 | "Expected the priority element to be a scalar, but " |
| 93 | "received shape: ", |
| 94 | tuple[0].shape().DebugString())); |
| 95 | return kComplete; |
| 96 | } |
| 97 | const int64 priority = tuple[0].scalar<int64>()(); |
| 98 | for (int i = 0; i < num_components(); ++i) { |
| 99 | queues_[i].emplace(priority, PersistentTensor(tuple[i])); |
| 100 | } |
| 101 | return kComplete; |
| 102 | } else { |
| 103 | return kNoProgress; |
| 104 | } |
| 105 | }); |
| 106 | } |
| 107 | } |
| 108 | if (!already_cancelled) { |
| 109 | FlushUnlocked(); |
| 110 | } else { |
| 111 | ctx->SetStatus(errors::Cancelled("Enqueue operation was cancelled")); |
| 112 | callback(); |
| 113 | } |
| 114 | } |
| 115 | |
| 116 | /* static */ |
| 117 | Status PriorityQueue::GetElementComponentFromBatch( |
nothing calls this directly
no test coverage detected