From f07e29ab56be325e91069c276e890cb7365a679b Mon Sep 17 00:00:00 2001
From: Abe Winter <abe-winter@users.noreply.github.com>
Date: Wed, 1 May 2024 18:16:30 -0400
Subject: [PATCH 1/2] add single_stage flag to graceful_exit

---
 grpclib/utils.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/grpclib/utils.py b/grpclib/utils.py
index 64ce881..c93728c 100644
--- a/grpclib/utils.py
+++ b/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:

From cf1a80df0e1d87fca5a0632ad059c34064f1f178 Mon Sep 17 00:00:00 2001
From: Abe Winter <abe-winter@users.noreply.github.com>
Date: Wed, 1 May 2024 18:21:48 -0400
Subject: [PATCH 2/2] test it

---
 tests/test_utils.py | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/tests/test_utils.py b/tests/test_utils.py
index 3e57175..738cee9 100644
--- a/tests/test_utils.py
+++ b/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