Add an on_write callback to the UDP server.

1 files changed, 19 insertions, 6 deletions

diff --git a/test/core/iomgr/udp_server_test.c b/test/core/iomgr/udp_server_test.c
index 9ec6d7bf20..ba7a52ea21 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,6 +79,15 @@ static void on_read(grpc_exec_ctx *exec_ctx, grpc_fd *emfd,
   gpr_mu_unlock(g_mu);
 }
 
+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_fd *emfd) {
   gpr_log(GPR_INFO, "gRPC FD about to be orphaned: %d",
           grpc_fd_wrapped_fd(emfd));
@@ -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);
@@ -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,