优化结构

2026-08-30 01:08:10 +08:00
parent 509df03d61
commit e0c35e4e9b
23 changed files with 1513 additions and 38 deletions
@@ -0,0 +1,592 @@
+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
+}