bobcat-ticketing / mqtt-recorder Goto Github PK

$ git diff
diff --git a/mqtt_recorder.py b/mqtt_recorder.py
index 0d30b26..8945456 100644
--- a/mqtt_recorder.py
+++ b/mqtt_recorder.py
@@ -8,13 +8,13 @@ import sys
 import base64
 import time
 from hbmqtt.client import MQTTClient, QOS_0, QOS_1
-
+import signal
+from concurrent.futures import CancelledError

 TOPICS = [("#", QOS_1)]

 logger = logging.getLogger('mqtt_recorder')

-
 async def mqtt_record(server: str, output: str = None) -> None:
     """Record MQTT messages"""
     mqtt = MQTTClient()
@@ -36,6 +36,9 @@ async def mqtt_record(server: str, output: str = None) -> None:
         print(json.dumps(record), file=output_file)


+
+
+
 async def mqtt_replay(server: str, input: str = None, delay: int = 0, realtime: bool = False, scale: float = 1) -> None:
     """Replay MQTT messages"""
     mqtt = MQTTClient()
@@ -73,6 +76,10 @@ async def mqtt_replay(server: str, input: str = None, delay: int = 0, realtime:
             await asyncio.sleep(delay_s)


+async def shutdown(sig, loop):
+    loop.stop()
+
+
 def main():
     """ Main function"""
     parser = argparse.ArgumentParser(description='MQTT recorder')
@@ -124,13 +131,22 @@ def main():
     else:
         logging.basicConfig(level=logging.INFO)

+
     if args.mode == 'replay':
         process = mqtt_replay(server=args.server, input=args.input, delay=args.delay, realtime=args.realtime, scale=1/args.speed)
     else:
         process = mqtt_record(server=args.server, output=args.output)
+    for signame in (signal.SIGINT, signal.SIGTERM):
+        asyncio.get_event_loop().add_signal_handler(signame,
+                            lambda: asyncio.ensure_future(shutdown(signame, asyncio.get_event_loop())))
+
+    try:
+        asyncio.get_event_loop().run_until_complete(process)
+    except:
+        pass

-    asyncio.get_event_loop().run_until_complete(process)


 if __name__ == "__main__":
     main()