FranklinWH-Toolbox/service/utils/mqttClient/mqttClient.go

package mqttClient

import (
	"crypto/tls"
	"fmt"
	"time"

	mqtt "github.com/eclipse/paho.mqtt.golang"
	"go.uber.org/zap"
)

type MqttClientOpt struct {
	Broker   string // 服务端地址
	ClientID string // (可选) 客户端名称
	TLS      *tls.Config
}

type MqttClient struct {
	client       mqtt.Client
	subscribeMap map[string]func(payload []byte)
	clientID     string // 客户端名称
}

const (
	connectTimeout      = 30 * time.Second // 30s
	subscribeMsgTimeout = 30 * time.Second // 30s
	sendMsgTimeout      = 30 * time.Second // 30s
)

// 创建客户端
func CreateMqttClient(opt MqttClientOpt) (*MqttClient, error) {
	// 创建客户端选项
	opts := mqtt.NewClientOptions()
	opts.AddBroker(opt.Broker)
	clientID := "go_mqtt_client"
	if len(opt.ClientID) > 0 {
		clientID = opt.ClientID
	}
	opts.SetClientID(clientID)

	opts.SetAutoReconnect(true) // 自动重连
	opts.SetCleanSession(true)  // 不保留会话
	opts.SetTLSConfig(opt.TLS)

	mqttClient := MqttClient{
		subscribeMap: map[string]func(payload []byte){},
	}

	// 连接成功回调
	opts.OnConnect = func(c mqtt.Client) {
		zap.L().Info("Connected to MQTT broker", zap.String("clientID", clientID))

		// 重新订阅丢失的 Subscribe
		mqttClient.ResumeSubscribe()
	}

	// 连接丢失回调
	opts.OnConnectionLost = func(c mqtt.Client, err error) {
		fmt.Printf("%s Connection lost: %s\n", clientID, err.Error())
		zap.L().Info("Connection lost", zap.String("clientID", clientID), zap.String("err", err.Error()))
	}

	// 创建客户端
	mqttClient.client = mqtt.NewClient(opts)

	return &mqttClient, nil
}

func (c *MqttClient) Connect() error {
	if c.client.IsConnected() {
		return nil
	}

	// 连接
	token := c.client.Connect()
	if !token.WaitTimeout(connectTimeout) {
		return fmt.Errorf("connect timeout (%s)", c.clientID)
	}

	return token.Error()
}

// 断开连接
func (c *MqttClient) Disconnect() {
	if !c.client.IsConnected() {
		return
	}

	// 清空订阅
	c.CleanSubscribe()

	// 安全断开，等待最多 2 秒让队列消息发送完
	c.client.Disconnect(2000)
}

// 检查是否连接
func (c *MqttClient) IsConnected() bool {
	return c.client.IsConnected()
}

// 恢复订阅消息
func (c *MqttClient) ResumeSubscribe() {
	for k, v := range c.subscribeMap {
		c.subscribeMessage(k, v)
	}
}

// 清空所有订阅
func (c *MqttClient) CleanSubscribe() {
	for k := range c.subscribeMap {
		delete(c.subscribeMap, k)
	}
}

// 订阅消息
func (c *MqttClient) subscribeMessage(topic string, callback func(payload []byte)) error {
	if !c.client.IsConnected() {
		return fmt.Errorf("subscribe failed. (client not connected)")
	}

	token := c.client.Subscribe(topic, 1, func(c mqtt.Client, m mqtt.Message) {
		callback(m.Payload())
	})
	if !token.WaitTimeout(subscribeMsgTimeout) {
		return fmt.Errorf("subscribe timeout (topic:%s)", topic)
	}

	err := token.Error()
	if err != nil {
		return err
	}

	return nil
}

// 订阅消息
func (c *MqttClient) SubscribeMessage(topic string, callback func(payload []byte)) error {
	err := c.subscribeMessage(topic, callback)
	if err != nil {
		return err
	}

	// 备份subscribe
	c.subscribeMap[topic] = callback

	return nil
}

// 取消订阅
func (c *MqttClient) UnSubscribeMessage(topic string) error {
	if !c.client.IsConnected() {
		return fmt.Errorf("unsubscribe failed. (client not connected)")
	}

	token := c.client.Unsubscribe(topic)
	if !token.WaitTimeout(subscribeMsgTimeout) {
		return fmt.Errorf("unsubscribe timeout (topic:%s)", topic)
	}

	err := token.Error()
	if err != nil {
		return err
	}

	delete(c.subscribeMap, topic)

	return nil
}

// 发送消息，只尝试一次，不管是否送达
func (c *MqttClient) SendMessageOnce(topic string, msg []byte) error {
	if !c.client.IsConnected() {
		return fmt.Errorf("send message once failed. (client not connected)")
	}

	token := c.client.Publish(topic, 0, false, msg)
	if !token.WaitTimeout(sendMsgTimeout) {
		return fmt.Errorf("send message once timeout (topic:%s, msg:%s)", topic, msg)
	}

	return token.Error()
}

// 发送消息，没收到 ACK 就重发，至少送到一次，可能收到多次
func (c *MqttClient) SendMessage(topic string, msg []byte, retained bool) error {
	if !c.client.IsConnected() {
		return fmt.Errorf("send message failed. (client not connected)")
	}

	token := c.client.Publish(topic, 1, retained, msg)
	if !token.WaitTimeout(sendMsgTimeout) {
		return fmt.Errorf("send message timeout (topic:%s, msg:%s)", topic, msg)
	}

	return token.Error()
}

// 发送消息，严格一次，不重复，不丢失
func (c *MqttClient) SendMessageMustOnce(topic string, msg []byte, retained bool) error {
	if !c.client.IsConnected() {
		return fmt.Errorf("send message must once failed. (client not connected)")
	}

	token := c.client.Publish(topic, 2, retained, msg)
	if !token.WaitTimeout(sendMsgTimeout) {
		return fmt.Errorf("send message must once timeout (topic:%s, msg:%s)", topic, msg)
	}

	return token.Error()
}

// 发送消息，没收到 ACK 就重发，至少送到一次，可能收到多次，并保留消息在Broker，新定义的设备会收到消息
func (c *MqttClient) SendMessageRetained(topic string, msg []byte) error {
	if !c.client.IsConnected() {
		return fmt.Errorf("send message retained failed. (client not connected)")
	}

	token := c.client.Publish(topic, 1, true, msg)
	if !token.WaitTimeout(sendMsgTimeout) {
		return fmt.Errorf("send message retained timeout (topic:%s, msg:%s)", topic, msg)
	}

	return token.Error()
}

// 清除保留在Broker的消息
func (c *MqttClient) CleanMessageRetained(topic string) error {
	if !c.client.IsConnected() {
		return fmt.Errorf("clean message retained failed. (client not connected)")
	}

	token := c.client.Publish(topic, 1, true, nil)
	if !token.WaitTimeout(sendMsgTimeout) {
		return fmt.Errorf("clean message retained timeout (topic:%s)", topic)
	}

	return token.Error()
}