From d8bdd717c75c91e30831b428775c4ff2150b96dd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?F=C3=BCleki=20F=C3=A1bi=C3=A1n?= <ffabi1997@gmail.com>
Date: Fri, 8 May 2020 21:24:37 +0200
Subject: [PATCH] Fixed up Redis super storage

---
 app.py | 83 ++++++++++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 64 insertions(+), 19 deletions(-)

diff --git a/app.py b/app.py
index 0a51413..4fb65e9 100644
--- a/app.py
+++ b/app.py
@@ -22,24 +22,67 @@ sentry_sdk.init("https://0a106e104e114bc9a3fa47f9cb0db2f4@sentry.kmlabz.com/10")
 
 
 class RedisSuperStorage:
-    def __init__(self, r):
-        self.r = r
+    def __init__(self, redis_url: str):
+        self.r = redis.from_url(redis_url)
 
-    def get_consumer_dictionary(self):
-        consumer_list_redis = json.loads((self.r.get('consumer_list') or b'{}').decode('utf-8'))
+    def get_consumer_list(self) -> dict:
+        keys = self.r.keys('consumer_*')
 
-        return consumer_list_redis
+        list_of_customers = {}
 
-    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
+        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, r):
-        self.r = r
+
+    def __init__(self):
+        pass
 
     def request_first_sync(self, ip_list, consumer_list_redis):
         temp_dict = {}
@@ -91,7 +134,6 @@ class Scheduler:
 
         elif old_ip != current_ip:
             logging.info(f'IP changed: {old_ip} -> {current_ip} Pushing updates...')
-            self.r.set('current_ip', current_ip.encode('utf-8'))
 
             # pushing updates...
             keys = self.r.keys('producer_*')
@@ -122,22 +164,25 @@ class Scheduler:
         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)
 
-    # Connect redis
-    r = redis.from_url(os.environ.get('REDIS_URL', "redis://localhost:6379/0"))
-
-    redis_storage = RedisSuperStorage(r)
-    scheduler = Scheduler(r)
+    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')
 
-    host_name = socket.gethostname()
     ip_list = redis.get_ip_list()
 
     scheduler.request_list(ip_list)