init repo

internal/storage/cache/formula.go

@@ -0,0 +1,100 @@
+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
+}