Merge branch 'idempotent_endpoint_shutdown' into error

1 files changed, 28 insertions, 13 deletions

diff --git a/test/core/iomgr/endpoint_tests.c b/test/core/iomgr/endpoint_tests.c
index fb5c4b1949..b79c22e42a 100644
--- a/test/core/iomgr/endpoint_tests.c
+++ b/test/core/iomgr/endpoint_tests.c
@@ -223,7 +223,7 @@ static void read_and_write_test(grpc_endpoint_test_config config,
      even when bytes_written is unsigned. */
   state.bytes_written -= state.current_write_size;
   read_and_write_test_write_handler(&exec_ctx, &state, GRPC_ERROR_NONE);
-  grpc_exec_ctx_finish(&exec_ctx);
+  grpc_exec_ctx_flush(&exec_ctx);
 
   grpc_endpoint_read(&exec_ctx, state.read_ep, &state.incoming,
                      &state.done_read);
@@ -234,7 +234,7 @@ static void read_and_write_test(grpc_endpoint_test_config config,
     gpr_log(GPR_DEBUG, "shutdown write");
     grpc_endpoint_shutdown(&exec_ctx, state.write_ep);
   }
-  grpc_exec_ctx_finish(&exec_ctx);
+  grpc_exec_ctx_flush(&exec_ctx);
 
   gpr_mu_lock(g_mu);
   while (!state.read_done || !state.write_done) {
@@ -246,7 +246,7 @@ static void read_and_write_test(grpc_endpoint_test_config config,
                           gpr_now(GPR_CLOCK_MONOTONIC), deadline)));
   }
   gpr_mu_unlock(g_mu);
-  grpc_exec_ctx_finish(&exec_ctx);
+  grpc_exec_ctx_flush(&exec_ctx);
 
   end_test(config);
   gpr_slice_buffer_destroy(&state.outgoing);
@@ -261,6 +261,24 @@ static void inc_on_failure(grpc_exec_ctx *exec_ctx, void *arg,
   *(int *)arg += (error != GRPC_ERROR_NONE);
 }
 
+static void wait_for_fail_count(grpc_exec_ctx *exec_ctx, int *fail_count,
+                                int want_fail_count) {
+  grpc_exec_ctx_flush(exec_ctx);
+  for (int i = 0; i < 5 && *fail_count < want_fail_count; i++) {
+    grpc_pollset_worker *worker = NULL;
+    gpr_timespec now = gpr_now(GPR_CLOCK_REALTIME);
+    gpr_timespec deadline =
+        gpr_time_add(now, gpr_time_from_seconds(1, GPR_TIMESPAN));
+    gpr_mu_lock(g_mu);
+    GPR_ASSERT(GRPC_LOG_IF_ERROR(
+        "pollset_work",
+        grpc_pollset_work(exec_ctx, g_pollset, &worker, now, deadline)));
+    gpr_mu_unlock(g_mu);
+    grpc_exec_ctx_flush(exec_ctx);
+  }
+  GPR_ASSERT(*fail_count == want_fail_count);
+}
+
 static void multiple_shutdown_test(grpc_endpoint_test_config config) {
   grpc_endpoint_test_fixture f =
       begin_test(config, "multiple_shutdown_test", 128);
@@ -270,25 +288,21 @@ static void multiple_shutdown_test(grpc_endpoint_test_config config) {
   gpr_slice_buffer_init(&slice_buffer);
 
   grpc_exec_ctx exec_ctx = GRPC_EXEC_CTX_INIT;
+  grpc_endpoint_add_to_pollset(&exec_ctx, f.client_ep, g_pollset);
   grpc_endpoint_read(&exec_ctx, f.client_ep, &slice_buffer,
                      grpc_closure_create(inc_on_failure, &fail_count));
-  grpc_exec_ctx_flush(&exec_ctx);
-  GPR_ASSERT(fail_count == 0);
+  wait_for_fail_count(&exec_ctx, &fail_count, 0);
   grpc_endpoint_shutdown(&exec_ctx, f.client_ep);
-  grpc_exec_ctx_flush(&exec_ctx);
-  GPR_ASSERT(fail_count == 1);
+  wait_for_fail_count(&exec_ctx, &fail_count, 1);
   grpc_endpoint_read(&exec_ctx, f.client_ep, &slice_buffer,
                      grpc_closure_create(inc_on_failure, &fail_count));
-  grpc_exec_ctx_flush(&exec_ctx);
-  GPR_ASSERT(fail_count == 2);
+  wait_for_fail_count(&exec_ctx, &fail_count, 2);
   gpr_slice_buffer_add(&slice_buffer, gpr_slice_from_copied_string("a"));
   grpc_endpoint_write(&exec_ctx, f.client_ep, &slice_buffer,
                       grpc_closure_create(inc_on_failure, &fail_count));
-  grpc_exec_ctx_flush(&exec_ctx);
-  GPR_ASSERT(fail_count == 3);
+  wait_for_fail_count(&exec_ctx, &fail_count, 3);
   grpc_endpoint_shutdown(&exec_ctx, f.client_ep);
-  grpc_exec_ctx_flush(&exec_ctx);
-  GPR_ASSERT(fail_count == 3);
+  wait_for_fail_count(&exec_ctx, &fail_count, 3);
 
   gpr_slice_buffer_destroy(&slice_buffer);
 
@@ -310,4 +324,5 @@ void grpc_endpoint_tests(grpc_endpoint_test_config config,
     read_and_write_test(config, 40320, i, i, false);
   }
   g_pollset = NULL;
+  g_mu = NULL;
 }