Allow specific pollers to be woken

Currently, if two threads call grpc_completion_queue_pluck on the same
completion queue for different tags, there is a 50% chance that we
deliver the completion wakeup to the wrong poller - forcing the correct
poller to wait until its polling times out before it can return an event
up to the application.

This change tweaks our polling interfaces so that we can indeed wake a
specific poller.

Nothing has been performance tuned yet. It's definitely sub-optimal in a
number of places. Wakeup file-descriptors should be recycled. We should
have a path that avoids calling poll() followed by epoll(). We can
probably live without it right at the second though.

This code will fail on Windows at least (I'll do that port when I'm in the office and have a Windows
machine).

10 files changed, 68 insertions, 171 deletions

diff --git a/test/core/httpcli/httpcli_test.c b/test/core/httpcli/httpcli_test.c
index 4801eb3e39..390afcdf63 100644
--- a/test/core/httpcli/httpcli_test.c
+++ b/test/core/httpcli/httpcli_test.c
@@ -64,7 +64,7 @@ static void on_finish(void *arg, const grpc_httpcli_response *response) {
   GPR_ASSERT(0 == memcmp(expect, response->body, response->body_length));
   gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
   g_done = 1;
-  grpc_pollset_kick(&g_pollset);
+  grpc_pollset_kick(&g_pollset, NULL);
   gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
 }
 
@@ -87,7 +87,8 @@ static void test_get(int use_ssl, int port) {
                    (void *)42);
   gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
   while (!g_done) {
-    grpc_pollset_work(&g_pollset, n_seconds_time(20));
+    grpc_pollset_worker worker;
+    grpc_pollset_work(&g_pollset, &worker, n_seconds_time(20));
   }
   gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
   gpr_free(host);
@@ -112,7 +113,8 @@ static void test_post(int use_ssl, int port) {
                     n_seconds_time(15), on_finish, (void *)42);
   gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
   while (!g_done) {
-    grpc_pollset_work(&g_pollset, n_seconds_time(20));
+    grpc_pollset_worker worker;
+    grpc_pollset_work(&g_pollset, &worker, n_seconds_time(20));
   }
   gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
   gpr_free(host);
diff --git a/test/core/iomgr/endpoint_tests.c b/test/core/iomgr/endpoint_tests.c
index cb6adc58cf..8186c96da1 100644
--- a/test/core/iomgr/endpoint_tests.c
+++ b/test/core/iomgr/endpoint_tests.c
@@ -132,7 +132,7 @@ static void read_and_write_test_read_handler(void *data, gpr_slice *slices,
     gpr_log(GPR_INFO, "Read handler shutdown");
     gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
     state->read_done = 1;
-    grpc_pollset_kick(g_pollset);
+    grpc_pollset_kick(g_pollset, NULL);
     gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
     return;
   }
