Merge pull request #1 from 116davinder/v2

create multithreaded kafka consumer ( Single Topic +  Multi Partition )

Author: 116davinder

README.md

@@ -4,7 +4,12 @@
 * It will auto resume from same point from where it died if given consumer group name is same before and after crash.
 * it will upload `current.bin` file to s3 which contains messages upto `NUMBER_OF_MESSAGE_PER_BACKUP_FILE`
 but will only upload with other backup files.
-* upload to s3 is background process and it depends on `RETRY_UPLOAD_SECONDS`.
+* `RETRY_UPLOAD_SECONDS` controls upload to s3 or other cloud storage.
+* `NUMBER_OF_KAFKA_THREADS` is used to parallelise reading from kafka topic.
+It should not be more than number of partitions.
+* `LOG_LEVEL` values can be found https://docs.python.org/3/library/logging.html#logging-levels
+* `NUMBER_OF_MESSAGE_PER_BACKUP_FILE` will try to keep this number consistent in file
+but if application got restarted then it may be vary for first back file.
 
 **Restore Application**
 * it will restore from backup dir into given topic.
@@ -25,50 +30,76 @@ python3 backup.py backup.json
 **Local Filesytem Backup.json**
 ```
 {
-  "BOOTSTRAP_SERVERS": "localhost:9092",
+  "BOOTSTRAP_SERVERS": "kafka01:9092,kafka02:9092,kafka03:9092",
   "TOPIC_NAMES": ["davinder.test"],
   "GROUP_ID": "Kafka-BackUp-Consumer-Group",
   "FILESYSTEM_TYPE": "LINUX",
   "FILESYSTEM_BACKUP_DIR": "/tmp/",
-  "NUMBER_OF_MESSAGE_PER_BACKUP_FILE": 50
+  "NUMBER_OF_MESSAGE_PER_BACKUP_FILE": 1000,
+  "RETRY_UPLOAD_SECONDS": 100,
+  "NUMBER_OF_KAFKA_THREADS": 3,
+  "LOG_LEVEL": 20
 }
 ```
 
 **S3 backup.json**
 ```
 {
-  "BOOTSTRAP_SERVERS": "localhost:9092",
+  "BOOTSTRAP_SERVERS": "kafka01:9092,kafka02:9092,kafka03:9092",
   "TOPIC_NAMES": ["davinder.test"],
   "GROUP_ID": "Kafka-BackUp-Consumer-Group",
   "FILESYSTEM_TYPE": "S3",
   "FILESYSTEM_BACKUP_DIR": "/tmp/",
-  "NUMBER_OF_MESSAGE_PER_BACKUP_FILE": 50,
-  "RETRY_UPLOAD_SECONDS": 100
+  "NUMBER_OF_MESSAGE_PER_BACKUP_FILE": 1000,
+  "RETRY_UPLOAD_SECONDS": 100,
+  "NUMBER_OF_KAFKA_THREADS": 3,
+  "LOG_LEVEL": 20
 }
 ```
