129 lines
3.6 KiB
Go
129 lines
3.6 KiB
Go
package tools
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"go-bot/config"
|
||
"log"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/go-redis/redis/v8"
|
||
)
|
||
|
||
var (
|
||
rdb *redis.Client
|
||
once sync.Once
|
||
ctx = context.Background()
|
||
)
|
||
|
||
// InitRedis 初始化 Redis 客户端（懒加载模式，确保只初始化一次）
|
||
func InitRedis() *redis.Client {
|
||
// 使用 sync.Once 确保客户端只初始化一次
|
||
once.Do(func() {
|
||
redisConfig := config.GetConfig()["REDIS"].(map[string]interface{})
|
||
rdb = redis.NewClient(&redis.Options{
|
||
Addr: redisConfig["REDIS_URL"].(string), // Redis 地址
|
||
Password: redisConfig["REDIS_PASSWORD"].(string), // Redis 密码，如果有设置
|
||
DB: int(redisConfig["REDIS_DB"].(int64)), // 使用的 Redis 数据库
|
||
PoolSize: 20, // 连接池大小
|
||
MinIdleConns: 10, // 最小空闲连接数
|
||
})
|
||
|
||
// 测试连接是否成功
|
||
_, err := rdb.Ping(ctx).Result()
|
||
if err != nil {
|
||
log.Println("Redis 连接失败:", err)
|
||
} else {
|
||
log.Println("Redis 连接成功")
|
||
}
|
||
})
|
||
return rdb
|
||
}
|
||
|
||
// GetRedisClient 获取 Redis 客户端实例
|
||
func GetRedisClient() *redis.Client {
|
||
if rdb == nil {
|
||
return InitRedis()
|
||
}
|
||
return rdb
|
||
}
|
||
|
||
// 添加对话到上下文，超过5条则删除所有上下文
|
||
func AddToContext(key string, message []byte, contextlenth int64) {
|
||
// key := fmt.Sprintf("context:%s:%s:%s", worker, groupID, qqID)
|
||
// 检查上下文持续时间，如果大于一小时就删除所有列表
|
||
lastRequestKey := fmt.Sprintf("last_request:%s", key)
|
||
now := time.Now().Unix()
|
||
exists, err := CheckKeyExists(lastRequestKey)
|
||
if err != nil || !exists {
|
||
log.Println("检查键是否存在时出错:", err)
|
||
rdb.Del(ctx, key)
|
||
rdb.Set(ctx, lastRequestKey, fmt.Sprintf("%d", now), 60*time.Minute)
|
||
}
|
||
|
||
// 如果上下文超过5条，删除所有上下文
|
||
listLength := GetListLength(key)
|
||
// log.Println("listLength:", listLength)
|
||
if contextlenth > 0 && listLength >= contextlenth*2 {
|
||
rdb.Del(ctx, key) // 删除该用户的所有上下文
|
||
}
|
||
rdb.RPush(ctx, key, message) // 添加新消息到列表
|
||
}
|
||
|
||
// 检查请求频率，10秒内只能请求一次
|
||
func CheckRequestFrequency(GID string, UID string, interval int64) bool {
|
||
key := fmt.Sprintf("time_interval:%s:%s", GID, UID)
|
||
// lastRequest, err := rdb.Get(ctx, key).Int64()
|
||
exists, err := CheckKeyExists(key)
|
||
if err != nil || !exists {
|
||
// log.Println("检查键是否存在时出错:", err)
|
||
|
||
rdb.Set(ctx, key, fmt.Sprintf("%d", time.Now().Unix()), time.Duration(interval)*time.Second)
|
||
return true
|
||
}
|
||
|
||
return false // 频率超限，拒绝请求
|
||
|
||
}
|
||
|
||
// getListLength 获取列表长度
|
||
func GetListLength(key string) int64 {
|
||
return rdb.LLen(ctx, key).Val()
|
||
}
|
||
|
||
func GetListValue(key string, index int64) (string, error) {
|
||
return rdb.LIndex(ctx, key, index).Result()
|
||
}
|
||
|
||
// 删除列表最后一个元素
|
||
func RemoveLastValueFromList(key string) (string, error) {
|
||
// RPOP 命令会移除并返回列表的最后一个元素
|
||
value, err := rdb.RPop(ctx, key).Result()
|
||
if err != nil {
|
||
if err == redis.Nil {
|
||
return "", fmt.Errorf("列表为空或不存在")
|
||
}
|
||
return "", fmt.Errorf("删除最后一个值时出错: %v", err)
|
||
}
|
||
return value, nil
|
||
}
|
||
|
||
// SetValue 设置键值对
|
||
func SetValue(key string, value string) error {
|
||
return rdb.Set(ctx, key, value, 0).Err()
|
||
}
|
||
|
||
// GetValue 获取键对应的值
|
||
func GetValue(key string) (string, error) {
|
||
return rdb.Get(ctx, key).Result()
|
||
}
|
||
|
||
// CheckKeyExists 检查键是否存在
|
||
func CheckKeyExists(key string) (bool, error) {
|
||
val, err := rdb.Exists(ctx, key).Result()
|
||
if err != nil {
|
||
return false, err
|
||
}
|
||
return val > 0, nil
|
||
}
|