chore: 并入 iotlcm-worker-svr 历史到 worker-svr/（原提交 96f0f4e）

2026-11-02 13:28:03 +08:00
parent 6de2e503cc 96f0f4eacd
commit 2923c78755
422 changed files with 80993 additions and 0 deletions
@@ -0,0 +1,22 @@
+package rabbitmq
+
+import "context"
+
+type DeliveryMetadata struct {
+	DeliveryTag uint64
+	Redelivered bool
+	Consumer    string
+	Queue       string
+	Headers     map[string]any
+}
+
+type deliveryMetadataContextKey struct{}
+
+func WithDeliveryMetadata(ctx context.Context, metadata DeliveryMetadata) context.Context {
+	return context.WithValue(ctx, deliveryMetadataContextKey{}, metadata)
+}
+
+func DeliveryMetadataFromContext(ctx context.Context) (DeliveryMetadata, bool) {
+	metadata, ok := ctx.Value(deliveryMetadataContextKey{}).(DeliveryMetadata)
+	return metadata, ok
+}
@@ -0,0 +1,445 @@
+package rabbitmq
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/gogf/gf/v2/errors/gerror"
+	"github.com/gogf/gf/v2/frame/g"
+	amqp "github.com/rabbitmq/amqp091-go"
+)
+
+type Config struct {
+	Enabled bool
+	Dsn     string
+}
+
+type channelLifecycle interface {
+	Cancel(consumer string, noWait bool) error
+	Close() error
+	IsClosed() bool
+}
+
+type managedConsumer struct {
+	channel channelLifecycle
+	tag     string
+	done    chan struct{}
+}
+
+// Handler 是 RabbitMQ 消费处理函数。
+// 返回 nil 时消息 Ack；返回 error 时消息 Nack 并重新入队，等待后续重试。
+type Handler func(ctx context.Context, body []byte) error
+
+// JSONHandler 是 RabbitMQ JSON 消息的业务处理函数。
+type JSONHandler[T any] func(ctx context.Context, message T) error
+
+var (
+	conn       *amqp.Connection
+	connMu     sync.Mutex
+	consumers  []*managedConsumer
+	consumerMu sync.Mutex
+	closeMu    sync.Mutex
+)
+
+func GetConfig(ctx context.Context) Config {
+	return Config{
+		Enabled: g.Cfg().MustGet(ctx, "rabbitmq.enabled", false).Bool(),
+		Dsn:     g.Cfg().MustGet(ctx, "rabbitmq.dsn", "").String(),
+	}
+}
+
+// Enabled 判断 RabbitMQ 是否启用且 DSN 是否已配置。
+func Enabled(ctx context.Context) bool {
+	cfg := GetConfig(ctx)
+	return cfg.Enabled && cfg.Dsn != ""
+}
+
+// Conn 懒加载 RabbitMQ 连接，供多个生产者/消费者复用。
+func Conn(ctx context.Context) (*amqp.Connection, error) {
+	cfg := GetConfig(ctx)
+	if !cfg.Enabled {
+		return nil, gerror.New("rabbitmq is disabled")
+	}
+	if cfg.Dsn == "" {
+		return nil, gerror.New("rabbitmq dsn is empty")
+	}
+
+	connMu.Lock()
+	defer connMu.Unlock()
+	if conn != nil && !conn.IsClosed() {
+		return conn, nil
+	}
+
+	c, err := amqp.Dial(cfg.Dsn)
+	if err != nil {
+		return nil, gerror.Wrap(err, "rabbitmq connect failed")
+	}
+	conn = c
+	return conn, nil
+}
+
+// Channel 基于共享连接创建独立 channel。
+func Channel(ctx context.Context) (*amqp.Channel, error) {
+	c, err := Conn(ctx)
+	if err != nil {
+		return nil, err
+	}
+	ch, err := c.Channel()
+	if err != nil {
+		return nil, gerror.Wrap(err, "rabbitmq open channel failed")
+	}
+	return ch, nil
+}
+
+// DeclareQueue 声明一个持久化队列。
+func DeclareQueue(ctx context.Context, queue string) error {
+	ch, err := Channel(ctx)
+	if err != nil {
+		return err
+	}
+	defer ch.Close()
+	return declareQueue(ch, queue)
+}
+
+// Publish 将结构体序列化为 JSON 后发布到指定队列。
+func Publish(ctx context.Context, queue string, data any) error {
+	body, err := json.Marshal(data)
+	if err != nil {
+		return gerror.Wrap(err, "rabbitmq marshal message failed")
+	}
+	return PublishBytes(ctx, queue, body)
+}
+
+// PublishDelayed 将结构体序列化为 JSON 后发布到延迟队列。
+// 消息在 delayQueue 中等待 delay 后，通过死信路由进入 targetQueue。
+func PublishDelayed(ctx context.Context, delayQueue, targetQueue string, delay time.Duration, data any) error {
+	body, err := json.Marshal(data)
+	if err != nil {
+		return gerror.Wrap(err, "rabbitmq marshal delayed message failed")
+	}
+	return PublishDelayedBytes(ctx, delayQueue, targetQueue, delay, body)
+}
+
+// PublishBytes 发布原始消息体到指定队列。
+func PublishBytes(ctx context.Context, queue string, body []byte) error {
+	ch, err := Channel(ctx)
+	if err != nil {
+		return err
+	}
+	defer ch.Close()
+
+	// 发布前也声明一次队列，避免生产者先启动时队列不存在。
+	if err = declareQueue(ch, queue); err != nil {
+		return err
+	}
+	return ch.PublishWithContext(ctx, "", queue, false, false, amqp.Publishing{
+		ContentType:  "application/json",
+		DeliveryMode: amqp.Persistent,
+		Timestamp:    time.Now(),
+		Body:         body,
+	})
+}
+
+// PublishDelayedBytes 发布原始消息体到延迟队列。
+func PublishDelayedBytes(ctx context.Context, delayQueue, targetQueue string, delay time.Duration, body []byte) error {
+	if delay <= 0 {
+		return gerror.New("rabbitmq delayed message delay must be positive")
+	}
+	ch, err := Channel(ctx)
+	if err != nil {
+		return err
+	}
+	defer ch.Close()
+
+	if err = declareQueue(ch, targetQueue); err != nil {
+		return err
+	}
+	if err = declareDelayQueue(ch, delayQueue, targetQueue); err != nil {
+		return err
+	}
+	return ch.PublishWithContext(ctx, "", delayQueue, false, false, amqp.Publishing{
+		ContentType:  "application/json",
+		DeliveryMode: amqp.Persistent,
+		Expiration:   fmt.Sprintf("%d", delay.Milliseconds()),
+		Timestamp:    time.Now(),
+		Body:         body,
+	})
+}
+
+// PublishDelayedBytesWithHeaders 发布带消息头的原始消息体到延迟队列。
+func PublishDelayedBytesWithHeaders(ctx context.Context, delayQueue, targetQueue string, delay time.Duration, body []byte, headers amqp.Table) error {
+	if delay <= 0 {
+		return gerror.New("rabbitmq delayed message delay must be positive")
+	}
+	ch, err := Channel(ctx)
+	if err != nil {
+		return err
+	}
+	defer ch.Close()
+
+	if err = declareQueue(ch, targetQueue); err != nil {
+		return err
+	}
+	if err = declareDelayQueue(ch, delayQueue, targetQueue); err != nil {
+		return err
+	}
+	return publishWithConfirm(ctx, ch, delayQueue, amqp.Publishing{
+		ContentType:  "application/json",
+		DeliveryMode: amqp.Persistent,
+		Expiration:   fmt.Sprintf("%d", delay.Milliseconds()),
+		Timestamp:    time.Now(),
+		Headers:      headers,
+		Body:         body,
+	})
+}
+
+// PublishBytesWithHeaders 发布带消息头的原始消息体到指定队列。
+func PublishBytesWithHeaders(ctx context.Context, queue string, body []byte, headers amqp.Table) error {
+	ch, err := Channel(ctx)
+	if err != nil {
+		return err
+	}
+	defer ch.Close()
+
+	if err = declareQueue(ch, queue); err != nil {
+		return err
+	}
+	return publishWithConfirm(ctx, ch, queue, amqp.Publishing{
+		ContentType:  "application/json",
+		DeliveryMode: amqp.Persistent,
+		Timestamp:    time.Now(),
+		Headers:      headers,
+		Body:         body,
+	})
+}
+
+func publishWithConfirm(ctx context.Context, ch *amqp.Channel, queue string, message amqp.Publishing) error {
+	if err := ch.Confirm(false); err != nil {
+		return gerror.Wrap(err, "rabbitmq enable publisher confirm failed")
+	}
+	confirms := ch.NotifyPublish(make(chan amqp.Confirmation, 1))
+	if err := ch.PublishWithContext(ctx, "", queue, false, false, message); err != nil {
+		return gerror.Wrapf(err, "rabbitmq publish failed, queue:%s", queue)
+	}
+	select {
+	case confirmation, ok := <-confirms:
+		if !ok {
+			return gerror.Newf("rabbitmq publisher confirm channel closed, queue:%s", queue)
+		}
+		if !confirmation.Ack {
+			return gerror.Newf("rabbitmq message was not confirmed, queue:%s", queue)
+		}
+		return nil
+	case <-ctx.Done():
+		return gerror.Wrapf(ctx.Err(), "wait rabbitmq publisher confirm failed, queue:%s", queue)
+	}
+}
+
+// Consume 启动消费者，使用手动 Ack，handler 成功才确认消息。
+func Consume(ctx context.Context, queue, consumer string, handler Handler) error {
+	_, err := ConsumeWithDone(ctx, queue, consumer, handler)
+	return err
+}
+
+func ConsumeWithDone(ctx context.Context, queue, consumer string, handler Handler) (<-chan struct{}, error) {
+	ch, err := Channel(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if err = declareQueue(ch, queue); err != nil {
+		_ = ch.Close()
+		return nil, err
+	}
+	if err = ch.Qos(1, 0, false); err != nil {
+		_ = ch.Close()
+		return nil, gerror.Wrap(err, "rabbitmq set qos failed")
+	}
+
+	// autoAck=false，业务处理成功后手动 Ack，处理失败时 Nack 并重新入队。
+	deliveries, err := ch.Consume(queue, consumer, false, false, false, false, nil)
+	if err != nil {
+		_ = ch.Close()
+		return nil, gerror.Wrap(err, "rabbitmq consume failed")
+	}
+
+	consumerState := &managedConsumer{
+		channel: ch,
+		tag:     consumer,
+		done:    make(chan struct{}),
+	}
+	registerConsumer(consumerState)
+	go func() {
+		defer close(consumerState.done)
+		defer ch.Close()
+		defer unregisterConsumer(consumerState)
+		for {
+			if ctx.Err() != nil {
+				return
+			}
+			select {
+			case <-ctx.Done():
+				return
+			case delivery, ok := <-deliveries:
+				if !ok {
+					return
+				}
+				consumerTag := delivery.ConsumerTag
+				if consumerTag == "" {
+					consumerTag = consumer
+				}
+				deliveryCtx := WithDeliveryMetadata(ctx, DeliveryMetadata{
+					DeliveryTag: delivery.DeliveryTag,
+					Redelivered: delivery.Redelivered,
+					Consumer:    consumerTag,
+					Queue:       queue,
+					Headers:     cloneHeaders(delivery.Headers),
+				})
+				if handlerErr := handler(deliveryCtx, delivery.Body); handlerErr != nil {
+					g.Log().Errorf(deliveryCtx, "rabbitmq message Nack, queue:%s, consumer:%s, deliveryTag:%d, redelivered:%t, err:%+v", queue, consumerTag, delivery.DeliveryTag, delivery.Redelivered, handlerErr)
+					_ = delivery.Nack(false, true)
+					continue
+				}
+				// handler 内部已经完成业务处理或重新发布下一次重试消息，这里只确认当前消息。
+				g.Log().Infof(deliveryCtx, "rabbitmq message Ack, queue:%s, consumer:%s, deliveryTag:%d, redelivered:%t", queue, consumerTag, delivery.DeliveryTag, delivery.Redelivered)
+				_ = delivery.Ack(false)
+			}
+		}
+	}()
+	return consumerState.done, nil
+}
+
+// ConsumeJSON 启动 JSON 消费者，统一完成反序列化和 Ack/Nack 处理。
+func ConsumeJSON[T any](ctx context.Context, queue, consumer string, handler JSONHandler[T]) error {
+	return Consume(ctx, queue, consumer, func(ctx context.Context, body []byte) error {
+		var message T
+		if err := json.Unmarshal(body, &message); err != nil {
+			g.Log().Errorf(ctx, "rabbitmq message dropped, invalid json, queue:%s, body:%s, err:%+v", queue, string(body), err)
+			return nil
+		}
+		return handler(ctx, message)
+	})
+}
+
+// Close 关闭消费者 channel 和共享连接。
+func Close(ctx context.Context) error {
+	closeMu.Lock()
+	defer closeMu.Unlock()
+
+	consumerMu.Lock()
+	registeredConsumers := consumers
+	consumers = nil
+	consumerMu.Unlock()
+
+	var closeErrors []error
+	for _, consumer := range registeredConsumers {
+		if consumer.channel.IsClosed() {
+			continue
+		}
+		if err := consumer.channel.Cancel(consumer.tag, false); !isClosedError(err) {
+			closeErrors = append(closeErrors, gerror.Wrapf(err, "cancel RabbitMQ consumer failed, consumer:%s", consumer.tag))
+		}
+	}
+
+	if err := waitForConsumers(ctx, registeredConsumers); err != nil {
+		closeErrors = append(closeErrors, err)
+	}
+
+	for _, consumer := range registeredConsumers {
+		if consumer.channel.IsClosed() {
+			continue
+		}
+		if err := consumer.channel.Close(); !isClosedError(err) {
+			closeErrors = append(closeErrors, gerror.Wrapf(err, "close RabbitMQ channel failed, consumer:%s", consumer.tag))
+		}
+	}
+
+	if err := closeConnection(); err != nil {
+		closeErrors = append(closeErrors, err)
+	}
+	return errors.Join(closeErrors...)
+}
+
+// declareQueue 统一声明持久化队列，生产者和消费者都会调用，保证队列存在。
+func declareQueue(ch *amqp.Channel, queue string) error {
+	_, err := ch.QueueDeclare(queue, true, false, false, false, nil)
+	if err != nil {
+		return gerror.Wrapf(err, "rabbitmq declare queue failed, queue:%s", queue)
+	}
+	return nil
+}
+
+// declareDelayQueue 声明 TTL 延迟队列，到期后死信转发到目标队列。
+func declareDelayQueue(ch *amqp.Channel, delayQueue, targetQueue string) error {
+	_, err := ch.QueueDeclare(delayQueue, true, false, false, false, amqp.Table{
+		"x-dead-letter-exchange":    "",
+		"x-dead-letter-routing-key": targetQueue,
+	})
+	if err != nil {
+		return gerror.Wrapf(err, "rabbitmq declare delay queue failed, queue:%s", delayQueue)
+	}
+	return nil
+}
+
+// registerCloser 记录消费者 channel，服务退出时统一关闭。
+func registerConsumer(consumer *managedConsumer) {
+	consumerMu.Lock()
+	defer consumerMu.Unlock()
+	consumers = append(consumers, consumer)
+}
+
+func waitForConsumers(ctx context.Context, registeredConsumers []*managedConsumer) error {
+	for _, consumer := range registeredConsumers {
+		select {
+		case <-consumer.done:
+		case <-ctx.Done():
+			return gerror.Wrap(ctx.Err(), "wait for RabbitMQ consumers failed")
+		}
+	}
+	return nil
+}
+
+func closeConnection() error {
+	connMu.Lock()
+	defer connMu.Unlock()
+	if conn == nil || conn.IsClosed() {
+		conn = nil
+		return nil
+	}
+	if err := conn.Close(); !isClosedError(err) {
+		return gerror.Wrap(err, "close RabbitMQ connection failed")
+	}
+	conn = nil
+	return nil
+}
+
+func cloneHeaders(headers amqp.Table) map[string]any {
+	if len(headers) == 0 {
+		return nil
+	}
+	cloned := make(map[string]any, len(headers))
+	for key, value := range headers {
+		cloned[key] = value
+	}
+	return cloned
+}
+
+func isClosedError(err error) bool {
+	return err == nil || errors.Is(err, amqp.ErrClosed)
+}
+
+func unregisterConsumer(target *managedConsumer) {
+	consumerMu.Lock()
+	defer consumerMu.Unlock()
+	for index, consumer := range consumers {
+		if consumer == target {
+			remaining := make([]*managedConsumer, 0, len(consumers)-1)
+			remaining = append(remaining, consumers[:index]...)
+			consumers = append(remaining, consumers[index+1:]...)
+			return
+		}
+	}
+}
@@ -0,0 +1,170 @@
+package rabbitmq
+
+import (
+	"context"
+	"reflect"
+	"sync"
+	"testing"
+)
+
+type fakeChannel struct {
+	mu          sync.Mutex
+	cancelCalls int
+	closeCalls  int
+	closed      bool
+}
+
+func (f *fakeChannel) Cancel(string, bool) error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.cancelCalls++
+	return nil
+}
+
+func (f *fakeChannel) Close() error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.closeCalls++
+	f.closed = true
+	return nil
+}
+
+func (f *fakeChannel) IsClosed() bool {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return f.closed
+}
+
+func TestCloseIsIdempotent(t *testing.T) {
+	resetRabbitMQState(t)
+
+	channel := &fakeChannel{}
+	done := make(chan struct{})
+	close(done)
+	registerConsumer(&managedConsumer{
+		channel: channel,
+		tag:     "test-consumer",
+		done:    done,
+	})
+
+	if err := Close(context.Background()); err != nil {
+		t.Fatalf("first Close() error = %v", err)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	if err := Close(ctx); err != nil {
+		t.Fatalf("second Close() error = %v", err)
+	}
+
+	channel.mu.Lock()
+	defer channel.mu.Unlock()
+	if channel.cancelCalls != 1 {
+		t.Fatalf("Cancel() calls = %d, want 1", channel.cancelCalls)
+	}
+	if channel.closeCalls != 1 {
+		t.Fatalf("Close() calls = %d, want 1", channel.closeCalls)
+	}
+}
+
+func TestCloseSkipsAlreadyClosedChannel(t *testing.T) {
+	resetRabbitMQState(t)
+
+	channel := &fakeChannel{closed: true}
+	done := make(chan struct{})
+	close(done)
+	registerConsumer(&managedConsumer{
+		channel: channel,
+		tag:     "closed-consumer",
+		done:    done,
+	})
+
+	if err := Close(context.Background()); err != nil {
+		t.Fatalf("Close() error = %v", err)
+	}
+
+	channel.mu.Lock()
+	defer channel.mu.Unlock()
+	if channel.cancelCalls != 0 || channel.closeCalls != 0 {
+		t.Fatalf("closed channel calls: Cancel=%d Close=%d, want 0", channel.cancelCalls, channel.closeCalls)
+	}
+}
+
+func TestCloseHonorsContextWhileWaiting(t *testing.T) {
+	resetRabbitMQState(t)
+
+	channel := &fakeChannel{}
+	registerConsumer(&managedConsumer{
+		channel: channel,
+		tag:     "blocked-consumer",
+		done:    make(chan struct{}),
+	})
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	if err := Close(ctx); err == nil {
+		t.Fatal("Close() error = nil, want context cancellation error")
+	}
+
+	channel.mu.Lock()
+	defer channel.mu.Unlock()
+	if channel.closeCalls != 1 {
+		t.Fatalf("Close() calls = %d, want forced cleanup", channel.closeCalls)
+	}
+}
+
+func TestDeliveryMetadataContext(t *testing.T) {
+	expected := DeliveryMetadata{
+		DeliveryTag: 42,
+		Redelivered: true,
+		Consumer:    "batch-consumer",
+		Queue:       "iotlcm.batch.execute",
+	}
+	ctx := WithDeliveryMetadata(context.Background(), expected)
+	actual, ok := DeliveryMetadataFromContext(ctx)
+	if !ok {
+		t.Fatal("delivery metadata not found in context")
+	}
+	if !reflect.DeepEqual(actual, expected) {
+		t.Fatalf("delivery metadata = %+v, want %+v", actual, expected)
+	}
+	if _, ok = DeliveryMetadataFromContext(context.Background()); ok {
+		t.Fatal("unexpected delivery metadata in empty context")
+	}
+}
+
+func resetRabbitMQState(t *testing.T) {
+	t.Helper()
+	consumerMu.Lock()
+	consumers = nil
+	consumerMu.Unlock()
+	connMu.Lock()
+	conn = nil
+	connMu.Unlock()
+	t.Cleanup(func() {
+		consumerMu.Lock()
+		consumers = nil
+		consumerMu.Unlock()
+		connMu.Lock()
+		conn = nil
+		connMu.Unlock()
+	})
+}
+
+func TestUnregisterConsumerPreservesShutdownSnapshot(t *testing.T) {
+	consumerMu.Lock()
+	original := consumers
+	first, second, third := &managedConsumer{}, &managedConsumer{}, &managedConsumer{}
+	consumers = []*managedConsumer{first, second, third}
+	snapshot := consumers
+	consumerMu.Unlock()
+	defer func() { consumerMu.Lock(); consumers = original; consumerMu.Unlock() }()
+	unregisterConsumer(second)
+	consumerMu.Lock()
+	defer consumerMu.Unlock()
+	if len(consumers) != 2 || consumers[0] != first || consumers[1] != third {
+		t.Fatal("consumer not removed")
+	}
+	if snapshot[1] != second || snapshot[2] != third {
+		t.Fatal("shutdown snapshot mutated")
+	}
+}