tuoheng_alg_move_20240715/util/KafkaUtils.py

# -*- coding: utf-8 -*-
import time

from kafka import KafkaProducer, KafkaConsumer, TopicPartition, OffsetAndMetadata
from kafka.errors import kafka_errors
import json
from loguru import logger

from common import YmlConstant, Constant


# 生产者
class CustomerKafkaProducer:

    def __init__(self, context):
        self.__context = context
        self.__configs = YmlConstant.get_kafka_producer_config(context)
        self.customerProducer = None
        self.get_producer()

    # 获取kafka生产者
    def get_producer(self):
        if self.customerProducer is None:
            logger.info("配置kafka生产者!")
            self.customerProducer = KafkaProducer(
                bootstrap_servers=YmlConstant.get_kafka_bootstrap_servers(self.__context),
                acks=self.__configs[YmlConstant.KAFKA_ACKS],
                retries=self.__configs[YmlConstant.KAFKA_RETRIES],
                linger_ms=self.__configs[YmlConstant.KAFKA_LINGER_MS],
                retry_backoff_ms=self.__configs[YmlConstant.KAFKA_RETRY_BACKOFF_MS],
                max_in_flight_requests_per_connection=self.__configs[YmlConstant.KAFKA_MAX_IN_FLIGHT_REQUESTS],
                key_serializer=lambda m: json.dumps(m).encode(Constant.UTF_8),
                value_serializer=lambda m: json.dumps(m).encode(Constant.UTF_8))

    # mode 模式1：异步发送 2：同步发送
    # def on_send_success(record_metadata): 成功回调
    # def on_send_error(excp): 失败回调
    def sender(self, topic, key, message, mode=1, customer_send_success=None, customer_send_error=None):
        retry_send_num = 0
        while True:
            try:
                self.get_producer()
                logger.info("kafka发送信息，topic:{}|key:{}|message:{}|mode:{}|requestId:{}", topic, key, message, mode,
                            message.get("request_id"))
                if mode == 1:
                    if not customer_send_success:
                        customer_send_success = CustomerKafkaProducer.on_send_success
                    if not customer_send_error:
                        customer_send_error = CustomerKafkaProducer.on_send_error
                    self.customerProducer.send(topic=topic, key=key, value=message) \
                        .add_callback(customer_send_success, message.get(YmlConstant.REQUEST_ID)) \
                        .add_errback(customer_send_error, message.get(YmlConstant.REQUEST_ID))
                if mode == 2:
                    try:
                        self.customerProducer.send(topic=topic, key=key, value=message).get(timeout=30)
                        logger.info("kafka同步发送信息成功, requestId:{}", message.get(YmlConstant.REQUEST_ID))
                    except kafka_errors as ke:
                        logger.exception("kafka同步发送消息异常: {}, requestId:{}", ke, message.get(YmlConstant.REQUEST_ID))
                        raise ke
                break
            except Exception as e:
                retry_send_num += 1
                logger.error("kafka发送消息异常, 开始重试, 当前重试次数：{} requestId:{}", retry_send_num,
                             message.get(YmlConstant.REQUEST_ID))
                self.customerProducer = None
                if retry_send_num > 3:
                    logger.exception("kafka发送消息重试失败: {}, requestId:{}", e, message.get(YmlConstant.REQUEST_ID))
                    raise e

    def close_producer(self):
        self.customerProducer.flush()
        self.customerProducer.close()
        logger.info("kafka生产者关闭完成!")

    @staticmethod
    def on_send_success(requestId, record_metadata):
        logger.info("kafka异步发送信息成功，topic:{}|partition:{}|offset:{}|requestId:{}", record_metadata.topic,
                    record_metadata.partition, record_metadata.offset, requestId)

    @staticmethod
    def on_send_error(requestId, excp):
        logger.exception("kafka异步发送消息异常: {}, requestId:{}", excp, requestId)


# 生产者
class CustomerKafkaConsumer:

    def __init__(self, context, topics=()):
        logger.info("初始化消费者")
        self.__context = context
        self.__configs = YmlConstant.get_kafka_consumer_config(context)
        self.customerConsumer = None
        self.__topics = topics
        self.subscribe()
        logger.info("初始化消费者完成")

    def subscribe(self):
        if self.customerConsumer is None:
            logger.info("获取消费者！")
            self.customerConsumer = KafkaConsumer(
                bootstrap_servers=YmlConstant.get_kafka_bootstrap_servers(self.__context),
                client_id=self.__configs[YmlConstant.KAFKA_CLIENT_ID],
                group_id=self.__configs[YmlConstant.KAFKA_GROUP_ID],
                auto_offset_reset=self.__configs[YmlConstant.KAFKA_AUTO_OFFSET_RESET],
                enable_auto_commit=self.__configs[YmlConstant.KAFKA_ENABLE_AUTO_COMMIT],
                max_poll_records=self.__configs[YmlConstant.KAFKA_MAX_POLL_RECORDS],
                value_deserializer=lambda m: json.loads(m.decode(Constant.UTF_8)))
            logger.info("kafka生产者订阅topic:{}", self.__topics)
            # if self.topics is None or len(self.topics) == 0:
            #     logger.error("消费者订阅topic不能为空！")
            #     raise Exception("消费者订阅topic不能为空！")
            # # 手动配置分区
            # customer_partition = []
            # for topic in self.topics:
            #     for p in self.content["kafka"][self.content["dsp"]["active"]][topic]["partition"]:
            #         customer_partition.append(TopicPartition(topic, p))
            # self.customerConsumer.assign(customer_partition)
            # 自动配置
            self.customerConsumer.subscribe(topics=self.__topics)
            logger.info("kafka生产者订阅topic完成")

    def poll(self):
        msg = None
        try:
            self.subscribe()
            msg = self.customerConsumer.poll()
        except Exception as e:
            self.customerConsumer = None
            logger.exception("消费者拉取消息异常: {}", e)
        return msg

    def commit_offset(self, message):
        retry_num = 1
        while True:
            try:
                self.subscribe()
                logger.info("消费者开始提交offset,topic:{}|offset:{}|partition:{}", message.topic, message.offset + 1,
                            message.partition)
                tp = TopicPartition(topic=message.topic, partition=message.partition)
                self.customerConsumer.commit(offsets={tp: (OffsetAndMetadata(message.offset + 1, None))})
                logger.info("消费者提交offset完成,topic:{}|offset:{}|partition:{}", message.topic, message.offset + 1,
                            message.partition)
                break
            except Exception as e:
                self.customerConsumer = None
                logger.exception("消费者提交offset异常: {}, 重试次数: {}", e, retry_num)
                time.sleep(1)
                retry_num += 1
                if retry_num > 3:
                    logger.exception("消费者提交offset重试失败: {}", e)
                    break

# if __name__=="__main__":
#     try:
#         1/0
#     except Exception as e:
#         logger.exception("aaaaa:{} {}", e, "11111")