Merge pull request #8707 from kpayson64/upstream_changes

Add configurable exit grace periods and shutdown handlers

4 files changed, 102 insertions, 40 deletions

diff --git a/src/python/grpcio/grpc/__init__.py b/src/python/grpcio/grpc/__init__.py
index 4e4062bafc..6087276d51 100644
--- a/src/python/grpcio/grpc/__init__.py
+++ b/src/python/grpcio/grpc/__init__.py
@@ -905,6 +905,21 @@ class Server(six.with_metaclass(abc.ABCMeta)):
     raise NotImplementedError()
 
   @abc.abstractmethod
+  def add_shutdown_handler(self, shutdown_handler):
+    """Adds a handler to be called on server shutdown.
+
+    Shutdown handlers are run on server stop() or in the event that a running
+    server is destroyed unexpectedly.  The handlers are run in series before
+    the stop grace period.
+
+    Args:
+      shutdown_handler:  A function taking a single arg, a time in seconds
+      within which the handler should complete.  None indicates the handler can
+      run for any duration.
+    """
+    raise NotImplementedError()
+
+  @abc.abstractmethod
   def start(self):
     """Starts this Server's service of RPCs.
 
@@ -914,7 +929,7 @@ class Server(six.with_metaclass(abc.ABCMeta)):
     raise NotImplementedError()
 
   @abc.abstractmethod
-  def stop(self, grace):
+  def stop(self, grace, shutdown_handler_grace=None):
     """Stops this Server's service of RPCs.
 
     All calls to this method immediately stop service of new RPCs. When existing
@@ -937,6 +952,8 @@ class Server(six.with_metaclass(abc.ABCMeta)):
         aborted by this Server's stopping. If None, all RPCs will be aborted
         immediately and this method will block until this Server is completely
         stopped.
+      shutdown_handler_grace:  A duration of time in seconds or None.  This
+        value is passed to all shutdown handlers.
 
     Returns:
       A threading.Event that will be set when this Server has completely
@@ -1231,7 +1248,8 @@ def secure_channel(target, credentials, options=None):
                           credentials._credentials)
 
 
-def server(thread_pool, handlers=None, options=None):
+def server(thread_pool, handlers=None, options=None, exit_grace=None,
+           exit_shutdown_handler_grace=None):
   """Creates a Server with which RPCs can be serviced.
 
   Args:
@@ -1244,13 +1262,19 @@ def server(thread_pool, handlers=None, options=None):
       returned Server is started.
     options: A sequence of string-value pairs according to which to configure
       the created server.
+    exit_grace:  The grace period to use when terminating
+      running servers at interpreter exit.  None indicates unspecified.
+    exit_shutdown_handler_grace:  The shutdown handler grace to use when
+      terminating running servers at interpreter exit.  None indicates
+      unspecified.
 
   Returns:
     A Server with which RPCs can be serviced.
   """
   from grpc import _server
   return _server.Server(thread_pool, () if handlers is None else handlers,
-                        () if options is None else options)
+                        () if options is None else options, exit_grace,
+                        exit_shutdown_handler_grace)
 
 
 ###################################  __all__  #################################
diff --git a/src/python/grpcio/grpc/_server.py b/src/python/grpcio/grpc/_server.py
index 5223712dfa..d83a2e6ded 100644
--- a/src/python/grpcio/grpc/_server.py
+++ b/src/python/grpcio/grpc/_server.py
@@ -60,7 +60,8 @@ _CANCELLED = 'cancelled'
 _EMPTY_FLAGS = 0
 _EMPTY_METADATA = cygrpc.Metadata(())
 
-_UNEXPECTED_EXIT_SERVER_GRACE = 1.0
+_DEFAULT_EXIT_GRACE = 1.0
+_DEFAULT_EXIT_SHUTDOWN_HANDLER_GRACE = 5.0
 
 
 def _serialized_request(request_event):
@@ -595,14 +596,18 @@ class _ServerStage(enum.Enum):
 
 class _ServerState(object):
 
