175 lines
5.7 KiB
Go
175 lines
5.7 KiB
Go
package message
|
|
|
|
import (
|
|
"context"
|
|
"github.com/gogf/gf/v2/database/gredis"
|
|
|
|
"github.com/gogf/gf/v2/errors/gerror"
|
|
)
|
|
|
|
func GetRedisClientTest(name string) *gredis.Redis {
|
|
return getRedisClientTest(name)
|
|
}
|
|
|
|
// GetLock 获取分布式锁
|
|
func GetLock(ctx context.Context, key string, expireSeconds int64, fn func(ctx context.Context) error) (success bool, err error) {
|
|
return lock(ctx, key, expireSeconds, fn)
|
|
}
|
|
|
|
// MessageConfig 消息配置接口
|
|
type MessageConfig interface {
|
|
start(ctx context.Context) error
|
|
publish(ctx context.Context, data interface{}, options ...map[string]interface{}) (messageID string, err error)
|
|
}
|
|
|
|
// RedisMessageConfig Redis Stream 消息配置
|
|
type RedisMessageConfig struct {
|
|
StreamKey string // Stream 键名
|
|
GroupName string // 消费者组名称
|
|
ConsumerName string // 消费者名称
|
|
BatchSize int64 // 最大并发数(信号量容量)
|
|
AutoAck bool // ACK确认,true自动确认,false手动确认
|
|
HandleFunc func(ctx context.Context, message map[string]interface{}) error
|
|
}
|
|
|
|
func (r *RedisMessageConfig) start(ctx context.Context) error {
|
|
return readFromStream(ctx, QueueMessage{
|
|
StreamKey: r.StreamKey,
|
|
GroupName: r.GroupName,
|
|
ConsumerName: r.ConsumerName,
|
|
BatchSize: r.BatchSize,
|
|
AutoAck: r.AutoAck,
|
|
HandleFunc: r.HandleFunc,
|
|
})
|
|
}
|
|
|
|
func (r *RedisMessageConfig) publish(ctx context.Context, data interface{}, options ...map[string]interface{}) (messageID string, err error) {
|
|
return publishToRedis(ctx, r.StreamKey, data)
|
|
}
|
|
|
|
// RabbitMQMessageConfig RabbitMQ 消息配置
|
|
type RabbitMQMessageConfig struct {
|
|
Queue string // 队列名称
|
|
Exchange string // 交换器名称
|
|
RoutingKey string // 路由键
|
|
PrefetchCount int // QoS: 预取数量(并发控制)
|
|
WorkerCount int // worker 数量
|
|
ConsumerTag string // 消费者标签
|
|
HandleFunc func(ctx context.Context, message map[string]interface{}) error
|
|
}
|
|
|
|
func (r *RabbitMQMessageConfig) start(ctx context.Context) error {
|
|
return startRabbitMQConsumer(ctx, QueueMessage{
|
|
Queue: r.Queue,
|
|
Exchange: r.Exchange,
|
|
RoutingKey: r.RoutingKey,
|
|
PrefetchCount: r.PrefetchCount,
|
|
WorkerCount: r.WorkerCount,
|
|
ConsumerTag: r.ConsumerTag,
|
|
AutoAck: true,
|
|
HandleFunc: r.HandleFunc,
|
|
})
|
|
}
|
|
|
|
func (r *RabbitMQMessageConfig) publish(ctx context.Context, data interface{}, options ...map[string]interface{}) (messageID string, err error) {
|
|
opts := make(map[string]interface{})
|
|
if len(options) > 0 {
|
|
opts = options[0]
|
|
}
|
|
exchange := r.Exchange
|
|
routingKey := r.RoutingKey
|
|
delay := 0
|
|
|
|
if v, ok := opts["exchange"].(string); ok {
|
|
exchange = v
|
|
}
|
|
if v, ok := opts["routingKey"].(string); ok {
|
|
routingKey = v
|
|
}
|
|
if v, ok := opts["delay"].(int); ok {
|
|
delay = v
|
|
}
|
|
|
|
if delay > 0 {
|
|
return publishDelayedToRabbitMQ(ctx, exchange, routingKey, data, delay)
|
|
}
|
|
return publishToRabbitMQ(ctx, exchange, routingKey, data)
|
|
}
|
|
|
|
// QueueMessage 统一消息队列配置结构体(内部使用)
|
|
type QueueMessage struct {
|
|
// Redis Stream 配置
|
|
StreamKey string
|
|
GroupName string
|
|
ConsumerName string
|
|
BatchSize int64
|
|
AutoAck bool
|
|
HandleFunc func(ctx context.Context, message map[string]interface{}) error
|
|
|
|
// RabbitMQ 配置
|
|
Queue string
|
|
Exchange string
|
|
RoutingKey string
|
|
PrefetchCount int
|
|
WorkerCount int
|
|
ConsumerTag string
|
|
}
|
|
|
|
// StartConsumers 启动消息消费者(统一入口)
|
|
// 支持同时启动多个消费者,包括 Redis Stream 和 RabbitMQ
|
|
func StartConsumers(ctx context.Context, configs ...MessageConfig) error {
|
|
for _, cfg := range configs {
|
|
if err := cfg.start(ctx); err != nil {
|
|
return gerror.Wrap(err, "启动消费者失败")
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// PublishMessage 发布消息(统一入口)
|
|
// 根据配置类型选择发布到 Redis Stream 或 RabbitMQ
|
|
func PublishMessage(ctx context.Context, cfg MessageConfig, data interface{}, options ...map[string]interface{}) (messageID string, err error) {
|
|
return cfg.publish(ctx, data, options...)
|
|
}
|
|
|
|
// ========== Redis Stream 公共方法(方便迁移) ==========
|
|
|
|
// AddToStream 将消息添加到 Redis Stream
|
|
//func AddToStream(ctx context.Context, streamKey string, msg interface{}) (messageID string, err error) {
|
|
// return addToStream(ctx, streamKey, msg)
|
|
//}
|
|
|
|
// ReadFromStream 从 Redis Stream 读取消息(已废弃)
|
|
// 请使用 RedisMessageConfig.StartConsumers 启动消费者
|
|
// 此方法保留用于向后兼容,但实际不会返回消息(异步消费模式)
|
|
func ReadFromStream(ctx context.Context, streamKey, groupName, consumerName string, count, blockMs int64) ([]StreamMessage, error) {
|
|
return nil, gerror.New("ReadFromStream 已废弃,请使用 RedisMessageConfig.StartConsumers 启动消费者")
|
|
}
|
|
|
|
// AckMessage 确认 Redis Stream 消息
|
|
func AckMessage(ctx context.Context, streamKey, groupName string, messageIDs ...string) error {
|
|
return ackMessage(ctx, streamKey, groupName, messageIDs...)
|
|
}
|
|
|
|
// InitStreamGroup 初始化 Redis Stream 消费者组
|
|
func InitStreamGroup(ctx context.Context, streamKey, groupName string) error {
|
|
return initStreamGroup(ctx, streamKey, groupName)
|
|
}
|
|
|
|
// ========== RabbitMQ 公共方法(方便迁移) ==========
|
|
|
|
// InitRabbitMQ 初始化 RabbitMQ 连接
|
|
func InitRabbitMQ(ctx context.Context) error {
|
|
return initRabbitMQ(ctx)
|
|
}
|
|
|
|
// PublishToRabbitMQ 发布消息到 RabbitMQ
|
|
//func PublishToRabbitMQ(ctx context.Context, exchange, routingKey string, message interface{}) error {
|
|
// return publishToRabbitMQ(ctx, exchange, routingKey, message)
|
|
//}
|
|
|
|
// PublishDelayedToRabbitMQ 发布延时消息到 RabbitMQ
|
|
//func PublishDelayedToRabbitMQ(ctx context.Context, exchange, routingKey string, message interface{}, delaySeconds int) error {
|
|
// return publishDelayedToRabbitMQ(ctx, exchange, routingKey, message, delaySeconds)
|
|
//}
|