101 lines
2.7 KiB
Go
101 lines
2.7 KiB
Go
package cache
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strconv"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
const (
|
|
FormulaRecognitionTaskCount = "formula_recognition_task"
|
|
FormulaRecognitionTaskQueue = "formula_recognition_queue"
|
|
FormulaRecognitionDistLock = "formula_recognition_dist_lock"
|
|
VLMFormulaCount = "vlm_formula_count:%s" // VLM公式识别次数 ip
|
|
VLMRecognitionTaskQueue = "vlm_recognition_queue"
|
|
DefaultLockTimeout = 60 * time.Second // 默认锁超时时间
|
|
)
|
|
|
|
// TODO the sigle queue not reliable, message maybe lost
|
|
func PushVLMRecognitionTask(ctx context.Context, taskID int64) (count int64, err error) {
|
|
count, err = RedisClient.LPush(ctx, VLMRecognitionTaskQueue, taskID).Result()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return count, nil
|
|
}
|
|
|
|
func PopVLMRecognitionTask(ctx context.Context) (int64, error) {
|
|
result, err := RedisClient.BRPop(ctx, 0, VLMRecognitionTaskQueue).Result()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return strconv.ParseInt(result[1], 10, 64)
|
|
}
|
|
|
|
func PushFormulaTask(ctx context.Context, taskID int64) (count int64, err error) {
|
|
count, err = RedisClient.LPush(ctx, FormulaRecognitionTaskQueue, taskID).Result()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return count, nil
|
|
}
|
|
|
|
func PopFormulaTask(ctx context.Context) (int64, error) {
|
|
result, err := RedisClient.BRPop(ctx, 0, FormulaRecognitionTaskQueue).Result()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return strconv.ParseInt(result[1], 10, 64)
|
|
}
|
|
|
|
func GetFormulaTaskCount(ctx context.Context) (int64, error) {
|
|
count, err := RedisClient.LLen(ctx, FormulaRecognitionTaskQueue).Result()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return count, nil
|
|
}
|
|
|
|
// GetDistributedLock 获取分布式锁
|
|
func GetDistributedLock(ctx context.Context) (bool, error) {
|
|
return RedisClient.SetNX(ctx, FormulaRecognitionDistLock, "locked", DefaultLockTimeout).Result()
|
|
}
|
|
|
|
// ReleaseLock 释放分布式锁
|
|
func ReleaseLock(ctx context.Context) error {
|
|
return RedisClient.Del(ctx, FormulaRecognitionDistLock).Err()
|
|
}
|
|
|
|
func GetVLMFormulaCount(ctx context.Context, ip string) (int64, error) {
|
|
count, err := RedisClient.Get(ctx, fmt.Sprintf(VLMFormulaCount, ip)).Result()
|
|
if err != nil {
|
|
if err == redis.Nil {
|
|
return 0, nil
|
|
}
|
|
return 0, err
|
|
}
|
|
return strconv.ParseInt(count, 10, 64)
|
|
}
|
|
|
|
func IncrVLMFormulaCount(ctx context.Context, ip string) (int64, error) {
|
|
key := fmt.Sprintf(VLMFormulaCount, ip)
|
|
count, err := RedisClient.Incr(ctx, key).Result()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
if count == 1 {
|
|
now := time.Now()
|
|
nextMidnight := time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location())
|
|
ttl := nextMidnight.Sub(now)
|
|
if err := RedisClient.Expire(ctx, key, ttl).Err(); err != nil {
|
|
return count, err
|
|
}
|
|
}
|
|
|
|
return count, nil
|
|
}
|