优化结构

2026-08-21 21:37:30 +08:00
parent a2ce3ae218
commit e4ba0dced9
41 changed files with 3646 additions and 364 deletions
@@ -0,0 +1,447 @@
+package mq
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+	"sync"
+	"time"
+
+	amqp "github.com/rabbitmq/amqp091-go"
+)
+
+type RabbitMQConfig struct {
+	Enabled        bool
+	Host           string
+	Port           int
+	Username       string
+	Password       string
+	VHost          string
+	Exchange       string
+	ExchangeType   string
+	Queue          string
+	RoutingKey     string
+	Durable        bool
+	AutoDelete     bool
+	PrefetchCount  int
+	Heartbeat      time.Duration
+	ConnectTimeout time.Duration
+	TLS            bool
+}
+
+type rabbitSubscription struct {
+	qos     byte
+	handler Handler
+}
+
+// RabbitMQ adapts AMQP exchanges and routing keys to the shared topic-based
+// Client contract. All subscriptions share the configured queue and a single
+// consumer; deliveries are dispatched to matching handlers locally.
+type RabbitMQ struct {
+	mu             sync.RWMutex
+	opMu           sync.Mutex
+	publishMu      sync.Mutex
+	consumeMu      sync.Mutex
+	connection     *amqp.Connection
+	publishChannel *amqp.Channel
+	consumeChannel *amqp.Channel
+	config         RabbitMQConfig
+	subscriptions  map[string]rabbitSubscription
+	consumerTag    string
+	consuming      bool
+	stop           chan struct{}
+	closed         bool
+}
+
+func NewRabbitMQ(config RabbitMQConfig) (*RabbitMQ, error) {
+	client := &RabbitMQ{subscriptions: make(map[string]rabbitSubscription), stop: make(chan struct{})}
+	if !config.Enabled {
+		return client, nil
+	}
+	config = defaultRabbitMQConfig(config)
+	if err := validateRabbitMQConfig(config); err != nil {
+		return nil, err
+	}
+
+	scheme := "amqp"
+	if config.TLS {
+		scheme = "amqps"
+	}
+	address := amqp.URI{
+		Scheme: scheme, Host: config.Host, Port: config.Port,
+		Username: config.Username, Password: config.Password, Vhost: config.VHost,
+		ConnectionTimeout: int(config.ConnectTimeout.Milliseconds()),
+	}.String()
+	connection, err := amqp.DialConfig(address, amqp.Config{
+		Heartbeat: config.Heartbeat,
+		Recovery:  &amqp.Recovery{},
+	})
+	if err != nil {
+		return nil, fmt.Errorf("connect rabbitmq: %w", err)
+	}
+
+	publishChannel, err := connection.Channel()
+	if err != nil {
+		_ = connection.Close()
+		return nil, fmt.Errorf("open rabbitmq publish channel: %w", err)
+	}
+	consumeChannel, err := connection.Channel()
+	if err != nil {
+		_ = publishChannel.Close()
+		_ = connection.Close()
+		return nil, fmt.Errorf("open rabbitmq consume channel: %w", err)
+	}
+	closeOnError := func() {
+		_ = consumeChannel.Close()
+		_ = publishChannel.Close()
+		_ = connection.Close()
+	}
+	if err = consumeChannel.ExchangeDeclare(config.Exchange, config.ExchangeType, config.Durable, config.AutoDelete, false, false, nil); err != nil {
+		closeOnError()
+		return nil, fmt.Errorf("declare rabbitmq exchange: %w", err)
+	}
+	if _, err = consumeChannel.QueueDeclare(config.Queue, config.Durable, config.AutoDelete, false, false, nil); err != nil {
+		closeOnError()
+		return nil, fmt.Errorf("declare rabbitmq queue: %w", err)
+	}
+	if config.PrefetchCount > 0 {
+		if err = consumeChannel.Qos(config.PrefetchCount, 0, false); err != nil {
+			closeOnError()
+			return nil, fmt.Errorf("configure rabbitmq qos: %w", err)
+		}
+	}
+	client.connection = connection
+	client.publishChannel = publishChannel
+	client.consumeChannel = consumeChannel
+	client.config = config
+	client.consumerTag = fmt.Sprintf("kra-%d", time.Now().UnixNano())
+	return client, nil
+}
+
+func defaultRabbitMQConfig(config RabbitMQConfig) RabbitMQConfig {
+	if config.Port <= 0 {
+		if config.TLS {
+			config.Port = 5671
+		} else {
+			config.Port = 5672
+		}
+	}
+	if config.Username == "" {
+		config.Username = "guest"
+	}
+	if config.Password == "" {
+		config.Password = "guest"
+	}
+	if config.VHost == "" {
+		config.VHost = "/"
+	}
+	if config.ExchangeType == "" {
+		config.ExchangeType = "topic"
+	}
+	if config.RoutingKey == "" {
+		config.RoutingKey = "#"
+	}
+	if config.Heartbeat <= 0 {
+		config.Heartbeat = 10 * time.Second
+	}
+	if config.ConnectTimeout <= 0 {
+		config.ConnectTimeout = 10 * time.Second
+	}
+	return config
+}
+
+func validateRabbitMQConfig(config RabbitMQConfig) error {
+	if strings.TrimSpace(config.Host) == "" {
+		return errors.New("rabbitmq host is empty")
+	}
+	if config.Port < 1 || config.Port > 65535 {
+		return fmt.Errorf("invalid rabbitmq port %d", config.Port)
+	}
+	if strings.TrimSpace(config.Exchange) == "" {
+		return errors.New("rabbitmq exchange is empty")
+	}
+	if strings.TrimSpace(config.Queue) == "" {
+		return errors.New("rabbitmq queue is empty")
+	}
+	switch strings.ToLower(config.ExchangeType) {
+	case "direct", "fanout", "topic":
+		return nil
+	default:
+		return fmt.Errorf("invalid rabbitmq exchange type %q", config.ExchangeType)
+	}
+}
+
+func (c *RabbitMQ) Publish(ctx context.Context, topic string, payload []byte, qos byte, _ bool) error {
+	if c == nil {
+		return ErrUnavailable
+	}
+	if c != nil && strings.TrimSpace(topic) == "" {
+		topic = c.config.RoutingKey
+	}
+	if strings.TrimSpace(topic) == "" {
+		return errors.New("rabbitmq routing key is empty")
+	}
+	if qos > AtLeastOnce {
+		return fmt.Errorf("rabbitmq supports qos 0 or 1, got %d", qos)
+	}
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
+	c.publishMu.Lock()
+	defer c.publishMu.Unlock()
+	c.mu.RLock()
+	if c.closed || c.publishChannel == nil || c.publishChannel.IsClosed() {
+		c.mu.RUnlock()
+		return ErrUnavailable
+	}
+	channel := c.publishChannel
+	exchange := c.config.Exchange
+	c.mu.RUnlock()
+	deliveryMode := amqp.Transient
+	if qos >= AtLeastOnce {
+		deliveryMode = amqp.Persistent
+	}
+	return channel.PublishWithContext(ctx, exchange, topic, false, false, amqp.Publishing{
+		ContentType:  "application/octet-stream",
+		DeliveryMode: deliveryMode,
+		Timestamp:    time.Now(),
+		Body:         append([]byte(nil), payload...),
+	})
+}
+
+func (c *RabbitMQ) Subscribe(ctx context.Context, topic string, qos byte, handler Handler) error {
+	if c == nil {
+		return ErrUnavailable
+	}
+	if c != nil && strings.TrimSpace(topic) == "" {
+		topic = c.config.RoutingKey
+	}
+	if strings.TrimSpace(topic) == "" {
+		return errors.New("rabbitmq routing key is empty")
+	}
+	if qos > AtLeastOnce {
+		return fmt.Errorf("rabbitmq supports qos 0 or 1, got %d", qos)
+	}
+	if handler == nil {
+		return errors.New("rabbitmq handler is nil")
+	}
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
+
+	c.opMu.Lock()
+	defer c.opMu.Unlock()
+	c.consumeMu.Lock()
+	defer c.consumeMu.Unlock()
+	c.mu.RLock()
+	if c.closed || c.consumeChannel == nil || c.consumeChannel.IsClosed() {
+		c.mu.RUnlock()
+		return ErrUnavailable
+	}
+	channel := c.consumeChannel
+	config := c.config
+	_, exists := c.subscriptions[topic]
+	c.mu.RUnlock()
+	if !exists {
+		if err := channel.QueueBind(config.Queue, topic, config.Exchange, false, nil); err != nil {
+			return fmt.Errorf("bind rabbitmq queue: %w", err)
+		}
+	}
+
+	c.mu.Lock()
+	c.subscriptions[topic] = rabbitSubscription{qos: qos, handler: handler}
+	shouldStart := !c.consuming
+	c.mu.Unlock()
+	if !shouldStart {
+		return nil
+	}
+	deliveries, err := channel.Consume(config.Queue, c.consumerTag, false, false, false, false, nil)
+	if err != nil {
+		c.mu.Lock()
+		delete(c.subscriptions, topic)
+		c.mu.Unlock()
+		if !exists {
+			_ = channel.QueueUnbind(config.Queue, topic, config.Exchange, nil)
+		}
+		return fmt.Errorf("consume rabbitmq queue: %w", err)
+	}
+	c.mu.Lock()
+	c.consuming = true
+	c.mu.Unlock()
+	go c.consume(deliveries)
+	return nil
+}
+
+func (c *RabbitMQ) consume(deliveries <-chan amqp.Delivery) {
+	for {
+		select {
+		case <-c.stop:
+			return
+		case delivery, ok := <-deliveries:
+			if !ok {
+				c.mu.Lock()
+				c.consuming = false
+				c.mu.Unlock()
+				return
+			}
+			subscriptions := c.subscriptionSnapshot()
+			for pattern, item := range subscriptions {
+				if rabbitRoutingKeyMatches(c.config.ExchangeType, pattern, delivery.RoutingKey) {
+					item.handler(context.Background(), Message{Topic: delivery.RoutingKey, Payload: append([]byte(nil), delivery.Body...), QoS: item.qos})
+				}
+			}
+			c.consumeMu.Lock()
+			_ = delivery.Ack(false)
+			c.consumeMu.Unlock()
+		}
+	}
+}
+
+func (c *RabbitMQ) subscriptionSnapshot() map[string]rabbitSubscription {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	result := make(map[string]rabbitSubscription, len(c.subscriptions))
+	for topic, item := range c.subscriptions {
+		result[topic] = item
+	}
+	return result
+}
+
+func (c *RabbitMQ) Unsubscribe(ctx context.Context, topics ...string) error {
+	if c == nil {
+		return ErrUnavailable
+	}
+	if len(topics) == 0 {
+		return errors.New("rabbitmq routing keys are empty")
+	}
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
+	c.opMu.Lock()
+	defer c.opMu.Unlock()
+	c.consumeMu.Lock()
+	defer c.consumeMu.Unlock()
+	c.mu.RLock()
+	if c.closed || c.consumeChannel == nil || c.consumeChannel.IsClosed() {
+		c.mu.RUnlock()
+		return ErrUnavailable
+	}
+	channel := c.consumeChannel
+	config := c.config
+	c.mu.RUnlock()
+	for _, topic := range topics {
+		c.mu.RLock()
+		_, exists := c.subscriptions[topic]
+		c.mu.RUnlock()
+		if !exists {
+			continue
+		}
+		if err := channel.QueueUnbind(config.Queue, topic, config.Exchange, nil); err != nil {
+			return fmt.Errorf("unbind rabbitmq queue: %w", err)
+		}
+		c.mu.Lock()
+		delete(c.subscriptions, topic)
+		c.mu.Unlock()
+	}
+	c.mu.RLock()
+	empty := len(c.subscriptions) == 0
+	consuming := c.consuming
+	c.mu.RUnlock()
+	if empty && consuming {
+		if err := channel.Cancel(c.consumerTag, false); err != nil {
+			return fmt.Errorf("cancel rabbitmq consumer: %w", err)
+		}
+		c.mu.Lock()
+		c.consuming = false
+		c.mu.Unlock()
+	}
+	return nil
+}
+
+func (c *RabbitMQ) Connected() bool {
+	if c == nil {
+		return false
+	}
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return !c.closed && c.connection != nil && !c.connection.IsClosed()
+}
+
+func (c *RabbitMQ) Close() error {
+	if c == nil {
+		return nil
+	}
+	c.opMu.Lock()
+	defer c.opMu.Unlock()
+	c.publishMu.Lock()
+	defer c.publishMu.Unlock()
+	c.consumeMu.Lock()
+	defer c.consumeMu.Unlock()
+	c.mu.Lock()
+	if c.closed {
+		c.mu.Unlock()
+		return nil
+	}
+	c.closed = true
+	close(c.stop)
+	publishChannel := c.publishChannel
+	consumeChannel := c.consumeChannel
+	connection := c.connection
+	c.publishChannel = nil
+	c.consumeChannel = nil
+	c.connection = nil
+	c.mu.Unlock()
+	var result error
+	if consumeChannel != nil {
+		result = errors.Join(result, consumeChannel.Close())
+	}
+	if publishChannel != nil {
+		result = errors.Join(result, publishChannel.Close())
+	}
+	if connection != nil {
+		result = errors.Join(result, connection.Close())
+	}
+	return result
+}
+
+func rabbitRoutingKeyMatches(exchangeType, pattern, routingKey string) bool {
+	switch strings.ToLower(exchangeType) {
+	case "fanout":
+		return true
+	case "direct":
+		return pattern == routingKey
+	}
+	patternParts := strings.Split(pattern, ".")
+	routingParts := strings.Split(routingKey, ".")
+	for len(patternParts) > 0 {
+		head := patternParts[0]
+		patternParts = patternParts[1:]
+		if head == "#" {
+			if len(patternParts) == 0 {
+				return true
+			}
+			for index := 0; index <= len(routingParts); index++ {
+				if rabbitTopicPartsMatch(patternParts, routingParts[index:]) {
+					return true
+				}
+			}
+			return false
+		}
+		if len(routingParts) == 0 || (head != "*" && head != routingParts[0]) {
+			return false
+		}
+		routingParts = routingParts[1:]
+	}
+	return len(routingParts) == 0
+}
+
+func rabbitTopicPartsMatch(patternParts, routingParts []string) bool {
+	return rabbitRoutingKeyMatches("topic", strings.Join(patternParts, "."), strings.Join(routingParts, "."))
+}
@@ -0,0 +1,48 @@
+package mq
+
+import (
+	"context"
+	"errors"
+	"testing"
+)
+
+func TestDisabledRabbitMQIsSafeAndUnavailable(t *testing.T) {
+	client, err := NewRabbitMQ(RabbitMQConfig{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if client.Connected() {
+		t.Fatal("disabled rabbitmq reported connected")
+	}
+	if err = client.Publish(context.Background(), "orders.paid", []byte("test"), AtLeastOnce, false); !errors.Is(err, ErrUnavailable) {
+		t.Fatalf("publish error = %v", err)
+	}
+	if err = client.Close(); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestEnabledRabbitMQRequiresTopology(t *testing.T) {
+	if _, err := NewRabbitMQ(RabbitMQConfig{Enabled: true}); err == nil {
+		t.Fatal("enabled rabbitmq without host and topology should fail")
+	}
+}
+
+func TestRabbitRoutingKeyMatches(t *testing.T) {
+	tests := []struct {
+		pattern string
+		key     string
+		want    bool
+	}{
+		{pattern: "orders.*.paid", key: "orders.cn.paid", want: true},
+		{pattern: "orders.#", key: "orders.cn.created", want: true},
+		{pattern: "#.paid", key: "orders.cn.paid", want: true},
+		{pattern: "orders.*", key: "orders.cn.paid", want: false},
+		{pattern: "orders.created", key: "orders.paid", want: false},
+	}
+	for _, test := range tests {
+		if got := rabbitRoutingKeyMatches("topic", test.pattern, test.key); got != test.want {
+			t.Fatalf("match(%q, %q) = %v, want %v", test.pattern, test.key, got, test.want)
+		}
+	}
+}