fix(go): batch durable kafka replay

This commit is contained in:
lingniu
2026-07-02 14:01:50 +08:00
parent c939cc6b0c
commit 5a958616a3
5 changed files with 176 additions and 5 deletions

View File

@@ -40,6 +40,10 @@ type durableRecordFile struct {
record durableRecord
}
type recordPublishingSink interface {
PublishRecords(context.Context, []durableRecord) error
}
const durableReplayOperationTimeout = 30 * time.Second
func NewDurableSink(delegate Sink, cfg DurableConfig) *DurableSink {
@@ -93,6 +97,24 @@ func (s *DurableSink) replay(ctx context.Context, limit int) error {
records = append(records, durableRecordFile{path: file, record: record})
}
sortDurableRecords(records)
if publisher, ok := s.delegate.(recordPublishingSink); ok {
durableRecords := make([]durableRecord, 0, len(records))
for _, item := range records {
durableRecords = append(durableRecords, item.record)
}
if err := publisher.PublishRecords(ctx, durableRecords); err != nil {
return err
}
for _, item := range records {
if err := os.Remove(item.path); err != nil {
return err
}
if item.record.Kind == "raw" {
s.clearRawPending(item.record.Envelope)
}
}
return nil
}
for _, item := range records {
if err := s.publishRecord(ctx, item.record); err != nil {
return err
@@ -163,10 +185,14 @@ func (s *DurableSink) publishRecord(ctx context.Context, record durableRecord) e
case "unified":
return s.delegate.PublishUnified(ctx, record.Envelope)
default:
return fmt.Errorf("unknown durable record kind %q", record.Kind)
return errUnknownRecordKind(record.Kind)
}
}
func errUnknownRecordKind(kind string) error {
return fmt.Errorf("unknown durable record kind %q", kind)
}
func (s *DurableSink) spool(kind string, env envelope.FrameEnvelope) error {
if s.dir == "" {
return fmt.Errorf("durable spool directory is empty")

View File

@@ -139,6 +139,34 @@ func TestDurableSinkReplayLoopTickRespectsBatchSize(t *testing.T) {
}
}
func TestDurableSinkReplayUsesBatchPublisherWhenAvailable(t *testing.T) {
dir := t.TempDir()
env := durableTestEnvelope()
writeDurableRecord(t, filepath.Join(dir, "0001-raw.json"), durableRecord{Kind: "raw", Envelope: env})
writeDurableRecord(t, filepath.Join(dir, "0002-unified.json"), durableRecord{Kind: "unified", Envelope: env})
delegate := &batchRecordingSink{}
sink := NewDurableSink(delegate, DurableConfig{Directory: dir})
if err := sink.ReplayOnce(context.Background()); err != nil {
t.Fatalf("ReplayOnce() error = %v", err)
}
if delegate.batchCalls != 1 {
t.Fatalf("batch calls = %d, want 1", delegate.batchCalls)
}
if delegate.rawCalls != 0 || delegate.unifiedCalls != 0 {
t.Fatalf("individual publish should not be used, raw=%d unified=%d", delegate.rawCalls, delegate.unifiedCalls)
}
if got, want := len(delegate.records), 2; got != want {
t.Fatalf("batch record count = %d, want %d", got, want)
}
if delegate.records[0].Kind != "raw" || delegate.records[1].Kind != "unified" {
t.Fatalf("batch order = %#v", delegate.records)
}
if files := spoolFiles(t, dir); len(files) != 0 {
t.Fatalf("spool files after replay = %#v, want none", files)
}
}
func TestDurableFilesFromReaderStopsAfterLimitedJSONBatch(t *testing.T) {
reader := &fakeNameReader{
batches: [][]string{
@@ -298,3 +326,30 @@ func (s *contextCheckingReplaySink) PublishUnified(ctx context.Context, _ envelo
func (s *contextCheckingReplaySink) Close() error {
return nil
}
type batchRecordingSink struct {
batchCalls int
records []durableRecord
rawCalls int
unifiedCalls int
}
func (s *batchRecordingSink) PublishRecords(_ context.Context, records []durableRecord) error {
s.batchCalls++
s.records = append(s.records, records...)
return nil
}
func (s *batchRecordingSink) PublishRaw(context.Context, envelope.FrameEnvelope) error {
s.rawCalls++
return nil
}
func (s *batchRecordingSink) PublishUnified(context.Context, envelope.FrameEnvelope) error {
s.unifiedCalls++
return nil
}
func (s *batchRecordingSink) Close() error {
return nil
}

View File

@@ -76,6 +76,25 @@ func (s *KafkaSink) PublishUnified(ctx context.Context, env envelope.FrameEnvelo
return s.publish(ctx, s.unifiedTopic, env)
}
func (s *KafkaSink) PublishRecords(ctx context.Context, records []durableRecord) error {
messages := make([]kafka.Message, 0, len(records))
for _, record := range records {
topic, err := s.topicForRecord(record)
if err != nil {
return err
}
message, err := kafkaMessage(topic, record.Envelope)
if err != nil {
return err
}
messages = append(messages, message)
}
if len(messages) == 0 {
return nil
}
return s.writer.WriteMessages(ctx, messages...)
}
func (s *KafkaSink) Close() error {
if s == nil || s.writer == nil {
return nil
@@ -84,13 +103,36 @@ func (s *KafkaSink) Close() error {
}
func (s *KafkaSink) publish(ctx context.Context, topic string, env envelope.FrameEnvelope) error {
payload, err := env.MarshalJSONBytes()
message, err := kafkaMessage(topic, env)
if err != nil {
return err
}
return s.writer.WriteMessages(ctx, kafka.Message{
return s.writer.WriteMessages(ctx, message)
}
func (s *KafkaSink) topicForRecord(record durableRecord) (string, error) {
switch record.Kind {
case "raw":
topic, ok := s.rawTopics[record.Envelope.Protocol]
if !ok || topic == "" {
return "", fmt.Errorf("raw topic not configured for protocol %s", record.Envelope.Protocol)
}
return topic, nil
case "unified":
return s.unifiedTopic, nil
default:
return "", fmt.Errorf("unknown durable record kind %q", record.Kind)
}
}
func kafkaMessage(topic string, env envelope.FrameEnvelope) (kafka.Message, error) {
payload, err := env.MarshalJSONBytes()
if err != nil {
return kafka.Message{}, err
}
return kafka.Message{
Topic: topic,
Key: env.KafkaKey(),
Value: payload,
})
}, nil
}

View File

@@ -64,6 +64,29 @@ func TestKafkaSinkRejectsUnknownProtocol(t *testing.T) {
}
}
func TestKafkaSinkPublishesDurableRecordsInSingleWriterCall(t *testing.T) {
writer := &recordingWriter{}
sink := newKafkaSinkWithWriter(writer, KafkaConfig{UnifiedTopic: "custom.unified"})
env := envelope.FrameEnvelope{Protocol: envelope.ProtocolJT808, Phone: "13307795425", MessageID: "0x0200"}
err := sink.PublishRecords(context.Background(), []durableRecord{
{Kind: "raw", Envelope: env},
{Kind: "unified", Envelope: env},
})
if err != nil {
t.Fatalf("PublishRecords() error = %v", err)
}
if writer.writeCalls != 1 {
t.Fatalf("WriteMessages calls = %d, want 1", writer.writeCalls)
}
if len(writer.messages) != 2 {
t.Fatalf("messages = %d, want 2", len(writer.messages))
}
if writer.messages[0].Topic != "vehicle.raw.jt808.v1" || writer.messages[1].Topic != "custom.unified" {
t.Fatalf("topics = %q, %q", writer.messages[0].Topic, writer.messages[1].Topic)
}
}
func TestNewKafkaSinkUsesProductionDeliveryGuarantees(t *testing.T) {
sink, err := NewKafkaSink(KafkaConfig{Brokers: []string{"127.0.0.1:9092"}})
if err != nil {
@@ -90,10 +113,12 @@ func TestNewKafkaSinkUsesProductionDeliveryGuarantees(t *testing.T) {
}
type recordingWriter struct {
messages []kafka.Message
messages []kafka.Message
writeCalls int
}
func (w *recordingWriter) WriteMessages(_ context.Context, messages ...kafka.Message) error {
w.writeCalls++
w.messages = append(w.messages, messages...)
return nil
}

View File

@@ -40,6 +40,29 @@ func (s *RetryingSink) PublishUnified(ctx context.Context, env envelope.FrameEnv
})
}
func (s *RetryingSink) PublishRecords(ctx context.Context, records []durableRecord) error {
if publisher, ok := s.delegate.(recordPublishingSink); ok {
return s.withRetry(ctx, func(attemptCtx context.Context) error {
return publisher.PublishRecords(attemptCtx, records)
})
}
for _, record := range records {
switch record.Kind {
case "raw":
if err := s.PublishRaw(ctx, record.Envelope); err != nil {
return err
}
case "unified":
if err := s.PublishUnified(ctx, record.Envelope); err != nil {
return err
}
default:
return errUnknownRecordKind(record.Kind)
}
}
return nil
}
func (s *RetryingSink) Close() error {
return s.delegate.Close()
}