fix(kafka): preserve partially fetched message batches

2026-10-04 23:37:10 +08:00
parent b05e8f60dd
commit 965f94f154
5 changed files with 527 additions and 9 deletions
@@ -161,11 +161,26 @@ type kafkaMessagePersistItem struct {
 	payload  kafkaMessagePersistPayload
 }

+type kafkaMessagePersistPreparer func(context.Context, kafka.Message, string) (*kafkaMessagePersistItem, bool, error)
+type kafkaMessagePersistGroupsWriter func(context.Context, map[string][]*kafkaMessagePersistItem) error
+
 func handleKafkaMessagePersistBatch(ctx context.Context, messages []kafka.Message, sourceTopicName string) error {
+	return processKafkaMessagePersistBatch(ctx, messages, sourceTopicName,
+		prepareKafkaMessagePersistItem, persistKafkaMessagePersistGroups)
+}
+
+func processKafkaMessagePersistBatch(ctx context.Context, messages []kafka.Message, sourceTopicName string,
+	prepare kafkaMessagePersistPreparer, persist kafkaMessagePersistGroupsWriter) error {
 	groups := make(map[string][]*kafkaMessagePersistItem)
 	for _, message := range messages {
-		item, terminal, err := prepareKafkaMessagePersistItem(ctx, message, sourceTopicName)
+		item, terminal, err := prepare(ctx, message, sourceTopicName)
 		if err != nil {
+			// Preparing an item acquires its idempotency claim. If a later item
+			// fails, finish all earlier claimed items before retrying the batch;
+			// otherwise retries can keep colliding with those processing claims.
+			if flushErr := persist(ctx, groups); flushErr != nil {
+				return fmt.Errorf("persist claimed Kafka messages after prepare failure: %w (prepare error: %v)", flushErr, err)
+			}
 			return err
 		}
 		if terminal {
@@ -177,7 +192,14 @@ func handleKafkaMessagePersistBatch(ctx context.Context, messages []kafka.Messag
 		}
 		groups[key] = append(groups[key], item)
 	}
+	return persist(ctx, groups)
+}
+
+func persistKafkaMessagePersistGroups(ctx context.Context, groups map[string][]*kafkaMessagePersistItem) error {
 	for _, items := range groups {
+		if len(items) == 0 {
+			continue
+		}
 		kind := items[0].payload.MessageKind
 		if kind == "" {
 			kind = "p2p"
@@ -301,6 +301,7 @@ 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
+type kafkaMessageFetcher func(context.Context) (kafka.Message, error)

 func startKafkaConsumer(topicName string, handler kafkaMessageHandler) {
 	startKafkaConsumerWithGroup(topicName, "", handler)
@@ -389,19 +390,23 @@ func startKafkaBatchConsumer(topicName string, handler kafkaBatchHandler) {
 		defer kafkaBus.wg.Done()
 		defer reader.Close()
 		for {
-			messages, err := fetchKafkaMessageBatch(kafkaBus.ctx, reader, kafkaBus.config.MessageBatchSize,
+			messages, fetchErr := 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) {
+			if fetchErr != nil && len(messages) == 0 {
+				if errors.Is(fetchErr, context.Canceled) || errors.Is(kafkaBus.ctx.Err(), context.Canceled) {
 					return
 				}
-				logs.Err.Println("Kafka batch fetch failed:", err)
+				logs.Err.Println("Kafka batch fetch failed:", fetchErr)
 				continue
 			}
+			if fetchErr != nil {
+				logs.Warn.Printf("Kafka batch fetch interrupted; processing retained messages: topic=%s size=%d err=%v",
+					topic, len(messages), fetchErr)
+			}
 			if len(messages) > 0 {
 				readerMetrics.observe(messages[0].Time)
 			}
-			if err = consumeKafkaMessageBatch(kafkaBus.ctx, messages, handler, reader.CommitMessages, waitKafkaRetry); err != nil {
+			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
 				}
@@ -468,7 +473,11 @@ func (m *kafkaReaderMetrics) update() {
 }

 func fetchKafkaMessageBatch(ctx context.Context, reader *kafka.Reader, maxSize int, maxWait time.Duration) ([]kafka.Message, error) {
-	first, err := reader.FetchMessage(ctx)
+	return fetchKafkaMessageBatchWith(ctx, maxSize, maxWait, reader.FetchMessage)
+}
+
+func fetchKafkaMessageBatchWith(ctx context.Context, maxSize int, maxWait time.Duration, fetch kafkaMessageFetcher) ([]kafka.Message, error) {
+	first, err := fetch(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -480,13 +489,16 @@ func fetchKafkaMessageBatch(ctx context.Context, reader *kafka.Reader, maxSize i
 			break
 		}
 		fetchCtx, cancel := context.WithTimeout(ctx, remaining)
-		message, fetchErr := reader.FetchMessage(fetchCtx)
+		message, fetchErr := fetch(fetchCtx)
 		cancel()
 		if errors.Is(fetchErr, context.DeadlineExceeded) && ctx.Err() == nil {
 			break
 		}
 		if fetchErr != nil {
-			return nil, fetchErr
+			// FetchMessage advances the reader's internal position before offsets are
+			// committed. Preserve a partial batch so a later higher-offset commit cannot
+			// skip messages already fetched into this batch.
+			return messages, fetchErr
 		}
 		messages = append(messages, message)
 	}
@@ -75,6 +75,103 @@ func TestConsumeKafkaMessageBatchDoesNotCommitWhileHandlerFails(t *testing.T) {
 	}
 }

+func TestFetchKafkaMessageBatchWithRetainsPartialBatchOnFetchError(t *testing.T) {
+	fetchCalls := 0
+	fetchErr := errors.New("broker connection reset")
+	messages, err := fetchKafkaMessageBatchWith(context.Background(), 10, time.Second,
+		func(context.Context) (kafka.Message, error) {
+			fetchCalls++
+			if fetchCalls == 1 {
+				return kafka.Message{Partition: 2, Offset: 41}, nil
+			}
+			return kafka.Message{}, fetchErr
+		})
+
+	if !errors.Is(err, fetchErr) {
+		t.Fatalf("fetchKafkaMessageBatchWith() error = %v, want %v", err, fetchErr)
+	}
+	if len(messages) != 1 || messages[0].Partition != 2 || messages[0].Offset != 41 {
+		t.Fatalf("fetchKafkaMessageBatchWith() messages = %#v, want retained offset 41", messages)
+	}
+}
+
+func TestFetchKafkaMessageBatchWithReturnsInitialFetchErrorWithoutMessages(t *testing.T) {
+	fetchErr := errors.New("broker unavailable")
+	messages, err := fetchKafkaMessageBatchWith(context.Background(), 10, time.Second,
+		func(context.Context) (kafka.Message, error) {
+			return kafka.Message{}, fetchErr
+		})
+
+	if !errors.Is(err, fetchErr) {
+		t.Fatalf("fetchKafkaMessageBatchWith() error = %v, want %v", err, fetchErr)
+	}
+	if len(messages) != 0 {
+		t.Fatalf("fetchKafkaMessageBatchWith() messages = %#v, want none", messages)
+	}
+}
+
+func TestProcessKafkaMessagePersistBatchFlushesClaimsBeforePrepareRetry(t *testing.T) {
+	prepareErr := errors.New("redis temporarily unavailable")
+	prepareCalls := 0
+	persistCalls := 0
+
+	err := processKafkaMessagePersistBatch(context.Background(), []kafka.Message{{Offset: 10}, {Offset: 11}}, "group_message_persist",
+		func(_ context.Context, message kafka.Message, topic string) (*kafkaMessagePersistItem, bool, error) {
+			prepareCalls++
+			if topic != "group_message_persist" {
+				t.Fatalf("prepare topic = %q", topic)
+			}
+			if message.Offset == 11 {
+				return nil, false, prepareErr
+			}
+			return &kafkaMessagePersistItem{
+				tenantID: types.TenantID(7),
+				payload: kafkaMessagePersistPayload{
+					MessageKind: "group",
+					ShardMonth:  "202610",
+					Message:     types.Message{ShardMonth: "202610"},
+				},
+			}, false, nil
+		},
+		func(_ context.Context, groups map[string][]*kafkaMessagePersistItem) error {
+			persistCalls++
+			items := groups["7:202610"]
+			if len(items) != 1 {
+				t.Fatalf("persisted groups = %#v, want one claimed item", groups)
+			}
+			return nil
+		})
+
+	if !errors.Is(err, prepareErr) {
+		t.Fatalf("processKafkaMessagePersistBatch() error = %v, want %v", err, prepareErr)
+	}
+	if prepareCalls != 2 || persistCalls != 1 {
+		t.Fatalf("calls prepare=%d persist=%d, want 2/1", prepareCalls, persistCalls)
+	}
+}
+
+func TestProcessKafkaMessagePersistBatchReportsFlushAndPrepareErrors(t *testing.T) {
+	prepareErr := errors.New("claim failed")
+	persistErr := errors.New("database unavailable")
+	prepareCalls := 0
+
+	err := processKafkaMessagePersistBatch(context.Background(), []kafka.Message{{Offset: 10}, {Offset: 11}}, "message_persist",
+		func(_ context.Context, _ kafka.Message, _ string) (*kafkaMessagePersistItem, bool, error) {
+			prepareCalls++
+			if prepareCalls == 2 {
+				return nil, false, prepareErr
+			}
+			return &kafkaMessagePersistItem{tenantID: types.TenantID(7)}, false, nil
+		},
+		func(context.Context, map[string][]*kafkaMessagePersistItem) error {
+			return persistErr
+		})
+
+	if !errors.Is(err, persistErr) || !strings.Contains(err.Error(), prepareErr.Error()) {
+		t.Fatalf("processKafkaMessagePersistBatch() error = %v, want both flush and prepare errors", err)
+	}
+}
+
 func TestConsumeKafkaMalformedMessageRetriesCommit(t *testing.T) {
 	commitCalls := 0
 	waitCalls := 0