From d6880dbf00cdb9db4bc815ec969fac4d6d47c856 Mon Sep 17 00:00:00 2001
From: keven1024 <keven@fudaoyuan.icu>
Date: Sat, 4 Apr 2026 21:20:17 +0800
Subject: [PATCH] refactor: improve context management in Redis stat updates by
 utilizing context in locking mechanism

---
 pkg/models/stat.go | 33 ++++++++++++++++++---------------
 1 file changed, 18 insertions(+), 15 deletions(-)

diff --git a/pkg/models/stat.go b/pkg/models/stat.go
index 6eab3dc..3dbc75d 100644
--- a/pkg/models/stat.go
+++ b/pkg/models/stat.go
@@ -1,6 +1,7 @@
 package models
 
 import (
+	"context"
 	"encoding/json"
 
 	"pkg/utils"
@@ -33,22 +34,24 @@ func GetRedisStat(key string) (*StatData, error) {
 }
 
 func SetRedisStat(key string, handler func(stat *StatData) *StatData) error {
-	rdb, ctx := utils.GetRedisClient()
-	old_stat, err := GetRedisStat(key)
-	if err != nil {
-		return err
-	}
-	if old_stat == nil {
-		old_stat = &StatData{
-			FileSize:    0,
-			FileNum:     0,
-			ShareNum:    0,
-			DownloadNum: 0,
+	return utils.WithLocker(context.Background(), "015:stat:"+key, 0, func(ctx context.Context) error {
+		rdb, _ := utils.GetRedisClient()
+		old_stat, err := GetRedisStat(key)
+		if err != nil {
+			return err
 		}
-	}
-	stat := handler(old_stat)
-	jsonData, _ := json.Marshal(stat)
-	return rdb.Do(ctx, rdb.B().Hset().Key("015:stat").FieldValue().FieldValue(key, string(jsonData)).Build()).Error()
+		if old_stat == nil {
+			old_stat = &StatData{
+				FileSize:    0,
+				FileNum:     0,
+				ShareNum:    0,
+				DownloadNum: 0,
+			}
+		}
+		stat := handler(old_stat)
+		jsonData, _ := json.Marshal(stat)
+		return rdb.Do(ctx, rdb.B().Hset().Key("015:stat").FieldValue().FieldValue(key, string(jsonData)).Build()).Error()
+	})
 }
 
 func GetRedisStatAll() (map[string]string, error) {