kra-new/pkg/mq/kafka.go

package mq

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	kafkago "github.com/segmentio/kafka-go"
	"github.com/segmentio/kafka-go/sasl"
	"github.com/segmentio/kafka-go/sasl/plain"
)

type KafkaConfig struct {
	Enabled                bool
	Brokers                []string
	ClientID               string
	GroupID                string
	Username               string
	Password               string
	TLS                    bool
	TLSSkipVerify          bool
	StartOffset            string
	MinBytes               int
	MaxBytes               int
	MaxWait                time.Duration
	ConnectTimeout         time.Duration
	ReconnectInterval      time.Duration
	AllowAutoTopicCreation bool
}

type kafkaSubscription struct {
	reader  *kafkago.Reader
	cancel  context.CancelFunc
	qos     byte
	handler Handler
}

type kafkaTopicMetadataError struct {
	err error
}

func (e *kafkaTopicMetadataError) Error() string { return e.err.Error() }
func (e *kafkaTopicMetadataError) Unwrap() error { return e.err }

// Kafka adapts Kafka topics and consumer groups to the shared messaging
// contract. QoS 0 commits before dispatch; QoS 1 commits after dispatch.
type Kafka struct {
	opMu          sync.Mutex
	mu            sync.RWMutex
	config        KafkaConfig
	dialer        *kafkago.Dialer
	metadata      *kafkago.Client
	transport     *kafkago.Transport
	bestEffort    *kafkago.Writer
	acknowledged  *kafkago.Writer
	subscriptions map[string]*kafkaSubscription
	connected     atomic.Bool
	reconnectAt   atomic.Int64
	closed        bool
}

func NewKafka(config KafkaConfig) (*Kafka, error) {
	client := &Kafka{subscriptions: make(map[string]*kafkaSubscription)}
	if !config.Enabled {
		return client, nil
	}
	config = defaultKafkaConfig(config)
	if err := validateKafkaConfig(config); err != nil {
		return nil, err
	}

	tlsConfig, mechanism := kafkaSecurity(config)
	dialer := &kafkago.Dialer{
		ClientID:      config.ClientID,
		Timeout:       config.ConnectTimeout,
		TLS:           tlsConfig,
		SASLMechanism: mechanism,
	}
	transport := &kafkago.Transport{
		DialTimeout: config.ConnectTimeout,
		ClientID:    config.ClientID,
		TLS:         tlsConfig,
		SASL:        mechanism,
	}
	metadata := &kafkago.Client{
		Addr:      kafkago.TCP(config.Brokers...),
		Timeout:   config.ConnectTimeout,
		Transport: transport,
	}
	if err := probeKafkaMetadata(context.Background(), metadata, ""); err != nil {
		transport.CloseIdleConnections()
		return nil, fmt.Errorf("connect kafka: %w", err)
	}
	newWriter := func(acks kafkago.RequiredAcks) *kafkago.Writer {
		return &kafkago.Writer{
			Addr:                   kafkago.TCP(config.Brokers...),
			Balancer:               &kafkago.LeastBytes{},
			RequiredAcks:           acks,
			ReadTimeout:            config.ConnectTimeout,
			WriteTimeout:           config.ConnectTimeout,
			Transport:              transport,
			AllowAutoTopicCreation: config.AllowAutoTopicCreation,
		}
	}
	client.config = config
	client.dialer = dialer
	client.metadata = metadata
	client.transport = transport
	client.bestEffort = newWriter(kafkago.RequireNone)
	client.acknowledged = newWriter(kafkago.RequireAll)
	client.markConnected()
	return client, nil
}

func probeKafkaMetadata(ctx context.Context, client *kafkago.Client, topic string) error {
	if client == nil {
		return ErrUnavailable
	}
	request := &kafkago.MetadataRequest{Topics: []string{}}
	if topic != "" {
		request.Topics = []string{topic}
	}
	response, err := client.Metadata(nonNilContext(ctx), request)
	if err != nil {
		return err
	}
	if response == nil || len(response.Brokers) == 0 {
		return errors.New("kafka metadata contains no brokers")
	}
	if topic == "" {
		return nil
	}
	for _, item := range response.Topics {
		if item.Name != topic {
			continue
		}
		if item.Error != nil {
			return &kafkaTopicMetadataError{err: item.Error}
		}
		if len(item.Partitions) == 0 {
			return &kafkaTopicMetadataError{err: errors.New("topic has no partitions")}
		}
		return nil
	}
	return &kafkaTopicMetadataError{err: errors.New("topic metadata was not returned")}
}

