red-future/common
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 19 Wiki Activity

重构消息队列模块，统一Redis/RabbitMQ/NATS客户端接口，新增消息代理抽象层

Browse Source
This commit is contained in:
曲皓东
2026-01-29 13:55:32 +08:00
committed by 张斌
parent 65b231b88f
commit a8993de6d5
8 changed files with 844 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
message/message.go
View File

@@ -7,11 +7,6 @@ import (
"github.com/gogf/gf/v2/errors/gerror"
)
// GetRedisClient 获取 Redis 客户端（供外部使用）
func GetRedisClient() *gredis.Redis {
return getRedisClient()
}
func GetRedisClientTest(name string) *gredis.Redis {
return getRedisClientTest(name)
}

152
message/msg_queue.go Normal file
View File

@@ -0,0 +1,152 @@
package message
import (
"context"
"fmt"
)
type RedisConfig struct {
// Stream 名称
Stream string
// 消费者组名称
Group string
// 消费者名称
Consumer string
// 每次消费数量
Count int64
// 是否自动 ACK
AutoAck bool
// 处理函数
HandleFunc func(ctx context.Context, message map[string]interface{}) error
}
// RabbitMQConfig RabbitMQ 队列配置
type RabbitMQConfig struct {
Mode string
Exchange string
Topic string
DelayMessage bool
// 队列名称（必需）
Name string
// 实际队列名（用于绑定）
Queue string
// 是否持久化
Durable bool
// QoS 预取数量（每次推送的消息数量，默认10）
PrefetchCount int
// 最大重试次数（默认3）
MaxRetry int
// 是否自动 ACK
AutoAck bool
// 处理函数
HandleFunc func(ctx context.Context, message map[string]interface{}) error
}
// NATSConfig NATS 队列配置
type NATSConfig struct {
DelayMessage bool
// Stream 名称
Stream string
// 消费者名称
Consumer string
// 是否持久化
Durable bool
// 副本数
Replicas int
// QoS 预取数量（每次推送的消息数量，默认10）
PrefetchCount int
// 是否自动 ACK
AutoAck bool
// 处理函数
HandleFunc func(ctx context.Context, message map[string]interface{}) error
}
// messageBroker 消息代理接口
type messageBroker interface {
// StreamGroup 创建消费组（支持单个配置或批量配置）
streamGroup(ctx context.Context, configs ...interface{}) error
// Publish 发布消息（支持单个配置或批量配置）
publish(ctx context.Context, config interface{}, data interface{}) error
// PublishDelayed 发布延迟消息（支持单个配置或批量配置）
publishDelayed(ctx context.Context, config interface{}, data interface{}, delay int) error
// Subscribe 订阅消息（支持单个配置或批量配置）
subscribe(ctx context.Context, configs ...interface{}) error
}
type messageClientType string
const (
ClientTypeRedis messageClientType = "redis"
ClientTypeRabbitMQ messageClientType = "rabbitmq"
ClientTypeNATS messageClientType = "nats"
)
// newMessageBroker 创建消息代理实例
func newMessageBroker(ctx context.Context, clientType messageClientType) (messageBroker, error) {
switch clientType {
case ClientTypeRedis:
return &redisMessageClient{clientType: clientType}, nil
case ClientTypeRabbitMQ:
return &rabbitMQMessageClient{clientType: clientType}, nil
case ClientTypeNATS:
return &natsMessageClient{clientType: clientType}, nil
default:
return nil, fmt.Errorf("unknown client type: %s", clientType)
}
}
// StreamGroup 直接创建消费组
func StreamGroup(ctx context.Context, clientType messageClientType, configs ...interface{}) error {
broker, err := newMessageBroker(ctx, clientType)
if err != nil {
return err
}
return broker.streamGroup(ctx, configs...)
}
// Publish 直接发布消息
func Publish(ctx context.Context, clientType messageClientType, config interface{}, data interface{}) error {
broker, err := newMessageBroker(ctx, clientType)
if err != nil {
return err
}
return broker.publish(ctx, config, data)
}
// PublishDelayed 直接发布延迟消息
func PublishDelayed(ctx context.Context, clientType messageClientType, config interface{}, data interface{}, delay int) error {
broker, err := newMessageBroker(ctx, clientType)
if err != nil {
return err
}
return broker.publishDelayed(ctx, config, data, delay)
}
// Subscribe 直接订阅消息
func Subscribe(ctx context.Context, clientType messageClientType, configs ...interface{}) error {
broker, err := newMessageBroker(ctx, clientType)
if err != nil {
return err
}
return broker.subscribe(ctx, configs...)
}

93
nats/connection.go → message/nats_client.go
View File