-
+**Example Local Backup Run Output**
+```
+{ "@timestamp": "2020-06-08 10:56:34,557","level": "INFO","thread": "Kafka Consumer 0","name": "root","message": "started polling on davinder.test" }
+{ "@timestamp": "2020-06-08 10:56:34,557","level": "INFO","thread": "Kafka Consumer 1","name": "root","message": "started polling on davinder.test" }
+{ "@timestamp": "2020-06-08 10:56:34,557","level": "INFO","thread": "Kafka Consumer 2","name": "root","message": "started polling on davinder.test" }
+{ "@timestamp": "2020-06-08 10:56:51,590","level": "INFO","thread": "Kafka Consumer 1","name": "root","message": "Created Successful Backupfile /tmp/davinder.test/1/20200608-105651.tar.gz" }
+{ "@timestamp": "2020-06-08 10:56:51,593","level": "INFO","thread": "Kafka Consumer 1","name": "root","message": "Created Successful Backup sha256 file of /tmp/davinder.test/1/20200608-105651.tar.gz.sha256" }
+{ "@timestamp": "2020-06-08 10:57:17,270","level": "INFO","thread": "Kafka Consumer 0","name": "root","message": "Created Successful Backupfile /tmp/davinder.test/0/20200608-105717.tar.gz" }
+{ "@timestamp": "2020-06-08 10:57:17,277","level": "INFO","thread": "Kafka Consumer 0","name": "root","message": "Created Successful Backup sha256 file of /tmp/davinder.test/0/20200608-105717.tar.gz.sha256" }
+{ "@timestamp": "2020-06-08 10:57:17,399","level": "INFO","thread": "Kafka Consumer 2","name": "root","message": "Created Successful Backupfile /tmp/davinder.test/2/20200608-105717.tar.gz" }
+{ "@timestamp": "2020-06-08 10:57:17,406","level": "INFO","thread": "Kafka Consumer 2","name": "root","message": "Created Successful Backup sha256 file of /tmp/davinder.test/2/20200608-105717.tar.gz.sha256" }
+...
+```
 **Example S3 Backup Run Output**
 ```
 $ python3 backup.py backup.json
-{ "@timestamp": "2020-06-01 10:37:00,168","level": "INFO","thread": "MainThread","name": "root","message": "Successful loading of config.json file" }
-{ "@timestamp": "2020-06-01 10:37:00,169","level": "INFO","thread": "MainThread","name": "root","message": "all required variables are successfully" }
-{ "@timestamp": "2020-06-01 10:37:00,187","level": "INFO","thread": "Kafka Consumer","name": "root","message": "starting polling on davinder.test" }
-{ "@timestamp": "2020-06-01 10:38:17,291","level": "INFO","thread": "Kafka Consumer","name": "root","message": "Created Successful Backupfile /tmp/davinder.test/20200601-103817.tar.gz" }
-{ "@timestamp": "2020-06-01 10:39:00,631","level": "INFO","thread": "S3-Upload","name": "root","message": "upload successful at s3://davinder-test-kafka-backup/davinder.test/20200601-103817.tar.gz" }
+{ "@timestamp": "2020-06-10 12:49:43,871","level": "INFO","thread": "S3 Upload","name": "botocore.credentials","message": "Found credentials in environment variables." }
+{ "@timestamp": "2020-06-10 12:49:43,912","level": "INFO","thread": "Kafka Consumer 1","name": "root","message": "started polling on davinder.test" }
+{ "@timestamp": "2020-06-10 12:49:43,915","level": "INFO","thread": "Kafka Consumer 0","name": "root","message": "started polling on davinder.test" }
+{ "@timestamp": "2020-06-10 12:49:43,916","level": "INFO","thread": "Kafka Consumer 2","name": "root","message": "started polling on davinder.test" }
+{ "@timestamp": "2020-06-10 12:49:44,307","level": "INFO","thread": "S3 Upload","name": "root","message": "upload successful at s3://davinder-test-kafka-backup/davinder.test/0/20200608-102909.tar.gz" }
+{ "@timestamp": "2020-06-10 12:49:45,996","level": "INFO","thread": "S3 Upload","name": "root","message": "waiting for new files to be generated" }
+{ "@timestamp": "2020-06-10 12:52:33,130","level": "INFO","thread": "Kafka Consumer 0","name": "root","message": "Created Successful Backupfile /tmp/davinder.test/0/20200610-125233.tar.gz" }
+{ "@timestamp": "2020-06-10 12:52:33,155","level": "INFO","thread": "Kafka Consumer 0","name": "root","message": "Created Successful Backup sha256 file of /tmp/davinder.test/0/20200610-125233.tar.gz.sha256" }
 ....
 ```
 
 # Backup Directory Structure
 ```
-$ tree davinder.test/
-davinder.test/
-├── 20204025-154046.tar.gz
-├── 20204025-154046.tar.gz.sha256
-├── 20204325-154344.tar.gz
-├── 20204325-154344.tar.gz.sha256
-└── current.bin
+/tmp/davinder.test/
+├── 0
+│   ├── 20200608-102909.tar.gz
+│   ├── 20200608-102909.tar.gz.sha256
+│   └── current.bin
+├── 1
+│   ├── 20200608-102909.tar.gz
+│   ├── 20200608-102909.tar.gz.sha256
+│   └── current.bin
+└── 2
+    ├── 20200608-102909.tar.gz
+    ├── 20200608-102909.tar.gz.sha256
+    └── current.bin
 
-0 directories, 5 files
+3 directories, 9 files
 ```
 
 # How to Run Kafka Restore Application

backup.json

@@ -5,5 +5,7 @@
   "FILESYSTEM_TYPE": "S3",
   "FILESYSTEM_BACKUP_DIR": "/tmp/",
   "NUMBER_OF_MESSAGE_PER_BACKUP_FILE": 1000,
-  "RETRY_UPLOAD_SECONDS": 100
+  "RETRY_UPLOAD_SECONDS": 100,
+  "NUMBER_OF_KAFKA_THREADS": 2,
+  "LOG_LEVEL": 20
 }

backup.py

@@ -3,7 +3,7 @@
 from common import Common
 import logging
 import threading
-from upload import Upload
+from cloud import Upload
 
 class KBackup:
     def __init__(self,config):
