| 76 | } |
| 77 | |
| 78 | int64 ReaderBase::ReadUpTo(const int64 num_records, QueueInterface* queue, |
| 79 | std::vector<string>* keys, |
| 80 | std::vector<string>* values, |
| 81 | OpKernelContext* context) { |
| 82 | mutex_lock lock(mu_); |
| 83 | int64 records_produced_this_call = 0; |
| 84 | while (true) { |
| 85 | // Records produced by this iteration of the ReadUpToLocked call. |
| 86 | int64 num_records_produced = 0; |
| 87 | int64 remaining = num_records - records_produced_this_call; |
| 88 | if (remaining == 0) { |
| 89 | return records_produced_this_call; |
| 90 | } |
| 91 | if (!work_in_progress()) { |
| 92 | work_ = GetNextWorkLocked(queue, context); |
| 93 | if (!context->status().ok()) { |
| 94 | return records_produced_this_call; |
| 95 | } |
| 96 | Status status = OnWorkStartedLocked(); |
| 97 | if (status.ok()) { |
| 98 | work_started_++; |
| 99 | } else { |
| 100 | context->SetStatus(status); |
| 101 | return records_produced_this_call; |
| 102 | } |
| 103 | } |
| 104 | bool at_end = false; |
| 105 | |
| 106 | Status status = |
| 107 | ReadUpToLocked(remaining, keys, values, &num_records_produced, &at_end); |
| 108 | // This call so far. |
| 109 | records_produced_this_call += num_records_produced; |
| 110 | |
| 111 | // In total, over the lifetime of the ReaderBase. |
| 112 | num_records_produced_ += num_records_produced; |
| 113 | |
| 114 | if (!at_end && status.ok() && num_records_produced == 0) { |
| 115 | status = errors::Internal( |
| 116 | "ReadManyLocked() for ", name(), |
| 117 | " must set *at_end=true, *num_produced > 0 or return an error."); |
| 118 | context->SetStatus(status); |
| 119 | return records_produced_this_call; |
| 120 | } |
| 121 | if (status.ok() && at_end) { |
| 122 | status = OnWorkFinishedLocked(); |
| 123 | work_finished_ = work_started_; |
| 124 | if (records_produced_this_call > 0) { |
| 125 | return records_produced_this_call; |
| 126 | } |
| 127 | } |
| 128 | if (!status.ok()) { |
| 129 | context->SetStatus(status); |
| 130 | return records_produced_this_call; |
| 131 | } |
| 132 | } |
| 133 | } |
| 134 | |
| 135 | // Default implementation just reads one record at a time. |