-  def __init__(self, completion_queue, server, generic_handlers, thread_pool):
+  def __init__(self, completion_queue, server, generic_handlers, thread_pool,
+               exit_grace, exit_shutdown_handler_grace):
     self.lock = threading.Lock()
     self.completion_queue = completion_queue
     self.server = server
     self.generic_handlers = list(generic_handlers)
     self.thread_pool = thread_pool
+    self.exit_grace = exit_grace
+    self.exit_shutdown_handler_grace = exit_shutdown_handler_grace
     self.stage = _ServerStage.STOPPED
     self.shutdown_events = None
+    self.shutdown_handlers = []
 
     # TODO(https://github.com/grpc/grpc/issues/6597): eliminate these fields.
     self.rpc_states = set()
@@ -672,41 +677,45 @@ def _serve(state):
             return
 
 
-def _stop(state, grace):
-  with state.lock:
-    if state.stage is _ServerStage.STOPPED:
-      shutdown_event = threading.Event()
-      shutdown_event.set()
-      return shutdown_event
-    else:
-      if state.stage is _ServerStage.STARTED:
-        state.server.shutdown(state.completion_queue, _SHUTDOWN_TAG)
+def _stop(state, grace, shutdown_handler_grace):
+  shutdown_event = threading.Event()
+
+  def cancel_all_calls_after_grace():
+    with state.lock:
+      if state.stage is _ServerStage.STOPPED:
+        shutdown_event.set()
+        return
+      elif state.stage is _ServerStage.STARTED:
+        do_shutdown = True
         state.stage = _ServerStage.GRACE
         state.shutdown_events = []
-        state.due.add(_SHUTDOWN_TAG)
-      shutdown_event = threading.Event()
+      else:
+        do_shutdown = False
       state.shutdown_events.append(shutdown_event)
-      if grace is None:
+
+    if do_shutdown:
+      # Run Shutdown Handlers without the lock
+      for handler in state.shutdown_handlers:
+        handler(shutdown_handler_grace)
+      with state.lock:
+        state.server.shutdown(state.completion_queue, _SHUTDOWN_TAG)
+        state.stage = _ServerStage.GRACE
+        state.due.add(_SHUTDOWN_TAG)
+
+    if not shutdown_event.wait(timeout=grace):
+      with state.lock:
         state.server.cancel_all_calls()
         # TODO(https://github.com/grpc/grpc/issues/6597): delete this loop.
         for rpc_state in state.rpc_states:
           with rpc_state.condition:
             rpc_state.client = _CANCELLED
             rpc_state.condition.notify_all()
-      else:
-        def cancel_all_calls_after_grace():
-          shutdown_event.wait(timeout=grace)
-          with state.lock:
-            state.server.cancel_all_calls()
-            # TODO(https://github.com/grpc/grpc/issues/6597): delete this loop.
-            for rpc_state in state.rpc_states:
-              with rpc_state.condition:
-                rpc_state.client = _CANCELLED
-                rpc_state.condition.notify_all()
-        thread = threading.Thread(target=cancel_all_calls_after_grace)
-        thread.start()
-        return shutdown_event
-  shutdown_event.wait()
+
+  if grace is None:
+    cancel_all_calls_after_grace()
+  else:
+    threading.Thread(target=cancel_all_calls_after_grace).start()
+
   return shutdown_event
 
 
@@ -716,12 +725,12 @@ def _start(state):
       raise ValueError('Cannot start already-started server!')
     state.server.start()
     state.stage = _ServerStage.STARTED
-    _request_call(state)    
+    _request_call(state)
     def cleanup_server(timeout):
       if timeout is None:
-        _stop(state, _UNEXPECTED_EXIT_SERVER_GRACE).wait()
+        _stop(state, state.exit_grace, state.exit_shutdown_handler_grace).wait()
       else:
-        _stop(state, timeout).wait()
+        _stop(state, timeout, 0).wait()
 
     thread = _common.CleanupThread(
         cleanup_server, target=_serve, args=(state,))
@@ -729,12 +738,16 @@ def _start(state):
 
 class Server(grpc.Server):
 
