From a3f217c9ae69bd4a0ea010f3d8cdd7f545e30745 Mon Sep 17 00:00:00 2001 From: Sumit Date: Thu, 23 Dec 2021 18:37:55 +0100 Subject: [PATCH] async sender with kafka --- journalpump/journalpump.py | 35 +++--- journalpump/senders/__init__.py | 1 + journalpump/senders/aws_cloudwatch.py | 1 - journalpump/senders/base.py | 44 ++++++++ journalpump/senders/kafka_async.py | 150 ++++++++++++++++++++++++++ 5 files changed, 218 insertions(+), 13 deletions(-) create mode 100644 journalpump/senders/kafka_async.py diff --git a/journalpump/journalpump.py b/journalpump/journalpump.py index 01ab427..ab80d2c 100644 --- a/journalpump/journalpump.py +++ b/journalpump/journalpump.py @@ -6,15 +6,17 @@ from .daemon import ServiceDaemon from .senders import ( AWSCloudWatchSender, ElasticsearchSender, FileSender, GoogleCloudLoggingSender, KafkaSender, LogplexSender, - RsyslogSender, WebsocketSender + RsyslogSender, WebsocketSender, AsyncKafkaSender ) -from .senders.base import MAX_KAFKA_MESSAGE_SIZE, SenderInitializationError, Tagged +from .senders.base import AsyncLogSender, MAX_KAFKA_MESSAGE_SIZE, SenderInitializationError, Tagged from .types import GeoIPProtocol from .util import atomic_replace_file, default_json_serialization from functools import reduce from systemd.journal import Reader +from threading import Thread from typing import Type, Union +import asyncio import copy import datetime import json @@ -114,6 +116,7 @@ class JournalReader(Tagged): "aws_cloudwatch": AWSCloudWatchSender, "google_cloud_logging": GoogleCloudLoggingSender, "websocket": WebsocketSender, + "kafka_async": AsyncKafkaSender, } def __init__( @@ -237,17 +240,25 @@ def initialize_senders(self): if not isinstance(extra_field_values, dict): self.log.warning("extra_field_values: %r not a dictionary object, ignoring", extra_field_values) extra_field_values = {} + kw = dict( + config=sender_config, + field_filter=field_filter, + extra_field_values=extra_field_values, + msg_buffer_max_length=self.msg_buffer_max_length, + name=sender_name, + reader=self, + stats=self.stats, + tags=self.make_tags(), + ) + if isinstance(sender_class, AsyncLogSender): + loop = asyncio.get_event_loop() + kw["aio_loop"] = loop + if not loop.is_running: + # to unbloock the journal reading process + # the same loop can be used for other senders once they become async + Thread(target=loop.run_forever).start() try: - sender = sender_class( - config=sender_config, - field_filter=field_filter, - extra_field_values=extra_field_values, - msg_buffer_max_length=self.msg_buffer_max_length, - name=sender_name, - reader=self, - stats=self.stats, - tags=self.make_tags(), - ) + sender = sender_class(**kw) except SenderInitializationError: # If sender init fails, log exception, don't start() the sender # and don't add it to self.senders dict. A metric about senders that diff --git a/journalpump/senders/__init__.py b/journalpump/senders/__init__.py index 94f2321..6e2fd2f 100644 --- a/journalpump/senders/__init__.py +++ b/journalpump/senders/__init__.py @@ -6,3 +6,4 @@ from .logplex import LogplexSender # noqa: F401 from .rsyslog import RsyslogSender # noqa: F401 from .websocket import WebsocketSender # noqa: F401 +from .kafka_async import AsyncKafkaSender # noqa: F401 diff --git a/journalpump/senders/aws_cloudwatch.py b/journalpump/senders/aws_cloudwatch.py index f46f280..9b5939f 100644 --- a/journalpump/senders/aws_cloudwatch.py +++ b/journalpump/senders/aws_cloudwatch.py @@ -1,6 +1,5 @@ from .base import ThreadedLogSender, SenderInitializationError - import boto3 import botocore import json diff --git a/journalpump/senders/base.py b/journalpump/senders/base.py index 0e452d0..91f2ade 100644 --- a/journalpump/senders/base.py +++ b/journalpump/senders/base.py @@ -1,5 +1,6 @@ from threading import Lock, Thread +import asyncio import logging import random import time @@ -265,3 +266,46 @@ def get_and_send_messages(self): # there is already a broad except handler in send_messages, so why this ? self.log.exception("Problem sending %r messages", msg_count) self._backoff() + +class AsyncLogSender(LogSender): + def __init__(self, **kw): + self.loop: asyncio.events.AbstractEventLoop = kw.pop('aio_loop') + LogSender.__init__(self, **kw) + + async def _backoff(self, *, base=0.5, cap=1800.0): + t = self._get_backoff_secs(base=base, cap=cap) + self.log.info("Sleeping for %.0f seconds", t) + await asyncio.sleep(t) + + def handle_maintenance_operations(self): + async def noop(): + pass + self.loop.create_task(noop()) + + async def run(self): + while self.running: + await self.handle_maintenance_operations() + if self.should_try_sending_messages(): + await self.get_and_send_messages() + else: + await asyncio.sleep(self._wait_for) + self.log.info("Stopping") + + async def get_and_send_messages(self): + batches = self.get_message_bodies_and_cursor() + msg_count = sum(len(batch[0]) for batch in batches) + self.log.debug("Got %d items from msg_buffer", msg_count) + start_time = time.monotonic() + try: + # pop to get free up memory as soon as the send was successful + while batches: + batch = batches.pop(0) + # die retrying, backoff is part of sending mechanism + while self.running and not await self.send_messages(messages=batch[0], cursor=batch[1]): + pass + self.log.debug("Sending %d msgs, took %.4fs", msg_count, time.monotonic() - start_time) + self.last_send_time = time.monotonic() + except Exception: # pylint: disable=broad-except + # there is already a broad except handler in send_messages, so why this ? + self.log.exception("Problem sending %r messages", msg_count) + await self._backoff() diff --git a/journalpump/senders/kafka_async.py b/journalpump/senders/kafka_async.py new file mode 100644 index 0000000..9a4913e --- /dev/null +++ b/journalpump/senders/kafka_async.py @@ -0,0 +1,150 @@ +from .base import AsyncLogSender +from aiokafka import AIOKafkaProducer +from kafka import errors + +import asyncio +import logging +import socket + + +try: + import snappy +except ImportError: + snappy = None + +try: + import zstandard as zstd +except ImportError: + zstd = None + +KAFKA_CONN_ERRORS = tuple(errors.RETRY_ERROR_TYPES) + ( + errors.UnknownError, + socket.timeout, +) + +logging.getLogger("kafka").setLevel(logging.CRITICAL) # remove client-internal tracebacks from logging output + + +class AsyncKafkaSender(AsyncLogSender): + def __init__(self, *, config, **kwargs): + super().__init__(config=config, max_send_interval=config.get("max_send_interval", 0.3), **kwargs) + self.kafka_producer = None + self.kafka_msg_key = self.config.get("kafka_msg_key") + if self.kafka_msg_key: + self.kafka_msg_key = self.kafka_msg_key.encode("utf8") + self.topic = self.config.get("kafka_topic") + + def _generate_client_config(self) -> dict: + config = { + "api_version": self.config.get("kafka_api_version"), + "bootstrap_servers": self.config.get("kafka_address"), + "reconnect_backoff_ms": 1000, # up from the default 50ms to reduce connection attempts + "reconnect_backoff_max_ms": 10000, # up the upper bound for backoff to 10 seconds + "loop": self.loop, + } + + if self.config.get("ssl"): + config["security_protocol"] = "SSL" + config["ssl_cafile"] = self.config.get("ca") + config["ssl_certfile"] = self.config.get("certfile") + config["ssl_keyfile"] = self.config.get("keyfile") + else: + config["security_protocol"] = "PLAINTEXT" + + return config + + def _generate_producer_config(self) -> dict: + producer_config = self._generate_client_config() + producer_config["linger_ms"] = 500 # wait up 500 ms to see if we can send msgs in a group + + # make sure the python client supports it as well + if zstd and "zstd" in AIOKafkaProducer._COMPRESSORS: # pylint: disable=protected-access + producer_config["compression_type"] = "zstd" + elif snappy: + producer_config["compression_type"] = "snappy" + else: + producer_config["compression_type"] = "gzip" + + if self.config.get("socks5_proxy"): + # Socks5_config is supported by Aiven fork of kafka-python for the time being + producer_config["socks5_proxy"] = self.config.get("socks5_proxy") + + return producer_config + + def start(self): + self.loop.create_task(self.run()) + + async def _init_kafka(self) -> None: + self.log.info("Initializing Kafka client, address: %r", self.config["kafka_address"]) + + if self.kafka_producer: + self.kafka_producer.close() + self.kafka_producer = None + + self.mark_disconnected() + + while self.running and not self._connected: + producer_config = self._generate_producer_config() + + try: + kafka_producer = AIOKafkaProducer(**producer_config) + await kafka_producer.start() + except KAFKA_CONN_ERRORS as ex: + self.mark_disconnected(ex) + self.log.warning("Retriable error during Kafka initialization: %s: %s", ex.__class__.__name__, ex) + await self._backoff() + else: + self.log.info("Initialized Kafka Client, address: %r", self.config["kafka_address"]) + self.kafka_producer = kafka_producer + self.mark_connected() + + # Assume that when the topic configuration is provided we should + # manually create it. This is useful for kafka clusters configured with + # `auto.create.topics.enable = false` +# topic_config = self.config.get("kafka_topic_config", dict()) +# num_partitions = topic_config.get("num_partitions") +# replication_factor = topic_config.get("replication_factor") +# if num_partitions is not None and replication_factor is not None: +# kafka_admin = KafkaAdminClient(**self._generate_client_config()) +# try: +# kafka_admin.create_topics([NewTopic(self.topic, num_partitions, replication_factor)]) +# except errors.TopicAlreadyExistsError: +# self.log.info("Kafka topic %r already exists", self.topic) +# else: +# self.log.info("Create Kafka topic, address: %r", self.topic) + +# def send_messages(self, *, messages, cursor): +# task = self.loop.(self._send_messages(messages=messages, cursor=cursor)) +# await task +# return task.result() + + async def send_messages(self, *, messages, cursor): + if not self.kafka_producer: + await self._init_kafka() + try: + # Collect return values of send(): + # FutureRecordMetadata which will trigger when message actually sent (during flush) + result_futures = asyncio.gather( + self.kafka_producer.send(topic=self.topic, value=msg, key=self.kafka_msg_key) for msg in messages + ) + await self.kafka_producer.flush() + for result_future in result_futures: + # get() throws error from future, catch below + # flush() above should have sent, getting with 1 sec timeout + result_future.get(timeout=1) + self.mark_sent(messages=messages, cursor=cursor) + return True + except KAFKA_CONN_ERRORS as ex: + self.mark_disconnected(ex) + self.log.info("Kafka retriable error during send: %s: %s, waiting", ex.__class__.__name__, ex) + if self.running: + await self._backoff() + await self._init_kafka() + except Exception as ex: # pylint: disable=broad-except + self.mark_disconnected(ex) + self.log.exception("Unexpected exception during send to kafka") + self.stats.unexpected_exception(ex=ex, where="sender", tags=self.make_tags({"app": "journalpump"})) + if self.running: + await self._backoff() + await self._init_kafka() + return False