Merge branch 'atomic-timers' of github.com:ctiller/grpc into atomic-timers

1 files changed, 20 insertions, 2 deletions

diff --git a/test/cpp/end2end/test_service_impl.cc b/test/cpp/end2end/test_service_impl.cc
index 59d36e9cb5..11729c425c 100644
--- a/test/cpp/end2end/test_service_impl.cc
+++ b/test/cpp/end2end/test_service_impl.cc
@@ -246,6 +246,9 @@ Status TestServiceImpl::ResponseStream(ServerContext* context,
   int server_try_cancel = GetIntValueFromMetadata(
       kServerTryCancelRequest, context->client_metadata(), DO_NOT_CANCEL);
 
+  int server_coalescing_api = GetIntValueFromMetadata(
+      kServerUseCoalescingApi, context->client_metadata(), 0);
+
   if (server_try_cancel == CANCEL_BEFORE_PROCESSING) {
     ServerTryCancel(context);
     return Status::CANCELLED;
@@ -260,7 +263,11 @@ Status TestServiceImpl::ResponseStream(ServerContext* context,
 
   for (int i = 0; i < kNumResponseStreamsMsgs; i++) {
     response.set_message(request->message() + grpc::to_string(i));
-    writer->Write(response);
+    if (i == kNumResponseStreamsMsgs - 1 && server_coalescing_api != 0) {
+      writer->WriteLast(response, WriteOptions());
+    } else {
+      writer->Write(response);
+    }
   }
 
   if (server_try_cancel_thd != nullptr) {
@@ -305,10 +312,21 @@ Status TestServiceImpl::BidiStream(
         new std::thread(&TestServiceImpl::ServerTryCancel, this, context);
   }
 
+  // kServerFinishAfterNReads suggests after how many reads, the server should
+  // write the last message and send status (coalesced using WriteLast)
+  int server_write_last = GetIntValueFromMetadata(
+      kServerFinishAfterNReads, context->client_metadata(), 0);
+
+  int read_counts = 0;
   while (stream->Read(&request)) {
+    read_counts++;
     gpr_log(GPR_INFO, "recv msg %s", request.message().c_str());
     response.set_message(request.message());
-    stream->Write(response);
+    if (read_counts == server_write_last) {
+      stream->WriteLast(response, WriteOptions());
+    } else {
+      stream->Write(response);
+    }
   }
 
   if (server_try_cancel_thd != nullptr) {