yvan/kra-new
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

优化结构

Browse Source
This commit is contained in:
Yvan
2026-08-22 08:25:25 +08:00
parent 626f505b5d
commit 2b78b195c1
10 changed files with 815 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/mq/mq.go
+9 -8
View File
@@ -83,14 +83,15 @@ type Registry interface {
}
type Config struct {
Enabled bool
Broker string
ClientID string
Username string
Password string
KeepAlive time.Duration
CleanSession bool
ConnectTimeout time.Duration
Enabled bool
Broker string
ClientID string
Username string
Password string
KeepAlive time.Duration
CleanSession bool
ConnectTimeout time.Duration
ReconnectInterval time.Duration
}
func PublishJSON(ctx context.Context, client Client, topic string, value any, qos byte, retain bool) error {
pkg/mq/mqtt.go
+4 -1
View File
@@ -28,9 +28,12 @@ func NewMQTT(cfg Config) (*MQTT, error) {
if cfg.ConnectTimeout <= 0 {
cfg.ConnectTimeout = 10 * time.Second
}
if cfg.ReconnectInterval <= 0 {
cfg.ReconnectInterval = 5 * time.Second
}
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.SetResumeSubs(true).SetOrderMatters(false)
opts.SetMaxReconnectInterval(cfg.ReconnectInterval).SetResumeSubs(true).SetOrderMatters(false)
c := &MQTT{}
token := paho.NewClient(opts)
connect := token.Connect()
pkg/mq/rabbitmq.go
+54 -18
View File
@@ -6,28 +6,30 @@ import (
"fmt"
"strings"
"sync"
"sync/atomic"
"time"
amqp "github.com/rabbitmq/amqp091-go"
)
type RabbitMQConfig struct {
Enabled bool
Host string
Port int
Username string
Password string
VHost string
Exchange string
ExchangeType string
Queue string
RoutingKey string
Durable bool
AutoDelete bool
PrefetchCount int
Heartbeat time.Duration
ConnectTimeout time.Duration
TLS bool
Enabled bool
Host string
Port int
Username string
Password string
VHost string
Exchange string
ExchangeType string
Queue string
RoutingKey string
Durable bool
AutoDelete bool
PrefetchCount int
Heartbeat time.Duration
ConnectTimeout time.Duration
ReconnectInterval time.Duration
TLS bool
}
type rabbitSubscription struct {
@@ -52,6 +54,7 @@ type RabbitMQ struct {
consuming bool
stop chan struct{}
closed bool
reconnecting atomic.Bool
}
func NewRabbitMQ(config RabbitMQConfig) (*RabbitMQ, error) {
@@ -75,7 +78,13 @@ func NewRabbitMQ(config RabbitMQConfig) (*RabbitMQ, error) {
}.String()
connection, err := amqp.DialConfig(address, amqp.Config{
Heartbeat: config.Heartbeat,
Recovery: &amqp.Recovery{},
Recovery: &amqp.Recovery{
ReconnectionConfig: &amqp.ReconnectionConfig{MaxRetryCount: 5, RetryInterval: config.ReconnectInterval},
OnTopologyEntityError: func(_ *amqp.Connection, entity amqp.TopologyRecoveryEntity) bool {
amqp.Logger.Printf("rabbitmq topology recovery failed: %s", entity.Error())
return false
},
},
})
if err != nil {
return nil, fmt.Errorf("connect rabbitmq: %w", err)
@@ -116,6 +125,9 @@ func NewRabbitMQ(config RabbitMQConfig) (*RabbitMQ, error) {
client.consumeChannel = consumeChannel
client.config = config
client.consumerTag = fmt.Sprintf("kra-%d", time.Now().UnixNano())
stateChanges := make(chan *amqp.StateChanged, 16)
connection.NotifyStateChange(stateChanges)
go client.watchState(stateChanges)
return client, nil
}
@@ -148,6 +160,9 @@ func defaultRabbitMQConfig(config RabbitMQConfig) RabbitMQConfig {
if config.ConnectTimeout <= 0 {
config.ConnectTimeout = 10 * time.Second
}
if config.ReconnectInterval <= 0 {
config.ReconnectInterval = 5 * time.Second
}
return config
}
@@ -371,7 +386,28 @@ func (c *RabbitMQ) Connected() bool {
}
c.mu.RLock()
defer c.mu.RUnlock()
return !c.closed && c.connection != nil && !c.connection.IsClosed()
return !c.closed && !c.reconnecting.Load() && c.connection != nil && !c.connection.IsClosed()
}
// Reconnecting reports whether the AMQP driver is actively recovering the
// current connection. The outer runtime waits for this state to settle before
// deciding whether a fresh client must be built.
func (c *RabbitMQ) Reconnecting() bool {
return c != nil && !c.closed && c.reconnecting.Load()
}
func (c *RabbitMQ) watchState(states <-chan *amqp.StateChanged) {
for state := range states {
if state == nil {
continue
}
switch state.To {
case amqp.StateReconnecting:
c.reconnecting.Store(true)
case amqp.StateOpen, amqp.StateClosed:
c.reconnecting.Store(false)
}
}
}
func (c *RabbitMQ) Close() error {
pkg/mq/subscription_test.go
+232
View File
@@ -0,0 +1,232 @@
package mq
import (
"context"
"errors"
"reflect"
"strings"
"testing"
)
func TestNormalizeSubscriptionSetValidDeclaration(t *testing.T) {
handler := func(context.Context, Message) {}
set := SubscriptionSet{
Owner: " orders ",
Provider: " RABBITMQ ",
Topics: []TopicSubscription{
{Topic: " orders.created ", QoS: AtMostOnce, Handler: handler},
{Topic: "orders.updated", QoS: AtLeastOnce, Handler: handler},
{Topic: "orders.deleted", QoS: ExactlyOnce, Handler: handler},
},
}
got, err := NormalizeSubscriptionSet(set)
if err != nil {
t.Fatalf("NormalizeSubscriptionSet() error = %v", err)
}
want := SubscriptionSet{
Owner: "orders",
Provider: ProviderRabbitMQ,
Topics: []TopicSubscription{
{Topic: "orders.created", QoS: AtMostOnce, Handler: handler},
{Topic: "orders.updated", QoS: AtLeastOnce, Handler: handler},
{Topic: "orders.deleted", QoS: ExactlyOnce, Handler: handler},
},
}
if got.Owner != want.Owner || got.Provider != want.Provider {
t.Fatalf("normalized identity = %#v, want owner=%q provider=%q", got, want.Owner, want.Provider)
}
if len(got.Topics) != len(want.Topics) {
t.Fatalf("normalized topic count = %d, want %d", len(got.Topics), len(want.Topics))
}
for index := range want.Topics {
if got.Topics[index].Topic != want.Topics[index].Topic {
t.Errorf("topic[%d] = %q, want %q", index, got.Topics[index].Topic, want.Topics[index].Topic)
}
if got.Topics[index].QoS != want.Topics[index].QoS {
t.Errorf("qos[%d] = %d, want %d", index, got.Topics[index].QoS, want.Topics[index].QoS)
}
if reflect.ValueOf(got.Topics[index].Handler).Pointer() != reflect.ValueOf(want.Topics[index].Handler).Pointer() {
t.Errorf("handler[%d] was changed", index)
}
}
}
func TestNormalizeSubscriptionSetRejectsInvalidDeclarations(t *testing.T) {
handler := func(context.Context, Message) {}
base := func() SubscriptionSet {
return SubscriptionSet{
Owner: "orders",
Provider: ProviderEMQX,
Topics: []TopicSubscription{
{Topic: "orders.created", QoS: AtLeastOnce, Handler: handler},
},
}
}
tests := []struct {
name string
set SubscriptionSet
wantErr string
}{
{
name: "empty owner",
set: func() SubscriptionSet { set := base(); set.Owner = " "; return set }(),
wantErr: "owner is empty",
},
{
name: "empty provider",
set: func() SubscriptionSet { set := base(); set.Provider = " "; return set }(),
wantErr: "unsupported mq provider",
},
{
name: "unsupported provider",
set: func() SubscriptionSet { set := base(); set.Provider = "kafka"; return set }(),
wantErr: "unsupported mq provider",
},
{
name: "empty topics",
set: func() SubscriptionSet { set := base(); set.Topics = nil; return set }(),
wantErr: "topics are empty",
},
{
name: "empty topic",
set: func() SubscriptionSet {
set := base()
set.Topics[0].Topic = " "
return set
}(),
wantErr: "topic at index 0 is empty",
},
{
name: "duplicate topic",
set: func() SubscriptionSet {
set := base()
set.Topics = append(set.Topics, TopicSubscription{
Topic: " orders.created ", QoS: AtMostOnce, Handler: handler,
})
return set
}(),
wantErr: "duplicate mq subscription topic",
},
{
name: "nil handler",
set: func() SubscriptionSet {
set := base()
set.Topics[0].Handler = nil
return set
}(),
wantErr: "handler is nil",
},
{
name: "invalid qos",
set: func() SubscriptionSet {
set := base()
set.Topics[0].QoS = ExactlyOnce + 1
return set
}(),
wantErr: "invalid mq qos",
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
_, err := NormalizeSubscriptionSet(test.set)
if err == nil {
t.Fatal("NormalizeSubscriptionSet() error = nil, want error")
}
if !strings.Contains(err.Error(), test.wantErr) {
t.Fatalf("error = %q, want substring %q", err, test.wantErr)
}
})
}
}
type subscriptionTestRegistrar struct {
sets []SubscriptionSet
}
func (r *subscriptionTestRegistrar) Register(set SubscriptionSet) error {
r.sets = append(r.sets, set)
return nil
}
func (*subscriptionTestRegistrar) Unregister(string) error { return nil }
type subscriptionTestContributor struct {
name string
set SubscriptionSet
order *[]string
err error
}
func (c *subscriptionTestContributor) RegisterSubscriptions(registrar SubscriptionRegistrar) error {
*c.order = append(*c.order, c.name)
if c.err != nil {
return c.err
}
return registrar.Register(c.set)
}
func TestApplySubscriptionsAppliesMultipleContributorsInOrder(t *testing.T) {
order := make([]string, 0, 2)
registrar := &subscriptionTestRegistrar{}
handler := func(context.Context, Message) {}
first := &subscriptionTestContributor{
name: "orders",
order: &order,
set: SubscriptionSet{
Owner: "orders",
Provider: ProviderRabbitMQ,
Topics: []TopicSubscription{{Topic: "orders.created", Handler: handler}},
},
}
second := &subscriptionTestContributor{
name: "notifications",
order: &order,
set: SubscriptionSet{
Owner: "notifications",
Provider: ProviderEMQX,
Topics: []TopicSubscription{{Topic: "notifications.sent", Handler: handler}},
},
}
if err := ApplySubscriptions(registrar, first, nil, second); err != nil {
t.Fatalf("ApplySubscriptions() error = %v", err)
}
if !reflect.DeepEqual(order, []string{"orders", "notifications"}) {
t.Fatalf("contributor order = %#v, want %#v", order, []string{"orders", "notifications"})
}
if len(registrar.sets) != 2 {
t.Fatalf("registered sets = %d, want 2", len(registrar.sets))
}
if registrar.sets[0].Owner != "orders" || registrar.sets[1].Owner != "notifications" {
t.Fatalf("registered owners = %q, %q", registrar.sets[0].Owner, registrar.sets[1].Owner)
}
}
func TestApplySubscriptionsPropagatesContributorError(t *testing.T) {
order := make([]string, 0, 2)
registrar := &subscriptionTestRegistrar{}
wantErr := errors.New("registration failed")
failing := &subscriptionTestContributor{name: "failing", order: &order, err: wantErr}
following := &subscriptionTestContributor{name: "following", order: &order}
err := ApplySubscriptions(registrar, failing, following)
if !errors.Is(err, wantErr) {
t.Fatalf("ApplySubscriptions() error = %v, want %v", err, wantErr)
}
if !reflect.DeepEqual(order, []string{"failing"}) {
t.Fatalf("contributors called = %#v, want %#v", order, []string{"failing"})
}
if len(registrar.sets) != 0 {
t.Fatalf("registered sets = %d, want 0", len(registrar.sets))
}
}
func TestApplySubscriptionsRejectsNilRegistrar(t *testing.T) {
if err := ApplySubscriptions(nil); err == nil || !strings.Contains(err.Error(), "registrar is nil") {
t.Fatalf("ApplySubscriptions(nil) error = %v, want nil-registrar error", err)
}
}