Clean shutdown on SIGTERM

webservice/job_monitor.py

@@ -3,6 +3,7 @@ import argparse
 import json
 import locale
 import logging
+import signal
 import time
 from datetime import datetime, timezone
 from pathlib import Path
@@ -93,6 +94,13 @@ class JobsMonitor:
         self.kafka_topic = config['kafka']['topic']
         self.time_interval = int(config['web-service']['job-update-interval'])
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.job_db.close()
+        self.kafka_prod.close(timeout=5)
+
     def run(self):
         while True:
             try:
@@ -293,6 +301,9 @@ class JobsMonitor:
             log.error("Failed to update MDC dark run id %s", dark_run_id)
             log.error(Errors.MDC_RESPONSE.format(response))
 
+def interrupted(signum, frame):
+    raise KeyboardInterrupt
+
 def main(argv=None):
     # Ensure files are opened as UTF-8 by default, regardless of environment.
     locale.setlocale(locale.LC_CTYPE, ('en_US', 'UTF-8'))
@@ -318,7 +329,15 @@ def main(argv=None):
     )
     # DEBUG logs from kafka-python are very verbose, so we'll turn them off
     logging.getLogger('kafka').setLevel(logging.INFO)
-    JobsMonitor(config).run()
+
+    # Treat SIGTERM like SIGINT (Ctrl-C) & do a clean shutdown
+    signal.signal(signal.SIGTERM, interrupted)
+
+    with JobsMonitor(config) as jm:
+        try:
+            jm.run()
+        except KeyboardInterrupt:
+            logging.info("Shutting down on SIGINT/SIGTERM")
 
 
 if __name__ == "__main__":