workaround for graceful_exit / wait_closed

grpclib/utils.py

@@ -154,12 +154,20 @@ def _exit_handler(
  flag.append(True)
 
 
+def _single_stage_exit_handler(sig_num: int, servers: Collection['IClosable']) -> None:
+ """More aggressive version of exit_handler which runs the two stages sequentially.
+ """
+ _first_stage(cast('signal.Signals', sig_num), servers)
+ _second_stage(cast('signal.Signals', sig_num))
+
+
 @contextmanager
 def graceful_exit(
  servers: Collection['IClosable'],
  *,
  loop: Optional[asyncio.AbstractEventLoop] = None,
  signals: Collection[int] = (signal.SIGINT, signal.SIGTERM),
+ single_stage: bool = False,
 ) -> Iterator[None]:
  """Utility context-manager to help properly shutdown server in response to
  the OS signals
@@ -203,6 +211,7 @@ async def main(...):
  :param servers: list of servers
  :param loop: (deprecated) asyncio-compatible event loop
  :param signals: set of the OS signals to handle
+ :param single_stage: call first and second stages sequentially
 
  .. note:: Not supported in Windows
  """
@@ -215,7 +224,10 @@ async def main(...):
  signals = set(signals)
  flag: 'List[bool]' = []
  for sig_num in signals:
- loop.add_signal_handler(sig_num, _exit_handler, sig_num, servers, flag)
+ if single_stage:
+ loop.add_signal_handler(sig_num, _single_stage_exit_handler, sig_num, servers)
+ else:
+ loop.add_signal_handler(sig_num, _exit_handler, sig_num, servers, flag)
  try:
  yield
  finally:

tests/test_utils.py

@@ -118,7 +118,7 @@ def test_graceful_exit_normal_server(sig_num):
 
 async def main():
  server = Server([])
- with graceful_exit([server]):
+ with graceful_exit([server]{extra_args}):
  await server.start('127.0.0.1')
  print("Started!")
  await server.wait_closed()
@@ -136,7 +136,7 @@ async def main():
  (signal.SIGTERM, signal.SIGINT),
 ])
 def test_graceful_exit_sluggish_server(sig1, sig2):
- cmd = [sys.executable, '-u', '-c', SLUGGISH_SERVER]
+ cmd = [sys.executable, '-u', '-c', SLUGGISH_SERVER.format(extra_args="")]
  with subprocess.Popen(cmd, stdout=subprocess.PIPE) as proc:
  try:
  assert proc.stdout.readline() == b'Started!\n'
@@ -151,6 +151,19 @@ def test_graceful_exit_sluggish_server(sig1, sig2):
  proc.kill()
 
 
+def test_graceful_exit_single_stage():
+ cmd = [sys.executable, '-u', '-c', SLUGGISH_SERVER.format(extra_args=", single_stage=True")]
+ with subprocess.Popen(cmd, stdout=subprocess.PIPE) as proc:
+ try:
+ assert proc.stdout.readline() == b'Started!\n'
+ time.sleep(0.001)
+ proc.send_signal(signal.SIGTERM)
+ assert proc.wait(1) == 128 + signal.SIGTERM
+ finally:
+ if proc.returncode is None:
+ proc.kill()
+
+
 def test_cached():
  def func():
  return 42