1 files changed, 24 insertions, 10 deletions

diff --git a/test/core/iomgr/udp_server_test.c b/test/core/iomgr/udp_server_test.c
index 9bea229466..ba572017a2 100644
--- a/test/core/iomgr/udp_server_test.c
+++ b/test/core/iomgr/udp_server_test.c
@@ -58,6 +58,7 @@
 static grpc_pollset *g_pollset;
 static gpr_mu *g_mu;
 static int g_number_of_reads = 0;
+static int g_number_of_writes = 0;
 static int g_number_of_bytes_read = 0;
 static int g_number_of_orphan_calls = 0;
 
@@ -78,7 +79,16 @@ static void on_read(grpc_exec_ctx *exec_ctx, grpc_fd *emfd,
   gpr_mu_unlock(g_mu);
 }
 
-static void on_fd_orphaned(grpc_fd *emfd) {
+static void on_write(grpc_exec_ctx *exec_ctx, grpc_fd *emfd) {
+  gpr_mu_lock(g_mu);
+  g_number_of_writes++;
+
+  GPR_ASSERT(
+      GRPC_LOG_IF_ERROR("pollset_kick", grpc_pollset_kick(g_pollset, NULL)));
+  gpr_mu_unlock(g_mu);
+}
+
+static void on_fd_orphaned(grpc_exec_ctx *exec_ctx, grpc_fd *emfd) {
   gpr_log(GPR_INFO, "gRPC FD about to be orphaned: %d",
           grpc_fd_wrapped_fd(emfd));
   g_number_of_orphan_calls++;
@@ -111,8 +121,8 @@ static void test_no_op_with_port(void) {
   memset(&resolved_addr, 0, sizeof(resolved_addr));
   resolved_addr.len = sizeof(struct sockaddr_in);
   addr->sin_family = AF_INET;
-  GPR_ASSERT(
-      grpc_udp_server_add_port(s, &resolved_addr, on_read, on_fd_orphaned));
+  GPR_ASSERT(grpc_udp_server_add_port(s, &resolved_addr, on_read, on_write,
+                                      on_fd_orphaned));
 
   grpc_udp_server_destroy(&exec_ctx, s, NULL);
   grpc_exec_ctx_finish(&exec_ctx);
@@ -132,8 +142,8 @@ static void test_no_op_with_port_and_start(void) {
   memset(&resolved_addr, 0, sizeof(resolved_addr));
   resolved_addr.len = sizeof(struct sockaddr_in);
   addr->sin_family = AF_INET;
-  GPR_ASSERT(
-      grpc_udp_server_add_port(s, &resolved_addr, on_read, on_fd_orphaned));
+  GPR_ASSERT(grpc_udp_server_add_port(s, &resolved_addr, on_read, on_write,
+                                      on_fd_orphaned));
 
   grpc_udp_server_start(&exec_ctx, s, NULL, 0, NULL);
 
@@ -164,8 +174,8 @@ static void test_receive(int number_of_clients) {
   memset(&resolved_addr, 0, sizeof(resolved_addr));
   resolved_addr.len = sizeof(struct sockaddr_storage);
   addr->ss_family = AF_INET;
-  GPR_ASSERT(
-      grpc_udp_server_add_port(s, &resolved_addr, on_read, on_fd_orphaned));
+  GPR_ASSERT(grpc_udp_server_add_port(s, &resolved_addr, on_read, on_write,
+                                      on_fd_orphaned));
 
   svrfd = grpc_udp_server_get_fd(s, 0);
   GPR_ASSERT(svrfd >= 0);
@@ -179,7 +189,7 @@ static void test_receive(int number_of_clients) {
   gpr_mu_lock(g_mu);
 
   for (i = 0; i < number_of_clients; i++) {
-    deadline = GRPC_TIMEOUT_SECONDS_TO_DEADLINE(10);
+    deadline = grpc_timeout_seconds_to_deadline(10);
 
     number_of_reads_before = g_number_of_reads;
     /* Create a socket, send a packet to the UDP server. */
@@ -212,6 +222,9 @@ static void test_receive(int number_of_clients) {
   /* The server had a single FD, which is orphaned once in *
    * deactivated_all_ports, and once in grpc_udp_server_destroy. */
   GPR_ASSERT(g_number_of_orphan_calls == 2);
+
+  /* The write callback should have fired a few times. */
+  GPR_ASSERT(g_number_of_writes > 0);
 }
 
 static void destroy_pollset(grpc_exec_ctx *exec_ctx, void *p,
@@ -234,11 +247,12 @@ int main(int argc, char **argv) {
   test_receive(1);
   test_receive(10);
 
-  grpc_closure_init(&destroyed, destroy_pollset, g_pollset);
+  grpc_closure_init(&destroyed, destroy_pollset, g_pollset,
+                    grpc_schedule_on_exec_ctx);
   grpc_pollset_shutdown(&exec_ctx, g_pollset, &destroyed);
   grpc_exec_ctx_finish(&exec_ctx);
   gpr_free(g_pollset);
-  grpc_iomgr_shutdown();
+  grpc_shutdown();
   return 0;
 }