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
@@ -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