From c871c55f79234c84103514b32bde063711017836 Mon Sep 17 00:00:00 2001
From: keven1024 <keven@fudaoyuan.icu>
Date: Sat, 4 Apr 2026 16:18:19 +0800
Subject: [PATCH] refactor: enhance Redis lock utility by adding expiration
 duration support and improving context management

---
 pkg/utils/redis_lock.go | 45 ++++++++++++++++++-----------------------
 1 file changed, 20 insertions(+), 25 deletions(-)

diff --git a/pkg/utils/redis_lock.go b/pkg/utils/redis_lock.go
index 82646d5..0815859 100644
--- a/pkg/utils/redis_lock.go
+++ b/pkg/utils/redis_lock.go
@@ -51,37 +51,32 @@ func GetRedisLocker() rueidislock.Locker {
 	return redisLock
 }
 
-func baseLocker(ctx context.Context, key string) (context.Context, context.CancelFunc, error) {
-	locker := GetRedisLocker()
-	return locker.WithContext(ctx, key)
-}
-
-func Locker(key string) (context.CancelFunc, error) {
-	_, cancel, err := baseLocker(context.Background(), key)
-	if err != nil {
-		return nil, err
+func baseLocker(ctx context.Context, key string, expired time.Duration) (context.Context, context.CancelFunc, error) {
+	if expired <= 0 {
+		expired = defaultRedisLockValidity
 	}
-	return cancel, nil
+	locker := GetRedisLocker()
+	timeoutCtx, timeoutCancel := context.WithTimeout(ctx, expired)
+	lockCtx, lockCancel, err := locker.WithContext(timeoutCtx, key)
+	if err != nil {
+		return nil, nil, err
+	}
+	return lockCtx, func() {
+		lockCancel()
+		timeoutCancel()
+	}, nil
 }
 
-func WithLocker(ctx context.Context, key string, fn func(context.Context) error) error {
-	lockCtx, cancel, err := baseLocker(ctx, key)
+func Locker(key string, expired time.Duration) (context.CancelFunc, error) {
+	_, cancel, err := baseLocker(context.Background(), key, expired)
+	return cancel, err
+}
+
+func WithLocker(ctx context.Context, key string, expired time.Duration, fn func(context.Context) error) error {
+	lockCtx, cancel, err := baseLocker(ctx, key, expired)
 	if err != nil {
 		return err
 	}
 	defer cancel()
 	return fn(lockCtx)
 }
-
-func TryWithLocker(ctx context.Context, key string, fn func(context.Context) error) error {
-	locker := GetRedisLocker()
-	lockCtx, cancel, err := locker.TryWithContext(ctx, key)
-	if err != nil {
-		if errors.Is(err, rueidislock.ErrNotLocked) {
-			return ErrRedisLockNotAcquired
-		}
-		return err
-	}
-	defer cancel()
-	return fn(lockCtx)
-}