-  def __init__(self, thread_pool, generic_handlers, options):
+  def __init__(self, thread_pool, generic_handlers, options, exit_grace,
+               exit_shutdown_handler_grace):
     completion_queue = cygrpc.CompletionQueue()
     server = cygrpc.Server(_common.channel_args(options))
     server.register_completion_queue(completion_queue)
     self._state = _ServerState(
-        completion_queue, server, generic_handlers, thread_pool)
+        completion_queue, server, generic_handlers, thread_pool,
+        _DEFAULT_EXIT_GRACE if exit_grace is None else exit_grace,
+        _DEFAULT_EXIT_SHUTDOWN_HANDLER_GRACE if exit_shutdown_handler_grace
+        is None else exit_shutdown_handler_grace)
 
   def add_generic_rpc_handlers(self, generic_rpc_handlers):
     _add_generic_handlers(self._state, generic_rpc_handlers)
@@ -745,11 +758,14 @@ class Server(grpc.Server):
   def add_secure_port(self, address, server_credentials):
     return _add_secure_port(self._state, _common.encode(address), server_credentials)
 
+  def add_shutdown_handler(self, handler):
+    self._state.shutdown_handlers.append(handler)
+
   def start(self):
     _start(self._state)
 
-  def stop(self, grace):
-    return _stop(self._state, grace)
+  def stop(self, grace, shutdown_handler_grace=None):
+    return _stop(self._state, grace, shutdown_handler_grace)
 
   def __del__(self):
-    _stop(self._state, None)
+    _stop(self._state, None, None)
diff --git a/src/python/grpcio_tests/tests/tests.json b/src/python/grpcio_tests/tests/tests.json
index dd4a0257f5..04a2e44178 100644
--- a/src/python/grpcio_tests/tests/tests.json
+++ b/src/python/grpcio_tests/tests/tests.json
@@ -27,6 +27,7 @@
   "unit._cython.cygrpc_test.TypeSmokeTest",
   "unit._empty_message_test.EmptyMessageTest",
   "unit._exit_test.ExitTest",
+  "unit._exit_test.ShutdownHandlerTest",
   "unit._metadata_code_details_test.MetadataCodeDetailsTest",
   "unit._metadata_test.MetadataTest",
   "unit._rpc_test.RPCTest",
diff --git a/src/python/grpcio_tests/tests/unit/_exit_test.py b/src/python/grpcio_tests/tests/unit/_exit_test.py
index 5a4a32887c..342f5fcc10 100644
--- a/src/python/grpcio_tests/tests/unit/_exit_test.py
+++ b/src/python/grpcio_tests/tests/unit/_exit_test.py
@@ -43,6 +43,8 @@ import threading
 import time
 import unittest
 
+import grpc
+from grpc.framework.foundation import logging_pool
 from tests.unit import _exit_scenarios
 
 SCENARIO_FILE = os.path.abspath(os.path.join(
@@ -52,7 +54,7 @@ BASE_COMMAND = [INTERPRETER, SCENARIO_FILE]
 BASE_SIGTERM_COMMAND = BASE_COMMAND + ['--wait_for_interrupt']
 
 INIT_TIME = 1.0
-
+SHUTDOWN_GRACE = 5.0
 
 processes = []
 process_lock = threading.Lock()
@@ -182,5 +184,24 @@ class ExitTest(unittest.TestCase):
     interrupt_and_wait(process)
 
 
+class _ShutDownHandler(object):
+
+  def __init__(self):
+    self.seen_handler_grace = None
+
+  def shutdown_handler(self, handler_grace):
+    self.seen_handler_grace = handler_grace
+
+  
+class ShutdownHandlerTest(unittest.TestCase):
+
+  def test_shutdown_handler(self):
+    server = grpc.server(logging_pool.pool(1))
+    handler = _ShutDownHandler()
+    server.add_shutdown_handler(handler.shutdown_handler)
+    server.start()
+    server.stop(0, shutdown_handler_grace=SHUTDOWN_GRACE).wait()
+    self.assertEqual(SHUTDOWN_GRACE, handler.seen_handler_grace)
+
 if __name__ == '__main__':
   unittest.main(verbosity=2)