Merge pull request #4526 from soltanmm/schrodingers-fix

Fix Python test timeouts.

src/python/grpcio/grpc/_cython/_cygrpc/completion_queue.pxd

@@ -37,3 +37,5 @@ cdef class CompletionQueue:
   cdef bint is_polling
   cdef bint is_shutting_down
   cdef bint is_shutdown
+
+  cdef _interpret_event(self, grpc.grpc_event event)

src/python/grpcio/grpc/_cython/_cygrpc/completion_queue.pyx

@@ -46,35 +46,13 @@ cdef class CompletionQueue:
     self.poll_condition = threading.Condition()
     self.is_polling = False
 
-  def poll(self, records.Timespec deadline=None):
-    # We name this 'poll' to avoid problems with CPython's expectations for
-    # 'special' methods (like next and __next__).
-    cdef grpc.gpr_timespec c_deadline = grpc.gpr_inf_future(
-        grpc.GPR_CLOCK_REALTIME)
+  cdef _interpret_event(self, grpc.grpc_event event):
     cdef records.OperationTag tag = None
     cdef object user_tag = None
     cdef call.Call operation_call = None
     cdef records.CallDetails request_call_details = None
     cdef records.Metadata request_metadata = None
     cdef records.Operations batch_operations = None
-    if deadline is not None:
-      c_deadline = deadline.c_time
-    cdef grpc.grpc_event event
-
-    # Poll within a critical section
-    # TODO consider making queue polling contention a hard error to enable
-    # easier bug discovery
-    with self.poll_condition:
-      while self.is_polling:
-        self.poll_condition.wait(float(deadline) - time.time())
-      self.is_polling = True
-    with nogil:
-      event = grpc.grpc_completion_queue_next(
-          self.c_completion_queue, c_deadline, NULL)
-    with self.poll_condition:
-      self.is_polling = False
-      self.poll_condition.notify()
-
     if event.type == grpc.GRPC_QUEUE_TIMEOUT:
       return records.Event(
           event.type, False, None, None, None, None, False, None)
@@ -104,6 +82,54 @@ cdef class CompletionQueue:
           request_call_details, request_metadata, tag.is_new_request,
           batch_operations)
 
+  def poll(self, records.Timespec deadline=None):
+    # We name this 'poll' to avoid problems with CPython's expectations for
+    # 'special' methods (like next and __next__).
+    cdef grpc.gpr_timespec c_deadline = grpc.gpr_inf_future(
+        grpc.GPR_CLOCK_REALTIME)
+    if deadline is not None:
+      c_deadline = deadline.c_time
+    cdef grpc.grpc_event event
+
+    # Poll within a critical section
+    # TODO(atash) consider making queue polling contention a hard error to
+    # enable easier bug discovery
+    with self.poll_condition:
+      while self.is_polling:
+        self.poll_condition.wait(float(deadline) - time.time())
+      self.is_polling = True
+    with nogil:
+      event = grpc.grpc_completion_queue_next(
+          self.c_completion_queue, c_deadline, NULL)
+    with self.poll_condition:
+      self.is_polling = False
+      self.poll_condition.notify()
+    return self._interpret_event(event)
+
+  def pluck(self, records.OperationTag tag, records.Timespec deadline=None):
+    # Plucking a 'None' tag is equivalent to passing control to GRPC core until
+    # the deadline.
+    cdef grpc.gpr_timespec c_deadline = grpc.gpr_inf_future(
+        grpc.GPR_CLOCK_REALTIME)
+    if deadline is not None:
+      c_deadline = deadline.c_time
+    cdef grpc.grpc_event event
+
+    # Poll within a critical section
+    # TODO(atash) consider making queue polling contention a hard error to
+    # enable easier bug discovery
+    with self.poll_condition:
+      while self.is_polling:
+        self.poll_condition.wait(float(deadline) - time.time())
+      self.is_polling = True
+    with nogil:
+      event = grpc.grpc_completion_queue_pluck(
+          self.c_completion_queue, <cpython.PyObject *>tag, c_deadline, NULL)
+    with self.poll_condition:
+      self.is_polling = False
+      self.poll_condition.notify()
+    return self._interpret_event(event)
+
   def shutdown(self):
     grpc.grpc_completion_queue_shutdown(self.c_completion_queue)
     self.is_shutting_down = True

src/python/grpcio/grpc/_cython/_cygrpc/grpc.pxd

@@ -292,6 +292,9 @@ cdef extern from "grpc/grpc.h":
   grpc_event grpc_completion_queue_next(grpc_completion_queue *cq,
                                         gpr_timespec deadline,
                                         void *reserved) nogil
+  grpc_event grpc_completion_queue_pluck(grpc_completion_queue *cq, void *tag,
+                                         gpr_timespec deadline,
+                                         void *reserved) nogil
   void grpc_completion_queue_shutdown(grpc_completion_queue *cq)
   void grpc_completion_queue_destroy(grpc_completion_queue *cq)
 

src/python/grpcio/grpc/_cython/_cygrpc/server.pyx

@@ -89,6 +89,8 @@ cdef class Server:
     self.register_completion_queue(self.backup_shutdown_queue)
     self.is_started = True
     grpc.grpc_server_start(self.c_server)
+    # Ensure the core has gotten a chance to do the start-up work
+    self.backup_shutdown_queue.pluck(None, records.Timespec(None))
 
   def add_http2_port(self, address,
                      credentials.ServerCredentials server_credentials=None):