@@ -12,30 +12,38 @@ def __init__(self,config):
         self.GROUP_ID = config['GROUP_ID']
         self.TOPIC_NAME_LIST = config['TOPIC_NAMES']
         self.BACKUP_DIR = os.path.join(config['FILESYSTEM_BACKUP_DIR'], self.TOPIC_NAME_LIST[0])
-        self.BACKUP_TMP_FILE = os.path.join(self.BACKUP_DIR, "current.bin")
         try:
             self.NUMBER_OF_MESSAGE_PER_BACKUP_FILE = int(config['NUMBER_OF_MESSAGE_PER_BACKUP_FILE'])
         except:
-            logging.info(
+            logging.error(
                 f"NUMBER_OF_MESSAGE_PER_BACKUP_FILE {str(config['NUMBER_OF_MESSAGE_PER_BACKUP_FILE'])} is not integer value"
             )
             self.NUMBER_OF_MESSAGE_PER_BACKUP_FILE = 50
-            logging.info(f"NUMBER_OF_MESSAGE_PER_BACKUP_FILE is set to default value 50")
-
+            logging.debug(f"NUMBER_OF_MESSAGE_PER_BACKUP_FILE is set to default value 50")
+        try:
+            self.NUMBER_OF_KAFKA_THREADS = config['NUMBER_OF_KAFKA_THREADS']
+        except:
+            self.NUMBER_OF_KAFKA_THREADS = 1
         self.CONSUMERCONFIG = {
             'bootstrap.servers': self.BOOTSTRAP_SERVERS,
             'group.id': self.GROUP_ID,
             'auto.offset.reset': 'earliest'
         }
-        logging.info(f"successful loading of all variables")
+        try:
+            self.LOG_LEVEL = config['LOG_LEVEL']
+        except:
+            self.LOG_LEVEL = logging.INFO
+
+        logging.debug(f"successful loading of all variables")
 
     def backup(self):
         _bt = confluent_kafka.Consumer(self.CONSUMERCONFIG)
         _bt.subscribe(self.TOPIC_NAME_LIST)
 
-        Common.createBackupTopicDir(self.BACKUP_DIR)
+        for p in Common.findNumberOfPartitionsInTopic(_bt.list_topics().topics[self.TOPIC_NAME_LIST[0]].partitions):
+            Common.createDir(os.path.join(self.BACKUP_DIR, str(p)))
 
-        count = Common.currentMessageCountInBinFile(self.BACKUP_TMP_FILE)
+        count = 0
         logging.info(f"started polling on {self.TOPIC_NAME_LIST[0]}")
         while True:
             msg = _bt.poll(timeout=1.0)
@@ -45,18 +53,22 @@ def backup(self):
             if msg.error():
                 logging.error(f"{msg.error()}")
                 continue
-            else:
+            if msg.partition() is not None:
+                _tmp_file = os.path.join(self.BACKUP_DIR, str(msg.partition()) ,"current.bin")
+                _tar_location = os.path.join(self.BACKUP_DIR, str(msg.partition()))
                 _msg = Common.decodeMsgToUtf8(msg)
                 if _msg is not None:
                     if count == 0:
-                        Common.writeDataToKafkaBinFile(self.BACKUP_TMP_FILE, _msg, "a+")
+                        Common.writeDataToKafkaBinFile(_tmp_file, _msg, "a+")
                     if count > 0:
                         if count % self.NUMBER_OF_MESSAGE_PER_BACKUP_FILE == 0:
-                            Common.createTarGz(self.BACKUP_DIR, self.BACKUP_TMP_FILE)
-                            Common.writeDataToKafkaBinFile(self.BACKUP_TMP_FILE, _msg, "w")
+                            Common.createTarGz(_tar_location, _tmp_file)
+                            Common.writeDataToKafkaBinFile(_tmp_file, _msg, "w")
                         else:
-                            Common.writeDataToKafkaBinFile(self.BACKUP_TMP_FILE, _msg, "a+")
-            
+                            Common.writeDataToKafkaBinFile(_tmp_file, _msg, "a+")
+            else:
+                logging.error(f"no partition found for message")
+
             count += 1
 
         _bt.close()
@@ -72,11 +84,15 @@ def main():
         exit(1)
 
     b = KBackup(config)
-    _r_thread = threading.Thread(
-        target=b.backup,
-        name="Kafka Consumer"
-    )
-    _r_thread.start()
+    Common.setLoggingFormat(b.LOG_LEVEL)
+
+    for _r_thread in range(b.NUMBER_OF_KAFKA_THREADS):
+        _r_thread = threading.Thread(
+            target=b.backup,
+            name="Kafka Consumer " + str(_r_thread)
+        )
+        _r_thread.start()
+
 
     if config['FILESYSTEM_TYPE'] == "S3":
         try:
@@ -89,12 +105,9 @@ def main():
             except:
                 logging.debug(f"setting RETRY_UPLOAD_SECONDS to default 60 ")
                 retry_upload_seconds = 60
-            _s3_upload_thread = threading.Thread(
-                target=Upload.s3_upload_files,
-                args=[bucket, tmp_dir, topic_name,retry_upload_seconds],
-                name="S3-Upload"
-            )
-            _s3_upload_thread.start()
+            
+            Upload.s3_upload(bucket,tmp_dir,topic_name,retry_upload_seconds,b.NUMBER_OF_KAFKA_THREADS + 1)
+
         except KeyError as e:
             logging.error(f"unable to set s3 required variables {e}")
 

cloud.py

@@ -0,0 +1,48 @@
+import logging
+import boto3
+from botocore.exceptions import ClientError
+import os
+import time
+from common import Common
+import pathlib
+import threading
+
+class Upload:
+
+    def findFiles(dir,pattern="*"):
+        _list = []
+        result = list(pathlib.Path(dir).rglob(pattern))
+        for path in result:
+            if os.path.isfile(path):
+                _list.append(path)
+        return _list
+
+    def s3_upload_file(s3_client,bucket,file_name,object_name):
+        try:
+            response = s3_client.upload_file(file_name,bucket,object_name)
+            logging.info(f"upload successful at s3://{bucket}/{object_name}")
+            if not file_name.endswith(".bin"):
+                logging.debug(f"deleting uploaded file {file_name}")
+                os.remove(file_name)
+        except ClientError as e:
+            logging.error(f"{file_path} upload failed error {e}")
+
+    def s3_upload(bucket,dir,topic_name,retry_upload_seconds,thread_count):
+        s3_client = boto3.client('s3')
+        while True:
+            _topic_dir = os.path.join(dir, topic_name)
+            _count_partition_dirs = len(Common.listDirs(_topic_dir))
+            _list = Upload.findFiles(_topic_dir)
+            if len(_list) > _count_partition_dirs and threading.active_count() <= thread_count:
+                for file_name in _list:
+                    file_name = str(file_name)
+                    if os.path.getsize(file_name) > 0:
+                        object_name = file_name.split(dir)[1]
+                        t = threading.Thread(
+                            target=Upload.s3_upload_file,
+                            args=[s3_client,bucket,file_name,object_name],
+                            name="S3 Upload Threads"
+                        ).start()
+            else:
+                logging.info(f"s3 upload retry for new files in {retry_upload_seconds} seconds")
+                time.sleep(retry_upload_seconds)

common.py

@@ -7,13 +7,13 @@
 
 class Common:
 
-    def setLoggingFormat():
+    def setLoggingFormat(level=20):
         logging.basicConfig(
             format='{ "@timestamp": "%(asctime)s","level": "%(levelname)s","thread": "%(threadName)s","name": "%(name)s","message": "%(message)s" }'
         )
-        logging.getLogger().setLevel(logging.INFO)
+        logging.getLogger().setLevel(level)
 
-    def listFiles(dir):
+    def listDirs(dir):
         try:
             return sorted(os.listdir(dir))
         except FileNotFoundError as e:
@@ -24,7 +24,7 @@ def listFiles(dir):
     def readJsonConfig(file):
         try:
             with open(file) as cf:
-                logging.info(f'loading {file} file')
+                logging.debug(f'loading {file} file')
                 return json.load(cf)
         except json.decoder.JSONDecodeError as e:
             logging.error(f'{e}')
@@ -45,11 +45,11 @@ def createSha256OfBackupFile(file,hash):
         except:
             logging.error(f'unable to write to {file}.sha256')
 
-    def createBackupTopicDir(dir):
+    def createDir(dir):
         try:
             os.mkdir(dir)
         except FileExistsError as e:
-            logging.info(e)
+            logging.debug(e)
         except:
             logging.error(f'unable to create folder {dir}')
             exit(1)
@@ -122,4 +122,10 @@ def extractBinFile(file,hashfile,extractDir):
                 return os.path.join(extractDir,_sname)
             except FileNotFoundError as e:
                 logging.error(e)
-        return None
+        return None
+
+    def findNumberOfPartitionsInTopic(list):
+        _lp = []
+        for i in list:
+            _lp.append(i)
+        return _lp

restore.py

@@ -33,7 +33,7 @@ def delivery_report(err, msg):
     def restore(self):
         _rt = confluent_kafka.Producer(self.PRODUCERCONFIG)
         while True:
-            _files_in_backup_dir = Common.listFiles(self.BACKUP_DIR)
+            _files_in_backup_dir = Common.listDirs(self.BACKUP_DIR)
             for file in _files_in_backup_dir:
                     file = os.path.join(self.BACKUP_DIR,file)
                     if file.endswith("tar.gz"):