Force the client to shutdown so that the server write can complete

2 files changed, 34 insertions, 15 deletions

diff --git a/test/core/bad_client/bad_client.cc b/test/core/bad_client/bad_client.cc
index f5fb82e5f0..6a35ffcb6d 100644
--- a/test/core/bad_client/bad_client.cc
+++ b/test/core/bad_client/bad_client.cc
@@ -75,6 +75,17 @@ static void set_read_done(void* arg, grpc_error* error) {
   gpr_event_set(read_done, (void*)1);
 }
 
+/* shutdown client */
+static void shutdown_client(grpc_endpoint** client_fd) {
+  if (*client_fd != nullptr) {
+    grpc_endpoint_shutdown(
+        *client_fd, GRPC_ERROR_CREATE_FROM_STATIC_STRING("Forced Disconnect"));
+    grpc_endpoint_destroy(*client_fd);
+    grpc_core::ExecCtx::Get()->Flush();
+    *client_fd = nullptr;
+  }
+}
+
 /* Runs client side validator */
 void grpc_run_client_side_validator(grpc_bad_client_arg* arg, uint32_t flags,
                                     grpc_endpoint_pair* sfd,
@@ -115,11 +126,7 @@ void grpc_run_client_side_validator(grpc_bad_client_arg* arg, uint32_t flags,
   }
 
   if (flags & GRPC_BAD_CLIENT_DISCONNECT) {
-    grpc_endpoint_shutdown(
-        sfd->client, GRPC_ERROR_CREATE_FROM_STATIC_STRING("Forced Disconnect"));
-    grpc_endpoint_destroy(sfd->client);
-    grpc_core::ExecCtx::Get()->Flush();
-    sfd->client = nullptr;
+    shutdown_client(&sfd->client);
   }
 
   if (sfd->client != nullptr) {
@@ -158,6 +165,19 @@ void grpc_run_client_side_validator(grpc_bad_client_arg* arg, uint32_t flags,
     grpc_core::ExecCtx::Get()->Flush();
   }
 
+  /* If flags is non-zero, it is time to shutdown the client */
+  if (flags) {
+    shutdown_client(&sfd->client);
+  }
+
+  /* Make sure that the server is done writing */
+  while (!gpr_event_get(&done_write)) {
+    GPR_ASSERT(
+        grpc_completion_queue_next(
+            client_cq, grpc_timeout_milliseconds_to_deadline(100), nullptr)
+            .type == GRPC_QUEUE_TIMEOUT);
+  }
+
   grpc_slice_buffer_destroy_internal(&outgoing);
   grpc_core::ExecCtx::Get()->Flush();
 }
@@ -206,18 +226,14 @@ void grpc_run_bad_client_test(
   /* Start validator */
   gpr_thd_new(&id, "grpc_bad_client", thd_func, &a, nullptr);
   for (int i = 0; i < num_args; i++) {
-    grpc_run_client_side_validator(&args[i], flags, &sfd, client_cq);
+    grpc_run_client_side_validator(&args[i], i == (num_args - 1) ? flags : 0,
+                                   &sfd, client_cq);
   }
   /* Wait for server thread to finish */
   GPR_ASSERT(gpr_event_wait(&a.done_thd, grpc_timeout_seconds_to_deadline(1)));
 
   /* Shutdown. */
-  if (sfd.client != nullptr) {
-    grpc_endpoint_shutdown(
-        sfd.client, GRPC_ERROR_CREATE_FROM_STATIC_STRING("Test Shutdown"));
-    grpc_endpoint_destroy(sfd.client);
-    grpc_core::ExecCtx::Get()->Flush();
-  }
+  shutdown_client(&sfd.client);
 
   shutdown_cq = grpc_completion_queue_create_for_pluck(nullptr);
   grpc_server_shutdown_and_notify(a.server, shutdown_cq, nullptr);
diff --git a/test/core/bad_client/bad_client.h b/test/core/bad_client/bad_client.h
index bf73a78886..46f25269ca 100644
--- a/test/core/bad_client/bad_client.h
+++ b/test/core/bad_client/bad_client.h
@@ -49,15 +49,18 @@ struct grpc_bad_client_arg {
 
 /* Test runner.
  *
- *  Create a server, and for each arg in \a args send client_payload. For each
- *  payload, run client_validator to make sure that the response is as expected.
- *  Also execute \a server_validator in a separate thread to assert that the
+ * Create a server, and for each arg in \a args send client_payload. For each
+ * payload, run client_validator to make sure that the response is as expected.
+ * Also execute \a server_validator in a separate thread to assert that the
  * bytes are handled as expected.
  */
 void grpc_run_bad_client_test(
     grpc_bad_client_server_side_validator server_validator,
     grpc_bad_client_arg args[], int num_args, uint32_t flags);
 
+/* A hack to let old tests work as before. In these tests, instead of an array,
+ * the tests provide a single client_validator and payload
+ */
 #define COMBINE1(X, Y) X##Y
 #define COMBINE(X, Y) COMBINE1(X, Y)