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

message/redis.go

@@ -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
 }