ipole_edge_server/common/mqtt/client_wrapper.go

package mqtt

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/sirupsen/logrus"
)

// BaseMqttOnline 上线消息和遗嘱消息接口
type BaseMqttOnline interface {
	GetOnlineMsg() (string, string)
	GetWillMsg() (string, string)
}

// EmptyMqttOnline 空实现：不上线、无遗嘱
type EmptyMqttOnline struct{}

func (o *EmptyMqttOnline) GetOnlineMsg() (string, string) { return "", "" }
func (o *EmptyMqttOnline) GetWillMsg() (string, string)   { return "", "" }

// MqttClient 封装 common/mqtt.Client，提供超时控制、主题记忆、断线自动重订阅
type MqttClient struct {
	mqtt       *Client
	mu         sync.Mutex
	mapTopics  map[string]QOS
	timeout    uint
	MqttOnline BaseMqttOnline
}

// NewMqttClient 创建并连接 MQTT 客户端
func NewMqttClient(server, clientID, user, password string, timeout uint, online BaseMqttOnline) *MqttClient {
	o := &MqttClient{
		mapTopics:  make(map[string]QOS),
		timeout:    timeout,
		MqttOnline: online,
	}
	client, err := NewClient(ClientOptions{
		Servers:       []string{server},
		ClientID:      clientID,
		Username:      user,
		Password:      password,
		AutoReconnect: true,
	}, o)
	if err != nil {
		panic(fmt.Errorf("MQTT错误: %w", err))
	}
	o.mqtt = client
	ctx, cancel := o.Ctx()
	defer cancel()
	_ = client.Connect(ctx)
	return o
}

// ---- 实现 ConnHandler 接口 ----

func (o *MqttClient) ConnectionLostHandler(err error) {
	logrus.Errorln("MqttClient.ConnectionLostHandler:MQTT连接已经断开,原因:", err)
}

func (o *MqttClient) OnConnectHandler() {
	logrus.Infoln("MqttClient.OnConnectHandler:MQTT连接成功")
	for k, v := range o.mapTopics {
		if err := o.Subscribe(k, v); err != nil {
			logrus.Errorf("重订阅主题 %s 失败: %v", k, err)
		}
	}
	topic, str := o.MqttOnline.GetOnlineMsg()
	if topic != "" {
		if err := o.PublishString(topic, str, 0); err != nil {
			logrus.Errorf("发布上线消息失败: %v", err)
		}
	}
}

func (o *MqttClient) GetWill() (string, string) {
	return o.MqttOnline.GetWillMsg()
}

// ---- 操作代理 ----

func (o *MqttClient) Connect() error {
	if !o.mqtt.IsConnected() {
		ctx, cancel := o.Ctx()
		defer cancel()
		return o.mqtt.Connect(ctx)
	}
	return nil
}

func (o *MqttClient) IsConnected() bool { return o.mqtt.IsConnected() }

func (o *MqttClient) Publish(topic string, payload []byte, qos QOS) error {
	ctx, cancel := o.Ctx()
	defer cancel()
	return o.mqtt.Publish(ctx, topic, payload, qos)
}

func (o *MqttClient) PublishString(topic string, payload string, qos QOS) error {
	ctx, cancel := o.Ctx()
	defer cancel()
	return o.mqtt.PublishString(ctx, topic, payload, qos)
}

func (o *MqttClient) PublishJSON(topic string, payload interface{}, qos QOS) error {
	ctx, cancel := o.Ctx()
	defer cancel()
	return o.mqtt.PublishJSON(ctx, topic, payload, qos)
}

func (o *MqttClient) Subscribe(topic string, qos QOS) error {
	o.mu.Lock()
	if _, ok := o.mapTopics[topic]; !ok {
		o.mapTopics[topic] = qos
	}
	o.mu.Unlock()
	ctx, cancel := o.Ctx()
	defer cancel()
	return o.mqtt.Subscribe(ctx, topic, qos)
}

func (o *MqttClient) Unsubscribe(topic string) error {
	o.mu.Lock()
	if _, ok := o.mapTopics[topic]; ok {
		delete(o.mapTopics, topic)
	}
	o.mu.Unlock()
	ctx, cancel := o.Ctx()
	defer cancel()
	return o.mqtt.Unsubscribe(ctx, topic)
}

func (o *MqttClient) Handle(topic string, handler MessageHandler) Route {
	return o.mqtt.Handle(topic, handler)
}

func (o *MqttClient) Ctx() (context.Context, context.CancelFunc) {
	return context.WithTimeout(context.Background(), time.Millisecond*time.Duration(o.timeout))
}