opensearch-python-handler/logzio/sender.py

# This class is responsible for handling all asynchronous Logz.io's
# communication
import sys
import json

from time import sleep
from datetime import datetime
from threading import Thread, enumerate

import requests

if sys.version[0] == '2':
    import Queue as queue
else:
    import queue as queue


MAX_BULK_SIZE_IN_BYTES = 1 * 1024 * 1024  # 1 MB


def backup_logs(logs):
    timestamp = datetime.now().strftime("%d%m%Y-%H%M%S")
    print("Backing up your logs to logzio-failures-{}.txt".format(timestamp))
    with open("logzio-failures-{}.txt".format(timestamp), "a") as f:
        f.writelines('\n'.join(logs))


class LogzioSender:

    def __init__(self,
                 token, url="https://listener.logz.io:8071",
                 logs_drain_timeout=5,
                 debug=False):
        self.token = token
        self.url = "{}/?token={}".format(url, token)
        self.logs_drain_timeout = logs_drain_timeout
        self.debug = debug

        # Function to see if the main thread is alive
        self.is_main_thread_active = lambda: any(
            (i.name == "MainThread") and i.is_alive() for i in enumerate())

        # Create a queue to hold logs
        self.queue = queue.Queue()

        self.sending_thread = Thread(target=self._drain_queue)
        self.sending_thread.daemon = False
        self.sending_thread.name = "logzio-sending-thread"
        self.sending_thread.start()

    def append(self, logs_message):
        # Queue lib is thread safe, no issue here
        self.queue.put(json.dumps(logs_message))

    def flush(self):
        self._flush_queue()

    def _debug(self, message):
        if self.debug:
            print(str(message))

    def _drain_queue(self):
        last_try = False

        while not last_try:
            # If main is exited, we should run one last time and try to remove
            # all logs
            if not self.is_main_thread_active():
                self._debug(
                    "Identified quit of main thread, sending logs one "
                    "last time")
                last_try = True

            try:
                self._flush_queue()

            # TODO: Which exception?
            except Exception as e:
                self._debug(
                    "Unexpected exception while draining queue to Logz.io, "
                    "swallowing. Exception: {}".format(e))

            if not last_try:
                sleep(self.logs_drain_timeout)

    def _flush_queue(self):
        # TODO: Break this down. This function is crazy.
        # Sending logs until queue is empty
        while not self.queue.empty():
            logs_list = self._get_messages_up_to_max_allowed_size()
            self._debug(
                "Starting to drain {} logs to Logz.io".format(len(logs_list)))

            # Not configurable from the outside
            sleep_between_retries = 2
            number_of_retries = 4

            should_backup_to_disk = True
            headers = {"Content-type": "text/plain"}

            for current_try in range(number_of_retries):
                should_retry = False
                try:
                    response = requests.post(
                        self.url, headers=headers, data='\n'.join(logs_list))
                    if response.status_code != 200:
                        if response.status_code == 400:
                            print("Got 400 code from Logz.io. This means that "
                                  "some of your logs are too big, or badly "
                                  "formatted. response: {}".format(
                                      response.text))
                            should_backup_to_disk = False
                            break

                        if response.status_code == 401:
                            print(
                                "You are not authorized with Logz.io! Token "
                                "OK? dropping logs...")
                            should_backup_to_disk = False
                            break
                        else:
                            print(
                                "Got {} while sending logs to Logz.io, "
                                "Try ({}/{}). Response: {}".format(
                                    response.status_code,
                                    current_try + 1,
                                    number_of_retries,
                                    response.text))
                            should_retry = True
                    else:
                        self._debug(
                            "Successfully sent bulk of {} logs to "
                            "Logz.io!".format(len(logs_list)))
                        should_backup_to_disk = False
                        break

                # TODO: Which exception?
                except Exception as e:
                    print("Got exception while sending logs to Logz.io, "
                          "Try ({}/{}). Message: {}".format(
                              current_try + 1, number_of_retries, e))
                    should_retry = True

                if should_retry:
                    sleep(sleep_between_retries)
                    sleep_between_retries *= 2

            if should_backup_to_disk:
                # Write to file
                print("Could not send logs to Logz.io after {} tries, "
                      "backing up to local file system".format(
                          number_of_retries))
                backup_logs(logs_list)

    def _get_messages_up_to_max_allowed_size(self):
        logs_list = []
        current_size = 0
        while not self.queue.empty():
            current_log = self.queue.get()
            current_size += sys.getsizeof(current_log)
            logs_list.append(current_log)
            if current_size >= MAX_BULK_SIZE_IN_BYTES:
                break
        return logs_list