feat: batch kafka message persistence

2026-09-24 22:45:07 +08:00
parent 7f2df1dd4f
commit f898a7171c
18 changed files with 682 additions and 98 deletions
@@ -51,6 +51,9 @@ type kafkaConfig struct {
 	SASLPassword     string            `json:"sasl_password"`
 	MaxRetryCount    int               `json:"max_retry_count"`
 	RetryBackoffMS   []int             `json:"retry_backoff_ms"`
+	MessageBatchSize int               `json:"message_batch_size"`
+	MessageBatchWait int               `json:"message_batch_wait_ms"`
+	MessageRetries   int               `json:"message_single_retry_count"`
 	DeadRetentionSec int               `json:"dead_letter_retention_seconds"`
 }

@@ -180,6 +183,17 @@ func normalizeKafkaConfig(cfg *kafkaConfig) {
 	if len(cfg.RetryBackoffMS) == 0 {
 		cfg.RetryBackoffMS = []int{10000, 30000, 60000, 300000, 900000}
 	}
+	if cfg.MessageBatchSize <= 0 {
+		cfg.MessageBatchSize = 100
+	} else if cfg.MessageBatchSize > 500 {
+		cfg.MessageBatchSize = 500
+	}
+	if cfg.MessageBatchWait <= 0 {
+		cfg.MessageBatchWait = 10
+	}
+	if cfg.MessageRetries <= 0 {
+		cfg.MessageRetries = 3
+	}
 	if cfg.DeadRetentionSec <= 0 {
 		cfg.DeadRetentionSec = 90 * 24 * 3600
 	}
@@ -276,6 +290,11 @@ func publishKafkaEnvelope(ctx context.Context, topicName string, key string, env

 type kafkaMessageHandler func(context.Context, kafkaEnvelope, kafka.Message) error

+type kafkaMessageCommitter func(context.Context, kafka.Message) error
+type kafkaRetryWaiter func(context.Context, time.Duration) error
+type kafkaBatchHandler func(context.Context, []kafka.Message) error
+type kafkaBatchCommitter func(context.Context, ...kafka.Message) error
+
 func startKafkaConsumer(topicName string, handler kafkaMessageHandler) {
 	startKafkaConsumerWithGroup(topicName, "", handler)
 }
@@ -327,24 +346,157 @@ func startKafkaConsumerWithGroup(topicName, groupID string, handler kafkaMessage
 				logs.Err.Println("Kafka fetch failed:", err)
 				continue
 			}
-			var envelope kafkaEnvelope
-			if err = json.Unmarshal(msg.Value, &envelope); err != nil {
-				logs.Err.Println("Kafka envelope decode failed:", err)
-				_ = reader.CommitMessages(context.Background(), msg)
-				continue
+			commit := func(ctx context.Context, message kafka.Message) error {
+				return reader.CommitMessages(ctx, message)
 			}
-			if err = handler(kafkaBus.ctx, envelope, msg); err != nil {
-				logs.Err.Println("Kafka handler failed:", envelope.EventID, envelope.EventType, err)
-				continue
-			}
-			if err = reader.CommitMessages(context.Background(), msg); err != nil {
-				logs.Err.Println("Kafka commit failed:", err)
+			if err = consumeKafkaMessage(kafkaBus.ctx, msg, handler, commit, waitKafkaRetry); err != nil {
+				if errors.Is(err, context.Canceled) || errors.Is(kafkaBus.ctx.Err(), context.Canceled) {
+					return
+				}
+				logs.Err.Printf("Kafka message stopped without commit: topic=%s partition=%d offset=%d err=%v", msg.Topic, msg.Partition, msg.Offset, err)
 			}
 		}
 	}()
 	logs.Info.Printf("Kafka consumer started: topic=%s group=%s", topic, groupID)
 }

+func startKafkaBatchConsumer(topicName string, handler kafkaBatchHandler) {
+	if kafkaBus == nil || handler == nil {
+		return
+	}
+	topic := kafkaBus.topics[topicName]
+	if topic == "" {
+		topic = topicName
+	}
+	groupID := kafkaBus.config.ConsumerGroup + "." + topicName
+	reader := kafka.NewReader(kafka.ReaderConfig{
+		Brokers: kafkaBus.config.Brokers, GroupID: groupID, Topic: topic, StartOffset: kafka.FirstOffset,
+		MinBytes: 1, MaxBytes: 10e6, MaxWait: time.Second, Dialer: kafkaBus.dialer(),
+	})
+	kafkaBus.wg.Add(1)
+	go func() {
+		defer kafkaBus.wg.Done()
+		defer reader.Close()
+		for {
+			messages, err := fetchKafkaMessageBatch(kafkaBus.ctx, reader, kafkaBus.config.MessageBatchSize,
+				time.Duration(kafkaBus.config.MessageBatchWait)*time.Millisecond)
+			if err != nil {
+				if errors.Is(err, context.Canceled) || errors.Is(kafkaBus.ctx.Err(), context.Canceled) {
+					return
+				}
+				logs.Err.Println("Kafka batch fetch failed:", err)
+				continue
+			}
+			if err = consumeKafkaMessageBatch(kafkaBus.ctx, messages, handler, reader.CommitMessages, waitKafkaRetry); err != nil {
+				if errors.Is(err, context.Canceled) || errors.Is(kafkaBus.ctx.Err(), context.Canceled) {
+					return
+				}
+				logs.Err.Printf("Kafka batch stopped without commit: topic=%s size=%d err=%v", topic, len(messages), err)
+			}
+		}
+	}()
+	logs.Info.Printf("Kafka batch consumer started: topic=%s group=%s size=%d wait_ms=%d", topic, groupID, kafkaBus.config.MessageBatchSize, kafkaBus.config.MessageBatchWait)
+}
+
+func fetchKafkaMessageBatch(ctx context.Context, reader *kafka.Reader, maxSize int, maxWait time.Duration) ([]kafka.Message, error) {
+	first, err := reader.FetchMessage(ctx)
+	if err != nil {
+		return nil, err
+	}
+	messages := []kafka.Message{first}
+	deadline := time.Now().Add(maxWait)
+	for len(messages) < maxSize {
+		remaining := time.Until(deadline)
+		if remaining <= 0 {
+			break
+		}
+		fetchCtx, cancel := context.WithTimeout(ctx, remaining)
+		message, fetchErr := reader.FetchMessage(fetchCtx)
+		cancel()
+		if errors.Is(fetchErr, context.DeadlineExceeded) && ctx.Err() == nil {
+			break
+		}
+		if fetchErr != nil {
+			return nil, fetchErr
+		}
+		messages = append(messages, message)
+	}
+	return messages, nil
+}
+
+func consumeKafkaMessageBatch(ctx context.Context, messages []kafka.Message, handler kafkaBatchHandler, commit kafkaBatchCommitter, wait kafkaRetryWaiter) error {
+	for attempt := 1; ; attempt++ {
+		if err := handler(ctx, messages); err == nil {
+			return commitKafkaMessageBatch(ctx, messages, commit, wait)
+		}
+		if err := wait(ctx, kafkaRetryBackoff(attempt)); err != nil {
+			return err
+		}
+	}
+}
+
+func consumeKafkaMessage(ctx context.Context, msg kafka.Message, handler kafkaMessageHandler, commit kafkaMessageCommitter, wait kafkaRetryWaiter) error {
+	var envelope kafkaEnvelope
+	if err := json.Unmarshal(msg.Value, &envelope); err != nil {
+		return commitKafkaMessage(ctx, msg, commit, wait)
+	}
+	for attempt := 1; ; attempt++ {
+		if err := handler(ctx, envelope, msg); err == nil {
+			return commitKafkaMessage(ctx, msg, commit, wait)
+		}
+		if err := wait(ctx, kafkaRetryBackoff(attempt)); err != nil {
+			return err
+		}
+	}
+}
+
+func commitKafkaMessage(ctx context.Context, msg kafka.Message, commit kafkaMessageCommitter, wait kafkaRetryWaiter) error {
+	for attempt := 1; ; attempt++ {
+		if err := commit(ctx, msg); err == nil {
+			return nil
+		}
+		if err := wait(ctx, kafkaRetryBackoff(attempt)); err != nil {
+			return err
+		}
+	}
+}
+
+func commitKafkaMessageBatch(ctx context.Context, messages []kafka.Message, commit kafkaBatchCommitter, wait kafkaRetryWaiter) error {
+	for attempt := 1; ; attempt++ {
+		if err := commit(ctx, messages...); err == nil {
+			return nil
+		}
+		if err := wait(ctx, kafkaRetryBackoff(attempt)); err != nil {
+			return err
+		}
+	}
+}
+
+func kafkaRetryBackoff(attempt int) time.Duration {
+	if kafkaBus == nil || len(kafkaBus.config.RetryBackoffMS) == 0 {
+		return time.Second
+	}
+	if attempt < 1 {
+		attempt = 1
+	}
+	index := attempt - 1
+	if index >= len(kafkaBus.config.RetryBackoffMS) {
+		index = len(kafkaBus.config.RetryBackoffMS) - 1
+	}
+	return time.Duration(kafkaBus.config.RetryBackoffMS[index]) * time.Millisecond
+}
+
+func waitKafkaRetry(ctx context.Context, delay time.Duration) error {
+	timer := time.NewTimer(delay)
+	defer timer.Stop()
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-timer.C:
+		return nil
+	}
+}
+
 func kafkaBroadcastConsumerGroup(topicName string) string {
 	if kafkaBus == nil {
 		return ""