kra-new/pkg/mq/mqtt.go

package mq

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
	reconnecting atomic.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
	}
	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)
	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 = client
	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)
	}
	ctx = nonNilContext(ctx)
	client, err := c.activeClient()
	if err != nil || client == nil || !client.IsConnected() {
		return ErrUnavailable
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
	}
	return waitToken(ctx, 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)
	}
	ctx = nonNilContext(ctx)
	client, err := c.activeClient()
	if err != nil || client == nil || !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, 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")
	}
	ctx = nonNilContext(ctx)
	client, err := c.activeClient()
	if err != nil || client == nil || !client.IsConnected() {
		return ErrUnavailable
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
	}
	return waitToken(ctx, client.Unsubscribe(topics...))
}

func waitToken(ctx context.Context, token paho.Token) error {
	ctx = nonNilContext(ctx)
	if token == nil {
		return ErrUnavailable
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-token.Done():
		return token.Error()
	}
}

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()
	if c.closed {
		c.mu.Unlock()
		return nil
	}
	c.closed = true
	client := c.client
	c.client = nil
	c.mu.Unlock()
	if client != nil {
		client.Disconnect(250)
	}
	return nil
}