@@ -143,7 +143,7 @@ static void read_and_write_test_read_handler(void *data, gpr_slice *slices,
     gpr_log(GPR_INFO, "Read handler done");
     gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
     state->read_done = 1;
-    grpc_pollset_kick(g_pollset);
+    grpc_pollset_kick(g_pollset, NULL);
     gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
   } else {
     grpc_endpoint_notify_on_read(state->read_ep,
@@ -167,7 +167,7 @@ static void read_and_write_test_write_handler(void *data,
     gpr_log(GPR_INFO, "Write handler shutdown");
     gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
     state->write_done = 1;
-    grpc_pollset_kick(g_pollset);
+    grpc_pollset_kick(g_pollset, NULL);
     gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
     return;
   }
@@ -201,7 +201,7 @@ static void read_and_write_test_write_handler(void *data,
   gpr_log(GPR_INFO, "Write handler done");
   gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
   state->write_done = 1;
-  grpc_pollset_kick(g_pollset);
+  grpc_pollset_kick(g_pollset, NULL);
   gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
 }
 
@@ -254,8 +254,9 @@ static void read_and_write_test(grpc_endpoint_test_config config,
 
   gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
   while (!state.read_done || !state.write_done) {
+    grpc_pollset_worker worker;
     GPR_ASSERT(gpr_time_cmp(gpr_now(GPR_CLOCK_MONOTONIC), deadline) < 0);
-    grpc_pollset_work(g_pollset, deadline);
+    grpc_pollset_work(g_pollset, &worker, deadline);
   }
   gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
 
@@ -287,7 +288,7 @@ static void shutdown_during_write_test_read_handler(
     grpc_endpoint_destroy(st->ep);
     gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
     st->done = error;
-    grpc_pollset_kick(g_pollset);
+    grpc_pollset_kick(g_pollset, NULL);
     gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
   } else {
     grpc_endpoint_notify_on_read(
@@ -309,7 +310,7 @@ static void shutdown_during_write_test_write_handler(
   }
   gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
   st->done = 1;
-  grpc_pollset_kick(g_pollset);
+  grpc_pollset_kick(g_pollset, NULL);
   gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
 }
 
@@ -350,15 +351,17 @@ static void shutdown_during_write_test(grpc_endpoint_test_config config,
         deadline = GRPC_TIMEOUT_SECONDS_TO_DEADLINE(10);
         gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
         while (!write_st.done) {
+          grpc_pollset_worker worker;
           GPR_ASSERT(gpr_time_cmp(gpr_now(deadline.clock_type), deadline) < 0);
-          grpc_pollset_work(g_pollset, deadline);
+          grpc_pollset_work(g_pollset, &worker, deadline);
         }
         gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
         grpc_endpoint_destroy(write_st.ep);
         gpr_mu_lock(GRPC_POLLSET_MU(g_pollset));
         while (!read_st.done) {
+          grpc_pollset_worker worker;
           GPR_ASSERT(gpr_time_cmp(gpr_now(deadline.clock_type), deadline) < 0);
-          grpc_pollset_work(g_pollset, deadline);
+          grpc_pollset_work(g_pollset, &worker, deadline);
         }
         gpr_mu_unlock(GRPC_POLLSET_MU(g_pollset));
         gpr_free(slices);
diff --git a/test/core/iomgr/fd_posix_test.c b/test/core/iomgr/fd_posix_test.c
index 7d00c098cc..adcbcafdbb 100644
--- a/test/core/iomgr/fd_posix_test.c
+++ b/test/core/iomgr/fd_posix_test.c
@@ -179,7 +179,7 @@ static void listen_shutdown_cb(void *arg /*server*/, int success) {
 
   gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
   sv->done = 1;
-  grpc_pollset_kick(&g_pollset);
+  grpc_pollset_kick(&g_pollset, NULL);
   gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
 }
 
@@ -249,7 +249,8 @@ static int server_start(server *sv) {
 static void server_wait_and_shutdown(server *sv) {
   gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
   while (!sv->done) {
-    grpc_pollset_work(&g_pollset, gpr_inf_future(GPR_CLOCK_MONOTONIC));
+    grpc_pollset_worker worker;
+    grpc_pollset_work(&g_pollset, &worker, gpr_inf_future(GPR_CLOCK_MONOTONIC));
   }
   gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
 }
@@ -286,7 +287,7 @@ static void client_session_shutdown_cb(void *arg /*client*/, int success) {
   client *cl = arg;
   grpc_fd_orphan(cl->em_fd, NULL, "c");
   cl->done = 1;
-  grpc_pollset_kick(&g_pollset);
+  grpc_pollset_kick(&g_pollset, NULL);
 }
 
 /* Write as much as possible, then register notify_on_write. */
@@ -356,7 +357,8 @@ static void client_start(client *cl, int port) {
 static void client_wait_and_shutdown(client *cl) {
   gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
   while (!cl->done) {
-    grpc_pollset_work(&g_pollset, gpr_inf_future(GPR_CLOCK_MONOTONIC));
+    grpc_pollset_worker worker;
+    grpc_pollset_work(&g_pollset, &worker, gpr_inf_future(GPR_CLOCK_MONOTONIC));
   }
   gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
 }
@@ -392,7 +394,7 @@ static void first_read_callback(void *arg /* fd_change_data */, int success) {
 
   gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
   fdc->cb_that_ran = first_read_callback;
-  grpc_pollset_kick(&g_pollset);
+  grpc_pollset_kick(&g_pollset, NULL);
   gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
 }
 
@@ -401,7 +403,7 @@ static void second_read_callback(void *arg /* fd_change_data */, int success) {
 
   gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
   fdc->cb_that_ran = second_read_callback;
-  grpc_pollset_kick(&g_pollset);
+  grpc_pollset_kick(&g_pollset, NULL);
   gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
 }
 
@@ -445,7 +447,8 @@ static void test_grpc_fd_change(void) {
   /* And now wait for it to run. */
   gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
   while (a.cb_that_ran == NULL) {
-    grpc_pollset_work(&g_pollset, gpr_inf_future(GPR_CLOCK_MONOTONIC));
+    grpc_pollset_worker worker;
+    grpc_pollset_work(&g_pollset, &worker, gpr_inf_future(GPR_CLOCK_MONOTONIC));
   }
   GPR_ASSERT(a.cb_that_ran == first_read_callback);
   gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
@@ -463,7 +466,8 @@ static void test_grpc_fd_change(void) {
 
   gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
   while (b.cb_that_ran == NULL) {
-    grpc_pollset_work(&g_pollset, gpr_inf_future(GPR_CLOCK_MONOTONIC));
+    grpc_pollset_worker worker;
+    grpc_pollset_work(&g_pollset, &worker, gpr_inf_future(GPR_CLOCK_MONOTONIC));
   }
   /* Except now we verify that second_read_callback ran instead */
   GPR_ASSERT(b.cb_that_ran == second_read_callback);
diff --git a/test/core/iomgr/poll_kick_posix_test.c b/test/core/iomgr/poll_kick_posix_test.c
deleted file mode 100644
index 3aa6807806..0000000000
--- a/test/core/iomgr/poll_kick_posix_test.c
+++ /dev/null
@@ -1,130 +0,0 @@
-/*
- *
- * Copyright 2015, Google Inc.
- * All rights reserved.
- *
- * Redistribution and use in source and binary forms, with or without
- * modification, are permitted provided that the following conditions are
- * met:
- *
- *     * Redistributions of source code must retain the above copyright
- * notice, this list of conditions and the following disclaimer.
- *     * Redistributions in binary form must reproduce the above
- * copyright notice, this list of conditions and the following disclaimer
- * in the documentation and/or other materials provided with the
- * distribution.
- *     * Neither the name of Google Inc. nor the names of its
- * contributors may be used to endorse or promote products derived from
- * this software without specific prior written permission.
- *
- * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
- * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
- * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
- * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
- * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
- * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
- * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
- * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
- * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
- * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
- *
- */
-
-#include "src/core/iomgr/pollset_kick_posix.h"
-
-#include <grpc/support/alloc.h>
-#include <grpc/support/log.h>
-#include "test/core/util/test_config.h"
-
-static void test_allocation(void) {
-  grpc_pollset_kick_state state;
-  grpc_pollset_kick_init(&state);
-  grpc_pollset_kick_destroy(&state);
-}
-
-static void test_non_kick(void) {
-  grpc_pollset_kick_state state;
-  grpc_kick_fd_info *kfd;
-
-  grpc_pollset_kick_init(&state);
-  kfd = grpc_pollset_kick_pre_poll(&state);
-  GPR_ASSERT(kfd != NULL);
-
-  grpc_pollset_kick_post_poll(&state, kfd);
-  grpc_pollset_kick_destroy(&state);
-}
-
-static void test_basic_kick(void) {
-  /* Kicked during poll */
-  grpc_pollset_kick_state state;
-  grpc_kick_fd_info *kfd;
-  grpc_pollset_kick_init(&state);
-
-  kfd = grpc_pollset_kick_pre_poll(&state);
-  GPR_ASSERT(kfd != NULL);
-
-  grpc_pollset_kick_kick(&state);
-
-  /* Now hypothetically we polled and found that we were kicked */
-  grpc_pollset_kick_consume(&state, kfd);
-
-  grpc_pollset_kick_post_poll(&state, kfd);
-
-  grpc_pollset_kick_destroy(&state);
-}
-
-static void test_non_poll_kick(void) {
-  /* Kick before entering poll */
-  grpc_pollset_kick_state state;
-  grpc_kick_fd_info *kfd;
-
-  grpc_pollset_kick_init(&state);
-
-  grpc_pollset_kick_kick(&state);
-  kfd = grpc_pollset_kick_pre_poll(&state);
-  GPR_ASSERT(kfd == NULL);
-  grpc_pollset_kick_destroy(&state);
-}
-
-#define GRPC_MAX_CACHED_PIPES 50
-
-static void test_over_free(void) {
-  /* Check high watermark pipe free logic */
-  int i;
-  grpc_kick_fd_info **kfds =
-      gpr_malloc(sizeof(grpc_kick_fd_info *) * GRPC_MAX_CACHED_PIPES);
-  grpc_pollset_kick_state state;
-  grpc_pollset_kick_init(&state);
-  for (i = 0; i < GRPC_MAX_CACHED_PIPES; ++i) {
-    kfds[i] = grpc_pollset_kick_pre_poll(&state);
-    GPR_ASSERT(kfds[i] != NULL);
-  }
-
-  for (i = 0; i < GRPC_MAX_CACHED_PIPES; ++i) {
-    grpc_pollset_kick_post_poll(&state, kfds[i]);
-  }
-  grpc_pollset_kick_destroy(&state);
-  gpr_free(kfds);
-}
-
-static void run_tests(void) {
-  test_allocation();
-  test_basic_kick();
-  test_non_poll_kick();
-  test_non_kick();
-  test_over_free();
-}
-
-int main(int argc, char **argv) {
-  grpc_test_init(argc, argv);
-
-  grpc_pollset_kick_global_init();
-  run_tests();
-  grpc_pollset_kick_global_destroy();
-
-  grpc_pollset_kick_global_init_fallback_fd();
-  run_tests();
-  grpc_pollset_kick_global_destroy();
-  return 0;
-}
diff --git a/test/core/iomgr/tcp_client_posix_test.c b/test/core/iomgr/tcp_client_posix_test.c
index 38b7b5909d..07bbe1f402 100644
--- a/test/core/iomgr/tcp_client_posix_test.c
+++ b/test/core/iomgr/tcp_client_posix_test.c
@@ -56,7 +56,7 @@ static gpr_timespec test_deadline(void) {
 static void finish_connection() {
   gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
   g_connections_complete++;
-  grpc_pollset_kick(&g_pollset);
+  grpc_pollset_kick(&g_pollset, NULL);
   gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
 }
 
@@ -111,7 +111,8 @@ void test_succeeds(void) {
   gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
 
   while (g_connections_complete == connections_complete_before) {
-    grpc_pollset_work(&g_pollset, GRPC_TIMEOUT_SECONDS_TO_DEADLINE(5));
+    grpc_pollset_worker worker;
+    grpc_pollset_work(&g_pollset, &worker, GRPC_TIMEOUT_SECONDS_TO_DEADLINE(5));
   }
 
   gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
@@ -140,7 +141,8 @@ void test_fails(void) {
 
   /* wait for the connection callback to finish */
   while (g_connections_complete == connections_complete_before) {
-    grpc_pollset_work(&g_pollset, test_deadline());
+    grpc_pollset_worker worker;
+    grpc_pollset_work(&g_pollset, &worker, test_deadline());
   }
 
   gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
@@ -199,6 +201,7 @@ void test_times_out(void) {
                       gpr_now(connect_deadline.clock_type)) > 0) {
     int is_after_deadline =
         gpr_time_cmp(connect_deadline, gpr_now(GPR_CLOCK_MONOTONIC)) <= 0;
+    grpc_pollset_worker worker;
     if (is_after_deadline &&
         gpr_time_cmp(gpr_time_add(connect_deadline,
                                   gpr_time_from_seconds(1, GPR_TIMESPAN)),
@@ -208,7 +211,7 @@ void test_times_out(void) {
       GPR_ASSERT(g_connections_complete ==
                  connections_complete_before + is_after_deadline);
     }
-    grpc_pollset_work(&g_pollset, GRPC_TIMEOUT_MILLIS_TO_DEADLINE(10));
+    grpc_pollset_work(&g_pollset, &worker, GRPC_TIMEOUT_MILLIS_TO_DEADLINE(10));
   }
   gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
 
diff --git a/test/core/iomgr/tcp_posix_test.c b/test/core/iomgr/tcp_posix_test.c
index 4e7fb446a7..17a85ceaec 100644
--- a/test/core/iomgr/tcp_posix_test.c
+++ b/test/core/iomgr/tcp_posix_test.c
@@ -186,7 +186,8 @@ static void read_test(ssize_t num_bytes, ssize_t slice_size) {
 
   gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
   while (state.read_bytes < state.target_read_bytes) {
-    grpc_pollset_work(&g_pollset, deadline);
+    grpc_pollset_worker worker;
+    grpc_pollset_work(&g_pollset, &worker, deadline);
   }
   GPR_ASSERT(state.read_bytes == state.target_read_bytes);
   gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
@@ -222,7 +223,8 @@ static void large_read_test(ssize_t slice_size) {
 
   gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
   while (state.read_bytes < state.target_read_bytes) {
-    grpc_pollset_work(&g_pollset, deadline);
+    grpc_pollset_worker worker;
+    grpc_pollset_work(&g_pollset, &worker, deadline);
   }
   GPR_ASSERT(state.read_bytes == state.target_read_bytes);
   gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
@@ -265,7 +267,7 @@ static void write_done(void *user_data /* write_socket_state */,
   gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
   gpr_log(GPR_INFO, "Signalling write done");
   state->write_done = 1;
-  grpc_pollset_kick(&g_pollset);
+  grpc_pollset_kick(&g_pollset, NULL);
   gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
 }
 
@@ -281,8 +283,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 (;;) {
+    grpc_pollset_worker worker;
     gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
-    grpc_pollset_work(&g_pollset, GRPC_TIMEOUT_MILLIS_TO_DEADLINE(10));
+    grpc_pollset_work(&g_pollset, &worker, GRPC_TIMEOUT_MILLIS_TO_DEADLINE(10));
     gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
     do {
       bytes_read =
@@ -358,10 +361,11 @@ static void write_test(ssize_t num_bytes, ssize_t slice_size) {
     drain_socket_blocking(sv[0], num_bytes, num_bytes);
     gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
     for (;;) {
+      grpc_pollset_worker worker;
       if (state.write_done) {
         break;
       }
-      grpc_pollset_work(&g_pollset, deadline);
+      grpc_pollset_work(&g_pollset, &worker, deadline);
     }
     gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
   }
@@ -387,6 +391,7 @@ static void write_error_test(ssize_t num_bytes, ssize_t slice_size) {
   size_t num_blocks;
   gpr_slice *slices;
   int current_data = 0;
+  grpc_pollset_worker worker;
   gpr_timespec deadline = GRPC_TIMEOUT_SECONDS_TO_DEADLINE(20);
 
   gpr_log(GPR_INFO, "Start write error test with %d bytes, slice size %d",
@@ -417,7 +422,7 @@ static void write_error_test(ssize_t num_bytes, ssize_t slice_size) {
         if (state.write_done) {
           break;
         }
-        grpc_pollset_work(&g_pollset, deadline);
+        grpc_pollset_work(&g_pollset, &worker, deadline);
       }
       gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
       break;
diff --git a/test/core/iomgr/tcp_server_posix_test.c b/test/core/iomgr/tcp_server_posix_test.c
index f8d0fe8217..b82d7c08b1 100644
--- a/test/core/iomgr/tcp_server_posix_test.c
+++ b/test/core/iomgr/tcp_server_posix_test.c
@@ -54,7 +54,7 @@ static void on_connect(void *arg, grpc_endpoint *tcp) {
 
   gpr_mu_lock(GRPC_POLLSET_MU(&g_pollset));
   g_nconnects++;
-  grpc_pollset_kick(&g_pollset);
+  grpc_pollset_kick(&g_pollset, NULL);
   gpr_mu_unlock(GRPC_POLLSET_MU(&g_pollset));
 }
 
@@ -136,7 +136,8 @@ static void test_connect(int n) {
     gpr_log(GPR_DEBUG, "wait");
     while (g_nconnects == nconnects_before &&
            gpr_time_cmp(deadline, gpr_now(deadline.clock_type)) > 0) {
-      grpc_pollset_work(&g_pollset, deadline);
+      grpc_pollset_worker worker;
+      grpc_pollset_work(&g_pollset, &worker, deadline);
     }
     gpr_log(GPR_DEBUG, "wait done");
 
diff --git a/test/core/security/oauth2_utils.c b/test/core/security/oauth2_utils.c
index ecd04fd9e1..bd4d44ac48 100644
--- a/test/core/security/oauth2_utils.c
+++ b/test/core/security/oauth2_utils.c
@@ -68,7 +68,7 @@ static void on_oauth2_response(void *user_data, grpc_credentials_md *md_elems,
   gpr_mu_lock(GRPC_POLLSET_MU(&request->pollset));
   request->is_done = 1;
   request->token = token;
-  grpc_pollset_kick(&request->pollset);
+  grpc_pollset_kick(&request->pollset, NULL);
   gpr_mu_unlock(GRPC_POLLSET_MU(&request->pollset));
 }
 
@@ -83,8 +83,11 @@ char *grpc_test_fetch_oauth2_token_with_credentials(grpc_credentials *creds) {
                                         on_oauth2_response, &request);
 
   gpr_mu_lock(GRPC_POLLSET_MU(&request.pollset));
-  while (!request.is_done)
-    grpc_pollset_work(&request.pollset, gpr_inf_future(GPR_CLOCK_REALTIME));
+  while (!request.is_done) {
+    grpc_pollset_worker worker;
+    grpc_pollset_work(&request.pollset, &worker,
+                      gpr_inf_future(GPR_CLOCK_REALTIME));
+  }
   gpr_mu_unlock(GRPC_POLLSET_MU(&request.pollset));
 
   grpc_pollset_shutdown(&request.pollset, do_nothing, NULL);
diff --git a/test/core/security/print_google_default_creds_token.c b/test/core/security/print_google_default_creds_token.c
index 0875cfb4fb..7238efbbfd 100644
--- a/test/core/security/print_google_default_creds_token.c
+++ b/test/core/security/print_google_default_creds_token.c
@@ -65,7 +65,7 @@ static void on_metadata_response(void *user_data,
   }
   gpr_mu_lock(GRPC_POLLSET_MU(&sync->pollset));
   sync->is_done = 1;
-  grpc_pollset_kick(&sync->pollset);
+  grpc_pollset_kick(&sync->pollset, NULL);
   gpr_mu_unlock(GRPC_POLLSET_MU(&sync->pollset));
 }
 
@@ -95,8 +95,11 @@ int main(int argc, char **argv) {
                                         on_metadata_response, &sync);
 
   gpr_mu_lock(GRPC_POLLSET_MU(&sync.pollset));
-  while (!sync.is_done)
-    grpc_pollset_work(&sync.pollset, gpr_inf_future(GPR_CLOCK_REALTIME));
+  while (!sync.is_done) {
+    grpc_pollset_worker worker;
+    grpc_pollset_work(&sync.pollset, &worker,
+                      gpr_inf_future(GPR_CLOCK_REALTIME));
+  }
   gpr_mu_unlock(GRPC_POLLSET_MU(&sync.pollset));
 
   grpc_credentials_release(creds);
diff --git a/test/core/security/verify_jwt.c b/test/core/security/verify_jwt.c
index cb073f19c7..9b334b3c3e 100644
--- a/test/core/security/verify_jwt.c
+++ b/test/core/security/verify_jwt.c
@@ -79,7 +79,7 @@ static void on_jwt_verification_done(void *user_data,
 
   gpr_mu_lock(GRPC_POLLSET_MU(&sync->pollset));
   sync->is_done = 1;
-  grpc_pollset_kick(&sync->pollset);
+  grpc_pollset_kick(&sync->pollset, NULL);
   gpr_mu_unlock(GRPC_POLLSET_MU(&sync->pollset));
 }
 
@@ -109,8 +109,11 @@ int main(int argc, char **argv) {
                            on_jwt_verification_done, &sync);
 
   gpr_mu_lock(GRPC_POLLSET_MU(&sync.pollset));
-  while (!sync.is_done)
-    grpc_pollset_work(&sync.pollset, gpr_inf_future(GPR_CLOCK_REALTIME));
+  while (!sync.is_done) {
+    grpc_pollset_worker worker;
+    grpc_pollset_work(&sync.pollset, &worker,
+                      gpr_inf_future(GPR_CLOCK_REALTIME));
+  }
   gpr_mu_unlock(GRPC_POLLSET_MU(&sync.pollset));
 
   grpc_jwt_verifier_destroy(verifier);