| 1265 | } |
| 1266 | |
| 1267 | Status KrpcDataStreamSender::FlushFinal(RuntimeState* state) { |
| 1268 | SCOPED_TIMER(profile()->total_time_counter()); |
| 1269 | DCHECK(!flushed_); |
| 1270 | DCHECK(!closed_); |
| 1271 | flushed_ = true; |
| 1272 | |
| 1273 | // Send out the final row batches and EOS signals on all channels in parallel. |
| 1274 | // If we hit an error here, we can return without closing the remaining channels as |
| 1275 | // the error is propagated back to the coordinator, which in turn cancels the query, |
| 1276 | // which will cause the remaining open channels to be closed. |
| 1277 | for (auto& [ch, ice_ch] : channel_to_ice_channel_) { |
| 1278 | RETURN_IF_ERROR(ice_ch->Flush()); |
| 1279 | } |
| 1280 | for (PartitionRowCollector& collector: partition_row_collectors_) { |
| 1281 | RETURN_IF_ERROR(collector.SendCurrentBatch()); |
| 1282 | } |
| 1283 | for (unique_ptr<Channel>& channel : channels_) { |
| 1284 | RETURN_IF_ERROR(channel->WaitForRpc()); |
| 1285 | } |
| 1286 | for (unique_ptr<Channel>& channel : channels_) { |
| 1287 | RETURN_IF_ERROR(channel->SendEosAsync()); |
| 1288 | } |
| 1289 | for (unique_ptr<Channel>& channel : channels_) { |
| 1290 | RETURN_IF_ERROR(channel->WaitForRpc()); |
| 1291 | } |
| 1292 | return Status::OK(); |
| 1293 | } |
| 1294 | |
| 1295 | void KrpcDataStreamSender::Close(RuntimeState* state) { |
| 1296 | SCOPED_TIMER(profile()->total_time_counter()); |