consumer-scheduler/app.py

#!/usr/bin/env python
import sentry_sdk
import time
import requests
import requests.exceptions
import os
import redis
import json
import logging
import socket

"""
Scheduler
"""

__author__ = "@kocsisr"
__copyright__ = "Copyright 2020, GoldenPogácsa Team"
__module_name__ = "app"
__version__text__ = "1"

sentry_sdk.init("https://0a106e104e114bc9a3fa47f9cb0db2f4@sentry.kmlabz.com/10")


class RedisSuperStorage:
    def __init__(self, redis_url: str):
        self.r = redis.from_url(redis_url)

    def get_consumer_list(self) -> dict:
        keys = self.r.keys('consumer_*')

        list_of_customers = {}

        for key in keys:
            info = json.loads((self.r.get(key) or b"{}").decode('utf-8'))

            if info:
                list_of_customers[info['uuid']] = info

        return list_of_customers

    def get_producer_list(self) -> list:
        keys = self.r.keys('producer_*')

        list_of_producer_ip = []

        for key in keys:
            ip = (self.r.get(key) or b"").decode('utf-8')

            if ip:
                list_of_producer_ip.append(ip)

        return list_of_producer_ip

    def is_ip_changed(self) -> bool:
        old_ip = self.r.get('current_ip')

        if old_ip:
            old_ip = old_ip.decode('utf-8')

        host_name = socket.gethostname()
        current_ip = socket.gethostbyname(host_name)

        if current_ip != old_ip:
            self.r.set('current_ip', current_ip.encode('utf-8'))

        return current_ip != old_ip

    def update_consumer(self, uuid: str, ip: str):

        cust_key = f"consumer_{uuid}"

        info = {
            "uuid": uuid,
            "ip": ip,
            "last_seen": time.time()
        }

        self.r.set(cust_key, json.dumps(info).encode('utf-8'))
        self.r.expire(cust_key, os.environ["CUSTOMER_TIMEOUT"])


class Scheduler:

    def __init__(self):
        pass

    def request_first_sync(self, ip_list, consumer_list_redis):
        temp_dict = {}
        for ip in ip_list:
            try:
                # request synchronization
                response = requests.post(f"http://{ip}/sync", json={'uuid': os.environ['LOCAL_UUID']}, timeout=5)
            except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e:
                logging.error(f"Error while syncing to {ip}: {str(e)}")
                continue

            if response.status_code == 200:
                temp_dict[response.json()['uuid']] = {'ip': ip}
        consumer_list_redis.update(temp_dict)
        self.r.set('consumer_list', json.dumps(consumer_list_redis).encode('utf-8'))
        logging.debug('Update redis consumers ip list from first answers: Done')

    def request_sync(self, consumer_list_redis):
        temp_dict = {}
        for uuid, info in consumer_list_redis.items():
            ip = info['ip']
            try:
                # request synchronization
                response = requests.post(f"http://{ip}/sync", json={'uuid': os.environ['LOCAL_UUID']}, timeout=5)
            except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e:
                logging.error(f"Error while syncing to {ip}: {str(e)}")
                continue

            if response.status_code == 200:
                temp_dict[response.json()['uuid']] = {'ip': ip}

        # update the dictionary of the currently available consumers
        consumer_list_redis.update(temp_dict)
        self.r.set('consumer_list', json.dumps(consumer_list_redis).encode('utf-8'))
        logging.debug('Update redis consumer ip list from answers: Done')

    def test_ip_change(self, host_name):

        old_ip = self.r.get('current_ip')

        if old_ip:
            old_ip = old_ip.decode('utf-8')

        current_ip = socket.gethostbyname(host_name)

        if not old_ip:  # Not set yet. I this case no update required
            self.r.set('current_ip', current_ip.encode('utf-8'))
            logging.debug(f"Previous info about the ip address could not be found! Current: {current_ip}")

        elif old_ip != current_ip:
            logging.info(f'IP changed: {old_ip} -> {current_ip} Pushing updates...')

            # pushing updates...
            keys = self.r.keys('producer_*')

            logging.debug(f'Pushing update to the following producers: ' + ', '.join(k.decode('utf-8') for k in keys))

            for key in keys:
                ip = self.r.get(key)
                if ip:
                    ip = ip.decode('utf-8')
                else:
                    continue

                try:
                    response = requests.post(
                        f"http://{ip}/ip",
                        json={'uuid': os.environ['LOCAL_UUID'], 'ip': current_ip},
                        timeout=5
                    )
                    logging.debug(f"Pushed update to {key.decode('utf-8')} at {ip}. Response: {response.status_code}")
                except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e:
                    logging.warning(f"Could not push update to {key.decode('utf-8')}: {str(e)}")
                    continue
        else:
            logging.debug(f'IP unchanged: {current_ip}')

        logging.debug('Waiting for next turn')
        time.sleep(os.environ.get("RUN_INTERVAL", 30))


def get_ip_list(self):
    # set initial consumer addresses
    ip_list = os.environ['INITIAL_SERVERS'].split(',')
    logging.debug('Get consumer list from environ at first: Done')
    return ip_list


def main():
    # set logging preferences
    logging.basicConfig(filename='', level=logging.DEBUG)

    redis_storage = RedisSuperStorage(os.environ.get('REDIS_URL', "redis://localhost:6379/0"))
    scheduler = Scheduler()

    # get the dictionary of the currently available consumers

    consumer_list_redis = redis_storage.get_consumer_dictionary()
    logging.debug('Get consumer list from redis at first: Done')

    ip_list = redis.get_ip_list()

    scheduler.request_list(ip_list)
    scheduler.request_first_sync(ip_list, consumer_list_redis)

    while True:
        logging.debug('Infinite Cycle start : Done')
        # get the dictionary of the currently available consumers
        consumer_list_redis = redis.get_consumer_dictionary()
        logging.debug('Get consumer list from redis: Done')

        temp_dict = {}

        scheduler.request_sync(consumer_list_redis)
        # Test ip change stuff
        scheduler.test_ip_change(host_name)


if __name__ == "__main__":
    try:

        main()
    except KeyboardInterrupt:
        pass