140 lines
3.0 KiB
Go
140 lines
3.0 KiB
Go
package main
|
||
|
||
import (
|
||
"encoding/json"
|
||
"fmt"
|
||
"time"
|
||
|
||
"github.com/redis/go-redis/v9"
|
||
"golang.org/x/net/context"
|
||
)
|
||
|
||
var ctx = context.Background()
|
||
|
||
type UserBenefits struct {
|
||
client *redis.Client
|
||
}
|
||
|
||
// NewUserBenefits 初始化 UserBenefits 实例
|
||
func NewUserBenefits(redisClient *redis.Client) *UserBenefits {
|
||
return &UserBenefits{
|
||
client: redisClient,
|
||
}
|
||
}
|
||
|
||
// CheckAndDecrement 检查用户是否可以访问并递减剩余次数，0表示可用，其他值表示异常
|
||
func (ub *UserBenefits) CheckAndDecrement(ID int, timeZone string, secondsFromGMT int) (int, error) {
|
||
loc, err := time.LoadLocation(timeZone)
|
||
if err != nil {
|
||
// 使用 secondsFromGMT 计算时区
|
||
loc = time.FixedZone("Custom", secondsFromGMT)
|
||
}
|
||
currentTime := time.Now().In(loc)
|
||
dateKey := currentTime.Format("20060102")
|
||
|
||
key := fmt.Sprintf("%d_%s", ID, dateKey)
|
||
|
||
// 获取当前计数
|
||
val, err := ub.client.Get(ctx, key).Int()
|
||
if err == redis.Nil {
|
||
// 如果没有设置，初始化为3
|
||
ub.client.Set(ctx, key, DAILY_FREE_COUNT, 24*time.Hour)
|
||
val = DAILY_FREE_COUNT
|
||
} else if err != nil {
|
||
return -1, err
|
||
}
|
||
|
||
// 检查次数是否已用完
|
||
if val <= 0 {
|
||
return 1, nil
|
||
}
|
||
|
||
// 递减次数
|
||
ub.client.Decr(ctx, key)
|
||
|
||
return 0, nil
|
||
}
|
||
|
||
// QueryUserBenefits 查询用户使用情况
|
||
func (ub *UserBenefits) QueryUserBenefits(ID int) (string, error) {
|
||
keys, err := ub.client.Keys(ctx, fmt.Sprintf("%d_*", ID)).Result()
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
|
||
result := make(map[string]int)
|
||
for _, key := range keys {
|
||
val, err := ub.client.Get(ctx, key).Int()
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
result[key] = val
|
||
}
|
||
|
||
// 将结果编码为 JSON 格式
|
||
jsonData, err := json.Marshal(result)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
|
||
return string(jsonData), nil
|
||
}
|
||
|
||
// ResetUserBenefits 删除指定日期或所有日期的用户权益，只限于内部调用
|
||
func (ub *UserBenefits) ResetUserBenefits(ID int, dateStr string) error {
|
||
var pattern string
|
||
if dateStr == "" {
|
||
// 删除该用户的所有权益
|
||
pattern = fmt.Sprintf("%d_*", ID)
|
||
} else {
|
||
// 删除特定日期的权益
|
||
pattern = fmt.Sprintf("%d_%s", ID, dateStr)
|
||
}
|
||
|
||
keys, err := ub.client.Keys(ctx, pattern).Result()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
if len(keys) == 0 {
|
||
fmt.Println("No keys to delete.")
|
||
return nil
|
||
}
|
||
|
||
_, err = ub.client.Del(ctx, keys...).Result()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
fmt.Printf("Deleted keys for pattern %s\n", pattern)
|
||
return nil
|
||
}
|
||
|
||
func test_smain() {
|
||
// Redis 客户端使用连接池初始化
|
||
rdb := redis.NewClient(&redis.Options{
|
||
Addr: "172.18.0.4:6379",
|
||
Password: "", // no password set
|
||
DB: 0, // use default DB
|
||
PoolSize: 10, // 连接池大小
|
||
})
|
||
|
||
ub := NewUserBenefits(rdb)
|
||
|
||
// 示例使用
|
||
status, err := ub.CheckAndDecrement(10001, "Europe/Berlin", 0)
|
||
if err != nil {
|
||
fmt.Println("Error:", err)
|
||
} else {
|
||
fmt.Println("Status Code:", status)
|
||
}
|
||
|
||
// 查询用户使用情况
|
||
userData, err := ub.QueryUserBenefits(10001)
|
||
if err != nil {
|
||
fmt.Println("Error:", err)
|
||
} else {
|
||
fmt.Println("User Data:", userData)
|
||
}
|
||
}
|