Merge pull request #11 from kusha/multiprocessing-support

kusha · web-flow · commit 1578726586df · 2019-05-02T13:00:29.000+02:00
Add multiprocessing support
diff --git a/examples/mt_mp.py b/examples/mt_mp.py
@@ -0,0 +1,179 @@
+"""Usage of kafka-logging-handler with multiprocessing and multithreading."""
+
+from concurrent.futures import ThreadPoolExecutor
+import logging
+import multiprocessing
+from multiprocessing import Process
+import os
+import sys
+import threading
+
+from kafka_logger.handlers import KafkaLoggingHandler
+
+REQUIRED_ENV_VARS = ['KAFKA_SERVER', 'KAFKA_CERT', 'KAFKA_TOPIC']
+
+MAIN_PROCESS_THREADS = 2
+CHILD_PROCESSES = 2
+THREAD_POOL_WORKERS = 2
+THREAD_POOL_SIZE = 3
+
+LOGGER = None
+
+
+def get_process_thread():
+    """Get string with process and thread names."""
+    # you can get PID and thread ID as well:
+    # os.getpid(), threading.current_thread().ident
+    return "(process: {}, thread: {})".format(
+        multiprocessing.current_process().name,
+        threading.current_thread().name
+    )
+
+
+def child_process(index):
+    """
+    Log a message.
+
+    Args:
+        index (int): index of the child process
+    """
+    LOGGER.info("Hi, I'm child process #%d %s", index, get_process_thread())
+
+
+def grandchild_process():
+    """Log a message."""
+    LOGGER.info("Hi, I'm sub sub process %s", get_process_thread())
+
+
+def child_process_with_grandchild():
+    """Spawn grandchild process."""
+    LOGGER.info("I'm going to spawn another child %s", get_process_thread())
+    subprocess = Process(
+        target=grandchild_process,
+        name="Grandchild process")
+    subprocess.start()
+    subprocess.join()
+
+
+def thread_worker(index):
+    """Log a message."""
+    LOGGER.info(
+        "Hi, I'm a thread worker #%d in the child process thread pool %s",
+        index, get_process_thread())
+
+
+def child_process_with_threads():
+    """Run thread executor pool."""
+    LOGGER.info("I'm going to spawn multiple threads %s", get_process_thread())
+    with ThreadPoolExecutor(max_workers=THREAD_POOL_WORKERS) as executor:
+        for thread_idx in range(THREAD_POOL_SIZE):
+            executor.submit(thread_worker, thread_idx)
+
+
+def main_process_thread(index):
+    """Log a message."""
+    LOGGER.info("Hi, I'm a thread #%d in the main process %s",
+                index, get_process_thread())
+
+
+def child_process_with_exception():
+    """Raise an exception after start."""
+    LOGGER.info("Hi, I'm child process that is going to raise exception %s",
+                get_process_thread())
+    raise Exception('This exception will not occur in Kafka!')
+
+
+def main():
+    """Setup logger and test logging."""
+    global LOGGER
+
+    # validate that Kafka configuration is available
+    assert all([(key in os.environ) for key in REQUIRED_ENV_VARS])
+
+    LOGGER = logging.getLogger("test.logger")
+    LOGGER.propagate = False
+    log_level = logging.DEBUG
+
+    log_format = logging.Formatter(
+        '%(asctime)s %(name)-12s %(levelname)-8s %(message)s',
+        '%Y-%m-%dT%H:%M:%S')
+
+    # create handler to show logs at stdout
+    stdout_handler = logging.StreamHandler(sys.stdout)
+    stdout_handler.setLevel(log_level)
+    stdout_handler.setFormatter(log_format)
+    LOGGER.addHandler(stdout_handler)
+
+    # create Kafka logging handler
+    kafka_handler = KafkaLoggingHandler(
+        os.environ['KAFKA_SERVER'],
+        os.environ['KAFKA_TOPIC'],
+        security_protocol='SSL',
+        ssl_cafile=os.environ['KAFKA_CERT'],
+        unhandled_exception_logger=LOGGER,
+        additional_fields={
+            "service": "test_service"
+        }
+    )
+    kafka_handler.setFormatter(log_format)
+    LOGGER.addHandler(kafka_handler)
+
+    LOGGER.setLevel(log_level)
+
+    LOGGER.info("Hi there, I'm the main process! %s", get_process_thread())
+
+    # test child processes
+    child_processes = []
+    for idx in range(CHILD_PROCESSES):
+        child = Process(
+            target=child_process,
+            name="Child process #{}".format(idx),
+            args=(idx,))
+        child_processes.append(child)
+        child.start()
+
+    # testing threads in the main process
+    threads = []
+    for idx in range(MAIN_PROCESS_THREADS):
+        thread = threading.Thread(
+            target=main_process_thread,
+            name="Thread of the main process #{}".format(idx),
+            args=(idx, ))
+        threads.append(thread)
+        thread.start()
+    # wait for threads to finish
+    for thread in threads:
+        thread.join()
+
+    # there is a chance of logs loss
+    # if the main process terminates without joining child processes
+    for child in child_processes:
+        child.join()
+
+    # test if a child of a child process logs correctly
+    child_with_subprocess = Process(
+        target=child_process_with_grandchild,
+        name="Child process that spawns another child")
+    child_with_subprocess.start()
+    child_with_subprocess.join()
+
+    # test threads in a child process
+    child_with_threads = Process(
+        target=child_process_with_threads,
+        name="Child process that has a thread pool")
+    child_with_threads.start()
+    child_with_threads.join()
+
+    # test unhandled exception in a child process
+    child_exception = Process(
+        target=child_process_with_exception,
+        name="Child process with an exception")
+    child_exception.start()
+    child_exception.join()
+
+    # top-level exception works only in the main process
+    raise Exception('Testing top-level exception!')
+
+
+if __name__ == '__main__':
+    main()
diff --git a/kafka_logger/handlers.py b/kafka_logger/handlers.py
@@ -3,9 +3,12 @@
 import atexit
 import json
 import logging
