| 51 | ~TensorBuf() { Cancel(); } |
| 52 | |
| 53 | Status Put(const std::vector<Tensor>& record, int64 timeout_millis) { |
| 54 | std::unique_lock<std::mutex> lock(mu_); |
| 55 | |
| 56 | bool should_retry = !put_cv_.wait_for( |
| 57 | lock, std::chrono::milliseconds(timeout_millis), |
| 58 | [this]() { return buffer_.size() < capacity_ || is_cancelled_; }); |
| 59 | if (should_retry) { |
| 60 | lock.unlock(); |
| 61 | LOG(WARNING) << "Prefetching was ignored since timeout."; |
| 62 | return Status::OK(); |
| 63 | } |
| 64 | |
| 65 | if (TF_PREDICT_FALSE(is_cancelled_)) { |
| 66 | lock.unlock(); |
| 67 | return Status(errors::Cancelled("Session was closed.")); |
| 68 | } |
| 69 | |
| 70 | buffer_.push_back(std::move(record)); |
| 71 | |
| 72 | lock.unlock(); |
| 73 | take_cv_.notify_all(); |
| 74 | |
| 75 | return Status::OK(); |
| 76 | } |
| 77 | |
| 78 | Status Take(std::vector<Tensor>* record) { |
| 79 | std::unique_lock<std::mutex> lock(mu_); |