less prone to crashing
This commit is contained in:
		
							
								
								
									
										79
									
								
								src/app.py
									
									
									
									
									
								
							
							
						
						
									
										79
									
								
								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
 | 
			
		||||
 | 
			
		||||
 
 | 
			
		||||
@@ -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))
 | 
			
		||||
 
 | 
			
		||||
		Reference in New Issue
	
	Block a user