fix cache lock context propagation

cache/lock.go

@@ -23,6 +23,8 @@ var (
 	ErrInvalidLockRetryInterval = errors.New("锁重试/续期间隔必须大于 0")
 	// ErrLockUnexpectedResult Lua 脚本返回了非预期的结果类型（C1b：裸类型断言防护）。
 	ErrLockUnexpectedResult = errors.New("锁脚本返回非预期结果")
+	// ErrLockFuncNil 表示分布式锁业务函数为空。
+	ErrLockFuncNil = errors.New("锁业务函数不能为空")
 )
 
 // toInt64 将 Lua 脚本返回值安全断言为 int64（C1b：禁止裸断言 panic）。
@@ -231,7 +233,10 @@ func TryLock(ctx context.Context, key string, ttl time.Duration, retryInterval t
 //
 // 解锁用独立 Background ctx（C1a 一致性修复）：fn 返回或 panic 后，原 ctx 可能已被
 // 调用方取消，用其解锁会失败导致锁泄漏到 TTL。fn panic 时 defer 也保证解锁执行。
-func WithLock(ctx context.Context, key string, ttl time.Duration, fn func() error) error {
+func WithLock(ctx context.Context, key string, ttl time.Duration, fn func(context.Context) error) error {
+	if fn == nil {
+		return ErrLockFuncNil
+	}
 	token, err := NewLock(ctx, key, ttl)
 	if err != nil {
 		return err
@@ -245,7 +250,7 @@ func WithLock(ctx context.Context, key string, ttl time.Duration, fn func() erro
 		_ = Unlock(unlockCtx, token)
 	}()
 
-	return fn()
+	return fn(ctx)
 }
 
 // WithLockAutoExtend 使用分布式锁执行函数（自动续期）。
@@ -256,7 +261,10 @@ func WithLock(ctx context.Context, key string, ttl time.Duration, fn func() erro
 // 避免旧实现 done 无缓冲 + 子 defer close(done) + 父 done<-struct{}{} 的 send-on-closed panic
 // （ctx 取消或 ExtendLock 失败时 done 已 closed，父再 send 即 panic，Unlock 不执行、锁泄漏到 TTL）。
 // Unlock 用 context.Background() 派生超时，避免原 ctx 已取消致 Unlock 失败再泄漏。
-func WithLockAutoExtend(ctx context.Context, key string, initialTTL time.Duration, extendInterval time.Duration, fn func() error) error {
+func WithLockAutoExtend(ctx context.Context, key string, initialTTL time.Duration, extendInterval time.Duration, fn func(context.Context) error) error {
+	if fn == nil {
+		return ErrLockFuncNil
+	}
 	if extendInterval <= 0 {
 		return ErrInvalidLockRetryInterval
 	}
@@ -302,8 +310,8 @@ func WithLockAutoExtend(ctx context.Context, key string, initialTTL time.Duratio
 		_ = Unlock(unlockCtx, token)
 	}()
 
-	// 执行业务函数。
-	err = fn()
+	// 执行业务函数。fn 必须接收 ctx，避免请求取消后业务 loader/DB/HTTP 调用继续运行。
+	err = fn(ctx)
 
 	return err
 }

cache/lock_concurrency_test.go

@@ -45,11 +45,11 @@ func TestWithLockAutoExtendCtxCancelNoPanic(t *testing.T) {
 	go func() {
 		defer wg.Done()
 		defer func() { panicked = recover() }()
-		err := cache.WithLockAutoExtend(ctx, key, 10*time.Second, 100*time.Millisecond, func() error {
+		err := cache.WithLockAutoExtend(ctx, key, 10*time.Second, 100*time.Millisecond, func(runCtx context.Context) error {
 			close(ran)
 			// 阻塞直到 ctx 被取消，模拟长任务。
-			<-ctx.Done()
-			return ctx.Err()
+			<-runCtx.Done()
+			return runCtx.Err()
 		})
 		// WithLockAutoExtend 返回 fn 的错误（ctx.Err），不应 panic。
 		if err != nil && !errors.Is(err, context.Canceled) {
@@ -83,7 +83,7 @@ func TestWithLockAutoExtendNormalRelease(t *testing.T) {
 	key := "c1a:normal"
 
 	called := false
-	err := cache.WithLockAutoExtend(ctx, key, 10*time.Second, 100*time.Millisecond, func() error {
+	err := cache.WithLockAutoExtend(ctx, key, 10*time.Second, 100*time.Millisecond, func(context.Context) error {
 		called = true
 		return nil
 	})
@@ -107,7 +107,7 @@ func TestWithLockAutoExtendExtendsLock(t *testing.T) {
 
 	// initialTTL=500ms，extendInterval=100ms。fn 执行 800ms，期间应多次续期。
 	// 若续期失效，锁会在 500ms 过期，另一 worker 可获取。
-	err := cache.WithLockAutoExtend(ctx, key, 500*time.Millisecond, 100*time.Millisecond, func() error {
+	err := cache.WithLockAutoExtend(ctx, key, 500*time.Millisecond, 100*time.Millisecond, func(context.Context) error {
 		time.Sleep(800 * time.Millisecond)
 		return nil
 	})
@@ -137,7 +137,7 @@ func TestWithLockAutoExtendBlocksContender(t *testing.T) {
 
 	go func() {
 		defer close(lockDone)
-		cache.WithLockAutoExtend(ctx, key, 2*time.Second, 100*time.Millisecond, func() error {
+		cache.WithLockAutoExtend(ctx, key, 2*time.Second, 100*time.Millisecond, func(context.Context) error {
 			close(fnStarted)
 			time.Sleep(600 * time.Millisecond)
 			close(fnDone)
@@ -173,7 +173,7 @@ func TestWithLockAutoExtendFnPanicReleasesLock(t *testing.T) {
 	var panicked any
 	func() {
 		defer func() { panicked = recover() }()
-		_ = cache.WithLockAutoExtend(ctx, key, 10*time.Second, 100*time.Millisecond, func() error {
+		_ = cache.WithLockAutoExtend(ctx, key, 10*time.Second, 100*time.Millisecond, func(context.Context) error {
 			time.Sleep(150 * time.Millisecond) // 让续期 ticker 至少触发一次
 			panic("boom")
 		})
@@ -201,10 +201,10 @@ func TestWithLockCtxCancelReleasesLock(t *testing.T) {
 	done := make(chan struct{})
 	go func() {
 		defer close(done)
-		_ = cache.WithLock(ctx, key, 10*time.Second, func() error {
+		_ = cache.WithLock(ctx, key, 10*time.Second, func(runCtx context.Context) error {
 			close(fnStarted)
-			<-ctx.Done()
-			return ctx.Err()
+			<-runCtx.Done()
+			return runCtx.Err()
 		})
 	}()
 
@@ -218,6 +218,39 @@ func TestWithLockCtxCancelReleasesLock(t *testing.T) {
 	}
 }
 
+func TestWithLockPassesContextToFunction(t *testing.T) {
+	setupMiniRedis(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	key := "m10:ctx-aware"
+
+	var seen context.Context
+	err := cache.WithLock(ctx, key, time.Second, func(runCtx context.Context) error {
+		seen = runCtx
+		cancel()
+		<-runCtx.Done()
+		return runCtx.Err()
+	})
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("WithLock ctx-aware err = %v, want context.Canceled", err)
+	}
+	if seen != ctx {
+		t.Fatal("WithLock 应把调用方 ctx 传入业务函数")
+	}
+}
+
+func TestWithLockRejectsNilFunction(t *testing.T) {
+	setupMiniRedis(t)
+	ctx := context.Background()
+
+	if err := cache.WithLock(ctx, "m10:nil-fn", time.Second, nil); !errors.Is(err, cache.ErrLockFuncNil) {
+		t.Fatalf("WithLock nil fn err = %v, want ErrLockFuncNil", err)
+	}
+	if err := cache.WithLockAutoExtend(ctx, "m10:nil-fn-auto", time.Second, 100*time.Millisecond, nil); !errors.Is(err, cache.ErrLockFuncNil) {
+		t.Fatalf("WithLockAutoExtend nil fn err = %v, want ErrLockFuncNil", err)
+	}
+}
+
 // ===== C1b：类型断言不 panic =====
 
 // 回归 C1b：NewLock/Unlock/ExtendLock 在正常路径返回正确结果（Lua 脚本返 int64）。
@@ -323,7 +356,7 @@ func TestWithLockAutoExtendRejectsNonPositiveInterval(t *testing.T) {
 	setupMiniRedis(t)
 	called := false
 
-	err := cache.WithLockAutoExtend(context.Background(), "m10:interval", time.Second, 0, func() error {
+	err := cache.WithLockAutoExtend(context.Background(), "m10:interval", time.Second, 0, func(context.Context) error {
 		called = true
 		return nil
 	})
@@ -347,7 +380,7 @@ func TestWithLockHeldReturnsErrLockNotAcquired(t *testing.T) {
 	defer cache.Unlock(ctx, holder)
 
 	called := false
-	err = cache.WithLock(ctx, key, time.Second, func() error {
+	err = cache.WithLock(ctx, key, time.Second, func(context.Context) error {
 		called = true
 		return nil
 	})