Port endpoint tests to pollsets

4 files changed, 96 insertions, 86 deletions

diff --git a/test/core/iomgr/endpoint_tests.c b/test/core/iomgr/endpoint_tests.c
index e63048e26c..d68f8a4f8c 100644
--- a/test/core/iomgr/endpoint_tests.c
+++ b/test/core/iomgr/endpoint_tests.c
@@ -57,6 +57,8 @@
 
 */
 
+static grpc_pollset *g_pollset;
+
 size_t count_and_unref_slices(gpr_slice *slices, size_t nslices,
                               int *current_data) {
   size_t num_bytes = 0;
@@ -111,8 +113,6 @@ static gpr_slice *allocate_blocks(size_t num_bytes, size_t slice_size,
 struct read_and_write_test_state {
   grpc_endpoint *read_ep;
   grpc_endpoint *write_ep;
-  gpr_mu mu;
-  gpr_cv cv;
   size_t target_bytes;
   size_t bytes_read;
   size_t current_write_size;
@@ -130,10 +130,10 @@ static void read_and_write_test_read_handler(void *data, gpr_slice *slices,
   GPR_ASSERT(error != GRPC_ENDPOINT_CB_ERROR);
   if (error == GRPC_ENDPOINT_CB_SHUTDOWN) {
     gpr_log(GPR_INFO, "Read handler shutdown");
-    gpr_mu_lock(&state->mu);
+    gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
     state->read_done = 1;
-    gpr_cv_signal(&state->cv);
-    gpr_mu_unlock(&state->mu);
+    grpc_pollset_kick(g_pollset);
+    gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
     return;
   }
 
@@ -141,10 +141,10 @@ static void read_and_write_test_read_handler(void *data, gpr_slice *slices,
       count_and_unref_slices(slices, nslices, &state->current_read_data);
   if (state->bytes_read == state->target_bytes) {
     gpr_log(GPR_INFO, "Read handler done");
-    gpr_mu_lock(&state->mu);
+    gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
     state->read_done = 1;
-    gpr_cv_signal(&state->cv);
-    gpr_mu_unlock(&state->mu);
+    grpc_pollset_kick(g_pollset);
+    gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
   } else {
     grpc_endpoint_notify_on_read(state->read_ep,
                                  read_and_write_test_read_handler, data);
@@ -164,10 +164,10 @@ static void read_and_write_test_write_handler(void *data,
 
   if (error == GRPC_ENDPOINT_CB_SHUTDOWN) {
     gpr_log(GPR_INFO, "Write handler shutdown");
-    gpr_mu_lock(&state->mu);
+    gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
     state->write_done = 1;
-    gpr_cv_signal(&state->cv);
-    gpr_mu_unlock(&state->mu);
+    grpc_pollset_kick(g_pollset);
+    gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
     return;
   }
 
@@ -198,10 +198,10 @@ static void read_and_write_test_write_handler(void *data,
   GPR_ASSERT(state->bytes_written == state->target_bytes);
 
   gpr_log(GPR_INFO, "Write handler done");
-  gpr_mu_lock(&state->mu);
+  gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
   state->write_done = 1;
-  gpr_cv_signal(&state->cv);
-  gpr_mu_unlock(&state->mu);
+  grpc_pollset_kick(g_pollset);
+  gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
 }
 
 /* Do both reading and writing using the grpc_endpoint API.
@@ -222,9 +222,6 @@ static void read_and_write_test(grpc_endpoint_test_config config,
             num_bytes, slice_size);
   }
 
-  gpr_mu_init(&state.mu);
-  gpr_cv_init(&state.cv);
-
   state.read_ep = f.client_ep;
   state.write_ep = f.server_ep;
   state.target_bytes = num_bytes;
@@ -253,29 +250,24 @@ static void read_and_write_test(grpc_endpoint_test_config config,
     grpc_endpoint_shutdown(state.write_ep);
   }
 
-  gpr_mu_lock(&state.mu);
+  gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
   while (!state.read_done || !state.write_done) {
-    if (gpr_cv_wait(&state.cv, &state.mu, deadline)) {
-      gpr_log(GPR_ERROR, "timeout: read_done=%d, write_done=%d",
-              state.read_done, state.write_done);
-      abort();
-    }
+    GPR_ASSERT(gpr_time_cmp(gpr_now(), deadline) < 0);
+    grpc_pollset_work(g_pollset, deadline);
   }
-  gpr_mu_unlock(&state.mu);
+  gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
 
   grpc_endpoint_destroy(state.read_ep);
   grpc_endpoint_destroy(state.write_ep);
-  gpr_mu_destroy(&state.mu);
-  gpr_cv_destroy(&state.cv);
   end_test(config);
 }
 
 struct timeout_test_state {
-  gpr_event io_done;
+  int io_done;
 };
 
 typedef struct {
-  gpr_event ev;
+  int done;
   grpc_endpoint *ep;
 } shutdown_during_write_test_state;
 
@@ -291,7 +283,10 @@ static void shutdown_during_write_test_read_handler(
 
   if (error != GRPC_ENDPOINT_CB_OK) {
     grpc_endpoint_destroy(st->ep);
-    gpr_event_set(&st->ev, (void *)(gpr_intptr) error);
+    gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
+    st->done = error;
+    grpc_pollset_kick(g_pollset);
+    gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
   } else {
     grpc_endpoint_notify_on_read(
         st->ep, shutdown_during_write_test_read_handler, user_data);
@@ -310,7 +305,10 @@ static void shutdown_during_write_test_write_handler(
     gpr_log(GPR_ERROR,
             "shutdown_during_write_test_write_handler completed unexpectedly");
   }
-  gpr_event_set(&st->ev, (void *)(gpr_intptr) 1);
+  gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
+  st->done = 1;
+  grpc_pollset_kick(g_pollset);
+  gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
 }
 
 static void shutdown_during_write_test(grpc_endpoint_test_config config,
@@ -329,8 +327,8 @@ static void shutdown_during_write_test(grpc_endpoint_test_config config,
 
   read_st.ep = f.client_ep;
   write_st.ep = f.server_ep;
-  gpr_event_init(&read_st.ev);
-  gpr_event_init(&write_st.ev);
+  read_st.done = 0;
+  write_st.done = 0;
 
   grpc_endpoint_notify_on_read(
       read_st.ep, shutdown_during_write_test_read_handler, &read_st);
@@ -347,9 +345,19 @@ static void shutdown_during_write_test(grpc_endpoint_test_config config,
       case GRPC_ENDPOINT_WRITE_PENDING:
         grpc_endpoint_shutdown(write_st.ep);
         deadline = GRPC_TIMEOUT_SECONDS_TO_DEADLINE(10);
-        GPR_ASSERT(gpr_event_wait(&write_st.ev, deadline));
+        gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
+        while (!write_st.done) {
+          GPR_ASSERT(gpr_time_cmp(gpr_now(), deadline) < 0);
+          grpc_pollset_work(g_pollset, deadline);
+        }
+        gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
         grpc_endpoint_destroy(write_st.ep);
-        GPR_ASSERT(gpr_event_wait(&read_st.ev, deadline));
+        gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
+        while (!read_st.done) {
+          GPR_ASSERT(gpr_time_cmp(gpr_now(), deadline) < 0);
+          grpc_pollset_work(g_pollset, deadline);
+        }
+        gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
         gpr_free(slices);
         end_test(config);
         return;
@@ -361,9 +369,11 @@ static void shutdown_during_write_test(grpc_endpoint_test_config config,
   abort();
 }
 
-void grpc_endpoint_tests(grpc_endpoint_test_config config) {
+void grpc_endpoint_tests(grpc_endpoint_test_config config, grpc_pollset *pollset) {
+  g_pollset = pollset;
   read_and_write_test(config, 10000000, 100000, 8192, 0);
   read_and_write_test(config, 1000000, 100000, 1, 0);
   read_and_write_test(config, 100000000, 100000, 1, 1);
   shutdown_during_write_test(config, 1000);
+  g_pollset = NULL;
 }
diff --git a/test/core/iomgr/endpoint_tests.h b/test/core/iomgr/endpoint_tests.h
index 1679d7bd4f..852e71d479 100644
--- a/test/core/iomgr/endpoint_tests.h
+++ b/test/core/iomgr/endpoint_tests.h
@@ -52,6 +52,6 @@ struct grpc_endpoint_test_config {
   void (*clean_up)();
 };
 
-void grpc_endpoint_tests(grpc_endpoint_test_config config);
+void grpc_endpoint_tests(grpc_endpoint_test_config config, grpc_pollset *pollset);
 
 #endif  /* GRPC_TEST_CORE_IOMGR_ENDPOINT_TESTS_H */
diff --git a/test/core/iomgr/tcp_posix_test.c b/test/core/iomgr/tcp_posix_test.c
index 40abed5f6e..04aed132f4 100644
--- a/test/core/iomgr/tcp_posix_test.c
+++ b/test/core/iomgr/tcp_posix_test.c
@@ -48,6 +48,8 @@
 #include "test/core/util/test_config.h"
 #include "test/core/iomgr/endpoint_tests.h"
 
+static grpc_pollset g_pollset;
+
 /*
    General test notes:
 
@@ -114,8 +116,6 @@ static size_t fill_socket_partial(int fd, size_t bytes) {
 
 struct read_socket_state {
   grpc_endpoint *ep;
-  gpr_mu mu;
-  gpr_cv cv;
   ssize_t read_bytes;
   ssize_t target_read_bytes;
 };
@@ -145,18 +145,18 @@ static void read_cb(void *user_data, gpr_slice *slices, size_t nslices,
 
   GPR_ASSERT(error == GRPC_ENDPOINT_CB_OK);
 
-  gpr_mu_lock(&state->mu);
+  gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
   current_data = state->read_bytes % 256;
   read_bytes = count_and_unref_slices(slices, nslices, &current_data);
   state->read_bytes += read_bytes;
   gpr_log(GPR_INFO, "Read %d bytes of %d", read_bytes,
           state->target_read_bytes);
   if (state->read_bytes >= state->target_read_bytes) {
-    gpr_cv_signal(&state->cv);
+    /* empty */
   } else {
     grpc_endpoint_notify_on_read(state->ep, read_cb, state);
   }
-  gpr_mu_unlock(&state->mu);
+  gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
 }
 
 /* Write to a socket, then read from it using the grpc_tcp API. */
@@ -173,31 +173,25 @@ static void read_test(ssize_t num_bytes, ssize_t slice_size) {
   create_sockets(sv);
 
   ep = grpc_tcp_create(grpc_fd_create(sv[1]), slice_size);
+  grpc_endpoint_add_to_pollset(ep, &g_pollset);
+
   written_bytes = fill_socket_partial(sv[0], num_bytes);
   gpr_log(GPR_INFO, "Wrote %d bytes", written_bytes);
 
-  gpr_mu_init(&state.mu);
-  gpr_cv_init(&state.cv);
   state.ep = ep;
   state.read_bytes = 0;
   state.target_read_bytes = written_bytes;
 
   grpc_endpoint_notify_on_read(ep, read_cb, &state);
 
-  gpr_mu_lock(&state.mu);
-  for (;;) {
-    GPR_ASSERT(gpr_cv_wait(&state.cv, &state.mu, deadline) == 0);
-    if (state.read_bytes >= state.target_read_bytes) {
-      break;
-    }
+  gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
+  while (state.read_bytes < state.target_read_bytes) {
+    grpc_pollset_work(&g_pollset, deadline);
   }
   GPR_ASSERT(state.read_bytes == state.target_read_bytes);
-  gpr_mu_unlock(&state.mu);
+  gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
 
   grpc_endpoint_destroy(ep);
-
-  gpr_mu_destroy(&state.mu);
-  gpr_cv_destroy(&state.cv);
 }
 
 /* Write to a socket until it fills up, then read from it using the grpc_tcp
@@ -214,37 +208,28 @@ static void large_read_test(ssize_t slice_size) {
   create_sockets(sv);
 
   ep = grpc_tcp_create(grpc_fd_create(sv[1]), slice_size);
+  grpc_endpoint_add_to_pollset(ep, &g_pollset);
   written_bytes = fill_socket(sv[0]);
   gpr_log(GPR_INFO, "Wrote %d bytes", written_bytes);
 
-  gpr_mu_init(&state.mu);
-  gpr_cv_init(&state.cv);
   state.ep = ep;
   state.read_bytes = 0;
   state.target_read_bytes = written_bytes;
 
   grpc_endpoint_notify_on_read(ep, read_cb, &state);
 
-  gpr_mu_lock(&state.mu);
-  for (;;) {
-    GPR_ASSERT(gpr_cv_wait(&state.cv, &state.mu, deadline) == 0);
-    if (state.read_bytes >= state.target_read_bytes) {
-      break;
-    }
+  gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
+  while (state.read_bytes < state.target_read_bytes) {
+    grpc_pollset_work(&g_pollset, deadline);
   }
   GPR_ASSERT(state.read_bytes == state.target_read_bytes);
-  gpr_mu_unlock(&state.mu);
+  gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
 
   grpc_endpoint_destroy(ep);
-
-  gpr_mu_destroy(&state.mu);
-  gpr_cv_destroy(&state.cv);
 }
 
 struct write_socket_state {
   grpc_endpoint *ep;
-  gpr_mu mu;
-  gpr_cv cv;
   int write_done;
 };
 
@@ -275,11 +260,11 @@ static void write_done(void *user_data /* write_socket_state */,
                        grpc_endpoint_cb_status error) {
   struct write_socket_state *state = (struct write_socket_state *)user_data;
   gpr_log(GPR_INFO, "Write done callback called");
-  gpr_mu_lock(&state->mu);
+  gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
   gpr_log(GPR_INFO, "Signalling write done");
   state->write_done = 1;
-  gpr_cv_signal(&state->cv);
-  gpr_mu_unlock(&state->mu);
+  grpc_pollset_kick(&g_pollset);
+  gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
 }
 
 void drain_socket_blocking(int fd, size_t num_bytes, size_t read_size) {
@@ -294,6 +279,9 @@ void drain_socket_blocking(int fd, size_t num_bytes, size_t read_size) {
   GPR_ASSERT(fcntl(fd, F_SETFL, flags & ~O_NONBLOCK) == 0);
 
   for (;;) {
+    gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
+    grpc_pollset_work(&g_pollset, GRPC_TIMEOUT_MILLIS_TO_DEADLINE(10));
+    gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
     do {
       bytes_read =
           read(fd, buf, bytes_left > read_size ? read_size : bytes_left);
@@ -351,9 +339,8 @@ static void write_test(ssize_t num_bytes, ssize_t slice_size) {
   create_sockets(sv);
 
   ep = grpc_tcp_create(grpc_fd_create(sv[1]), GRPC_TCP_DEFAULT_READ_SLICE_SIZE);
+  grpc_endpoint_add_to_pollset(ep, &g_pollset);
 
-  gpr_mu_init(&state.mu);
-  gpr_cv_init(&state.cv);
   state.ep = ep;
   state.write_done = 0;
 
@@ -366,19 +353,17 @@ static void write_test(ssize_t num_bytes, ssize_t slice_size) {
     GPR_ASSERT(read_bytes == num_bytes);
   } else {
     drain_socket_blocking(sv[0], num_bytes, num_bytes);
-    gpr_mu_lock(&state.mu);
+    gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
     for (;;) {
       if (state.write_done) {
         break;
       }
-      GPR_ASSERT(gpr_cv_wait(&state.cv, &state.mu, deadline) == 0);
+      grpc_pollset_work(&g_pollset, deadline);
     }
-    gpr_mu_unlock(&state.mu);
+    gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
   }
 
   grpc_endpoint_destroy(ep);
-  gpr_mu_destroy(&state.mu);
-  gpr_cv_destroy(&state.cv);
   gpr_free(slices);
 }
 
@@ -407,10 +392,9 @@ static void write_error_test(ssize_t num_bytes, ssize_t slice_size) {
   create_sockets(sv);
 
   ep = grpc_tcp_create(grpc_fd_create(sv[1]), GRPC_TCP_DEFAULT_READ_SLICE_SIZE);
+  grpc_endpoint_add_to_pollset(ep, &g_pollset);
   close(sv[0]);
 
-  gpr_mu_init(&state.mu);
-  gpr_cv_init(&state.cv);
   state.ep = ep;
   state.write_done = 0;
 
@@ -423,20 +407,18 @@ static void write_error_test(ssize_t num_bytes, ssize_t slice_size) {
       break;
     case GRPC_ENDPOINT_WRITE_PENDING:
       grpc_endpoint_notify_on_read(ep, read_done_for_write_error, NULL);
-      gpr_mu_lock(&state.mu);
+      gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
       for (;;) {
         if (state.write_done) {
           break;
         }
-        GPR_ASSERT(gpr_cv_wait(&state.cv, &state.mu, deadline) == 0);
+        grpc_pollset_work(&g_pollset, deadline);
       }
-      gpr_mu_unlock(&state.mu);
+      gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
       break;
   }
 
   grpc_endpoint_destroy(ep);
-  gpr_mu_destroy(&state.mu);
-  gpr_cv_destroy(&state.cv);
   free(slices);
 }
 
@@ -475,6 +457,8 @@ static grpc_endpoint_test_fixture create_fixture_tcp_socketpair(
   create_sockets(sv);
   f.client_ep = grpc_tcp_create(grpc_fd_create(sv[0]), slice_size);
   f.server_ep = grpc_tcp_create(grpc_fd_create(sv[1]), slice_size);
+  grpc_endpoint_add_to_pollset(f.client_ep, &g_pollset);
+  grpc_endpoint_add_to_pollset(f.server_ep, &g_pollset);
 
   return f;
 }
@@ -483,11 +467,17 @@ static grpc_endpoint_test_config configs[] = {
     {"tcp/tcp_socketpair", create_fixture_tcp_socketpair, clean_up},
 };
 
+static void destroy_pollset(void *p) {
+  grpc_pollset_destroy(p);
+}
+
 int main(int argc, char **argv) {
   grpc_test_init(argc, argv);
   grpc_init();
+  grpc_pollset_init(&g_pollset);
   run_tests();
-  grpc_endpoint_tests(configs[0]);
+  grpc_endpoint_tests(configs[0], &g_pollset);
+  grpc_pollset_shutdown(&g_pollset, destroy_pollset, &g_pollset);
   grpc_shutdown();
 
   return 0;
diff --git a/test/core/security/secure_endpoint_test.c b/test/core/security/secure_endpoint_test.c
index 6477454e8a..28c909f34d 100644
--- a/test/core/security/secure_endpoint_test.c
+++ b/test/core/security/secure_endpoint_test.c
@@ -44,6 +44,8 @@
 #include "test/core/util/test_config.h"
 #include "src/core/tsi/fake_transport_security.h"
 
+static grpc_pollset g_pollset;
+
 static grpc_endpoint_test_fixture secure_endpoint_create_fixture_tcp_socketpair(
     size_t slice_size, gpr_slice *leftover_slices, size_t leftover_nslices) {
   tsi_frame_protector *fake_read_protector = tsi_create_fake_protector(NULL);
@@ -52,6 +54,8 @@ static grpc_endpoint_test_fixture secure_endpoint_create_fixture_tcp_socketpair(
   grpc_endpoint_pair tcp;
 
   tcp = grpc_iomgr_create_endpoint_pair(slice_size);
+  grpc_endpoint_add_to_pollset(tcp.client, &g_pollset);
+  grpc_endpoint_add_to_pollset(tcp.server, &g_pollset);
 
   if (leftover_nslices == 0) {
     f.client_ep =
@@ -190,13 +194,19 @@ static void test_destroy_ep_early(grpc_endpoint_test_config config,
   clean_up();
 }
 
+static void destroy_pollset(void *p) {
+  grpc_pollset_destroy(p);
+}
+
 int main(int argc, char **argv) {
   grpc_test_init(argc, argv);
 
   grpc_iomgr_init();
-  grpc_endpoint_tests(configs[0]);
+  grpc_pollset_init(&g_pollset);
+  grpc_endpoint_tests(configs[0], &g_pollset);
   test_leftover(configs[1], 1);
   test_destroy_ep_early(configs[1], 1);
+  grpc_pollset_shutdown(&g_pollset, destroy_pollset, &g_pollset);
   grpc_iomgr_shutdown();
 
   return 0;