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
@@ -0,0 +1,143 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/segmentio/kafka-go"
+	"github.com/tinode/chat/server/store/types"
+)
+
+func TestConsumeKafkaMessageBatchRetriesHandlerAndCommit(t *testing.T) {
+	handlerCalls := 0
+	commitCalls := 0
+	waitCalls := 0
+	messages := []kafka.Message{{Offset: 1}, {Offset: 2}}
+
+	err := consumeKafkaMessageBatch(context.Background(), messages,
+		func(context.Context, []kafka.Message) error {
+			handlerCalls++
+			if handlerCalls == 1 {
+				return errors.New("batch failed")
+			}
+			return nil
+		},
+		func(_ context.Context, committed ...kafka.Message) error {
+			commitCalls++
+			if len(committed) != len(messages) {
+				t.Fatalf("committed %d messages, want %d", len(committed), len(messages))
+			}
+			if commitCalls == 1 {
+				return errors.New("commit failed")
+			}
+			return nil
+		},
+		func(context.Context, time.Duration) error {
+			waitCalls++
+			return nil
+		})
+
+	if err != nil {
+		t.Fatalf("consumeKafkaMessageBatch() error = %v", err)
+	}
+	if handlerCalls != 2 || commitCalls != 2 || waitCalls != 2 {
+		t.Fatalf("calls handler=%d commit=%d wait=%d, want 2/2/2", handlerCalls, commitCalls, waitCalls)
+	}
+}
+
+func TestConsumeKafkaMalformedMessageRetriesCommit(t *testing.T) {
+	commitCalls := 0
+	waitCalls := 0
+	handlerCalls := 0
+	err := consumeKafkaMessage(context.Background(), kafka.Message{Value: []byte("{")},
+		func(context.Context, kafkaEnvelope, kafka.Message) error {
+			handlerCalls++
+			return nil
+		},
+		func(context.Context, kafka.Message) error {
+			commitCalls++
+			if commitCalls == 1 {
+				return errors.New("commit failed")
+			}
+			return nil
+		},
+		func(context.Context, time.Duration) error {
+			waitCalls++
+			return nil
+		})
+
+	if err != nil {
+		t.Fatalf("consumeKafkaMessage() error = %v", err)
+	}
+	if handlerCalls != 0 || commitCalls != 2 || waitCalls != 1 {
+		t.Fatalf("calls handler=%d commit=%d wait=%d, want 0/2/1", handlerCalls, commitCalls, waitCalls)
+	}
+}
+
+func TestRunKafkaMessageSingleRetriesThenDeadLetters(t *testing.T) {
+	operationCalls := 0
+	markFailedCalls := 0
+	markDeadCalls := 0
+	waitCalls := 0
+
+	err := runKafkaMessageSingleRetries(context.Background(), 3,
+		func() error {
+			operationCalls++
+			return errors.New("write failed")
+		},
+		func(error) (*types.KafkaConsumeRecord, error) {
+			markFailedCalls++
+			return &types.KafkaConsumeRecord{RetryCount: markFailedCalls}, nil
+		},
+		func(record *types.KafkaConsumeRecord, err error) error {
+			markDeadCalls++
+			if record == nil || record.RetryCount != 4 || err == nil {
+				t.Fatalf("unexpected dead-letter input: record=%+v err=%v", record, err)
+			}
+			return nil
+		},
+		func(context.Context, time.Duration) error {
+			waitCalls++
+			return nil
+		})
+
+	if err != nil {
+		t.Fatalf("runKafkaMessageSingleRetries() error = %v", err)
+	}
+	if operationCalls != 4 || markFailedCalls != 4 || waitCalls != 3 || markDeadCalls != 1 {
+		t.Fatalf("calls operation=%d failed=%d wait=%d dead=%d, want 4/4/3/1",
+			operationCalls, markFailedCalls, waitCalls, markDeadCalls)
+	}
+}
+
+func TestRunKafkaMessageSingleRetriesStopsOnSuccess(t *testing.T) {
+	operationCalls := 0
+	markFailedCalls := 0
+	markDeadCalls := 0
+	err := runKafkaMessageSingleRetries(context.Background(), 3,
+		func() error {
+			operationCalls++
+			if operationCalls < 3 {
+				return errors.New("temporary failure")
+			}
+			return nil
+		},
+		func(error) (*types.KafkaConsumeRecord, error) {
+			markFailedCalls++
+			return &types.KafkaConsumeRecord{RetryCount: markFailedCalls}, nil
+		},
+		func(*types.KafkaConsumeRecord, error) error {
+			markDeadCalls++
+			return nil
+		},
+		func(context.Context, time.Duration) error { return nil })
+
+	if err != nil {
+		t.Fatalf("runKafkaMessageSingleRetries() error = %v", err)
+	}
+	if operationCalls != 3 || markFailedCalls != 2 || markDeadCalls != 0 {
+		t.Fatalf("calls operation=%d failed=%d dead=%d, want 3/2/0", operationCalls, markFailedCalls, markDeadCalls)
+	}
+}