From 87cfdecc2de3224e4ceb331ed792b212a81be9b3 Mon Sep 17 00:00:00 2001
From: marcsello <punkosdmarcell@rocketmail.com>
Date: Wed, 28 Jul 2021 14:10:59 +0200
Subject: [PATCH] less prone to crashing

---
 src/app.py    | 79 +++++++++++++++++++++++++++++++++++----------------
 src/config.py |  1 +
 2 files changed, 56 insertions(+), 24 deletions(-)

diff --git a/src/app.py b/src/app.py
index 33fd089..e884226 100644
--- a/src/app.py
+++ b/src/app.py
@@ -40,18 +40,22 @@ if config.SENTRY_DSN:
 def setup_rabbit(mqtt_: MQTT) -> None:
     logging.info("Connecting to RabbitMQ...")
     credentials = pika.PlainCredentials(config.RABBIT_USERNAME, config.RABBIT_PASSWORD)
-    connection = pika.BlockingConnection(pika.ConnectionParameters(host=config.RABBIT_HOSTNAME,
-                                                                   credentials=credentials,
-                                                                   heartbeat=30,
-                                                                   socket_timeout=45))
-    channel = connection.channel()
-    channel.exchange_declare(exchange=config.RABBIT_EXCHANGE, exchange_type='fanout')
-    queue = channel.queue_declare(durable=True, auto_delete=True, queue=uuid.uuid4().urn.split(':')[2],
-                                  exclusive=True).method.queue
-    channel.queue_bind(exchange=config.RABBIT_EXCHANGE, queue=queue)
-    channel.basic_consume(queue=queue, on_message_callback=on_message_creator(mqtt_), auto_ack=True)
-    logging.debug("Starting consumption...")
-    channel.start_consuming()
+    while True:
+        connection = pika.BlockingConnection(pika.ConnectionParameters(host=config.RABBIT_HOSTNAME,
+                                                                       credentials=credentials,
+                                                                       heartbeat=30,
+                                                                       socket_timeout=45))
+        channel = connection.channel()
+        channel.exchange_declare(exchange=config.RABBIT_EXCHANGE, exchange_type='fanout')
+        queue = channel.queue_declare(durable=True, auto_delete=True, queue=uuid.uuid4().urn.split(':')[2],
+                                      exclusive=True).method.queue
+        channel.queue_bind(exchange=config.RABBIT_EXCHANGE, queue=queue)
+        channel.basic_consume(queue=queue, on_message_callback=on_message_creator(mqtt_), auto_ack=False)
+        logging.debug("Starting consumption...")
+        try:
+            channel.start_consuming()  # this automagically responds to heartbeats
+        except pika.exceptions.AMQPConnectionError as e:
+            logging.warning(f"AMQP Error happened: {e}; Reconnecting...")
 
 
 def on_message_creator(mqtt_: MQTT):
@@ -59,27 +63,54 @@ def on_message_creator(mqtt_: MQTT):
     This generator is used, so that the mqtt object can be injected just when the callback is registered
     """
 
-    def on_message(channel, method_frame, header_frame, body):
-        msg_json = json.loads(body)
-        if 'probability' not in msg_json:
-            logging.error("Malformed message from classifier")
+    def on_message(
+            channel: pika.channel.Channel,
+            method: pika.spec.Basic.Deliver,
+            properties: pika.spec.BasicProperties,
+            body: bytes
+    ):
+        try:
+            msg_json = json.loads(body)
+        except (json.JSONDecodeError, UnicodeDecodeError) as e:
+            logging.error(f"Malformed message from classifier: {e}")
+            channel.basic_ack(delivery_tag=method.delivery_tag)
             return
 
-        if 'class' not in msg_json:
-            logging.error("Malformed message from classifier")
+        if ('probability' not in msg_json) or ('class' not in msg_json):
+            logging.error("Malformed message from classifier: Missing fields")
+            channel.basic_ack(delivery_tag=method.delivery_tag)
             return
 
         # TODO: strurnus should not be hardcoded here
         if (msg_json['class'] == 'sturnus') and (msg_json['probability'] > config.TRIGGER_LEVEL):
-            r = requests.get(f"http://{config.INPUT_HOSTNAME}/sample/{msg_json['tag']}")
-            r.raise_for_status()
+            try:
+                r = requests.get(
+                    f"http://{config.INPUT_HOSTNAME}/sample/{msg_json['tag']}",
+                    timeout=config.INPUT_TIMEOUT
+                )
+            except requests.exceptions.Timeout:
+                logging.error(f"Input-service timed out! (Timeout: {config.INPUT_TIMEOUT} sec)")
+                return  # no ack
+
+            if r.status_code != 200:
+                logging.error(f"Input-service status code is not 200: {r.status_code}")
+                return  # no ack
+
             if 'device_id' not in r.json():
                 logging.error("Input-service response invalid")
-                return
+                return  # no ack
 
-            logging.info(f"Sending alert command to device {r.json()['device_id']}")
-            mqtt_.publish(subtopic=r.json()['device_id'],
-                          message=json.dumps({"command": "doAlert"}))
+            logging.info(f"Sending alert command to device {r.json()['device_id']}...")
+            mqtt_.publish(
+                subtopic=r.json()['device_id'],
+                message=json.dumps({"command": "doAlert"})
+            )
+
+        else:
+            logging.debug(f"Probability is either bellow trigger level, or not the target class. Nothing to do.")
+
+        # This concludes the job
+        channel.basic_ack(delivery_tag=method.delivery_tag)
 
     return on_message
 
diff --git a/src/config.py b/src/config.py
index b5552cd..d432a05 100644
--- a/src/config.py
+++ b/src/config.py
@@ -30,4 +30,5 @@ MQTT_PASSWORD = os.getenv("GUARD_MQTT_PASSWORD", "guard-service")
 MQTT_TOPIC = os.getenv("GUARD_MQTT_TOPIC", "guard-service")
 
 INPUT_HOSTNAME = os.getenv("INPUT_SVC_HOSTNAME", "input-service")
+INPUT_TIMEOUT = int(os.environ.get("INPUT_SVC_TIMEOUT", 5))
 TRIGGER_LEVEL = float(os.environ.get("TRIGGER_LEVEL", 0.51))