commit

2025-11-10 15:30:21 +08:00
parent 4716b82f29
commit df39693dff
43 changed files with 6864 additions and 5 deletions
--- a/main.go
+++ b/main.go
@@ -0,0 +1,526 @@
+package main
+
+import (
+	"bufio"
+	"context"
+	"crypto/sha512"
+	"fmt"
+	"html/template"
+	"io"
+	"log"
+	"net/http"
+	"os"
+	"os/signal"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/IBM/sarama"
+	"github.com/xdg/scram"
+)
+
+// SCRAMClientGeneratorFunc 必须指向一个返回 sarama.SCRAMClient 接口的函数
+// 这里的 func() sarama.SCRAMClient { ... } 就是所需的实现。
+func createScramClient() sarama.SCRAMClient {
+	// 传入用于 SCRAM 认证的 Hash 算法，这里是 SHA-512
+	return &XDGSCRAMClient{HashGeneratorFcn: sha512.New}
+}
+
+// XDGSCRAMClient 是对 github.com/xdg/scram/client 库的包装
+// 使得它符合 Sarama 的 sarama.SCRAMClient 接口
+type XDGSCRAMClient struct {
+	*scram.Client
+	*scram.ClientConversation
+	scram.HashGeneratorFcn
+}
+
+// Begin 用于启动 SCRAM 认证流程
+func (x *XDGSCRAMClient) Begin(userName, password, salt string) (err error) {
+	x.Client, err = x.HashGeneratorFcn.NewClient(userName, password, salt)
+	if err != nil {
+		return err
+	}
+	x.ClientConversation = x.Client.NewConversation()
+	return nil
+}
+
+// Step 用于处理 SCRAM 认证中的每一步挑战和响应
+func (x *XDGSCRAMClient) Step(challenge string) (response string, err error) {
+	response, err = x.ClientConversation.Step(challenge)
+	return
+}
+
+// Done 检查 SCRAM 认证是否完成
+func (x *XDGSCRAMClient) Done() bool {
+	return x.ClientConversation.Done()
+}
+
+// --- 配置常量 ---
+const (
+	// 假设 Kafka 运行在宿主机的 9092 端口
+	// BrokerList = "localhost:9092"
+	// Topic      = "test-topic"
+	BrokerList = "k8s-kafka-kafkaclu-ea0121c7a5-0f29d8348027ca52.elb.us-west-2.amazonaws.com:9094"
+	// prd
+	// BrokerList = "k8s-middlewa-kafkakaf-8362de7238-60b500b1ed68bc69.elb.us-west-2.amazonaws.com:9094"
+	Topic   = "message-users"
+	GroupID = "go-sarama-consumer-group-v1" // 消费组 ID
+)
+
+// --- Web 服务相关常量 ---
+const (
+	WebServerAddress = ":8080" // Web 服务监听的地址和端口
+)
+
+// 初始化 Sarama 配置
+func newConfig() *sarama.Config {
+	config := sarama.NewConfig()
+	// 假设您使用的 Kafka 版本为 4.0.0 或更高 (Sarama 库会兼容)
+	config.Version = sarama.V4_0_0_0 // 请根据您的 Kafka 版本进行调整
+	// 生产者配置
+	config.Producer.RequiredAcks = sarama.WaitForAll // 等待所有副本确认
+	config.Producer.Retry.Max = 5
+	config.Producer.Return.Successes = true // 必须开启，以便同步生产者可以等待确认
+	// 消费者配置：使用 Consumer Group 模式
+	if Topic == "message-users" {
+		// 2. 启用 SASL
+		config.Net.SASL.Enable = true
+		config.Net.SASL.SCRAMClientGeneratorFunc = func() sarama.SCRAMClient {
+			return createScramClient()
+		}
+
+		config.Net.SASL.Mechanism = sarama.SASLTypeSCRAMSHA512
+
+		// 4. 设置用户名和密码
+		config.Net.SASL.User = "kafka-user"
+		config.Net.SASL.Password = "Nrp0vwxgFJXfydV8zISPxjp8RSETjOEq"
+		// prd
+		// config.Net.SASL.Password = "oP0rACG2ka5EBVWfuAyR1fslH4yAlLZX"
+	}
+
+	config.Consumer.Return.Errors = true
+	config.Consumer.Group.Session.Timeout = 20 * time.Second
+	config.Metadata.Timeout = 15 * time.Second
+	return config
+}
+
+// --------------------------------------------------------------------------------------
+// 生产者逻辑
+// --------------------------------------------------------------------------------------
+func startProducer(config *sarama.Config) {
+	// 使用同步生产者 (SyncProducer)
+	producer, err := sarama.NewSyncProducer([]string{BrokerList}, config)
+	if err != nil {
+		log.Printf("Error creating sync producer: %v", err) // 改为 Printf 以避免程序退出
+		return
+	}
+	defer producer.Close()
+	log.Println("Producer started. Sending messages from long1.txt...")
+
+	f, err := os.Open("./long1.txt")
+	if err != nil {
+		log.Printf("failed to read file: %v", err) // 改为 Printf
+		return
+	}
+	defer f.Close()
+
+	bs, _ := io.ReadAll(f)
+	str := string(bs)
+	lines := strings.Split(str, "\n")
+
+	messageID := 0
+	// 循环发送消息
+	for _, line := range lines {
+		// 忽略空行
+		if strings.TrimSpace(line) == "" {
+			continue
+		}
+
+		// 1. 构造 ProducerRecord
+		msg := &sarama.ProducerMessage{
+			Topic: Topic,
+			Key:   sarama.StringEncoder(fmt.Sprintf("Key-%d", messageID)),
+			Value: sarama.StringEncoder(line),
+		}
+
+		// 2. 发送消息
+		partition, offset, err := producer.SendMessage(msg)
+		if err != nil {
+			log.Printf("Failed to send message: %v", err)
+		} else {
+			log.Printf("Producer sent: Topic=%s | Partition=%d | Offset=%d | Value='%s'", Topic, partition, offset, line)
+			messageID++
+		}
+
+		// 3. 暂停 (根据需求调整)
+		time.Sleep(10 * time.Millisecond)
+	}
+	log.Println("Producer finished sending all messages.")
+}
+
+// fetchPartitionMessages connects to a specific partition starting from the given offset,
+// writes up to limit messages into outputPath (one per line), then returns.
+// If filterKey is not empty, only messages with matching key are written.
+// If withTimestamp is true, prepends formatted timestamp to each message.
+func fetchPartitionMessages(config *sarama.Config, partition int32, offset int64, limit int, outputPath string, filterKey string, withTimestamp bool) error {
+	if limit <= 0 {
+		return fmt.Errorf("limit must be > 0")
+	}
+
+	client, err := sarama.NewClient([]string{BrokerList}, config)
+	if err != nil {
+		return err
+	}
+	defer client.Close()
+
+	consumer, err := sarama.NewConsumerFromClient(client)
+	if err != nil {
+		return err
+	}
+	defer consumer.Close()
+
+	pc, err := consumer.ConsumePartition(Topic, partition, offset)
+	if err != nil {
+		return err
+	}
+	defer pc.Close()
+
+	outFile, err := os.Create(outputPath)
+	if err != nil {
+		return err
+	}
+	defer outFile.Close()
+
+	writer := bufio.NewWriter(outFile)
+	defer writer.Flush()
+
+	received := 0
+	scanned := 0
+	filterEnabled := filterKey != ""
+
+	// 设置超时和最大扫描数
+	timeout := time.After(15 * time.Second)
+	maxScan := limit * 1000 // 如果过滤模式下扫描了这么多消息还没找到足够的匹配，就放弃
+	if !filterEnabled {
+		maxScan = limit // 非过滤模式下，扫描数就是 limit
+	}
+
+	for received < limit && scanned < maxScan {
+		select {
+		case <-timeout:
+			log.Printf("Fetch timeout after 30s. Retrieved %d messages (scanned %d)", received, scanned)
+			return nil // 超时不算错误，返回已获取的消息
+		case err, ok := <-pc.Errors():
+			if !ok {
+				return fmt.Errorf("partition consumer error channel closed")
+			}
+			if err != nil {
+				return err.Err
+			}
+		case msg, ok := <-pc.Messages():
+			if !ok {
+				log.Printf("Partition consumer closed. Retrieved %d messages (scanned %d)", received, scanned)
+				return nil // channel 关闭说明没有更多消息了
+			}
+			if msg == nil {
+				continue
+			}
+
+			scanned++
+
+			// 如果启用了 key 过滤，检查 key 是否匹配
+			if filterEnabled {
+				msgKey := string(msg.Key)
+				if msgKey != filterKey {
+					continue // 跳过不匹配的消息
+				}
+				log.Printf("Found matching message %d/%d (scanned %d total)", received+1, limit, scanned)
+			}
+
+			// 构造要写入的内容
+			var line string
+			if withTimestamp {
+				// 格式化时间戳为易读格式: 2006-01-02 15:04:05.000
+				timestampStr := msg.Timestamp.Format("2006-01-02 15:04:05.000")
+				line = fmt.Sprintf("%s %s\n", timestampStr, string(msg.Value))
+			} else {
+				line = string(msg.Value) + "\n"
+			}
+
+			if _, err := writer.WriteString(line); err != nil {
+				return err
+			}
+			received++
+		}
+	}
+
+	if scanned >= maxScan && received < limit {
+		log.Printf("Reached max scan limit (%d messages scanned). Retrieved %d matching messages", scanned, received)
+	}
+
+	return nil
+}
+
+// --------------------------------------------------------------------------------------
+// 消费者逻辑 (保持不变)
+// --------------------------------------------------------------------------------------
+// ConsumerHandler 结构体实现了 sarama.ConsumerGroupHandler 接口
+type ConsumerHandler struct{}
+
+// Setup 是在新会话开始时运行的，在 ConsumeClaim 之前
+func (consumer *ConsumerHandler) Setup(sarama.ConsumerGroupSession) error {
+	log.Println("Consumer Group session setup complete.")
+	return nil
+}
+
+// Cleanup 是在会话结束时运行的，在所有 ConsumeClaim 协程退出后
+func (consumer *ConsumerHandler) Cleanup(sarama.ConsumerGroupSession) error {
+	log.Println("Consumer Group session cleanup complete.")
+	return nil
+}
+
+// ConsumeClaim 必须启动一个消费循环，处理分配给当前 consumer 的分区消息
+func (consumer *ConsumerHandler) ConsumeClaim(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
+	for message := range claim.Messages() {
+		log.Printf("Consumer Received: Topic=%s | Partition=%d | Offset=%d | Value='%s' (Timestamp: %v)",
+			message.Topic, message.Partition, message.Offset, string(message.Value), message.Timestamp)
+		session.MarkMessage(message, "")
+	}
+	log.Printf("Consumer for Partition %d stopped or rebalancing...", claim.Partition())
+	return nil
+}
+
+func startConsumer(config *sarama.Config) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	group, err := sarama.NewConsumerGroup([]string{BrokerList}, GroupID, config)
+	if err != nil {
+		log.Fatalf("Error creating consumer group client: %v", err)
+	}
+	defer group.Close()
+
+	log.Printf("Consumer Group '%s' started. Listening on topic '%s'...", GroupID, Topic)
+	wg := &sync.WaitGroup{}
+	wg.Add(1)
+
+	go func() {
+		defer wg.Done()
+		handler := ConsumerHandler{}
+		for {
+			if err := group.Consume(ctx, []string{Topic}, &handler); err != nil {
+				if ctx.Err() != nil {
+					return
+				}
+				log.Printf("Error from consumer group: %v", err)
+				time.Sleep(5 * time.Second)
+			}
+		}
+	}()
+
+	signals := make(chan os.Signal, 1)
+	signal.Notify(signals, os.Interrupt)
+	if _, ok := <-signals; ok {
+		log.Println("Interrupt signal received. Shutting down consumer...")
+		cancel()
+	}
+
+	wg.Wait()
+	log.Println("Consumer shutdown complete.")
+}
+
+// --------------------------------------------------------------------------------------
+// Web 服务逻辑
+// --------------------------------------------------------------------------------------
+
+// handleRoot 处理根路径请求，显示操作页面
+func handleRoot(w http.ResponseWriter, r *http.Request) {
+	// 定义一个简单的 HTML 页面模板
+	const htmlTemplate = `
+	<!DOCTYPE html>
+	<html lang="en">
+	<head>
+		<meta charset="UTF-8">
+		<title>Kafka Web Producer</title>
+		<style>
+			body { font-family: Arial, sans-serif; display: flex; justify-content: center; align-items: center; height: 100vh; margin: 0; background-color: #f4f4f4; }
+			.container { text-align: center; }
+			#startButton {
+				padding: 15px 30px;
+				font-size: 18px;
+				cursor: pointer;
+				background-color: #4CAF50;
+				color: white;
+				border: none;
+				border-radius: 5px;
+				transition: background-color 0.3s;
+			}
+			#startButton:hover { background-color: #45a049; }
+			#status { margin-top: 20px; font-size: 16px; color: #333; }
+		</style>
+	</head>
+	<body>
+		<div class="container">
+			<h1>Kafka 消息生产者</h1>
+			<button id="startButton">开始发送消息</button>
+			<p id="status"></p>
+		</div>
+		<script>
+			document.getElementById('startButton').addEventListener('click', function() {
+				const statusElem = document.getElementById('status');
+				statusElem.textContent = '正在请求服务器开始发送消息...';
+
+				fetch('/start-producer', { method: 'POST' })
+					.then(response => response.json())
+					.then(data => {
+						console.log('Success:', data);
+						statusElem.textContent = data.message;
+					})
+					.catch((error) => {
+						console.error('Error:', error);
+						statusElem.textContent = '请求失败: ' + error;
+					});
+			});
+		</script>
+	</body>
+	</html>
+	`
+	// 解析并执行模板
+	tmpl, err := template.New("index").Parse(htmlTemplate)
+	if err != nil {
+		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
+		return
+	}
+	tmpl.Execute(w, nil)
+}
+
+// handleStartProducer 处理开始生产消息的请求
+func handleStartProducer(config *sarama.Config) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodPost {
+			http.Error(w, "Invalid request method", http.StatusMethodNotAllowed)
+			return
+		}
+
+		log.Println("Received request to start producer...")
+
+		// 使用 goroutine 在后台运行生产者，避免阻塞 HTTP 响应
+		go startProducer(config)
+
+		// 立即返回响应给前端
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+		fmt.Fprintln(w, `{"message": "生产者已在后台启动，请查看服务器日志确认消息发送状态。"}`)
+	}
+}
+
+// startWebServer 启动 Web 服务
+func startWebServer(config *sarama.Config) {
+	// 注册 HTTP 路由处理器
+	http.HandleFunc("/", handleRoot)
+	http.HandleFunc("/start-producer", handleStartProducer(config))
+
+	log.Printf("Web server starting on %s...", WebServerAddress)
+	log.Println("Visit http://localhost:8080 in your browser.")
+
+	// 启动 HTTP 服务器
+	if err := http.ListenAndServe(WebServerAddress, nil); err != nil {
+		log.Fatalf("Failed to start web server: %v", err)
+	}
+}
+
+// --------------------------------------------------------------------------------------
+// 主函数与命令行处理
+// --------------------------------------------------------------------------------------
+func main() {
+	if len(os.Args) < 2 {
+		fmt.Println("Usage: go run main.go <web|consumer|fetch>")
+		fmt.Println("  web:      Starts the web server to trigger the producer.")
+		fmt.Println("  consumer: Starts the Kafka consumer.")
+		fmt.Println("  fetch <partition> <offset|oldest|latest> <count> <output-file> [key] [--with-timestamp]: Fetches messages to a file.")
+		fmt.Println("           If [key] is provided, only messages with that key will be fetched.")
+		fmt.Println("           If --with-timestamp is provided, timestamp will be prepended to each message.")
+		return
+	}
+
+	// 初始化配置
+	config := newConfig()
+
+	role := strings.ToLower(os.Args[1])
+	switch role {
+	case "web":
+		startWebServer(config) // 启动 Web 服务
+	case "consumer":
+		startConsumer(config) // 启动消费者
+	case "fetch":
+		if len(os.Args) < 6 {
+			fmt.Println("Usage: go run main.go fetch <partition> <offset|oldest|latest> <count> <output-file> [key] [--with-timestamp]")
+			return
+		}
+
+		partitionVal, err := strconv.Atoi(os.Args[2])
+		if err != nil {
+			log.Fatalf("invalid partition: %v", err)
+		}
+
+		offsetArg := strings.ToLower(os.Args[3])
+		var offset int64
+		switch offsetArg {
+		case "oldest":
+			offset = sarama.OffsetOldest
+		case "latest":
+			offset = sarama.OffsetNewest
+		default:
+			offset, err = strconv.ParseInt(offsetArg, 10, 64)
+			if err != nil {
+				log.Fatalf("invalid offset: %v", err)
+			}
+		}
+
+		limit, err := strconv.Atoi(os.Args[4])
+		if err != nil {
+			log.Fatalf("invalid count: %v", err)
+		}
+
+		outputPath := os.Args[5]
+
+		// 可选的 key 参数和 timestamp 标志
+		var filterKey string
+		withTimestamp := false
+
+		// 检查剩余的参数
+		for i := 6; i < len(os.Args); i++ {
+			arg := os.Args[i]
+			if arg == "--with-timestamp" || arg == "-t" {
+				withTimestamp = true
+			} else if filterKey == "" {
+				// 第一个非标志参数作为 filterKey
+				filterKey = arg
+			}
+		}
+
+		if filterKey != "" {
+			log.Printf("Filtering by key: %s", filterKey)
+		}
+		if withTimestamp {
+			log.Printf("Will include timestamps in output")
+		}
+
+		if err := fetchPartitionMessages(config, int32(partitionVal), offset, limit, outputPath, filterKey, withTimestamp); err != nil {
+			log.Fatalf("fetch failed: %v", err)
+		}
+
+		offsetDesc := offsetArg
+		if offsetArg != "oldest" && offsetArg != "latest" {
+			offsetDesc = fmt.Sprintf("%d", offset)
+		}
+		if filterKey != "" {
+			log.Printf("Fetched %d messages with key '%s' from partition %d starting at %s into %s", limit, filterKey, partitionVal, offsetDesc, outputPath)
+		} else {
+			log.Printf("Fetched %d messages from partition %d starting at %s into %s", limit, partitionVal, offsetDesc, outputPath)
+		}
+	default:
+		fmt.Println("Invalid argument. Usage: go run main.go <web|consumer|fetch>")
+	}
+}