Merge pull request #19299 from lidizheng/wait

Add wait_for_termination method to grpc.Server

src/python/grpcio/grpc/__init__.py

@@ -1444,6 +1444,29 @@ class Server(six.with_metaclass(abc.ABCMeta)):
         """
         raise NotImplementedError()
 
+    def wait_for_termination(self, timeout=None):
+        """Block current thread until the server stops.
+
+        This is an EXPERIMENTAL API.
+
+        The wait will not consume computational resources during blocking, and
+        it will block until one of the two following conditions are met:
+
+        1) The server is stopped or terminated;
+        2) A timeout occurs if timeout is not `None`.
+
+        The timeout argument works in the same way as `threading.Event.wait()`.
+        https://docs.python.org/3/library/threading.html#threading.Event.wait
+
+        Args:
+          timeout: A floating point number specifying a timeout for the
+            operation in seconds.
+
+        Returns:
+          A bool indicates if the operation times out.
+        """
+        raise NotImplementedError()
+
 
 #################################  Functions    ################################
 

src/python/grpcio/grpc/_server.py

@@ -50,6 +50,7 @@ _CANCELLED = 'cancelled'
 _EMPTY_FLAGS = 0
 
 _DEALLOCATED_SERVER_CHECK_PERIOD_S = 1.0
+_INF_TIMEOUT = 1e9
 
 
 def _serialized_request(request_event):
@@ -764,7 +765,8 @@ class _ServerState(object):
         self.interceptor_pipeline = interceptor_pipeline
         self.thread_pool = thread_pool
         self.stage = _ServerStage.STOPPED
-        self.shutdown_events = None
+        self.termination_event = threading.Event()
+        self.shutdown_events = [self.termination_event]
         self.maximum_concurrent_rpcs = maximum_concurrent_rpcs
         self.active_rpc_count = 0
 
@@ -876,7 +878,6 @@ def _begin_shutdown_once(state):
         if state.stage is _ServerStage.STARTED:
             state.server.shutdown(state.completion_queue, _SHUTDOWN_TAG)
             state.stage = _ServerStage.GRACE
-            state.shutdown_events = []
             state.due.add(_SHUTDOWN_TAG)
 
 
@@ -959,6 +960,15 @@ class _Server(grpc.Server):
     def start(self):
         _start(self._state)
 
+    def wait_for_termination(self, timeout=None):
+        # NOTE(https://bugs.python.org/issue35935)
+        # Remove this workaround once threading.Event.wait() is working with
+        # CTRL+C across platforms.
+        return _common.wait(
+            self._state.termination_event.wait,
+            self._state.termination_event.is_set,
+            timeout=timeout)
+
     def stop(self, grace):
         return _stop(self._state, grace)
 

src/python/grpcio_tests/tests/tests.json

@@ -66,6 +66,7 @@
   "unit._server_ssl_cert_config_test.ServerSSLCertReloadTestWithClientAuth",
   "unit._server_ssl_cert_config_test.ServerSSLCertReloadTestWithoutClientAuth",
   "unit._server_test.ServerTest",
+  "unit._server_wait_for_termination_test.ServerWaitForTerminationTest",
   "unit._session_cache_test.SSLSessionCacheTest",
   "unit._signal_handling_test.SignalHandlingTest",
   "unit._version_test.VersionTest",

src/python/grpcio_tests/tests/unit/BUILD.bazel

@@ -33,6 +33,7 @@ GRPCIO_TESTS_UNIT = [
     # "_server_ssl_cert_config_test.py",
     "_server_test.py",
     "_server_shutdown_test.py",
+    "_server_wait_for_termination_test.py",
     "_session_cache_test.py",
 ]
 

src/python/grpcio_tests/tests/unit/_server_wait_for_termination_test.py

@@ -0,0 +1,92 @@
+# Copyright 2019 The gRPC Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import division
+
+import datetime
+from concurrent import futures
+import unittest
+import time
+import threading
+import six
+
+import grpc
+from tests.unit.framework.common import test_constants
+
+_WAIT_FOR_BLOCKING = datetime.timedelta(seconds=1)
+
+
+def _block_on_waiting(server, termination_event, timeout=None):
+    server.start()
+    server.wait_for_termination(timeout=timeout)
+    termination_event.set()
+
+
+class ServerWaitForTerminationTest(unittest.TestCase):
+
+    def test_unblock_by_invoking_stop(self):
+        termination_event = threading.Event()
+        server = grpc.server(futures.ThreadPoolExecutor())
+
+        wait_thread = threading.Thread(
+            target=_block_on_waiting, args=(
+                server,
+                termination_event,
+            ))
+        wait_thread.daemon = True
+        wait_thread.start()
+        time.sleep(_WAIT_FOR_BLOCKING.total_seconds())
+
+        server.stop(None)
+        termination_event.wait(timeout=test_constants.SHORT_TIMEOUT)
+        self.assertTrue(termination_event.is_set())
+
+    def test_unblock_by_del(self):
+        termination_event = threading.Event()
+        server = grpc.server(futures.ThreadPoolExecutor())
+
+        wait_thread = threading.Thread(
+            target=_block_on_waiting, args=(
+                server,
+                termination_event,
+            ))
+        wait_thread.daemon = True
+        wait_thread.start()
+        time.sleep(_WAIT_FOR_BLOCKING.total_seconds())
+
+        # Invoke manually here, in Python 2 it will be invoked by GC sometime.
+        server.__del__()
+        termination_event.wait(timeout=test_constants.SHORT_TIMEOUT)
+        self.assertTrue(termination_event.is_set())
+
+    def test_unblock_by_timeout(self):
+        termination_event = threading.Event()
+        server = grpc.server(futures.ThreadPoolExecutor())
+
+        wait_thread = threading.Thread(
+            target=_block_on_waiting,
+            args=(
+                server,
+                termination_event,
+                test_constants.SHORT_TIMEOUT / 2,
+            ))
+        wait_thread.daemon = True
+        wait_thread.start()
+
+        termination_event.wait(timeout=test_constants.SHORT_TIMEOUT)
+        self.assertTrue(termination_event.is_set())
+
+
+if __name__ == '__main__':
+    unittest.main(verbosity=2)