| 262 | num_remaining_senders_(num_senders) { } |
| 263 | |
| 264 | Status KrpcDataStreamRecvr::SenderQueue::GetBatch(RowBatch** next_batch) { |
| 265 | SCOPED_TIMER(recvr_->queue_get_batch_timer_); |
| 266 | DCHECK(TestInfo::is_test() || FragmentInstanceState::IsFragmentExecThread()); |
| 267 | DCHECK(!recvr_->closed_); |
| 268 | int num_to_dequeue = 0; |
| 269 | // The sender id is set below when we decide to dequeue entries from 'deferred_rpcs_'. |
| 270 | int sender_id = -1; |
| 271 | { |
| 272 | unique_lock<SpinLock> l(lock_); |
| 273 | // current_batch_ must be replaced with the returned batch. |
| 274 | current_batch_.reset(); |
| 275 | *next_batch = nullptr; |
| 276 | |
| 277 | // Wait until something shows up or we know we're done |
| 278 | while (batch_queue_.empty() && status_.ok() && !is_cancelled_ && |
| 279 | num_remaining_senders_ > 0) { |
| 280 | DCHECK_EQ(pending_deferred_rpcs_.size(), num_deserialize_tasks_pending_); |
| 281 | // Verify before waiting on 'data_arrival_cv_' that if there are any deferred |
| 282 | // batches, either there is outstanding deserialization request queued or there |
| 283 | // is pending insertion so this thread is guaranteed to wake up at some point. |
| 284 | DCHECK(deferred_rpcs_.empty() || |
| 285 | (num_deserialize_tasks_pending_ + num_pending_enqueue_) > 0); |
| 286 | VLOG_ROW << "wait arrival fragment_instance_id=" |
| 287 | << PrintId(recvr_->fragment_instance_id()) |
| 288 | << " node=" << recvr_->dest_node_id(); |
| 289 | // Don't count time spent waiting on the sender as active time. |
| 290 | CANCEL_SAFE_SCOPED_TIMER3(recvr_->data_wait_timer_, recvr_->inactive_timer_, |
| 291 | received_first_batch_ ? nullptr : recvr_->first_batch_wait_total_timer_, |
| 292 | &is_cancelled_); |
| 293 | data_arrival_cv_.wait(l); |
| 294 | } |
| 295 | |
| 296 | // Return early if there is any error when inserting row batches. |
| 297 | RETURN_IF_ERROR(status_); |
| 298 | |
| 299 | if (UNLIKELY(is_cancelled_)) { |
| 300 | // Cancellation should have drained the entire 'deferred_rpcs_' queue. |
| 301 | // Make sure the senders were replied to or they may be stuck waiting for a reply. |
| 302 | DCHECK(!HasDeferredRpcs()); |
| 303 | return Status::CANCELLED; |
| 304 | } |
| 305 | |
| 306 | // All senders have sent their row batches. Nothing to do. |
| 307 | if (num_remaining_senders_ == 0 && batch_queue_.empty()) { |
| 308 | // Note that it's an invariant that a sender cannot send the EOS RPC until all |
| 309 | // outstanding TransmitData() RPCs have been replied to. Therefore, it should be |
| 310 | // impossible for num_remaining_senders_ to reach 0 before all RPCs in |
| 311 | // 'deferred_rpcs_' have been replied to. |
| 312 | DCHECK(!HasDeferredRpcs()); |
| 313 | DCHECK_EQ(num_pending_enqueue_, 0); |
| 314 | return Status::OK(); |
| 315 | } |
| 316 | |
| 317 | DCHECK(!batch_queue_.empty()); |
| 318 | received_first_batch_ = true; |
| 319 | RowBatch* result = batch_queue_.front().second.release(); |
| 320 | int64_t batch_size = batch_queue_.front().first; |
| 321 | COUNTER_ADD(recvr_->bytes_dequeued_counter_, batch_size); |