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
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"strings"
 	"time"

 	"github.com/segmentio/kafka-go"
@@ -24,8 +25,12 @@ type kafkaMessagePersistPayload struct {
 var publishKafkaMessagePersist = publishKafkaEnvelope

 func startKafkaMessagePersistConsumer() {
-	startKafkaConsumer("message_persist", handleKafkaMessagePersist)
-	startKafkaConsumer("group_message_persist", handleKafkaGroupMessagePersist)
+	startKafkaBatchConsumer("message_persist", func(ctx context.Context, messages []kafka.Message) error {
+		return handleKafkaMessagePersistBatch(ctx, messages, "message_persist")
+	})
+	startKafkaBatchConsumer("group_message_persist", func(ctx context.Context, messages []kafka.Message) error {
+		return handleKafkaMessagePersistBatch(ctx, messages, "group_message_persist")
+	})
 }

 func saveMessageMaybeKafka(ctx context.Context, tenantID types.TenantID, msg *types.Message, attachmentURLs []string, readBySender bool) (*types.MessageSaveResult, error) {
@@ -132,70 +137,155 @@ func enqueueKafkaMessagePersist(ctx context.Context, tenantID types.TenantID, ms
 	return publishKafkaMessagePersist(ctx, kafkaMessagePersistTopicName(msg.Topic), msg.Topic, envelope)
 }

-func handleKafkaMessagePersist(ctx context.Context, envelope kafkaEnvelope, _ kafka.Message) error {
-	return handleKafkaMessagePersistFromTopic(ctx, envelope, "message_persist")
+type kafkaMessagePersistItem struct {
+	envelope kafkaEnvelope
+	tenantID types.TenantID
+	payload  kafkaMessagePersistPayload
 }

-func handleKafkaGroupMessagePersist(ctx context.Context, envelope kafkaEnvelope, _ kafka.Message) error {
-	return handleKafkaMessagePersistFromTopic(ctx, envelope, "group_message_persist")
+func handleKafkaMessagePersistBatch(ctx context.Context, messages []kafka.Message, sourceTopicName string) error {
+	groups := make(map[string][]*kafkaMessagePersistItem)
+	for _, message := range messages {
+		item, terminal, err := prepareKafkaMessagePersistItem(ctx, message, sourceTopicName)
+		if err != nil {
+			return err
+		}
+		if terminal {
+			continue
+		}
+		key := fmt.Sprintf("%d", item.tenantID)
+		if item.payload.MessageKind == "group" {
+			key += ":" + item.payload.Message.ShardMonth
+		}
+		groups[key] = append(groups[key], item)
+	}
+	for _, items := range groups {
+		batch := make([]store.PreparedMessageBatchItem, 0, len(items))
+		for _, item := range items {
+			batch = append(batch, store.PreparedMessageBatchItem{Message: &item.payload.Message, ReadBySender: item.payload.ReadBySender})
+		}
+		if err := store.PersistPreparedMessageBatch(items[0].tenantID, batch); err == nil {
+			for _, item := range items {
+				if err = finishKafkaMessagePersist(ctx, item); err != nil {
+					if err = persistKafkaMessageIndividually(ctx, item); err != nil {
+						return err
+					}
+				}
+			}
+			continue
+		} else {
+			logs.Warn.Printf("kafka message batch persist failed, degrading to singles: tenant=%d kind=%s shard=%s size=%d err=%v",
+				items[0].tenantID, items[0].payload.MessageKind, items[0].payload.Message.ShardMonth, len(items), err)
+		}
+		for _, item := range items {
+			if err := persistKafkaMessageIndividually(ctx, item); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
 }

-func handleKafkaMessagePersistFromTopic(ctx context.Context, envelope kafkaEnvelope, sourceTopicName string) error {
+func prepareKafkaMessagePersistItem(ctx context.Context, message kafka.Message, sourceTopicName string) (*kafkaMessagePersistItem, bool, error) {
+	var envelope kafkaEnvelope
+	if err := json.Unmarshal(message.Value, &envelope); err != nil {
+		logs.Err.Printf("kafka message malformed envelope skipped: partition=%d offset=%d err=%v", message.Partition, message.Offset, err)
+		return nil, true, nil
+	}
 	tenantID := envelope.TenantID
 	var err error
 	if tenantID.IsZero() {
 		tenantID, err = kafkaTenantID(envelope.TenantCode)
 		if err != nil {
-			return err
+			return nil, false, err
 		}
 	}
 	consume, claimed, err := beginKafkaConsumeForTopic(tenantID, envelope, sourceTopicName)
-	if err != nil || !claimed {
-		return err
+	if err != nil {
+		return nil, false, err
 	}
-	var payload kafkaMessagePersistPayload
-	if err = json.Unmarshal(envelope.Payload, &payload); err != nil {
-		return markKafkaDeadLetter(ctx, tenantID, envelope, consume, "MALFORMED_PAYLOAD", err.Error())
+	if !claimed {
+		return nil, true, nil
 	}
-	if payload.Message.TenantID == types.ZeroTenantID {
-		payload.Message.TenantID = tenantID
+	item := &kafkaMessagePersistItem{envelope: envelope, tenantID: tenantID}
+	if err = json.Unmarshal(envelope.Payload, &item.payload); err != nil {
+		return nil, true, markKafkaDeadLetter(ctx, tenantID, envelope, consume, "MALFORMED_PAYLOAD", err.Error())
 	}
-	if payload.Message.ShardMonth == "" {
-		payload.Message.ShardMonth = payload.ShardMonth
+	if item.payload.Message.TenantID == types.ZeroTenantID {
+		item.payload.Message.TenantID = tenantID
 	}
-	if payload.Message.TenantID != tenantID {
-		return markKafkaDeadLetter(ctx, tenantID, envelope, consume, "TENANT_MISMATCH", "message tenant does not match envelope tenant")
+	if item.payload.Message.ShardMonth == "" {
+		item.payload.Message.ShardMonth = item.payload.ShardMonth
 	}
-	if err = validateKafkaMessageKind(payload); err != nil {
-		return markKafkaDeadLetter(ctx, tenantID, envelope, consume, "MALFORMED_PAYLOAD", err.Error())
+	if item.payload.Message.TenantID != tenantID {
+		return nil, true, markKafkaDeadLetter(ctx, tenantID, envelope, consume, "TENANT_MISMATCH", "message tenant does not match envelope tenant")
 	}
-	if _, err = store.PersistPreparedMessage(tenantID, &payload.Message, nil, payload.ReadBySender); err != nil && !errors.Is(err, types.ErrDuplicate) {
-		failed, markErr := store.HTTPInfrastructure.MarkKafkaConsumeFailed(tenantID, envelope.EventID, kafkaBus.config.ConsumerGroup,
-			kafkaCommandErrorReason(err), err.Error(), time.Now().UTC().Round(time.Millisecond))
-		if markErr != nil {
-			return markErr
+	if err = validateKafkaMessageKind(item.payload); err != nil {
+		return nil, true, markKafkaDeadLetter(ctx, tenantID, envelope, consume, "MALFORMED_PAYLOAD", err.Error())
+	}
+	return item, false, nil
+}
+
+func persistKafkaMessageIndividually(ctx context.Context, item *kafkaMessagePersistItem) error {
+	operation := func() error {
+		err := store.PersistPreparedMessageBatch(item.tenantID, []store.PreparedMessageBatchItem{{
+			Message: &item.payload.Message, ReadBySender: item.payload.ReadBySender,
+		}})
+		if errors.Is(err, types.ErrUnsupported) {
+			_, err = store.PersistPreparedMessage(item.tenantID, &item.payload.Message, nil, item.payload.ReadBySender)
 		}
-		if failed != nil && failed.RetryCount >= kafkaBus.config.MaxRetryCount {
-			return markKafkaDeadLetter(ctx, tenantID, envelope, failed, kafkaCommandErrorReason(err), err.Error())
-		}
-		sleepKafkaRetryBackoff(failed.RetryCount)
-		return err
-	}
-	if len(payload.AttachmentURLs) > 0 {
-		if err = enqueueKafkaFileAttachmentsBind(ctx, tenantID, &payload.Message, payload.AttachmentURLs); err != nil {
-			failed, markErr := store.HTTPInfrastructure.MarkKafkaConsumeFailed(tenantID, envelope.EventID, kafkaBus.config.ConsumerGroup,
-				"KAFKA_FILE_LIFECYCLE_FAILED", err.Error(), time.Now().UTC().Round(time.Millisecond))
-			if markErr != nil {
-				return markErr
-			}
-			if failed != nil && failed.RetryCount >= kafkaBus.config.MaxRetryCount {
-				return markKafkaDeadLetter(ctx, tenantID, envelope, failed, "KAFKA_FILE_LIFECYCLE_FAILED", err.Error())
-			}
-			sleepKafkaRetryBackoff(failed.RetryCount)
+		if err != nil && !errors.Is(err, types.ErrDuplicate) {
 			return err
 		}
+		return finishKafkaMessagePersist(ctx, item)
 	}
-	return store.HTTPInfrastructure.MarkKafkaConsumeSucceeded(tenantID, envelope.EventID, kafkaBus.config.ConsumerGroup, time.Now().UTC().Round(time.Millisecond))
+	markFailed := func(err error) (*types.KafkaConsumeRecord, error) {
+		return store.HTTPInfrastructure.MarkKafkaConsumeFailed(item.tenantID, item.envelope.EventID, kafkaBus.config.ConsumerGroup,
+			kafkaMessagePersistFailureReason(err), err.Error(), time.Now().UTC().Round(time.Millisecond))
+	}
+	markDead := func(failed *types.KafkaConsumeRecord, err error) error {
+		return markKafkaDeadLetter(ctx, item.tenantID, item.envelope, failed, kafkaMessagePersistFailureReason(err), err.Error())
+	}
+	return runKafkaMessageSingleRetries(ctx, kafkaBus.config.MessageRetries, operation, markFailed, markDead, waitKafkaRetry)
+}
+
+func runKafkaMessageSingleRetries(ctx context.Context, retryCount int, operation func() error,
+	markFailed func(error) (*types.KafkaConsumeRecord, error), markDead func(*types.KafkaConsumeRecord, error) error,
+	wait kafkaRetryWaiter) error {
+	var lastErr error
+	var failed *types.KafkaConsumeRecord
+	for attempt := 1; attempt <= retryCount+1; attempt++ {
+		if lastErr = operation(); lastErr == nil {
+			return nil
+		}
+		var err error
+		if failed, err = markFailed(lastErr); err != nil {
+			return err
+		}
+		if attempt <= retryCount {
+			if err = wait(ctx, kafkaRetryBackoff(attempt)); err != nil {
+				return err
+			}
+		}
+	}
+	return markDead(failed, lastErr)
+}
+
+func finishKafkaMessagePersist(ctx context.Context, item *kafkaMessagePersistItem) error {
+	if len(item.payload.AttachmentURLs) > 0 {
+		if err := enqueueKafkaFileAttachmentsBind(ctx, item.tenantID, &item.payload.Message, item.payload.AttachmentURLs); err != nil {
+			return fmt.Errorf("file lifecycle enqueue failed: %w", err)
+		}
+	}
+	return store.HTTPInfrastructure.MarkKafkaConsumeSucceeded(item.tenantID, item.envelope.EventID,
+		kafkaBus.config.ConsumerGroup, time.Now().UTC().Round(time.Millisecond))
+}
+
+func kafkaMessagePersistFailureReason(err error) string {
+	if err != nil && strings.Contains(err.Error(), "file lifecycle") {
+		return "KAFKA_FILE_LIFECYCLE_FAILED"
+	}
+	return kafkaCommandErrorReason(err)
 }

 func kafkaMessagePersistTopicName(topic string) string {