diff --git a/webservice/job_monitor.py b/webservice/job_monitor.py index 5b1d2140b4a7b1334cffd4a446a007c3a2148031..b63383cfa93a4a3fca4abba5b4f5352e1461eff2 100644 --- a/webservice/job_monitor.py +++ b/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__":