优化结构

pkg/mq/rabbitmq.go

@@ -6,28 +6,30 @@ import (
 	"fmt"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"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
+	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
+	ReconnectInterval time.Duration
+	TLS               bool
 }
 
 type rabbitSubscription struct {
@@ -52,6 +54,7 @@ type RabbitMQ struct {
 	consuming      bool
 	stop           chan struct{}
 	closed         bool
+	reconnecting   atomic.Bool
 }
 
 func NewRabbitMQ(config RabbitMQConfig) (*RabbitMQ, error) {
@@ -75,7 +78,13 @@ func NewRabbitMQ(config RabbitMQConfig) (*RabbitMQ, error) {
 	}.String()
 	connection, err := amqp.DialConfig(address, amqp.Config{
 		Heartbeat: config.Heartbeat,
-		Recovery:  &amqp.Recovery{},
+		Recovery: &amqp.Recovery{
+			ReconnectionConfig: &amqp.ReconnectionConfig{MaxRetryCount: 5, RetryInterval: config.ReconnectInterval},
+			OnTopologyEntityError: func(_ *amqp.Connection, entity amqp.TopologyRecoveryEntity) bool {
+				amqp.Logger.Printf("rabbitmq topology recovery failed: %s", entity.Error())
+				return false
+			},
+		},
 	})
 	if err != nil {
 		return nil, fmt.Errorf("connect rabbitmq: %w", err)
@@ -116,6 +125,9 @@ func NewRabbitMQ(config RabbitMQConfig) (*RabbitMQ, error) {
 	client.consumeChannel = consumeChannel
 	client.config = config
 	client.consumerTag = fmt.Sprintf("kra-%d", time.Now().UnixNano())
+	stateChanges := make(chan *amqp.StateChanged, 16)
+	connection.NotifyStateChange(stateChanges)
+	go client.watchState(stateChanges)
 	return client, nil
 }
 
@@ -148,6 +160,9 @@ func defaultRabbitMQConfig(config RabbitMQConfig) RabbitMQConfig {
 	if config.ConnectTimeout <= 0 {
 		config.ConnectTimeout = 10 * time.Second
 	}
+	if config.ReconnectInterval <= 0 {
+		config.ReconnectInterval = 5 * time.Second
+	}
 	return config
 }
 
@@ -371,7 +386,28 @@ func (c *RabbitMQ) Connected() bool {
 	}
 	c.mu.RLock()
 	defer c.mu.RUnlock()
-	return !c.closed && c.connection != nil && !c.connection.IsClosed()
+	return !c.closed && !c.reconnecting.Load() && c.connection != nil && !c.connection.IsClosed()
+}
+
+// Reconnecting reports whether the AMQP driver is actively recovering the
+// current connection. The outer runtime waits for this state to settle before
+// deciding whether a fresh client must be built.
+func (c *RabbitMQ) Reconnecting() bool {
+	return c != nil && !c.closed && c.reconnecting.Load()
+}
+
+func (c *RabbitMQ) watchState(states <-chan *amqp.StateChanged) {
+	for state := range states {
+		if state == nil {
+			continue
+		}
+		switch state.To {
+		case amqp.StateReconnecting:
+			c.reconnecting.Store(true)
+		case amqp.StateOpen, amqp.StateClosed:
+			c.reconnecting.Store(false)
+		}
+	}
 }
 
 func (c *RabbitMQ) Close() error {