| 151 | } |
| 152 | |
| 153 | void ReaderBase::Read(QueueInterface* queue, string* key, string* value, |
| 154 | OpKernelContext* context) { |
| 155 | mutex_lock lock(mu_); |
| 156 | while (true) { |
| 157 | if (!work_in_progress()) { |
| 158 | work_ = GetNextWorkLocked(queue, context); |
| 159 | if (!context->status().ok()) { |
| 160 | return; |
| 161 | } |
| 162 | Status status = OnWorkStartedLocked(); |
| 163 | if (status.ok()) { |
| 164 | work_started_++; |
| 165 | } else { |
| 166 | context->SetStatus(status); |
| 167 | return; |
| 168 | } |
| 169 | } |
| 170 | |
| 171 | bool produced = false; |
| 172 | bool at_end = false; |
| 173 | Status status = ReadLocked(key, value, &produced, &at_end); |
| 174 | |
| 175 | if (!at_end && status.ok() && !produced) { |
| 176 | status = errors::Internal( |
| 177 | "ReadLocked() for ", name(), |
| 178 | " must set *at_end=true, *produced=true, or return an error."); |
| 179 | } |
| 180 | if (!status.ok() && produced) { |
| 181 | status = errors::Internal( |
| 182 | "ReadLocked() for ", name(), |
| 183 | " set *produced=true *and* returned an error: ", status.ToString()); |
| 184 | } |
| 185 | if (status.ok() && at_end) { |
| 186 | status = OnWorkFinishedLocked(); |
| 187 | work_finished_ = work_started_; |
| 188 | } |
| 189 | if (!status.ok()) { |
| 190 | context->SetStatus(status); |
| 191 | return; |
| 192 | } |
| 193 | if (produced) { |
| 194 | ++num_records_produced_; |
| 195 | return; |
| 196 | } |
| 197 | } |
| 198 | } |
| 199 | |
| 200 | string ReaderBase::GetNextWorkLocked(QueueInterface* queue, |
| 201 | OpKernelContext* context) const { |