init repo

internal/storage/cache/engine.go

@@ -0,0 +1,33 @@
+package cache
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"gitea.com/bitwsd/document_ai/config"
+	"github.com/redis/go-redis/v9"
+)
+
+var RedisClient *redis.Client
+
+func InitRedisClient(config config.RedisConfig) {
+	fmt.Println("Initializing Redis client...")
+	RedisClient = redis.NewClient(&redis.Options{
+		Addr:         config.Addr,
+		Password:     config.Password,
+		DB:           config.DB,
+		DialTimeout:  10 * time.Second,
+		ReadTimeout:  10 * time.Second,
+		WriteTimeout: 10 * time.Second,
+	})
+
+	fmt.Println("Pinging Redis server...")
+	_, err := RedisClient.Ping(context.Background()).Result()
+	if err != nil {
+		fmt.Printf("Init redis client failed, err: %v\n", err)
+		panic(err)
+	}
+
+	fmt.Println("Redis client initialized successfully.")
+}

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

internal/storage/cache/url.go

@@ -0,0 +1,12 @@
+package cache
+
+import "context"
+
+func IncrURLCount(ctx context.Context) (int64, error) {
+	key := "formula_recognition:url_count"
+	count, err := RedisClient.Incr(ctx, key).Result()
+	if err != nil {
+		return 0, err
+	}
+	return count, nil
+}

internal/storage/cache/user.go

@@ -0,0 +1,63 @@
+package cache
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strconv"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+)
+
+const (
+	UserSmsCodeTTL        = 10 * time.Minute
+	UserSendSmsLimitTTL   = 24 * time.Hour
+	UserSendSmsLimitCount = 5
+)
+
+const (
+	UserSmsCodePrefix = "user:sms_code:%s"
+	UserSendSmsLimit  = "user:send_sms_limit:%s"
+)
+
+func GetUserSmsCode(ctx context.Context, phone string) (string, error) {
+	code, err := RedisClient.Get(ctx, fmt.Sprintf(UserSmsCodePrefix, phone)).Result()
+	if err != nil {
+		if err == redis.Nil {
+			return "", nil
+		}
+		return "", err
+	}
+	return code, nil
+}
+
+func SetUserSmsCode(ctx context.Context, phone, code string) error {
+	return RedisClient.Set(ctx, fmt.Sprintf(UserSmsCodePrefix, phone), code, UserSmsCodeTTL).Err()
+}
+
+func GetUserSendSmsLimit(ctx context.Context, phone string) (int, error) {
+	limit, err := RedisClient.Get(ctx, fmt.Sprintf(UserSendSmsLimit, phone)).Result()
+	if err != nil {
+		if err == redis.Nil {
+			return 0, nil
+		}
+		return 0, err
+	}
+	return strconv.Atoi(limit)
+}
+
+func SetUserSendSmsLimit(ctx context.Context, phone string) error {
+	count, err := RedisClient.Incr(ctx, fmt.Sprintf(UserSendSmsLimit, phone)).Result()
+	if err != nil {
+		return err
+	}
+	if count > UserSendSmsLimitCount {
+		return errors.New("send sms limit")
+	}
+	return RedisClient.Expire(ctx, fmt.Sprintf(UserSendSmsLimit, phone), UserSendSmsLimitTTL).Err()
+}
+
+func DeleteUserSmsCode(ctx context.Context, phone string) error {
+	return RedisClient.Del(ctx, fmt.Sprintf(UserSmsCodePrefix, phone)).Err()
+}