func defaultKafkaConfig(config KafkaConfig) KafkaConfig {
	brokers := make([]string, 0, len(config.Brokers))
	for _, broker := range config.Brokers {
		if broker = strings.TrimSpace(broker); broker != "" {
			brokers = append(brokers, broker)
		}
	}
	config.Brokers = brokers
	config.ClientID = strings.TrimSpace(config.ClientID)
	config.GroupID = strings.TrimSpace(config.GroupID)
	config.StartOffset = strings.ToLower(strings.TrimSpace(config.StartOffset))
	if config.StartOffset == "" {
		config.StartOffset = "earliest"
	}
	if config.MinBytes <= 0 {
		config.MinBytes = 1
	}
	if config.MaxBytes <= 0 {
		config.MaxBytes = 10 << 20
	}
	if config.MaxWait <= 0 {
		config.MaxWait = time.Second
	}
	if config.ConnectTimeout <= 0 {
		config.ConnectTimeout = 10 * time.Second
	}
	if config.ReconnectInterval <= 0 {
		config.ReconnectInterval = 5 * time.Second
	}
	return config
}

func validateKafkaConfig(config KafkaConfig) error {
	if len(config.Brokers) == 0 {
		return errors.New("kafka brokers are empty")
	}
	for _, broker := range config.Brokers {
		host, portText, err := net.SplitHostPort(broker)
		port, parseErr := strconv.Atoi(portText)
		if err != nil || parseErr != nil || strings.TrimSpace(host) == "" || port < 1 || port > 65535 {
			return fmt.Errorf("invalid kafka broker %q", broker)
		}
	}
	if config.ClientID == "" {
		return errors.New("kafka client id is empty")
	}
	if config.GroupID == "" {
		return errors.New("kafka group id is empty")
	}
	if (config.Username == "") != (config.Password == "") {
		return errors.New("kafka username and password must be configured together")
	}
	if config.TLSSkipVerify && !config.TLS {
		return errors.New("kafka tls skip verify requires tls")
	}
	if config.StartOffset != "earliest" && config.StartOffset != "latest" {
		return fmt.Errorf("invalid kafka start offset %q", config.StartOffset)
	}
	if config.MinBytes < 1 || config.MaxBytes < config.MinBytes {
		return errors.New("kafka byte limits are invalid")
	}
	return nil
}

func kafkaSecurity(config KafkaConfig) (*tls.Config, sasl.Mechanism) {
	var tlsConfig *tls.Config
	if config.TLS {
		tlsConfig = &tls.Config{MinVersion: tls.VersionTLS12, InsecureSkipVerify: config.TLSSkipVerify} //nolint:gosec // explicit operator setting
	}
	var mechanism sasl.Mechanism
	if config.Username != "" {
		mechanism = plain.Mechanism{Username: config.Username, Password: config.Password}
	}
	return tlsConfig, mechanism
}

func (c *Kafka) Publish(ctx context.Context, topic string, payload []byte, qos byte, retain bool) error {
	topic = strings.TrimSpace(topic)
	if topic == "" {
		return errors.New("kafka topic is empty")
	}
	if retain {
		return errors.New("kafka does not support retained messages")
	}
	if qos > AtLeastOnce {
		return fmt.Errorf("kafka supports qos 0 or 1, got %d", qos)
	}
	if c == nil {
		return ErrUnavailable
	}
	ctx = nonNilContext(ctx)
	c.mu.RLock()
	if c.closed {
		c.mu.RUnlock()
		return ErrUnavailable
	}
	writer := c.bestEffort
	if qos == AtLeastOnce {
		writer = c.acknowledged
	}
	config := c.config
	c.mu.RUnlock()
	if writer == nil {
		return ErrUnavailable
	}
	err := writeKafkaMessage(ctx, writer, kafkago.Message{Topic: topic, Value: append([]byte(nil), payload...)}, config)
	if err == nil {
		c.markConnected()
	}
	return err
}

