优化结构

2026-08-21 08:54:29 +08:00
parent c968f9f24e
commit 77e192b6c0
167 changed files with 3657 additions and 1528 deletions
@@ -0,0 +1,126 @@
+package mq
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+
+	paho "github.com/eclipse/paho.mqtt.golang"
+)
+
+type MQTT struct {
+	client paho.Client
+	mu     sync.RWMutex
+	closed bool
+}
+
+func NewMQTT(cfg Config) (*MQTT, error) {
+	if !cfg.Enabled {
+		return &MQTT{}, nil
+	}
+	if cfg.Broker == "" {
+		return nil, fmt.Errorf("mqtt broker is empty")
+	}
+	if cfg.KeepAlive <= 0 {
+		cfg.KeepAlive = 30 * time.Second
+	}
+	if cfg.ConnectTimeout <= 0 {
+		cfg.ConnectTimeout = 10 * time.Second
+	}
+	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.SetResumeSubs(true).SetOrderMatters(false)
+	c := &MQTT{}
+	token := paho.NewClient(opts)
+	connect := token.Connect()
+	if !connect.WaitTimeout(cfg.ConnectTimeout) {
+		return nil, fmt.Errorf("connect mqtt: %w", ErrUnavailable)
+	}
+	if err := connect.Error(); err != nil {
+		return nil, err
+	}
+	c.client = token
+	return c, nil
+}
+
+func (c *MQTT) Publish(ctx context.Context, topic string, payload []byte, qos byte, retain bool) error {
+	if topic == "" {
+		return fmt.Errorf("mqtt topic is empty")
+	}
+	if qos > ExactlyOnce {
+		return fmt.Errorf("invalid mqtt qos %d", qos)
+	}
+	if c == nil || c.client == nil || !c.client.IsConnected() {
+		return ErrUnavailable
+	}
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
+	return waitToken(ctx, c.client.Publish(topic, qos, retain, payload))
+}
+
+func (c *MQTT) Subscribe(ctx context.Context, topic string, qos byte, handler Handler) error {
+	if topic == "" {
+		return fmt.Errorf("mqtt topic is empty")
+	}
+	if qos > ExactlyOnce {
+		return fmt.Errorf("invalid mqtt qos %d", qos)
+	}
+	if c == nil || c.client == nil || !c.client.IsConnected() {
+		return ErrUnavailable
+	}
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
+	if handler == nil {
+		return fmt.Errorf("mqtt handler is nil")
+	}
+	return waitToken(ctx, c.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()})
+	}))
+}
+
+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() {
+		return ErrUnavailable
+	}
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
+	return waitToken(ctx, c.client.Unsubscribe(topics...))
+}
+
+func waitToken(ctx context.Context, token paho.Token) error {
+	if token == nil {
+		return ErrUnavailable
+	}
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-token.Done():
+		return token.Error()
+	}
+}
+func (c *MQTT) Connected() bool { return c != nil && c.client != nil && c.client.IsConnected() }
+func (c *MQTT) Close() error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.closed {
+		return nil
+	}
+	c.closed = true
+	if c.client != nil && c.client.IsConnected() {
+		c.client.Disconnect(250)
+	}
+	return nil
+}