优化结构

2026-08-21 08:54:29 +08:00
parent c968f9f24e
commit 77e192b6c0
167 changed files with 3657 additions and 1528 deletions
@@ -0,0 +1,154 @@
+package mq
+
+import (
+	"context"
+	"log/slog"
+	"sync"
+	"time"
+
+	"kra/app/system/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
+}
@@ -0,0 +1,56 @@
+package mq
+
+import (
+	"context"
+	"testing"
+
+	platformmq "kra/pkg/mq"
+)
+
+type fakeClient struct {
+	subscribed   []string
+	unsubscribed []string
+}
+
+func (*fakeClient) Publish(context.Context, string, []byte, byte, bool) error { return nil }
+func (f *fakeClient) Subscribe(_ context.Context, topic string, _ byte, _ platformmq.Handler) error {
+	f.subscribed = append(f.subscribed, topic)
+	return nil
+}
+func (f *fakeClient) Unsubscribe(_ context.Context, topics ...string) error {
+	f.unsubscribed = append(f.unsubscribed, topics...)
+	return nil
+}
+func (*fakeClient) Connected() bool { return true }
+func (*fakeClient) Close() error    { return nil }
+
+func TestReloadableTracksSubscriptions(t *testing.T) {
+	client := &fakeClient{}
+	r := &Reloadable{current: client, subscriptions: make(map[string]subscription)}
+	handler := func(context.Context, platformmq.Message) {}
+	if err := r.Subscribe(context.Background(), "orders/+/paid", platformmq.AtLeastOnce, handler); err != nil {
+		t.Fatal(err)
+	}
+	if _, ok := r.subscriptions["orders/+/paid"]; !ok {
+		t.Fatal("subscription was not retained for configuration reload")
+	}
+	if err := r.Unsubscribe(context.Background(), "orders/+/paid"); err != nil {
+		t.Fatal(err)
+	}
+	if _, ok := r.subscriptions["orders/+/paid"]; ok {
+		t.Fatal("unsubscribed topic remained in the reload registry")
+	}
+}
+
+func TestReloadableRestoresSubscriptions(t *testing.T) {
+	client := &fakeClient{}
+	r := &Reloadable{subscriptions: map[string]subscription{
+		"orders/+/paid": {qos: platformmq.AtLeastOnce, handler: func(context.Context, platformmq.Message) {}},
+	}}
+	if err := r.restoreSubscriptions(context.Background(), client); err != nil {
+		t.Fatal(err)
+	}
+	if len(client.subscribed) != 1 || client.subscribed[0] != "orders/+/paid" {
+		t.Fatalf("restored subscriptions = %v", client.subscribed)
+	}
+}