kra-new/internal/integration/mq/emqx.go

package mq

import (
	"context"
	"log/slog"
	"sync"
	"time"

	"kra/internal/conf"
	"kra/pkg/mq"
)

// Reloadable follows the runtime snapshot and keeps a single shared EMQX
// connection for all modules in this process.
type Reloadable struct {
	mu            sync.RWMutex
	opMu          sync.Mutex
	current       mq.Client
	subscriptions map[string]subscription
	stop          func()
	logger        *slog.Logger
	closed        bool
}

type subscription struct {
	qos     byte
	handler mq.Handler
}

func New(runtime *conf.Runtime, logger *slog.Logger) (*Reloadable, func(), error) {
	if logger == nil {
		logger = slog.Default()
	}
	r := &Reloadable{logger: logger, subscriptions: make(map[string]subscription)}
	if runtime != nil {
		var config *conf.AdminBackend_MQ
		if admin := runtime.Admin(); admin != nil {
			config = admin.GetMq()
		}
		r.replace(config)
		r.stop = runtime.Subscribe(func(_ *conf.Data, admin *conf.AdminBackend) {
			if admin != nil {
				r.replace(admin.GetMq())
			}
		})
	}
	cleanup := func() {
		if r.stop != nil {
			r.stop()
		}
		_ = r.Close()
	}
	return r, cleanup, nil
}

func (r *Reloadable) replace(config *conf.AdminBackend_MQ) {
	r.opMu.Lock()
	defer r.opMu.Unlock()
	if r.closed {
		return
	}
	if config == nil {
		config = &conf.AdminBackend_MQ{}
	}
	cfg := mq.Config{Enabled: config.Enabled, Broker: config.Broker, ClientID: config.ClientId, Username: config.Username, Password: config.Password, CleanSession: config.CleanSession}
	if config.KeepAlive > 0 {
		cfg.KeepAlive = time.Duration(config.KeepAlive) * time.Second
	}
	if config.ConnectTimeout > 0 {
		cfg.ConnectTimeout = time.Duration(config.ConnectTimeout) * time.Second
	}
	client, err := mq.NewMQTT(cfg)
	if err != nil {
		r.logger.Warn("emqx unavailable", "mod", "mq", "error", err)
		return
	}
	if config.Enabled {
		if err = r.restoreSubscriptions(context.Background(), client); err != nil {
			_ = client.Close()
			r.logger.Warn("restore emqx subscriptions failed", "mod", "mq", "error", err)
			return
		}
	}
	r.mu.Lock()
	old := r.current
	r.current = client
	r.mu.Unlock()
	if old != nil {
		_ = old.Close()
	}
}

func (r *Reloadable) restoreSubscriptions(ctx context.Context, client mq.Client) error {
	for topic, item := range r.subscriptions {
		if err := client.Subscribe(ctx, topic, item.qos, item.handler); err != nil {
			return err
		}
	}
	return nil
}

func (r *Reloadable) client() mq.Client { r.mu.RLock(); defer r.mu.RUnlock(); return r.current }
func (r *Reloadable) Publish(ctx context.Context, topic string, payload []byte, qos byte, retain bool) error {
	c := r.client()
	if c == nil {
		return mq.ErrUnavailable
	}
	return c.Publish(ctx, topic, payload, qos, retain)
}
func (r *Reloadable) Subscribe(ctx context.Context, topic string, qos byte, handler mq.Handler) error {
	r.opMu.Lock()
	defer r.opMu.Unlock()
	c := r.client()
	if c == nil {
		return mq.ErrUnavailable
	}
	if err := c.Subscribe(ctx, topic, qos, handler); err != nil {
		return err
	}
	r.subscriptions[topic] = subscription{qos: qos, handler: handler}
	return nil
}
func (r *Reloadable) Unsubscribe(ctx context.Context, topics ...string) error {
	r.opMu.Lock()
	defer r.opMu.Unlock()
	c := r.client()
	if c == nil {
		return mq.ErrUnavailable
	}
	if err := c.Unsubscribe(ctx, topics...); err != nil {
		return err
	}
	for _, topic := range topics {
		delete(r.subscriptions, topic)
	}
	return nil
}
func (r *Reloadable) Connected() bool { c := r.client(); return c != nil && c.Connected() }
func (r *Reloadable) Close() error {
	r.opMu.Lock()
	defer r.opMu.Unlock()
	if r.closed {
		return nil
	}
	r.closed = true
	r.mu.Lock()
	old := r.current
	r.current = nil
	r.mu.Unlock()
	if old != nil {
		return old.Close()
	}
	return nil
}