+from multiprocessing import Queue
+import os
 import socket
 import sys
-from threading import Lock, Timer
+from threading import Lock, Thread, Timer
+import time
 
 from kafka import KafkaProducer  # pylint: disable=import-error
 
@@ -21,11 +24,18 @@ class KafkaLoggingHandler(logging.Handler):
     This handler enables the user to forward logs to Kafka.
 
     Attributes:
+        additional_fields (dict): extra fields attached to logs
         buffer (list): logs (dict) waiting for a flush to Kafka
         buffer_lock (threading.Lock): multithreading lock for a buffer access
         flush_interval (float): scheduled flush interval in seconds
         kafka_topic_name (str): topic name
+        main_process_pid (int): pid of the process which initialized logger
         max_buffer_size (int): flush if buffer > max size
+        mp_log_handler_flush_lock (threading.Lock):
+            locked when mp_log_handler_thread flushes logs
+        mp_log_handler_thread (threading.Thread): thread that flushed mp queue
+        mp_log_queue (multiprocessing.Queue):
+            queue used to redirect logs of child processes
         producer (kafka.KafkaProducer): producer object
         timer (threading.Timer): thread with a scheduled logs flush
         unhandled_exception_logger (logging.Logger):
@@ -37,6 +47,7 @@ class KafkaLoggingHandler(logging.Handler):
                                'relativeCreated',
                                'levelno',
                                'created']