@@ -1,4 +1,4 @@
package nats
package message
import (
"context"
@@ -16,7 +16,7 @@ var (
nc *nats.Conn
js jetstream.JetStream
inited bool
mu sync.RWMutex
natsMu sync.RWMutex
natsURL string
healthCtx context.Context
healthCancel context.CancelFunc
@@ -24,15 +24,15 @@ var (
reconnectChan chan struct{}
// 连接状态变化监听器
connStateListeners []ConnStateListener
connStateListeners []connStateListener
connListenersMu sync.RWMutex
// 监控指标
metrics Metrics
metrics metricsCounter
)
// Metrics 监控指标
type Metrics struct {
type metricsCounter struct {
PublishCount atomic.Int64
PublishError atomic.Int64
SubscribeCount atomic.Int64
@@ -43,33 +43,33 @@ type Metrics struct {
}
// ConnState 连接状态
type ConnState int
type connState int
const (
ConnStateDisconnected ConnState = iota
ConnStateConnecting
ConnStateConnected
ConnStateReconnecting
ConnStateClosed
connStateDisconnected connState = iota
connStateConnecting
connStateConnected
connStateReconnecting
connStateClosed
)
// ConnStateListener 连接状态监听器
type ConnStateListener func(state ConnState, err error)
type connStateListener func(state connState, err error)
// GetMetrics 获取监控指标
func GetMetrics() Metrics {
func getMetrics() metricsCounter {
return metrics
}
// RegisterConnStateListener 注册连接状态监听器
func RegisterConnStateListener(listener ConnStateListener) {
// registerConnStateListener 注册连接状态监听器
func registerConnStateListener(listener connStateListener) {
connListenersMu.Lock()
defer connListenersMu.Unlock()
connStateListeners = append(connStateListeners, listener)
}
// UnregisterConnStateListener 取消注册连接状态监听器
func UnregisterConnStateListener(listener ConnStateListener) {
// unregisterConnStateListener 取消注册连接状态监听器
func unregisterConnStateListener(listener connStateListener) {
connListenersMu.Lock()
defer connListenersMu.Unlock()
for i, l := range connStateListeners {
@@ -81,9 +81,9 @@ func UnregisterConnStateListener(listener ConnStateListener) {
}
// notifyConnState 通知所有监听器连接状态变化
func notifyConnState(state ConnState, err error) {
func notifyConnState(state connState, err error) {
connListenersMu.RLock()
listeners := make([]ConnStateListener, len(connStateListeners))
listeners := make([]connStateListener, len(connStateListeners))
copy(listeners, connStateListeners)
connListenersMu.RUnlock()
@@ -119,17 +119,17 @@ func init() {
// initConnection 初始化连接
func initConnection() {
ctx := context.Background()
notifyConnState(ConnStateConnecting, nil)
notifyConnState(connStateConnecting, nil)
if err := connect(ctx); err != nil {
g.Log().Errorf(ctx, "NATS 初始连接失败: %v", err)
notifyConnState(ConnStateDisconnected, err)
notifyConnState(connStateDisconnected, err)
}
}
// connect 建立 NATS 连接
func connect(ctx context.Context) error {
mu.Lock()
defer mu.Unlock()
natsMu.Lock()
defer natsMu.Unlock()
if nc != nil && !nc.IsClosed() {
nc.Close()
@@ -152,7 +152,7 @@ func connect(ctx context.Context) error {
}
// 通知重连成功
notifyConnState(ConnStateConnected, nil)
notifyConnState(connStateConnected, nil)
// 使用非阻塞发送避免阻塞
select {
@@ -164,12 +164,12 @@ func connect(ctx context.Context) error {
nats.DisconnectErrHandler(func(nc *nats.Conn, err error) {
g.Log().Warningf(ctx, "⚠️ NATS 连接断开: %v, 准备重连...", err)
connected = false
notifyConnState(ConnStateReconnecting, err)
notifyConnState(connStateReconnecting, err)
}),
nats.ClosedHandler(func(nc *nats.Conn) {
g.Log().Infof(ctx, "NATS 连接已关闭: %s", nc.ConnectedUrl())
connected = false
notifyConnState(ConnStateClosed, nil)
notifyConnState(connStateClosed, nil)
}),
nats.ErrorHandler(func(nc *nats.Conn, sub *nats.Subscription, err error) {
g.Log().Errorf(ctx, "NATS 错误: %v", err)
@@ -186,7 +186,7 @@ func connect(ctx context.Context) error {
if nc.Status() != nats.CONNECTED {
select {
case <-time.After(5 * time.Second):
notifyConnState(ConnStateDisconnected, fmt.Errorf("连接超时"))
notifyConnState(connStateDisconnected, fmt.Errorf("连接超时"))
return fmt.Errorf("NATS 连接超时")
case <-nc.StatusChanged(nats.CONNECTED):
}
@@ -201,7 +201,7 @@ func connect(ctx context.Context) error {
connected = true
inited = true
g.Log().Infof(ctx, "✅ NATS 连接成功: %s", nc.ConnectedUrl())
notifyConnState(ConnStateConnected, nil)
notifyConnState(connStateConnected, nil)
return nil
}
@@ -215,10 +215,10 @@ func healthCheck() {
case <-healthCtx.Done():
return
case <-ticker.C:
mu.RLock()
natsMu.RLock()
currentConnected := connected
currentConn := nc
mu.RUnlock()
natsMu.RUnlock()
if !currentConnected || currentConn == nil || currentConn.IsClosed() {
// 仅记录日志，不尝试重连（NATS 已有自动重连机制）
@@ -233,38 +233,33 @@ func healthCheck() {
// checkConnected 检查连接状态
func checkConnected() bool {
mu.RLock()
defer mu.RUnlock()
natsMu.RLock()
defer natsMu.RUnlock()
return connected && nc != nil && !nc.IsClosed()
}
// IsConnected 检查 NATS 是否已连接
func IsConnected() bool {
return checkConnected()
}
// GetConnState 获取当前连接状态
func GetConnState() ConnState {
mu.RLock()
defer mu.RUnlock()
// getConnState 获取当前连接状态
func getConnState() connState {
natsMu.RLock()
defer natsMu.RUnlock()
if nc == nil {
return ConnStateDisconnected
return connStateDisconnected
}
if nc.IsClosed() {
return ConnStateClosed
return connStateClosed
}
if connected {
return ConnStateConnected
return connStateConnected
}
return ConnStateDisconnected
return connStateDisconnected
}
// Shutdown 优雅关闭：自动注销所有已注册的服务并关闭 NATS 连接
func Shutdown() error {
// shutdown 优雅关闭：自动注销所有已注册的服务并关闭 NATS 连接
func shutdown() error {
ctx := context.Background()
g.Log().Info(ctx, "开始优雅关闭 NATS RPC 服务...")
@@ -299,8 +294,8 @@ func Shutdown() error {
g.Log().Infof(ctx, "已注销 %d 个单实例服务和 %d 个队列服务", singleServiceCount, queueServiceCount)
mu.Lock()
defer mu.Unlock()
natsMu.Lock()
defer natsMu.Unlock()
// 停止健康检查协程
if healthCancel != nil {

93
nats/nats_rpc.go → message/nats_rpc.go
View File

@@ -1,4 +1,4 @@
package nats
package message
import (
"context"
@@ -7,6 +7,7 @@ import (
"fmt"
"github.com/gogf/gf/v2/frame/g"
"github.com/nats-io/nats.go"
"go.opentelemetry.io/otel/trace"
"reflect"
"sync"
)
@@ -509,12 +510,12 @@ func AutoRegisterServices(ctx context.Context, serviceInstances map[string]inter
return fmt.Errorf("未能注册任何服务")
}
// 设置取消监听器（监听基于 TraceID 的取消请求）
//if _, err := setupCancelListener(ctx); err != nil {
// g.Log().Errorf(ctx, "设置取消监听器失败: %v", err)
//} else {
// g.Log().Infof(ctx, "✅ 取消监听器已自动设置")
//}
//g.Log().Infof(ctx, "✅ 共自动注册了 %d 个服务", totalRegistered)
if _, err := setupCancelListener(ctx); err != nil {
g.Log().Errorf(ctx, "设置取消监听器失败: %v", err)
} else {
g.Log().Infof(ctx, "✅ 取消监听器已自动设置")
}
g.Log().Infof(ctx, "✅ 共自动注册了 %d 个服务", totalRegistered)
return nil
}
@@ -671,3 +672,81 @@ func registerService(service interface{}, serviceNamePrefix string, options ...R
g.Log().Infof(context.Background(), "✅ Service %v 共注册了 %d 个 RPC 方法", serviceNamePrefix, registeredCount)
return nil
}
// ============ 上下文元数据工具函数 ============
// 以下函数用于在 context 和 NATS 消息头之间互转元数据
// 定义常见的上下文元数据 key
const (
TraceIDKey = "trace_id"
TokenKey = "token"
)
func getTraceID(ctx context.Context) (traceID string, err error) {
// 提取 traceId：首先尝试从 OpenTelemetry Span 中提取,从 context 中提取 TraceID
span := trace.SpanFromContext(ctx)
if span != nil && span.SpanContext().HasTraceID() {
traceID = span.SpanContext().TraceID().String()
} else if tid := ctx.Value(TraceIDKey); tid != nil {
traceID = fmt.Sprintf("%v", tid)
}
if traceID == "" {
return traceID, fmt.Errorf("context 中没有 TraceID")
}
return
}
// contextToHeaders 将 context 中的元数据转换为 NATS 消息头
// 支持提取 user_id、tenant_id、trace_id、token 等常见字段
func contextToHeaders(ctx context.Context) (nats.Header, error) {
headers := make(nats.Header)
// 提取 traceId：首先尝试从 OpenTelemetry Span 中提取
if traceID, err := getTraceID(ctx); err != nil {
return headers, err
} else {
headers.Set(TraceIDKey, traceID)
}
// 提取 token（优先级：context value > HTTP Authorization header）
token := ""
if t := ctx.Value(TokenKey); t != nil {
token = fmt.Sprintf("%v", t)
} else if r := g.RequestFromCtx(ctx); r != nil {
// 从 HTTP 请求的 Authorization header 中提取 token
auth := r.GetHeader("Authorization")
if auth != "" {
// 移除 "Bearer " 前缀
if len(auth) > 7 && auth[:7] == "Bearer " {
token = auth[7:]
} else {
token = auth
}
}
}
if token != "" {
headers.Set(TokenKey, token)
}
return headers, nil
}
// headersToContext 从 NATS 消息头重建 context
// 支持还原 user_id、tenant_id、trace_id、token 等字段
func headersToContext(ctx context.Context, headers nats.Header) context.Context {
if headers == nil {
return ctx
}
// 恢复 trace_id
if traceID := headers.Get(TraceIDKey); traceID != "" {
ctx = context.WithValue(ctx, TraceIDKey, traceID)
}
// 恢复 token
if token := headers.Get(TokenKey); token != "" {
ctx = context.WithValue(ctx, TokenKey, token)
}
return ctx
}

6
message/rabbit.go
View File

@@ -22,7 +22,7 @@ var (
)
// Config RabbitMQ 配置
type RabbitMQConfig struct {
type RabbitMQConfig1 struct {
Host string
Port int
Username string
@@ -31,8 +31,8 @@ type RabbitMQConfig struct {
}
// rabbitMQConfig 默认配置
func getRabbitMQConfig() *RabbitMQConfig {
return &RabbitMQConfig{
func getRabbitMQConfig() *RabbitMQConfig1 {
return &RabbitMQConfig1{
Host: g.Cfg().MustGet(context.Background(), "rabbitmq.host").String(),
Port: g.Cfg().MustGet(context.Background(), "rabbitmq.port").Int(),
Username: g.Cfg().MustGet(context.Background(), "rabbitmq.username").String(),

28
rabbitmq/client.go → message/rabbitmq_client.go
View File

@@ -1,4 +1,4 @@
package rabbitmq
package message
import (
"context"
@@ -14,8 +14,8 @@ import (
var (
conn *amqp.Connection
channel *amqp.Channel
once sync.Once
mu sync.RWMutex
rabbitmqOnce sync.Once
rabbitmqMu sync.RWMutex
closeWatcher chan struct{} // 用于停止监听 goroutine
watcherStarted bool // 防止重复启动监听
)
@@ -32,7 +32,7 @@ type Config struct {
// Init 初始化 RabbitMQ 连接
func Init(ctx context.Context, cfg *Config) error {
var err error
once.Do(func() {
rabbitmqOnce.Do(func() {
// 构建连接字符串
url := "amqp://" + cfg.Username + ":" + cfg.Password + "@" + cfg.Host + ":" + gconv.String(cfg.Port) + "/" + cfg.VHost
@@ -80,8 +80,8 @@ func InitFromConfig(ctx context.Context) error {
// GetChannel 获取 Channel
func GetChannel() (*amqp.Channel, error) {
mu.RLock()
defer mu.RUnlock()
rabbitmqMu.RLock()
defer rabbitmqMu.RUnlock()
if channel == nil || channel.IsClosed() {
return nil, gerror.New("RabbitMQ Channel 未初始化或已关闭")
@@ -92,8 +92,8 @@ func GetChannel() (*amqp.Channel, error) {
// GetConnection 获取连接
func GetConnection() (*amqp.Connection, error) {
mu.RLock()
defer mu.RUnlock()
rabbitmqMu.RLock()
defer rabbitmqMu.RUnlock()
if conn == nil || conn.IsClosed() {
return nil, gerror.New("RabbitMQ 连接未初始化或已关闭")
@@ -113,9 +113,9 @@ func handleConnectionClose(ctx context.Context) {
default:
}
mu.RLock()
rabbitmqMu.RLock()
currentConn := conn
mu.RUnlock()
rabbitmqMu.RUnlock()
if currentConn == nil {
return
@@ -141,8 +141,8 @@ func handleConnectionClose(ctx context.Context) {
// reconnect 重新连接
func reconnect(ctx context.Context) {
mu.Lock()
defer mu.Unlock()
rabbitmqMu.Lock()
defer rabbitmqMu.Unlock()
for i := 0; i < 10; i++ {
time.Sleep(time.Duration(i+1) * time.Second)
@@ -180,8 +180,8 @@ func reconnect(ctx context.Context) {
// Close 关闭连接
func Close(ctx context.Context) (err error) {
mu.Lock()
defer mu.Unlock()
rabbitmqMu.Lock()
defer rabbitmqMu.Unlock()
// 停止监听 goroutine
if closeWatcher != nil {

111
message/redis.go
View File

@@ -3,6 +3,7 @@ package message
import (
"context"
"errors"
"fmt"
"strings"
"time"
@@ -18,11 +19,6 @@ type StreamMessage struct {
Values map[string]interface{} // 消息内容
}
// getClient 获取 Redis 客户端
func getRedisClient() *gredis.Redis {
return g.Redis()
}
// getClient 获取 Redis 客户端
func getRedisClientTest(name string) *gredis.Redis {
return g.Redis(name)
@@ -47,46 +43,66 @@ func getRedisClientByDB(db int) *gredis.Redis {
// lock 分布式锁
func lock(ctx context.Context, key string, expireSeconds int64, fn func(ctx context.Context) error) (success bool, err error) {
limit := 3
LOOP:
if limit < 0 {
return false, errors.New("锁重试次数耗尽")
ds, err := GetManager().GetDefaultDataSource()
if err != nil {
return false, fmt.Errorf("获取默认数据源失败: %w", err)
}
limit--
if val, err := getRedisClient().Set(ctx, key, true, gredis.SetOption{
TTLOption: gredis.TTLOption{
EX: &expireSeconds,
},
NX: true,
}); err != nil {
return false, err
} else {
if val.Bool() {
defer func(RedisClient *gredis.Redis, ctx context.Context, key string) {
if _, err = RedisClient.Del(ctx, key); err != nil {
glog.Errorf(ctx, "RedisClient.Del error: %v", err)
}
}(getRedisClient(), ctx, key)
if err = fn(ctx); err != nil {
return false, err
}
return true, nil
maxRetries := 3
for i := 0; i < maxRetries; i++ {
if val, err := ds.Redis().Set(ctx, key, true, gredis.SetOption{
TTLOption: gredis.TTLOption{
EX: &expireSeconds,
},
NX: true,
}); err != nil {
return false, err
} else {
time.Sleep(time.Second)
goto LOOP
if val.Bool() {
defer func(redisClient *gredis.Redis, ctx context.Context, key string) {
if _, err = redisClient.Del(ctx, key); err != nil {
glog.Errorf(ctx, "RedisClient.Del error: %v", err)
}
}(ds.Redis(), ctx, key)
if err = fn(ctx); err != nil {
return false, err
}
return true, nil
} else {
// 检查上下文是否已取消
if ctx.Err() != nil {
return false, ctx.Err()
}
// 非最后一次重试时才等待
if i < maxRetries-1 {
time.Sleep(time.Second)
}
}
}
}
return false, errors.New("锁重试次数耗尽")
}
// publishToRedis 将消息添加到 Redis Stream
func publishToRedis(ctx context.Context, streamKey string, msg interface{}) (messageID string, err error) {
ds, err := GetManager().GetDefaultDataSource()
if err != nil {
return "", fmt.Errorf("获取默认数据源失败: %w", err)
}
if !ds.IsConnected() {
if err := ds.Reconnect(ctx); err != nil {
return "", fmt.Errorf("redis重连失败: %w", err)
}
}
values := gconv.Map(msg)
args := make([]interface{}, 0, len(values)*2+2)
args = append(args, streamKey, "*")
for key, val := range values {
args = append(args, key, val)
}
result, err := getRedisClient().Do(ctx, "XADD", args...)
result, err := ds.Redis().Do(ctx, "XADD", args...)
if err != nil {
return
}
@@ -96,7 +112,18 @@ func publishToRedis(ctx context.Context, streamKey string, msg interface{}) (mes
// initStreamGroup 初始化 Stream 和消费者组
func initStreamGroup(ctx context.Context, streamKey, groupName string) error {
_, err := getRedisClient().Do(ctx, "XGROUP", "CREATE", streamKey, groupName, "0", "MKSTREAM")
ds, err := GetManager().GetDefaultDataSource()
if err != nil {
return fmt.Errorf("获取默认数据源失败: %w", err)
}
if !ds.IsConnected() {
if err := ds.Reconnect(ctx); err != nil {
return fmt.Errorf("redis重连失败: %w", err)
}
}
_, err = ds.Redis().Do(ctx, "XGROUP", "CREATE", streamKey, groupName, "0", "MKSTREAM")
if err != nil {
// 如果组已存在，忽略错误
errStr := err.Error()
@@ -113,6 +140,11 @@ func initStreamGroup(ctx context.Context, streamKey, groupName string) error {
// readFromStream 从 Stream 读取消息
func readFromStream(ctx context.Context, msg QueueMessage) error {
ds, err := GetManager().GetDefaultDataSource()
if err != nil {
return fmt.Errorf("获取默认数据源失败: %w", err)
}
// 初始化 Stream 和消费者组
if err := initStreamGroup(ctx, msg.StreamKey, msg.GroupName); err != nil {
return err
@@ -120,7 +152,7 @@ func readFromStream(ctx context.Context, msg QueueMessage) error {
go func() {
RECONNECT:
for {
result, err := getRedisClient().Do(ctx, "XREADGROUP", "GROUP", msg.GroupName, msg.ConsumerName, "COUNT", msg.BatchSize, "BLOCK", 0, "STREAMS", msg.StreamKey, ">")
result, err := ds.Redis().Do(ctx, "XREADGROUP", "GROUP", msg.GroupName, msg.ConsumerName, "COUNT", msg.BatchSize, "BLOCK", 0, "STREAMS", msg.StreamKey, ">")
if err != nil {
//select {
//case <-ctx.Done():
@@ -222,11 +254,22 @@ func readFromStream(ctx context.Context, msg QueueMessage) error {
// ackMessage 确认消息已处理
func ackMessage(ctx context.Context, streamKey, groupName string, messageIDs ...string) error {
ds, err := GetManager().GetDefaultDataSource()
if err != nil {
return fmt.Errorf("获取默认数据源失败: %w", err)
}
if !ds.IsConnected() {
if err := ds.Reconnect(ctx); err != nil {
return fmt.Errorf("redis重连失败: %w", err)
}
}
args := make([]interface{}, 0, len(messageIDs)+2)
args = append(args, streamKey, groupName)
for _, id := range messageIDs {
args = append(args, id)
}
_, err := getRedisClient().Do(ctx, "XACK", args...)
_, err = ds.Redis().Do(ctx, "XACK", args...)
return err
}

468
message/redis_client.go Normal file
View File

@@ -0,0 +1,468 @@
// =============================================================================
// Redis 数据源连接管理
// 使用 GoFrame 框架自带的 Redis 客户端,负责数据源的连接、重连、健康检查和优雅关闭
// =============================================================================
package message
import (
"context"
"fmt"
"os"
"os/signal"
"sync"
"sync/atomic"
"syscall"
"time"
"github.com/gogf/gf/v2/database/gredis"
"github.com/gogf/gf/v2/frame/g"
"github.com/gogf/gf/v2/os/glog"
"github.com/gogf/gf/v2/util/gconv"
)
// =============================================================================
// 数据源配置结构
// =============================================================================
type RedisDataSourceConfig struct {
Name string `json:"name"` // 数据源名称
Address string `json:"address"` // Redis 地址,如: 127.0.0.1:6379
Db int `json:"db"` // 数据库编号
Pass string `json:"pass"` // 密码
Timeout time.Duration `json:"timeout"` // 连接超时
MaxIdle int `json:"maxIdle"` // 最大空闲连接数
MaxOpen int `json:"maxOpen"` // 最大活跃连接数
}
// =============================================================================
// 单个数据源接口
// =============================================================================
type DataSource interface {
Name() string
Redis() *gredis.Redis
IsConnected() bool
Connect(ctx context.Context) error
Reconnect(ctx context.Context) error
Close(ctx context.Context) error
}
// =============================================================================
// 数据源实现
// =============================================================================
type BaseDataSource struct {
config *RedisDataSourceConfig
redis *gredis.Redis
isConnected bool
mu sync.RWMutex
lastError error
lastErrorTime time.Time
metrics RedisMetrics
}
func NewBaseDataSource(config *RedisDataSourceConfig) *BaseDataSource {
return &BaseDataSource{
config: config,
isConnected: false,
}
}
func (d *BaseDataSource) Name() string {
return d.config.Name
}
func (d *BaseDataSource) Redis() *gredis.Redis {
d.mu.RLock()
defer d.mu.RUnlock()
return d.redis
}
func (d *BaseDataSource) IsConnected() bool {
d.mu.RLock()
defer d.mu.RUnlock()
return d.isConnected && d.redis != nil
}
func (d *BaseDataSource) Connect(ctx context.Context) error {
d.mu.Lock()
defer d.mu.Unlock()
// 设置默认值
config := d.config
if config.Timeout == 0 {
config.Timeout = 10 * time.Second
}
if config.MaxIdle == 0 {
config.MaxIdle = 10
}
if config.MaxOpen == 0 {
config.MaxOpen = 100
}
// 构建 GoFrame Redis 配置
redisConfig := &gredis.Config{
Address: config.Address,
Db: config.Db,
Pass: config.Pass,
}
// 使用 GoFrame 的 Redis 连接
redisObj, err := gredis.New(redisConfig)
if err != nil {
d.isConnected = false
d.lastError = err
d.lastErrorTime = time.Now()
d.metrics.PingError.Add(1)
return fmt.Errorf("datasource [%s] connection failed: %w", d.config.Name, err)
}
d.redis = redisObj
// 测试连接
if err := d.Ping(ctx); err != nil {
d.isConnected = false
d.lastError = err
d.lastErrorTime = time.Now()
return fmt.Errorf("datasource [%s] ping failed: %w", d.config.Name, err)
}
d.isConnected = true
d.lastError = nil
glog.Infof(ctx, "✅ datasource [%s] connected successfully", d.config.Name)
return nil
}
func (d *BaseDataSource) Ping(ctx context.Context) error {
defer func() {
if r := recover(); r != nil {
d.metrics.PingError.Add(1)
glog.Errorf(ctx, "❌ datasource [%s] ping panic: %v", d.config.Name, r)
}
}()
if d.redis == nil {
d.metrics.PingError.Add(1)
return fmt.Errorf("redis client is nil")
}
_, err := d.redis.Do(ctx, "PING")
if err != nil {
d.metrics.PingError.Add(1)
return err
}
d.metrics.PingCount.Add(1)
return nil
}
func (d *BaseDataSource) Reconnect(ctx context.Context) error {
glog.Infof(ctx, "🔄 reconnecting datasource [%s]", d.config.Name)
return d.Connect(ctx)
}
func (d *BaseDataSource) Close(ctx context.Context) error {
d.mu.Lock()
defer d.mu.Unlock()
if d.redis != nil {
if err := d.redis.Close(ctx); err != nil {
return fmt.Errorf("datasource [%s] close failed: %w", d.config.Name, err)
}
}
d.isConnected = false
d.redis = nil
glog.Infof(ctx, "datasource [%s] closed", d.config.Name)
return nil
}
func (d *BaseDataSource) GetMetrics() RedisMetrics {
return d.metrics
}
// =============================================================================
// 监控指标
// =============================================================================
type RedisMetrics struct {
PingCount atomic.Int64
PingError atomic.Int64
CommandCount atomic.Int64
CommandError atomic.Int64
}
// GetPingMetrics 获取 Ping 相关指标
func (m *RedisMetrics) GetPingMetrics() (int64, int64) {
return m.PingCount.Load(), m.PingError.Load()
}
// GetCommandMetrics 获取命令相关指标
func (m *RedisMetrics) GetCommandMetrics() (int64, int64) {
return m.CommandCount.Load(), m.CommandError.Load()
}
// =============================================================================
// 多数据源管理器
// =============================================================================
type DataSourceManager struct {
sources map[string]DataSource
mu sync.RWMutex
ctx context.Context
cancel context.CancelFunc
started bool
maxRetries int
metrics RedisMetrics
}
var (
manager *DataSourceManager
once sync.Once
)
// GetManager 获取全局管理器
func GetManager() *DataSourceManager {
once.Do(func() {
ctx, cancel := context.WithCancel(context.Background())
manager = &DataSourceManager{
sources: make(map[string]DataSource),
ctx: ctx,
cancel: cancel,
started: false,
maxRetries: 3,
}
})
return manager
}
// RegisterDataSource 注册数据源
func (m *DataSourceManager) RegisterDataSource(config *RedisDataSourceConfig) error {
m.mu.Lock()
defer m.mu.Unlock()
if _, exists := m.sources[config.Name]; exists {
return fmt.Errorf("datasource [%s] already exists", config.Name)
}
source := NewBaseDataSource(config)
m.sources[config.Name] = source
return nil
}
// GetDataSource 获取数据源
func (m *DataSourceManager) GetDataSource(name string) (DataSource, error) {
m.mu.RLock()
defer m.mu.RUnlock()
source, exists := m.sources[name]
if !exists {
return nil, fmt.Errorf("datasource [%s] not found", name)
}
return source, nil
}
// GetAllDataSourceNames 获取所有数据源名称
func (m *DataSourceManager) GetAllDataSourceNames() []string {
m.mu.RLock()
defer m.mu.RUnlock()
names := make([]string, 0, len(m.sources))
for name := range m.sources {
names = append(names, name)
}
return names
}
// GetDefaultDataSource 获取默认数据源(第一个注册的数据源)
func (m *DataSourceManager) GetDefaultDataSource() (DataSource, error) {
m.mu.RLock()
defer m.mu.RUnlock()
for _, source := range m.sources {
return source, nil
}
return nil, fmt.Errorf("no datasource available")
}
// GetMetrics 获取全局监控指标
func (m *DataSourceManager) GetMetrics() RedisMetrics {
return m.metrics
}
// init 初始化多数据源
func init() {
ctx := context.Background()
// 从配置初始化多数据源
if err := GetManager().InitializeFromConfig(ctx); err != nil {
glog.Errorf(ctx, "❌ Failed to initialize Redis datasources: %v", err)
} else {
glog.Infof(ctx, "✅ Redis datasources initialized: %v", GetManager().GetAllDataSourceNames())
}
// 启动健康检查
GetManager().StartHealthCheck()
// 设置优雅关闭
setupGracefulShutdown()
}
// InitializeFromConfig 从配置初始化数据源
// 动态读取 config.yml 中 redis 下的所有配置项
func (m *DataSourceManager) InitializeFromConfig(ctx context.Context) error {
var firstErr error
// 获取 redis 配置下的所有子键
redisConfig := g.Cfg().MustGet(ctx, "redis")
if redisConfig.IsNil() {
glog.Warningf(ctx, "no redis configuration found in config.yml")
return nil
}
// 将配置转换为 map
configMap := redisConfig.Map()
if configMap == nil {
glog.Warningf(ctx, "redis configuration is not a map")
return nil
}
// 遍历所有 redis 子配置
for name, subConfig := range configMap {
// 跳过非对象类型的配置
subMap, ok := subConfig.(map[string]interface{})
if !ok {
continue
}
// 检查是否有 address 配置
address, hasAddress := subMap["address"]
if !hasAddress || gconv.String(address) == "" {
continue
}
// 构建数据源配置
config := &RedisDataSourceConfig{
Name: name,
Address: gconv.String(address),
Db: gconv.Int(subMap["db"]),
Pass: gconv.String(subMap["pass"]),
}
// 设置默认值
if config.Db == 0 {
config.Db = 0
}
if config.Timeout == 0 {
config.Timeout = 10 * time.Second
}
if config.MaxIdle == 0 {
config.MaxIdle = 10
}
if config.MaxOpen == 0 {
config.MaxOpen = 100
}
// 注册数据源
if err := m.RegisterDataSource(config); err != nil {
glog.Errorf(ctx, "failed to register datasource [%s]: %v", name, err)
if firstErr == nil {
firstErr = err
}
continue
}
// 连接数据源
source, _ := m.GetDataSource(name)
if err := source.Connect(ctx); err != nil {
glog.Errorf(ctx, "failed to initialize datasource [%s]: %v", name, err)
if firstErr == nil {
firstErr = err
}
}
}
return firstErr
}
// StartHealthCheck 启动健康检查
func (m *DataSourceManager) StartHealthCheck() {
if m.started {
return
}
m.started = true
go m.healthCheckLoop()
}
// healthCheckLoop 健康检查循环
func (m *DataSourceManager) healthCheckLoop() {
ticker := time.NewTicker(30 * time.Second)
defer ticker.Stop()
for {
select {
case <-m.ctx.Done():
return
case <-ticker.C:
m.checkAndReconnect()
}
}
}
// checkAndReconnect 检查并重新连接
func (m *DataSourceManager) checkAndReconnect() {
m.mu.RLock()
defer m.mu.RUnlock()
for name, source := range m.sources {
if !source.IsConnected() {
glog.Warningf(context.Background(), "datasource [%s] disconnected, attempting reconnect", name)
reconnectCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
if err := source.Reconnect(reconnectCtx); err != nil {
glog.Errorf(reconnectCtx, "datasource [%s] reconnect failed: %v", name, err)
} else {
glog.Infof(reconnectCtx, "✅ datasource [%s] reconnected successfully", name)
}
}
}
}
// CloseAll 关闭所有数据源
func (m *DataSourceManager) CloseAll(ctx context.Context) error {
m.cancel()
m.mu.RLock()
defer m.mu.RUnlock()
var lastErr error
for name, source := range m.sources {
if err := source.Close(ctx); err != nil {
glog.Errorf(ctx, "failed to close datasource [%s]: %v", name, err)
lastErr = err
}
}
return lastErr
}
// setupGracefulShutdown 设置优雅关闭
func setupGracefulShutdown() {
go func() {
sigCh := make(chan os.Signal, 1)
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
<-sigCh
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
glog.Info(ctx, "🔄 Shutting down Redis connections...")
if err := GetManager().CloseAll(ctx); err != nil {
glog.Errorf(ctx, "❌ Failed to close Redis connections: %v", err)
} else {
glog.Info(ctx, "✅ Redis connections closed successfully")
}
}()
}