lingniu-vehicle-ingest/go/vehicle-gateway/internal/eventbus/kafka_sink.go

package eventbus

import (
	"context"
	"errors"
	"fmt"

	"github.com/segmentio/kafka-go"

	"lingniu-vehicle-ingest/go/vehicle-gateway/internal/envelope"
	"lingniu-vehicle-ingest/go/vehicle-gateway/internal/topics"
)

type Sink interface {
	PublishRaw(context.Context, envelope.FrameEnvelope) error
	PublishUnified(context.Context, envelope.FrameEnvelope) error
	PublishFields(context.Context, envelope.FrameEnvelope) error
	Close() error
}

type KafkaSink struct {
	writer       kafkaWriter
	rawTopics    map[envelope.Protocol]string
	fieldsTopics map[envelope.Protocol]string
	unifiedTopic string
}

type kafkaWriter interface {
	WriteMessages(context.Context, ...kafka.Message) error
	Close() error
}

type KafkaConfig struct {
	Brokers      []string
	RawTopics    map[envelope.Protocol]string
	FieldsTopics map[envelope.Protocol]string
	UnifiedTopic string
}

func NewKafkaSink(cfg KafkaConfig) (*KafkaSink, error) {
	if len(cfg.Brokers) == 0 {
		return nil, errors.New("kafka brokers are required")
	}
	if err := ValidateKafkaConfig(cfg); err != nil {
		return nil, err
	}
	rawTopics, fieldsTopics := kafkaTopicMaps(cfg)
	return newKafkaSinkWithWriter(&kafka.Writer{
		Addr:                   kafka.TCP(cfg.Brokers...),
		Balancer:               &kafka.Hash{},
		AllowAutoTopicCreation: false,
		RequiredAcks:           kafka.RequireAll,
		Async:                  false,
	}, KafkaConfig{
		RawTopics:    rawTopics,
		FieldsTopics: fieldsTopics,
		UnifiedTopic: cfg.UnifiedTopic,
	}), nil
}

func ValidateKafkaConfig(cfg KafkaConfig) error {
	rawTopics, fieldsTopics := kafkaTopicMaps(cfg)
	return topics.ValidateKafkaRawFields(protocolTopicLabels(rawTopics), protocolTopicLabels(fieldsTopics))
}

func newKafkaSinkWithWriter(writer kafkaWriter, cfg KafkaConfig) *KafkaSink {
	rawTopics, fieldsTopics := kafkaTopicMaps(cfg)
	unifiedTopic := cfg.UnifiedTopic
	if unifiedTopic == "" {
		unifiedTopic = topics.Unified
	}
	return &KafkaSink{writer: writer, rawTopics: rawTopics, fieldsTopics: fieldsTopics, unifiedTopic: unifiedTopic}
}

func kafkaTopicMaps(cfg KafkaConfig) (map[envelope.Protocol]string, map[envelope.Protocol]string) {
	rawTopics := map[envelope.Protocol]string{
		envelope.ProtocolGB32960:    topics.RawGB32960,
		envelope.ProtocolJT808:      topics.RawJT808,
		envelope.ProtocolYutongMQTT: topics.RawYutongMQTT,
	}
	for protocol, topic := range cfg.RawTopics {
		if topic != "" {
			rawTopics[protocol] = topic
		}
	}
	fieldsTopics := map[envelope.Protocol]string{
		envelope.ProtocolGB32960:    topics.FieldsGB32960,
		envelope.ProtocolJT808:      topics.FieldsJT808,
		envelope.ProtocolYutongMQTT: topics.FieldsYutongMQTT,
	}
	for protocol, topic := range cfg.FieldsTopics {
		if topic != "" {
			fieldsTopics[protocol] = topic
		}
	}
	return rawTopics, fieldsTopics
}

func protocolTopicLabels(values map[envelope.Protocol]string) map[string]string {
	out := make(map[string]string, len(values))
	for protocol, topic := range values {
		out[string(protocol)] = topic
	}
	return out
}

func (s *KafkaSink) PublishRaw(ctx context.Context, env envelope.FrameEnvelope) error {
	topic, ok := s.rawTopics[env.Protocol]
	if !ok || topic == "" {
		return fmt.Errorf("raw topic not configured for protocol %s", env.Protocol)
	}
	return s.publish(ctx, topic, env)
}

func (s *KafkaSink) PublishUnified(ctx context.Context, env envelope.FrameEnvelope) error {
	return s.publish(ctx, s.unifiedTopic, env)
}

func (s *KafkaSink) PublishFields(ctx context.Context, env envelope.FrameEnvelope) error {
	topic, ok := s.fieldsTopics[env.Protocol]
	if !ok || topic == "" {
		return fmt.Errorf("fields topic not configured for protocol %s", env.Protocol)
	}
	return s.publish(ctx, topic, 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
	}
	return s.writer.Close()
}

func (s *KafkaSink) publish(ctx context.Context, topic string, env envelope.FrameEnvelope) error {
	message, err := kafkaMessage(topic, env)
	if err != nil {
		return err
	}
	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
	case "fields":
		topic, ok := s.fieldsTopics[record.Envelope.Protocol]
		if !ok || topic == "" {
			return "", fmt.Errorf("fields topic not configured for protocol %s", record.Envelope.Protocol)
		}
		return topic, 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
}