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

package eventbus

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/nats-io/nats.go"

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

type NATSConfig struct {
	URL             string
	Name            string
	RawSubjects     map[envelope.Protocol]string
	FieldsSubjects  map[envelope.Protocol]string
	UnifiedSubject  string
	AsyncMaxPending int
	AsyncAckTimeout time.Duration
}

type NATSSink struct {
	conn           *nats.Conn
	publisher      natsPublisher
	asyncPublisher natsAsyncPublisher
	rawSubjects    map[envelope.Protocol]string
	fieldsSubjects map[envelope.Protocol]string
	unifiedSubject string
}

type NATSPublishOption = nats.PubOpt

type natsPublisher interface {
	Publish(context.Context, string, []byte, ...NATSPublishOption) error
}

type natsAsyncPublisher interface {
	PublishAsync(string, []byte, ...NATSPublishOption) (nats.PubAckFuture, error)
}

func NewNATSSink(cfg NATSConfig) (*NATSSink, error) {
	if cfg.URL == "" {
		return nil, errors.New("nats url is required")
	}
	if err := ValidateNATSConfig(cfg); err != nil {
		return nil, err
	}
	rawSubjects, fieldsSubjects := natsSubjectMaps(cfg)
	name := cfg.Name
	if name == "" {
		name = "lingniu-vehicle-gateway"
	}
	conn, err := nats.Connect(cfg.URL, nats.Name(name), nats.Timeout(5*time.Second))
	if err != nil {
		return nil, err
	}
	var jsOptions []nats.JSOpt
	if cfg.AsyncMaxPending > 0 {
		jsOptions = append(jsOptions, nats.PublishAsyncMaxPending(cfg.AsyncMaxPending))
	}
	if cfg.AsyncAckTimeout > 0 {
		jsOptions = append(jsOptions, nats.PublishAsyncTimeout(cfg.AsyncAckTimeout))
	}
	js, err := conn.JetStream(jsOptions...)
	if err != nil {
		conn.Close()
		return nil, err
	}
	publisher := natsJetStreamPublisher{js: js}
	sink := newNATSSinkWithPublishers(publisher, publisher, NATSConfig{
		RawSubjects:    rawSubjects,
		FieldsSubjects: fieldsSubjects,
		UnifiedSubject: cfg.UnifiedSubject,
	})
	sink.conn = conn
	return sink, nil
}

func ValidateNATSConfig(cfg NATSConfig) error {
	rawSubjects, fieldsSubjects := natsSubjectMaps(cfg)
	raw := protocolTopicLabels(rawSubjects)
	fields := protocolTopicLabels(fieldsSubjects)
	if err := topics.ValidateKnownRawFieldsProtocols(raw, fields, "nats subject"); err != nil {
		return err
	}
	return topics.ValidateRawFieldsDisjoint(raw, fields, "nats subject")
}

func newNATSSinkWithPublisher(publisher natsPublisher, cfg NATSConfig) *NATSSink {
	return newNATSSinkWithPublishers(publisher, nil, cfg)
}

func newNATSSinkWithPublishers(publisher natsPublisher, asyncPublisher natsAsyncPublisher, cfg NATSConfig) *NATSSink {
	rawSubjects, fieldsSubjects := natsSubjectMaps(cfg)
	unifiedSubject := cfg.UnifiedSubject
	if unifiedSubject == "" {
		unifiedSubject = topics.Unified
	}
	return &NATSSink{
		publisher:      publisher,
		asyncPublisher: asyncPublisher,
		rawSubjects:    rawSubjects,
		fieldsSubjects: fieldsSubjects,
		unifiedSubject: unifiedSubject,
	}
}

func natsSubjectMaps(cfg NATSConfig) (map[envelope.Protocol]string, map[envelope.Protocol]string) {
	rawSubjects := map[envelope.Protocol]string{
		envelope.ProtocolGB32960:    topics.RawGB32960,
		envelope.ProtocolJT808:      topics.RawJT808,
		envelope.ProtocolYutongMQTT: topics.RawYutongMQTT,
	}
	for protocol, subject := range cfg.RawSubjects {
		if subject != "" {
			rawSubjects[protocol] = subject
		}
	}
	fieldsSubjects := map[envelope.Protocol]string{
		envelope.ProtocolGB32960:    topics.FieldsGB32960,
		envelope.ProtocolJT808:      topics.FieldsJT808,
		envelope.ProtocolYutongMQTT: topics.FieldsYutongMQTT,
	}
	for protocol, subject := range cfg.FieldsSubjects {
		if subject != "" {
			fieldsSubjects[protocol] = subject
		}
	}
	return rawSubjects, fieldsSubjects
}

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

