diff --git a/test/cpp/end2end/nonblocking_test.cc b/test/cpp/end2end/nonblocking_test.cc index eb651df21df..4be070ec717 100644 --- a/test/cpp/end2end/nonblocking_test.cc +++ b/test/cpp/end2end/nonblocking_test.cc @@ -39,14 +39,18 @@ #ifdef GRPC_POSIX_SOCKET // Thread-local variable to so that only polls from this test assert -// non-blocking (not polls from resolver, timer thread, etc) -GPR_TLS_DECL(g_is_nonblocking_test); +// non-blocking (not polls from resolver, timer thread, etc), and only when the +// thread is waiting on polls caused by CompletionQueue::AsyncNext (not for +// picking a port or other reasons). +GPR_TLS_DECL(g_is_nonblocking_poll); namespace { int maybe_assert_non_blocking_poll(struct pollfd* pfds, nfds_t nfds, int timeout) { - if (gpr_tls_get(&g_is_nonblocking_test)) { + // Only assert that this poll should have zero timeout if we're in the + // middle of a zero-timeout CQ Next. + if (gpr_tls_get(&g_is_nonblocking_poll)) { GPR_ASSERT(timeout == 0); } return poll(pfds, nfds, timeout); @@ -74,11 +78,17 @@ class NonblockingTest : public ::testing::Test { } bool LoopForTag(void** tag, bool* ok) { + // Temporarily set the thread-local nonblocking poll flag so that the polls + // caused by this loop are indeed sent by the library with zero timeout. + intptr_t orig_val = gpr_tls_get(&g_is_nonblocking_poll); + gpr_tls_set(&g_is_nonblocking_poll, static_cast(true)); for (;;) { auto r = cq_->AsyncNext(tag, ok, gpr_time_0(GPR_CLOCK_REALTIME)); if (r == CompletionQueue::SHUTDOWN) { + gpr_tls_set(&g_is_nonblocking_poll, orig_val); return false; } else if (r == CompletionQueue::GOT_EVENT) { + gpr_tls_set(&g_is_nonblocking_poll, orig_val); return true; } } @@ -185,10 +195,20 @@ int main(int argc, char** argv) { #ifdef GRPC_POSIX_SOCKET // Override the poll function before anything else can happen grpc_poll_function = maybe_assert_non_blocking_poll; -#endif // GRPC_POSIX_SOCKET grpc::testing::TestEnvironment env(argc, argv); ::testing::InitGoogleTest(&argc, argv); + gpr_tls_init(&g_is_nonblocking_poll); + + // Start the nonblocking poll thread-local variable as false because the + // thread that issues RPCs starts by picking a port (which has non-zero + // timeout). + gpr_tls_set(&g_is_nonblocking_poll, static_cast(false)); + int ret = RUN_ALL_TESTS(); + gpr_tls_destroy(&g_is_nonblocking_poll); return ret; +#else // GRPC_POSIX_SOCKET + return 0; +#endif // GRPC_POSIX_SOCKET }