[wip] Shut down gracefully on SIGTERM

snekomatic/app.py

@@ -1,5 +1,7 @@
+from functools import partial
 import sys
 import os
+import signal
 import trio
 from glom import glom
 import hypercorn
@@ -144,6 +146,18 @@ async def _member_state(gh_client, org, member):
         return glom(response, "state")
 
 
+# This is used for testing; it never actually happens
+@github_app.route("pull_request", action="rotated")
+async def pull_request_rotated(event_type, payload, gh_client):
+    print("PR rotated (I guess you're running the test suite)")
+    try:
+        await trio.sleep(glom(payload, "rotation_time"))
+    except BaseException as exc:
+        print(f"rotation interrupted by {exc!r}")
+        raise
+    print("rotation complete")
+
+
 # There's no "merged" event; instead you get action=closed + merged=True
 @github_app.route("pull_request", action="closed")
 async def pull_request_merged(event_type, payload, gh_client):
@@ -184,6 +198,16 @@ async def pull_request_merged(event_type, payload, gh_client):
 
 
 async def main(*, task_status=trio.TASK_STATUS_IGNORED):
+    shutdown_event = trio.Event()
+
+    async def listen_for_sigterm(*, task_status=trio.TASK_STATUS_IGNORED):
+        with trio.open_signal_receiver(signal.SIGTERM) as signal_aiter:
+            task_status.started()
+            async for signum in signal_aiter:
+                print(f"shutdown on signal signum: {signum}")
+                shutdown_event.set()
+                return
+
     print("~~~ Starting up! ~~~")
     # On Heroku, have to bind to whatever $PORT says:
     # https://devcenter.heroku.com/articles/dynos#local-environment-variables
@@ -197,6 +221,12 @@ async def main(*, task_status=trio.TASK_STATUS_IGNORED):
             # Setting this just silences a warning:
             worker_class="trio",
         )
-        urls = await nursery.start(hypercorn.trio.serve, quart_app, config)
+        await nursery.start(listen_for_sigterm)
+        urls = await nursery.start(partial(
+            hypercorn.trio.serve,
+            quart_app,
+            config,
+            shutdown_trigger=shutdown_event.wait
+        ))
         print("Accepting HTTP requests at:", urls)
         task_status.started(urls)

tests/test_app.py

@@ -5,24 +5,31 @@
 import attr
 import urllib.parse
 import json
+import signal
+import sys
 
+import trio
 from snekomatic.app import main, SENT_INVITATION
 from snekomatic.gh import GithubApp, BaseGithubClient
 from .util import fake_webhook, save_environ
 from .credentials import *
 
 
 @pytest.fixture
-async def our_app_url(nursery, heroku_style_pg):
+def our_app_environment(heroku_style_pg):
     with save_environ():
         os.environ["GITHUB_USER_AGENT"] = TEST_USER_AGENT
         os.environ["GITHUB_APP_ID"] = TEST_APP_ID
         os.environ["GITHUB_PRIVATE_KEY"] = TEST_PRIVATE_KEY
         os.environ["GITHUB_WEBHOOK_SECRET"] = TEST_WEBHOOK_SECRET
         os.environ["PORT"] = "0"  # let the OS pick an unused port
+        yield
 
-        urls = await nursery.start(main)
-        yield urls[0]
+
+@pytest.fixture
+async def our_app_url(nursery, our_app_environment):
+    urls = await nursery.start(main)
+    return urls[0]
 
 
 async def test_main_smoke(our_app_url):
@@ -161,3 +168,76 @@ async def fake_request(self, method, url, headers, body):
     assert did_invite == did_comment
     assert did_invite == s.expect_invite
     assert (PR_CREATOR in SENT_INVITATION) == s.expect_in_db_after
+
+
+@pytest.mark.skipif(sys.platform.startswith("win"), reason="requires Unix")
+async def test_exits_cleanly_on_sigterm(our_app_environment, mock_clock):
+    # wait 1 second to make sure all I/O settles before jumping clock
+    mock_clock.autojump_threshold = 1.0
+
+    # Simple test: make sure SIGTERM does cause a clean shutdown
+    async with trio.open_nursery() as nursery:
+        await nursery.start(main)
+        os.kill(os.getpid(), signal.SIGTERM)
+
+    # Harder test: make sure running webhook handler completes before shutdown
+
+    async with trio.open_nursery() as nursery:
+        our_app_urls = await nursery.start(main)
+        our_app_url = our_app_urls[0]
+        # 1. send webhook
+        headers, body = fake_webhook(
+            "pull_request",
+            {
+                "action": "rotated",
+                "rotation_time": 10, # seconds
+                "installation": {
+                    "id": "000000",
+                },
+            },
+            TEST_WEBHOOK_SECRET,
+        )
+
+        post_finished = trio.Event()
+        async def post_then_notify():
+            try:
+                await asks.post(
+                    urllib.parse.urljoin(our_app_url, "webhook/github"),
+                    headers=headers,
+                    data=body,
+                )
+                post_finished.set()
+            except BaseException as exc:
+                print(f"post_then_notify crashing with {exc!r}")
+
+        nursery.start_soon(post_then_notify)
+
+        # 2. give webhook chance to be delivered; it should sleep for 10
+        # seconds or so
+        try:
+            print("first sleep")
+            await trio.sleep(3)
+            print("first wake")
+
+            # 3. send SIGTERM
+            print("killing")
+            os.kill(os.getpid(), signal.SIGTERM)
+
+            # 4. test sleeps for "5 seconds", to give SIGTERM time to work
+            print("second sleep")
+            await trio.sleep(3)
+            print("second wake")
+
+            # 5. make sure app has not actually exited
+            print(f"post_finished {post_finished!r}")
+            assert not post_finished.is_set()
+            # 5b. but it is refusing new requests
+            with pytest.raises(asks.errors.BadHttpResponse):
+                response = await asks.get(our_app_url)
+            assert "Hi!" in response.text
+
+            # 6. test sleep for "6 seconds"
+            # 7. make sure app has exited (and webhook got response)
+        except BaseException as exc:
+            print(f"test is failing with {exc!r}")
+            raise