func writeKafkaMessage(ctx context.Context, writer *kafkago.Writer, message kafkago.Message, config KafkaConfig) error {
	if writer == nil {
		return ErrUnavailable
	}
	ctx = nonNilContext(ctx)
	if !config.AllowAutoTopicCreation {
		return writer.WriteMessages(ctx, message)
	}
	attemptCtx, cancel := context.WithTimeout(ctx, config.ConnectTimeout)
	defer cancel()
	delay := 100 * time.Millisecond
	var lastErr error
	for {
		lastErr = writer.WriteMessages(attemptCtx, message)
		if lastErr == nil || !kafkaTopicNotReady(lastErr) {
			return lastErr
		}
		timer := time.NewTimer(delay)
		select {
		case <-attemptCtx.Done():
			timer.Stop()
			if err := ctx.Err(); err != nil {
				return err
			}
			return lastErr
		case <-timer.C:
		}
		if delay < time.Second {
			delay *= 2
			if delay > time.Second {
				delay = time.Second
			}
		}
	}
}

func kafkaTopicNotReady(err error) bool {
	if err == nil {
		return false
	}
	if writeErrors, ok := err.(kafkago.WriteErrors); ok {
		found := false
		for _, item := range writeErrors {
			if item == nil {
				continue
			}
			found = true
			if !kafkaTopicNotReady(item) {
				return false
			}
		}
		return found
	}
	var kafkaErr kafkago.Error
	if !errors.As(err, &kafkaErr) {
		return false
	}
	return kafkaErr == kafkago.UnknownTopicOrPartition || kafkaErr == kafkago.LeaderNotAvailable || kafkaErr == kafkago.NotLeaderForPartition
}

func (c *Kafka) Subscribe(ctx context.Context, topic string, qos byte, handler Handler) error {
	topic = strings.TrimSpace(topic)
	if topic == "" {
		return errors.New("kafka topic is empty")
	}
	if qos > AtLeastOnce {
		return fmt.Errorf("kafka supports qos 0 or 1, got %d", qos)
	}
	if handler == nil {
		return errors.New("kafka handler is nil")
	}
	if c == nil {
		return ErrUnavailable
	}
	ctx = nonNilContext(ctx)
	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
	}

	c.opMu.Lock()
	defer c.opMu.Unlock()
	c.mu.RLock()
	if c.closed || c.dialer == nil || c.metadata == nil {
		c.mu.RUnlock()
		return ErrUnavailable
	}
	config := c.config
	dialer := c.dialer
	metadata := c.metadata
	c.mu.RUnlock()
	if err := probeKafkaMetadata(ctx, metadata, topic); err != nil {
		var topicErr *kafkaTopicMetadataError
		if errors.As(err, &topicErr) {
			c.markConnected()
		} else if ctx.Err() == nil {
			c.markUnavailable()
		}
		return fmt.Errorf("lookup kafka topic %q: %w", topic, err)
	}
	if old := c.removeSubscription(topic); old != nil {
		closeKafkaSubscription(old)
	}
	reader := kafkago.NewReader(kafkago.ReaderConfig{
		Brokers:          append([]string(nil), config.Brokers...),
		GroupID:          config.GroupID,
		Topic:            topic,
		Dialer:           dialer,
		MinBytes:         config.MinBytes,
		MaxBytes:         config.MaxBytes,
		MaxWait:          config.MaxWait,
		JoinGroupBackoff: config.ReconnectInterval,
		ReadLagInterval:  -1,
		StartOffset:      kafkaStartOffset(config.StartOffset),
	})
	consumeCtx, cancel := context.WithCancel(context.Background())
	subscription := &kafkaSubscription{reader: reader, cancel: cancel, qos: qos, handler: handler}
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		cancel()
		_ = reader.Close()
		return ErrUnavailable
	}
	c.subscriptions[topic] = subscription
	c.mu.Unlock()
	c.markConnected()
	go c.consume(consumeCtx, subscription)
	return nil
}

func kafkaStartOffset(value string) int64 {
	if value == "latest" {
		return kafkago.LastOffset
	}
	return kafkago.FirstOffset
}