func (s *NATSSink) PublishUnified(ctx context.Context, env envelope.FrameEnvelope) error {
	if s.unifiedSubject == "" {
		return errors.New("unified subject is empty")
	}
	return s.publish(ctx, s.unifiedSubject, "unified", env)
}

func (s *NATSSink) PublishFields(ctx context.Context, env envelope.FrameEnvelope) error {
	subject, ok := s.fieldsSubjects[env.Protocol]
	if !ok || subject == "" {
		return fmt.Errorf("fields subject not configured for protocol %s", env.Protocol)
	}
	return s.publish(ctx, subject, "fields", env)
}

func (s *NATSSink) PublishRecords(ctx context.Context, records []durableRecord) error {
	for _, record := range records {
		subject, err := s.subjectForRecord(record)
		if err != nil {
			return err
		}
		if err := s.publish(ctx, subject, record.Kind, record.Envelope); err != nil {
			return err
		}
	}
	return nil
}

func (s *NATSSink) PublishRecordAsync(record durableRecord, complete func(error)) error {
	if s == nil || s.asyncPublisher == nil {
		return errors.New("nats async publisher is not configured")
	}
	if complete == nil {
		return errors.New("nats async publish completion callback is required")
	}
	subject, err := s.subjectForRecord(record)
	if err != nil {
		return err
	}
	payload, err := record.Envelope.MarshalJSONBytes()
	if err != nil {
		return err
	}
	future, err := s.asyncPublisher.PublishAsync(
		subject,
		payload,
		nats.MsgId(natsMessageID(record.Kind, subject, record.Envelope)),
	)
	if err != nil {
		return err
	}
	go func() {
		select {
		case <-future.Ok():
			complete(nil)
		case asyncErr := <-future.Err():
			if asyncErr == nil {
				asyncErr = errors.New("nats async publish failed without error detail")
			}
			complete(asyncErr)
		}
	}()
	return nil
}

func (s *NATSSink) ValidateRecord(record durableRecord) error {
	_, err := s.subjectForRecord(record)
	return err
}

func (s *NATSSink) Close() error {
	if s == nil || s.conn == nil {
		return nil
	}
	_ = s.conn.Drain()
	s.conn.Close()
	return nil
}

func (s *NATSSink) publish(ctx context.Context, subject string, kind string, env envelope.FrameEnvelope) error {
	payload, err := env.MarshalJSONBytes()
	if err != nil {
		return err
	}
	return s.publisher.Publish(ctx, subject, payload, nats.MsgId(natsMessageID(kind, subject, env)))
}

func natsMessageID(kind string, subject string, env envelope.FrameEnvelope) string {
	kind = strings.TrimSpace(kind)
	if kind == "" {
		kind = "unknown"
	}
	subject = strings.TrimSpace(subject)
	if subject == "" {
		subject = "unknown"
	}
	return kind + ":" + subject + ":" + env.StableEventID()
}

func (s *NATSSink) subjectForRecord(record durableRecord) (string, error) {
	switch record.Kind {
	case "raw":
		subject, ok := s.rawSubjects[record.Envelope.Protocol]
		if !ok || subject == "" {
			return "", fmt.Errorf("raw subject not configured for protocol %s", record.Envelope.Protocol)
		}
		return subject, nil
	case "unified":
		if s.unifiedSubject == "" {
			return "", errors.New("unified subject is empty")
		}
		return s.unifiedSubject, nil
	case "fields":
		subject, ok := s.fieldsSubjects[record.Envelope.Protocol]
		if !ok || subject == "" {
			return "", fmt.Errorf("fields subject not configured for protocol %s", record.Envelope.Protocol)
		}
		return subject, nil
	default:
		return "", errUnknownRecordKind(record.Kind)
	}
}

type natsJetStreamPublisher struct {
	js nats.JetStreamContext
}

func (p natsJetStreamPublisher) Publish(ctx context.Context, subject string, data []byte, opts ...NATSPublishOption) error {
	_, err := p.js.Publish(subject, data, append(opts, nats.Context(ctx))...)
	return err
}

func (p natsJetStreamPublisher) PublishAsync(subject string, data []byte, opts ...NATSPublishOption) (nats.PubAckFuture, error) {
	return p.js.PublishAsync(subject, data, opts...)
}