Merge pull request #22486 from vjpai/nonblocking

Fix nonblocking_test by properly setting thread-local variable for test

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<intptr_t>(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<intptr_t>(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
 }