优化结构

2026-08-22 10:49:48 +08:00
parent 2b78b195c1
commit 54f76bd913
11 changed files with 501 additions and 98 deletions
@@ -6,6 +6,7 @@ package module
 import (
 	"github.com/gin-gonic/gin"
 	"kra/pkg/database/migration"
+	"kra/pkg/mq"
 	"kra/pkg/task"
 )

@@ -29,11 +30,12 @@ type TimedTask struct {
 }

 type Definition struct {
-	Name       string
-	Migrations []migration.Step
-	Surface    Surface
-	TimedTasks []TimedTask
-	Tasks      []task.Method
+	Name                     string
+	Migrations               []migration.Step
+	Surface                  Surface
+	TimedTasks               []TimedTask
+	Tasks                    []task.Method
+	SubscriptionContributors []mq.SubscriptionContributor
 }

 type Catalog struct {
@@ -73,6 +75,16 @@ func (c Catalog) TaskMethods() []task.Method {
 	return methods
 }

+// SubscriptionContributors returns the module declarations that are applied
+// to the process-wide MQ runtime during application startup.
+func (c Catalog) SubscriptionContributors() []mq.SubscriptionContributor {
+	var contributors []mq.SubscriptionContributor
+	for _, item := range c.Definitions {
+		contributors = append(contributors, item.SubscriptionContributors...)
+	}
+	return contributors
+}
+
 type RouteRegistrar interface {
 	RegisterRoutes(public, private *gin.RouterGroup, engine *gin.Engine)
 }
@@ -4,15 +4,17 @@ import (
 	"context"
 	"fmt"
 	"sync"
+	"sync/atomic"
 	"time"

 	paho "github.com/eclipse/paho.mqtt.golang"
 )

 type MQTT struct {
-	client paho.Client
-	mu     sync.RWMutex
-	closed bool
+	client       paho.Client
+	mu           sync.RWMutex
+	closed       bool
+	reconnecting atomic.Bool
 }

 func NewMQTT(cfg Config) (*MQTT, error) {
@@ -31,19 +33,29 @@ func NewMQTT(cfg Config) (*MQTT, error) {
 	if cfg.ReconnectInterval <= 0 {
 		cfg.ReconnectInterval = 5 * time.Second
 	}
+	c := &MQTT{}
 	opts := paho.NewClientOptions().AddBroker(cfg.Broker).SetClientID(cfg.ClientID).SetUsername(cfg.Username).SetPassword(cfg.Password)
 	opts.SetKeepAlive(cfg.KeepAlive).SetCleanSession(cfg.CleanSession).SetConnectTimeout(cfg.ConnectTimeout).SetAutoReconnect(true)
 	opts.SetMaxReconnectInterval(cfg.ReconnectInterval).SetResumeSubs(true).SetOrderMatters(false)
-	c := &MQTT{}
-	token := paho.NewClient(opts)
-	connect := token.Connect()
+	opts.SetConnectionNotificationHandler(func(_ paho.Client, notification paho.ConnectionNotification) {
+		switch notification.Type() {
+		case paho.ConnectionNotificationTypeConnecting, paho.ConnectionNotificationTypeLost:
+			c.reconnecting.Store(true)
+		case paho.ConnectionNotificationTypeConnected:
+			c.reconnecting.Store(false)
+		}
+	})
+	client := paho.NewClient(opts)
+	connect := client.Connect()
 	if !connect.WaitTimeout(cfg.ConnectTimeout) {
+		client.Disconnect(0)
 		return nil, fmt.Errorf("connect mqtt: %w", ErrUnavailable)
 	}
 	if err := connect.Error(); err != nil {
+		client.Disconnect(0)
 		return nil, err
 	}
-	c.client = token
+	c.client = client
 	return c, nil
 }

@@ -54,7 +66,9 @@ func (c *MQTT) Publish(ctx context.Context, topic string, payload []byte, qos by
 	if qos > ExactlyOnce {
 		return fmt.Errorf("invalid mqtt qos %d", qos)
 	}
-	if c == nil || c.client == nil || !c.client.IsConnected() {
+	ctx = nonNilContext(ctx)
+	client, err := c.activeClient()
+	if err != nil || client == nil || !client.IsConnected() {
 		return ErrUnavailable
 	}
 	select {
@@ -62,7 +76,7 @@ func (c *MQTT) Publish(ctx context.Context, topic string, payload []byte, qos by
 		return ctx.Err()
 	default:
 	}
-	return waitToken(ctx, c.client.Publish(topic, qos, retain, payload))
+	return waitToken(ctx, client.Publish(topic, qos, retain, payload))
 }

 func (c *MQTT) Subscribe(ctx context.Context, topic string, qos byte, handler Handler) error {
@@ -72,7 +86,9 @@ func (c *MQTT) Subscribe(ctx context.Context, topic string, qos byte, handler Ha
 	if qos > ExactlyOnce {
 		return fmt.Errorf("invalid mqtt qos %d", qos)
 	}
-	if c == nil || c.client == nil || !c.client.IsConnected() {
+	ctx = nonNilContext(ctx)
+	client, err := c.activeClient()
+	if err != nil || client == nil || !client.IsConnected() {
 		return ErrUnavailable
 	}
 	select {
@@ -83,7 +99,7 @@ func (c *MQTT) Subscribe(ctx context.Context, topic string, qos byte, handler Ha
 	if handler == nil {
 		return fmt.Errorf("mqtt handler is nil")
 	}
-	return waitToken(ctx, c.client.Subscribe(topic, qos, func(_ paho.Client, msg paho.Message) {
+	return waitToken(ctx, client.Subscribe(topic, qos, func(_ paho.Client, msg paho.Message) {
 		handler(context.Background(), Message{Topic: msg.Topic(), Payload: append([]byte(nil), msg.Payload()...), QoS: msg.Qos(), Retain: msg.Retained()})
 	}))
 }
@@ -92,7 +108,9 @@ func (c *MQTT) Unsubscribe(ctx context.Context, topics ...string) error {
 	if len(topics) == 0 {
 		return fmt.Errorf("mqtt topics are empty")
 	}
-	if c == nil || c.client == nil || !c.client.IsConnected() {
+	ctx = nonNilContext(ctx)
+	client, err := c.activeClient()
+	if err != nil || client == nil || !client.IsConnected() {
 		return ErrUnavailable
 	}
 	select {
@@ -100,10 +118,11 @@ func (c *MQTT) Unsubscribe(ctx context.Context, topics ...string) error {
 		return ctx.Err()
 	default:
 	}
-	return waitToken(ctx, c.client.Unsubscribe(topics...))
+	return waitToken(ctx, client.Unsubscribe(topics...))
 }

 func waitToken(ctx context.Context, token paho.Token) error {
+	ctx = nonNilContext(ctx)
 	if token == nil {
 		return ErrUnavailable
 	}
@@ -114,16 +133,58 @@ func waitToken(ctx context.Context, token paho.Token) error {
 		return token.Error()
 	}
 }
-func (c *MQTT) Connected() bool { return c != nil && c.client != nil && c.client.IsConnected() }
+
+func nonNilContext(ctx context.Context) context.Context {
+	if ctx == nil {
+		return context.Background()
+	}
+	return ctx
+}
+
+func (c *MQTT) activeClient() (paho.Client, error) {
+	if c == nil {
+		return nil, ErrUnavailable
+	}
+	c.mu.RLock()
+	client := c.client
+	closed := c.closed
+	c.mu.RUnlock()
+	if closed || client == nil {
+		return nil, ErrUnavailable
+	}
+	return client, nil
+}
+
+func (c *MQTT) Connected() bool {
+	client, err := c.activeClient()
+	return err == nil && client.IsConnected()
+}
+
+// Reconnecting reports whether the MQTT driver is recovering its connection.
+// The reloadable integration waits for this state instead of replacing a
+// client while driver-level recovery is still in progress.
+func (c *MQTT) Reconnecting() bool {
+	if c == nil {
+		return false
+	}
+	c.mu.RLock()
+	closed := c.closed
+	c.mu.RUnlock()
+	return !closed && c.reconnecting.Load()
+}
+
 func (c *MQTT) Close() error {
 	c.mu.Lock()
-	defer c.mu.Unlock()
 	if c.closed {
+		c.mu.Unlock()
 		return nil
 	}
 	c.closed = true
-	if c.client != nil && c.client.IsConnected() {
-		c.client.Disconnect(250)
+	client := c.client
+	c.client = nil
+	c.mu.Unlock()
+	if client != nil {
+		client.Disconnect(250)
 	}
 	return nil
 }