基于kafka-go写的生产者和消费者

在这一篇《kafka-go使用：以及kafka一些基本概念说明》文章中简单梳理了kafka的相关知识点。

设计思想说明

这一篇文章，主要是基于kafka-go写的生产者和消费者的代码。我自定义了一个kafka消息体，消息体中，有一个actionType的概念。为什么要有这么一个定义呢，主要是考虑到kafka的topic的颗粒度不能太细，太细会导致kafka的性能下降，所以，我们的消费需要基于actionType进行细分，以提升消费的并发能力，这里指的是从业务的角度来提升消费能力，而不是通过kafka自身的调整。每一个actionType我会起一个协程来处理，如此，我们可以从业务维度来提升消费的速度，避免消息的堆积。比如，考虑付款，退款流程，我们可以用一个actionType，如此，当有此类消息，我们走到付款或者退款的函数中，考虑到，付款或者退款会涉及到先后顺序，所以，付款、退款用同一个actionType进行管理。

实现逻辑

实现逻辑我就不细说了，代码中都加了注释。我直接贴上代码。 

消息体定义

package customerkafka

type CustomKafkaMsg struct {
	ActionType  string      `json:"action_type"`  // 消息类型
	MsgId       string      `json:"msg_id"`       // 消息ID
	BalancerKey string      `json:"balancer_key"` // 分区键 //如果需要考虑消息的顺序性，那么就需要设置相同的分区键。并且采用Hash算法，即设置writerConfig.Balancer = &kafka.Hash{}，保证同一个分区键的消息会被发送到同一个分区中。
	Data        interface{} `json:"data"`         // 消息内容
	Timestamp   int64       `json:"timestamp"`    // 消息时间戳
}

生产者

package customerkafka

import (
	"context"
	"encoding/json"
	"fmt"

	"github.com/segmentio/kafka-go"
	"gorm.io/gorm/logger"
)

type Producer interface {
	Send(ctx context.Context, messages ...*CustomKafkaMsg) error
}

type ProduceClient struct {
	writer *kafka.Writer
	logger logger.Interface
}

func NewProduceClient(writerConfig kafka.WriterConfig, log logger.Interface) *ProduceClient {
	writer := kafka.NewWriter(writerConfig)
	if log == nil {
		log = logger.Default.LogMode(logger.Info)
	}
	return &ProduceClient{
		writer: writer,
		logger: log,
	}
}

func (c *ProduceClient) Close() error {
	return c.writer.Close()
}

func (c *ProduceClient) Send(ctx context.Context, messages ...*CustomKafkaMsg) error {
	kafkaMsgs := make([]kafka.Message, 0, len(messages))
	for _, msg := range messages {
		if msg == nil {
			c.logger.Warn(ctx, "Nil message detected and skipped")
			continue
		}
		kafkaValue, err := json.Marshal(msg)
		if err != nil {
			c.logger.Error(ctx, fmt.Sprintf("Failed to serialize message,error:%v", err))
			return err
		}
		c.logger.Info(ctx, fmt.Sprintf("kafka message:%s", string(kafkaValue)))
		kafkaMsgs = append(kafkaMsgs, kafka.Message{
			Key:   []byte(msg.BalancerKey),
			Value: kafkaValue,
		})
	}
	if err := c.writer.WriteMessages(ctx, kafkaMsgs...); err != nil {
		c.logger.Error(ctx, fmt.Sprintf("Failed to serialize message,error:%v", err))
		return err
	}
	return nil
}

消费者

package customerkafka

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"

	"github.com/segmentio/kafka-go"
	"gorm.io/gorm/logger"
)

// ActionHandler 是一个函数类型，用于处理不同类型的 Kafka 消息
type ActionHandler func(ctx context.Context, msg *CustomKafkaMsg) error

// Consumer 接口定义了消费者需要实现的方法
type Consumer interface {
	Start(ctx context.Context, handlers map[string]ActionHandler) error
	Close() error
}

// ConsumerClient 结构体表示一个Kafka消费者
type ConsumerClient struct {
	reader    *kafka.Reader
	logger    logger.Interface
	wg        sync.WaitGroup
	workerMap sync.Map // 使用 sync.Map 来存储每个 actionType 的 worker 信息
}

// NewConsumeClient 函数用于创建并初始化一个新的ConsumeClient实例
func NewConsumerClient(readerConfig kafka.ReaderConfig, log logger.Interface) *ConsumerClient {
	reader := kafka.NewReader(readerConfig)
	if log == nil {
		log = logger.Default.LogMode(logger.Info)
	}
	return &ConsumerClient{
		reader: reader,
		logger: log,
	}
}

// Start 方法启动消费者并开始读取消息，根据actionType调用不同的处理函数
func (c *ConsumerClient) Start(ctx context.Context, handlers map[string]ActionHandler) error {
	for {
		msg, err := c.reader.ReadMessage(ctx)
		if err != nil {
			c.logger.Error(ctx, "Failed to read message from Kafka", "error", err)
			continue
		}
		c.logger.Info(ctx, fmt.Sprintf("Message on topic: %s value: %s partion:%d offset:%d", msg.Topic, string(msg.Value), msg.Partition, msg.Offset))
		var kafkaMsg CustomKafkaMsg
		if err := json.Unmarshal(msg.Value, &kafkaMsg); err != nil {
			c.logger.Error(ctx, "Failed to unmarshal Kafka message", "error", err)
			continue
		}
		channels := make(map[string]chan *CustomKafkaMsg)
		channels[kafkaMsg.ActionType] = make(chan *CustomKafkaMsg)
		// 使用 sync.Map 来管理 worker
		worker, loaded := c.workerMap.LoadOrStore(kafkaMsg.ActionType, channels)
		if !loaded {
			c.wg.Add(1) // 增加 WaitGroup 计数
			if ch, ok := worker.(map[string]chan *CustomKafkaMsg)[kafkaMsg.ActionType]; ok {
				go c.startWorker(ctx, kafkaMsg.ActionType, handlers, ch)
			}
		}

		// 发送消息到对应的通道,避免阻塞其他消息消费
		if ch, ok := worker.(map[string]chan *CustomKafkaMsg)[kafkaMsg.ActionType]; ok {
			ch <- &kafkaMsg
		}
	}
}

// startWorker 启动处理特定 actionType 的 goroutine
func (c *ConsumerClient) startWorker(ctx context.Context, actionType string, handlers map[string]ActionHandler, ch chan *CustomKafkaMsg) {
	defer c.wg.Done() // 确保在退出时调用

	handlerFunc, ok := handlers[actionType]
	if !ok {
		//不存在对应的handler报错，读到该消息后，不处理，直接丢弃，避免阻塞其他消息消费
		c.logger.Error(ctx, fmt.Sprintf("No handler found for action type: %s", actionType))
		<-ch
		return
	}

	for {
		select {
		case msg, ok := <-ch:
			if !ok {
				return // 通道关闭，退出
			}
			if err := handlerFunc(ctx, msg); err != nil {
				c.logger.Error(ctx, fmt.Sprintf("Failed to handle Kafka message for action type %s, error: %+v", actionType, err))
			}
		case <-ctx.Done():
			return // 上下文取消，退出
		}
	}
}

// Close 方法关闭消费者并释放相关资源
func (c *ConsumerClient) Close() error {
	c.workerMap.Range(func(key, value interface{}) bool {
		close(value.(chan *CustomKafkaMsg)) // 关闭通道
		return true
	})
	c.wg.Wait() // 等待所有 goroutine 完成
	return c.reader.Close()
}

限于个人能力，如有不妥之处，可以给我留言。