func (c *Kafka) consume(ctx context.Context, subscription *kafkaSubscription) {
	for {
		var message kafkago.Message
		var err error
		if subscription.qos == AtMostOnce {
			message, err = subscription.reader.ReadMessage(ctx)
		} else {
			message, err = subscription.reader.FetchMessage(ctx)
		}
		if err != nil {
			if ctx.Err() != nil || errors.Is(err, io.EOF) {
				return
			}
			c.markUnavailable()
			if !sleepContext(ctx, c.config.ReconnectInterval) {
				return
			}
			continue
		}
		c.markConnected()
		subscription.handler(ctx, Message{Topic: message.Topic, Payload: append([]byte(nil), message.Value...), QoS: subscription.qos})
		if ctx.Err() != nil {
			return
		}
		if subscription.qos == AtLeastOnce {
			if err = subscription.reader.CommitMessages(ctx, message); err != nil {
				c.markUnavailable()
			}
		}
	}
}

func sleepContext(ctx context.Context, duration time.Duration) bool {
	timer := time.NewTimer(duration)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return false
	case <-timer.C:
		return true
	}
}

func (c *Kafka) Unsubscribe(ctx context.Context, topics ...string) error {
	if len(topics) == 0 {
		return errors.New("kafka topics are empty")
	}
	if c == nil {
		return ErrUnavailable
	}
	ctx = nonNilContext(ctx)
	normalized := make([]string, len(topics))
	for index, topic := range topics {
		normalized[index] = strings.TrimSpace(topic)
		if normalized[index] == "" {
			return errors.New("kafka topic is empty")
		}
	}
	c.opMu.Lock()
	defer c.opMu.Unlock()
	for _, topic := range normalized {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}
		if subscription := c.removeSubscription(topic); subscription != nil {
			closeKafkaSubscription(subscription)
		}
	}
	return nil
}

func (c *Kafka) removeSubscription(topic string) *kafkaSubscription {
	c.mu.Lock()
	defer c.mu.Unlock()
	subscription := c.subscriptions[topic]
	delete(c.subscriptions, topic)
	return subscription
}

func closeKafkaSubscription(subscription *kafkaSubscription) {
	if subscription == nil {
		return
	}
	if subscription.cancel != nil {
		subscription.cancel()
	}
	if subscription.reader != nil {
		_ = subscription.reader.Close()
	}
}

func (c *Kafka) Connected() bool {
	if c == nil {
		return false
	}
	c.mu.RLock()
	closed := c.closed
	c.mu.RUnlock()
	return !closed && c.connected.Load()
}

// Reconnecting lets the outer reloadable honor the configured retry delay
// before replacing a client after a runtime broker failure.
func (c *Kafka) Reconnecting() bool {
	if c == nil {
		return false
	}
	c.mu.RLock()
	closed := c.closed
	c.mu.RUnlock()
	return !closed && time.Now().UnixNano() < c.reconnectAt.Load()
}

func (c *Kafka) markConnected() {
	if c == nil {
		return
	}
	c.mu.RLock()
	closed := c.closed
	c.mu.RUnlock()
	if closed {
		return
	}
	c.reconnectAt.Store(0)
	c.connected.Store(true)
}

func (c *Kafka) markUnavailable() {
	if c == nil {
		return
	}
	c.mu.RLock()
	closed := c.closed
	c.mu.RUnlock()
	if closed {
		return
	}
	c.connected.Store(false)
	c.reconnectAt.Store(time.Now().Add(c.config.ReconnectInterval).UnixNano())
}

func (c *Kafka) Close() error {
	if c == nil {
		return nil
	}
	c.opMu.Lock()
	defer c.opMu.Unlock()
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return nil
	}
	c.closed = true
	c.connected.Store(false)
	c.reconnectAt.Store(0)
	subscriptions := make([]*kafkaSubscription, 0, len(c.subscriptions))
	for topic, subscription := range c.subscriptions {
		subscriptions = append(subscriptions, subscription)
		delete(c.subscriptions, topic)
	}
	bestEffort := c.bestEffort
	acknowledged := c.acknowledged
	transport := c.transport
	c.bestEffort = nil
	c.acknowledged = nil
	c.transport = nil
	c.dialer = nil
	c.metadata = nil
	c.mu.Unlock()
	for _, subscription := range subscriptions {
		closeKafkaSubscription(subscription)
	}
	var result error
	if bestEffort != nil {
		result = errors.Join(result, bestEffort.Close())
	}
	if acknowledged != nil {
		result = errors.Join(result, acknowledged.Close())
	}
	if transport != nil {
		transport.CloseIdleConnections()
	}
	return result
}