GO消费kafka，动态扩展协程并发数，消费海量数据

1. 初始化项目

go mod init kafkaConsumer

2. 安装kafka-go

go get github.com/segmentio/kafka-go

3. 代码实现

• main.go

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/segmentio/kafka-go"
)

// consumeAndProcess 消费Kafka数据并处理
func consumeAndProcess(brokers []string, group, topic string, maxGoroutines int) {
	// make a new reader that consumes from topic-A
	r := kafka.NewReader(kafka.ReaderConfig{
		Brokers:     brokers,
		GroupID:     group,
		Topic:       topic,
		MaxBytes:    10e6,             // 10MB
		StartOffset: kafka.LastOffset, // 默认从最新的数据开始消费
	})
	// 创建一个带缓冲的通道来控制并发量
	semaphore := make(chan struct{}, maxGoroutines)

	// 创建一个通道来监听系统信号
	sigchan := make(chan os.Signal, 1)
	signal.Notify(sigchan, syscall.SIGINT, syscall.SIGTERM)

	run := true

	for run {
		select {
		case sig := <-sigchan:
			fmt.Printf("捕获到信号 %v: 开始优雅关闭\n", sig)
			run = false
		default:
			semaphore <- struct{}{}

			msg, err := r.ReadMessage(context.Background())
			if err != nil {
				log.Fatal("failed to read msg, err:", err)
			}
			go func(msg kafka.Message) {
				defer func() { <-semaphore }()
				processMessage(msg)
			}(msg)
		}
	}
	if err := r.Close(); err != nil {
		log.Fatal("failed to close reader:", err)
	}
}

// processMessage 数据处理
func processMessage(msg kafka.Message) {
	//fmt.Printf("message at topic/partition/offset %v/%v/%v: %s = %s\n", msg.Topic, msg.Partition, msg.Offset, string(msg.Key), string(msg.Value))
	fmt.Printf("key:%s, value:%s \n", string(msg.Key), string(msg.Value))
	time.Sleep(2 * time.Second) // 假设处理消息需要2秒
}

func main() {
	brokers := []string{"127.0.0.1：9092"}
	group := "go-group"
	topic := "dlp-client"
	//maxGoroutines := 10
	maxGoroutines := flag.Int("max-goroutines", 2, "Maximum number of goroutines")

	// 解析命令行参数
	flag.Parse()

	consumeAndProcess(brokers, group, topic, *maxGoroutines)
}

• go.mod

module kafkaConsumer

go 1.22.5

require (
	github.com/klauspost/compress v1.15.9 // indirect
	github.com/pierrec/lz4/v4 v4.1.15 // indirect
	github.com/segmentio/kafka-go v0.4.47 // indirect
)

4. 动态传递最大协程数

• 设置为5，默认不设置就是2

go run main.go -max-goroutines 5

5. 代码分析

• 通过协程并发处理这些消息，同时通过通道控制并发量
• 创建一个带缓冲的通道 semaphore 用于控制最大并发量
• 创建一个通道 sigchan 用于监听系统信号（如 SIGINT 和 SIGTERM），以便在收到终止信号时优雅地关闭消费者
• 运行分析
  • 启动消费者：从指定的Kafka broker中消费指定主题的消息。
  • 并发控制：通过带缓冲的通道 semaphore 控制最大并发量，避免过多的协程同时运行。
  • 优雅关闭：通过监听系统信号，实现优雅关闭消费者。
  • 消息处理：使用协程并发处理每个消费到的消息，并模拟消息处理的时间延迟。
• 错误处理
  • 如果读取消息失败，会记录错误日志并终止程序。
  • 如果关闭消费者失败，也会记录错误日志并终止程序。

6. 总结

这段代码展示了如何使用 kafka-go 库实现一个高效的Kafka消费者，并通过协程并发处理消息，同时使用通道控制并发量，确保在处理过程中不会创建过多的协程，从而实现资源的高效利用。