127 lines
4.0 KiB
Go
127 lines
4.0 KiB
Go
package redisdao
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"time"
|
||
|
||
"git.lzcwy.cn/zhengchiliu/zt-closet-backend/global"
|
||
"git.lzcwy.cn/zhengchiliu/zt-closet-backend/internal/redismodel"
|
||
"github.com/redis/go-redis/v9"
|
||
)
|
||
|
||
// aiRecommendTaskTTL 任务快照保留时长，需覆盖前端轮询窗口。
|
||
const aiRecommendTaskTTL = 15 * time.Minute
|
||
|
||
// aiRecommendReleaseScript 仅当活跃位仍属于本次任务时才删除，
|
||
// 避免任务超时释放后又被后续任务占用，却被旧任务误删。
|
||
var aiRecommendReleaseScript = redis.NewScript(`
|
||
if redis.call('GET', KEYS[1]) == ARGV[1] then
|
||
return redis.call('DEL', KEYS[1])
|
||
else
|
||
return 0
|
||
end
|
||
`)
|
||
|
||
// AIRecommendTask AI 推荐异步任务快照。
|
||
type AIRecommendTask struct {
|
||
TaskID string `json:"task_id"`
|
||
UserID int64 `json:"user_id"`
|
||
Status int8 `json:"status"` // 1=进行中 2=成功 3=失败
|
||
Result json.RawMessage `json:"result,omitempty"`
|
||
ErrorCode int `json:"error_code,omitempty"`
|
||
ErrorMsg string `json:"error_msg,omitempty"`
|
||
CreateTime int64 `json:"create_time"`
|
||
}
|
||
|
||
// AIRecommendTaskAcquire 抢占用户的 AI 推荐活跃位。
|
||
// - 抢占成功：(true, "", nil)；
|
||
// - 已有任务在执行：(false, 进行中的 task_id, nil)，调用方应把它当作幂等结果返回；
|
||
// - Redis 异常：(false, "", err)。
|
||
func AIRecommendTaskAcquire(ctx context.Context, userID int64, taskID string, ttl time.Duration) (bool, string, error) {
|
||
ctx, span := global.Tracer.Start(ctx, "redisdao.AIRecommendTaskAcquire")
|
||
defer span.End()
|
||
|
||
key := fmt.Sprintf(redismodel.KeyLockAIRecommend, userID)
|
||
|
||
ok, err := global.RedisEngine.SetNX(ctx, key, taskID, ttl).Result()
|
||
if err != nil {
|
||
span.RecordError(err)
|
||
return false, "", err
|
||
}
|
||
if ok {
|
||
return true, "", nil
|
||
}
|
||
|
||
// 已有任务在跑：读回它的 task_id，让前端继续轮询同一个任务。
|
||
running, gerr := global.RedisEngine.Get(ctx, key).Result()
|
||
if gerr == redis.Nil {
|
||
// 活跃位正好在两次调用之间过期，交由上层按「稍后重试」处理。
|
||
return false, "", nil
|
||
}
|
||
if gerr != nil {
|
||
span.RecordError(gerr)
|
||
return false, "", gerr
|
||
}
|
||
return false, running, nil
|
||
}
|
||
|
||
// AIRecommendTaskRelease 释放活跃位，仅当值仍等于本次 task_id 时才删除。
|
||
// 失败仅由调用方记录，不阻断主流程。
|
||
func AIRecommendTaskRelease(ctx context.Context, userID int64, taskID string) error {
|
||
if global.RedisEngine == nil {
|
||
return nil
|
||
}
|
||
ctx, span := global.Tracer.Start(ctx, "redisdao.AIRecommendTaskRelease")
|
||
defer span.End()
|
||
|
||
key := fmt.Sprintf(redismodel.KeyLockAIRecommend, userID)
|
||
if err := aiRecommendReleaseScript.Run(ctx, global.RedisEngine, []string{key}, taskID).Err(); err != nil {
|
||
span.RecordError(err)
|
||
return err
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// AIRecommendTaskSave 写入任务快照（覆盖式），TTL 由 aiRecommendTaskTTL 决定。
|
||
func AIRecommendTaskSave(ctx context.Context, task *AIRecommendTask) error {
|
||
ctx, span := global.Tracer.Start(ctx, "redisdao.AIRecommendTaskSave")
|
||
defer span.End()
|
||
|
||
data, err := json.Marshal(task)
|
||
if err != nil {
|
||
span.RecordError(err)
|
||
return fmt.Errorf("marshal ai recommend task: %w", err)
|
||
}
|
||
key := fmt.Sprintf(redismodel.KeyAIRecommendTask, task.TaskID)
|
||
if err := global.RedisEngine.Set(ctx, key, data, aiRecommendTaskTTL).Err(); err != nil {
|
||
span.RecordError(err)
|
||
return err
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// AIRecommendTaskGet 读取任务快照，任务不存在或已过期返回 (nil, nil)。
|
||
func AIRecommendTaskGet(ctx context.Context, taskID string) (*AIRecommendTask, error) {
|
||
ctx, span := global.Tracer.Start(ctx, "redisdao.AIRecommendTaskGet")
|
||
defer span.End()
|
||
|
||
key := fmt.Sprintf(redismodel.KeyAIRecommendTask, taskID)
|
||
data, err := global.RedisEngine.Get(ctx, key).Bytes()
|
||
if err == redis.Nil {
|
||
return nil, nil
|
||
}
|
||
if err != nil {
|
||
span.RecordError(err)
|
||
return nil, err
|
||
}
|
||
|
||
var task AIRecommendTask
|
||
if err := json.Unmarshal(data, &task); err != nil {
|
||
span.RecordError(err)
|
||
return nil, fmt.Errorf("unmarshal ai recommend task: %w", err)
|
||
}
|
||
return &task, nil
|
||
}
|