ipole_edge_server/common/mqtt/mgr.go

package mqtt

import (
	"runtime/debug"
	"time"

	"github.com/sirupsen/logrus"

	"lc/common/util"
)

// OptType 消息发布/订阅方向
type OptType uint8

const (
	ToAll   OptType = 0 // 同时发布到云端和边缘端
	ToCloud OptType = 1 // 仅云端
	ToEdge  OptType = 2 // 仅边缘端
)

// MQTTMessage 队列中的消息
type MQTTMessage struct {
	topic   string
	payload string
	qos     QOS
	tp      OptType
}

// RestartFunc 协程 panic 后重建回调
type RestartFunc func(fn func(args ...interface{}) interface{}, args ...interface{})

// MQTTMgr 管理 Cloud/Edge 双客户端，支持异步队列发布
type MQTTMgr struct {
	Cloud     *MqttClient
	Edge      *MqttClient
	Queue     *util.MlQueue
	restartFn RestartFunc
}

// MQTTMgrConfig 管理器构造参数
type MQTTMgrConfig struct {
	CloudServer   string
	CloudClientID string
	CloudUser     string
	CloudPassword string
	CloudTimeout  uint
	CloudOnline   BaseMqttOnline // nil 则使用 EmptyMqttOnline

	EdgeServer   string
	EdgeClientID string
	EdgeUser     string
	EdgePassword string
	EdgeTimeout  uint
	EdgeOnline   BaseMqttOnline // nil 则使用 EmptyMqttOnline
}

// NewMQTTMgr 创建 MQTT 管理器
func NewMQTTMgr(cfg MQTTMgrConfig) *MQTTMgr {
	mgr := &MQTTMgr{
		Queue: util.NewQueue(2000),
	}
	if cfg.CloudOnline == nil {
		cfg.CloudOnline = &EmptyMqttOnline{}
	}
	if cfg.EdgeOnline == nil {
		cfg.EdgeOnline = &EmptyMqttOnline{}
	}
	if cfg.CloudServer != "" {
		mgr.Cloud = NewMqttClient(cfg.CloudServer, cfg.CloudClientID,
			cfg.CloudUser, cfg.CloudPassword, cfg.CloudTimeout, cfg.CloudOnline)
	}
	if cfg.EdgeServer != "" {
		mgr.Edge = NewMqttClient(cfg.EdgeServer, cfg.EdgeClientID,
			cfg.EdgeUser, cfg.EdgePassword, cfg.EdgeTimeout, cfg.EdgeOnline)
	}
	return mgr
}

// SetRestartFn 设置 panic 恢复后的协程重建回调
func (o *MQTTMgr) SetRestartFn(fn RestartFunc) {
	o.restartFn = fn
}

// Subscribe 订阅主题，支持方向选择
func (o *MQTTMgr) Subscribe(topic string, qos QOS, handler MessageHandler, tp OptType) {
	switch tp {
	case ToAll:
		if o.Cloud != nil {
			o.Cloud.Handle(topic, handler)
			_ = o.Cloud.Subscribe(topic, qos)
		}
		if o.Edge != nil {
			o.Edge.Handle(topic, handler)
			_ = o.Edge.Subscribe(topic, qos)
		}
	case ToCloud:
		if o.Cloud != nil {
			o.Cloud.Handle(topic, handler)
			_ = o.Cloud.Subscribe(topic, qos)
		}
	case ToEdge:
		if o.Edge != nil {
			o.Edge.Handle(topic, handler)
			_ = o.Edge.Subscribe(topic, qos)
		}
	}
}

// UnSubscribe 退订主题
func (o *MQTTMgr) UnSubscribe(topic string, tp OptType) {
	switch tp {
	case ToAll:
		if o.Cloud != nil {
			_ = o.Cloud.Unsubscribe(topic)
		}
		if o.Edge != nil {
			_ = o.Edge.Unsubscribe(topic)
		}
	case ToCloud:
		if o.Cloud != nil {
			_ = o.Cloud.Unsubscribe(topic)
		}
	case ToEdge:
		if o.Edge != nil {
			_ = o.Edge.Unsubscribe(topic)
		}
	}
}

// Publish 异步发布（入队列）
func (o *MQTTMgr) Publish(topic, payload string, qos QOS, tp OptType) {
	o.Queue.Put(&MQTTMessage{topic: topic, payload: payload, qos: qos, tp: tp})
}

// doPublish 底层同步发布
func (o *MQTTMgr) doPublish(msg *MQTTMessage) error {
	var err error
	switch msg.tp {
	case ToAll:
		if o.Cloud != nil {
			err = o.Cloud.PublishString(msg.topic, msg.payload, msg.qos)
		}
		if o.Edge != nil {
			_ = o.Edge.PublishString(msg.topic, msg.payload, msg.qos)
		}
	case ToCloud:
		if o.Cloud != nil {
			err = o.Cloud.PublishString(msg.topic, msg.payload, msg.qos)
		}
	case ToEdge:
		if o.Edge != nil {
			err = o.Edge.PublishString(msg.topic, msg.payload, msg.qos)
		}
	}
	return err
}

// MQTTConnectMgr 连接保持协程（每 10 秒重连）
func (o *MQTTMgr) MQTTConnectMgr(args ...interface{}) interface{} {
	for {
		time.Sleep(10 * time.Second)
		if o.Cloud != nil {
			o.Cloud.Connect()
		}
		if o.Edge != nil {
			o.Edge.Connect()
		}
	}
}

// MQTTMessageHandle 队列消费协程（带 RETRY）
func (o *MQTTMgr) MQTTMessageHandle(args ...interface{}) interface{} {
	defer func() {
		if err := recover(); err != nil {
			logrus.Errorf("MQTTMgr.MQTTMessageHandle发生异常:%v", err)
			logrus.Errorf("MQTTMgr.MQTTMessageHandle发生异常,堆栈信息:%s", string(debug.Stack()))
			time.Sleep(time.Second)
			if o.restartFn != nil {
				o.restartFn(o.MQTTMessageHandle, args)
			}
		}
	}()
	for {
		if m, ok, _ := o.Queue.Get(); ok {
			if msg, ok := m.(*MQTTMessage); ok {
				for {
					if err := o.doPublish(msg); err != nil {
						logrus.Errorf("发布主题为%s的消息失败，原因:%s", msg.topic, err.Error())
						time.Sleep(time.Second)
						continue
					}
					break
				}
			}
		} else {
			time.Sleep(200 * time.Millisecond)
		}
	}
}