+    __MULTIPROCESSING_QUEUE_FLUSH_DELAY = 0.2
 
     def __init__(self,
                  hosts_list,
@@ -102,20 +113,35 @@ def __init__(self,
             **kafka_producer_args)
 
         # setup exit hooks
+        # exit hooks work only in main process
+        # termination of child processes uses os.exit() and ignore any hooks
         atexit.register(self.at_exit)
         sys.excepthook = self.unhandled_exception
 
-    def emit(self, record):
+        # multiprocessing support
+        self.main_process_pid = os.getpid()
+        self.mp_log_queue = Queue()
+        # main process thread that will flush mp queue
+        self.mp_log_handler_flush_lock = Lock()
+        self.mp_log_handler_thread = Thread(
+            target=self.mp_log_handler,
+            name="Kafka Logger Multiprocessing Handler")
+        # daemon will terminate with the main process
+        self.mp_log_handler_thread.setDaemon(True)
+        self.mp_log_handler_thread.start()
+
+    def prepare_record_dict(self, record):
         """
-        Add a new log to the buffer.
+        Prepare a dictionary log item.
+
+        Format a log record and extend dictionary with default values.
 
         Args:
-            record: Logging message
-        """
-        # drop Kafka logging to avoid infinite recursion.
-        if record.name == 'kafka.client':
-            return
+            record (logging.LogRecord): log record
 
+        Returns:
+            dict: log item ready for Kafka
+        """
         # use default formatting
         # Update the msg dict to include all of the message attributes
         self.format(record)
@@ -137,8 +163,37 @@ def emit(self, record):
                     value = tuple(arg.__repr__() for arg in value)
                 rec[key] = "" if value is None else value
 
+        return rec
+
+    def emit(self, record):
+        """
+        Add log to the buffer or forward to the main process.
+
+        Args:
+            record: Logging message
+        """
+        # drop Kafka logging to avoid infinite recursion.
+        if record.name == 'kafka.client':
+            return
+
+        record_dict = self.prepare_record_dict(record)
+
+        if os.getpid() == self.main_process_pid:
+            self.append_to_buffer(record_dict)
+        else:  # if forked
+            self.mp_log_queue.put(record_dict)
+
+    def append_to_buffer(self, record_dict):
+        """
+        Place log dictionary to the buffer.
+
+        Triggers/schedules a flush of the buffer.
+
+        Args:
+            record_dict (dict): log item
+        """
         with self.buffer_lock:
-            self.buffer.append(rec)
+            self.buffer.append(record_dict)
 
         # schedule a flush
         if len(self.buffer) >= self.max_buffer_size:
@@ -175,6 +230,16 @@ def schedule_flush(self):
             self.timer.setDaemon(True)
             self.timer.start()
 
+    def mp_log_handler(self):
+        """Emit logs from multiprocessing queue."""
+        while True:
+            if self.mp_log_handler_flush_lock.locked():
+                # second+ iteration
+                self.mp_log_handler_flush_lock.release()
+            record_dict = self.mp_log_queue.get(block=True)  # wait for logs
+            self.mp_log_handler_flush_lock.acquire()
+            self.append_to_buffer(record_dict)
+
     def at_exit(self):
         """
         Flush logs at exit, close the producer.
@@ -183,6 +248,27 @@ def at_exit(self):
         Kafka raises RecordAccumulator in case of flushing in close method.
         """
         # Kafka's RecordAccumulator is still alive here
+        if self.unhandled_exception_logger is not None:
+            # check if there are running subprocesses and log a warning
+            try:
+                import psutil
+                main_process = psutil.Process(pid=self.main_process_pid)
+            except ImportError:
+                pass
+            except psutil.NoSuchProcess:
+                pass
+            else:
+                children = main_process.children(recursive=True)
+                if children:
+                    self.unhandled_exception_logger.warning(
+                        "There are %d child process(es) at the moment of the "
+                        "main process termination. This may cause logs loss.",
+                        len(children))
+        while self.mp_log_queue.qsize() != 0:
+            time.sleep(KafkaLoggingHandler.__MULTIPROCESSING_QUEUE_FLUSH_DELAY)
+        # wait until everything in multiprocessing queue will be buffered
+        self.mp_log_handler_flush_lock.acquire()
+
         if self.timer is not None:
             self.flush()
         self.producer.close()
