feat: init project with distributed lock implementation

- Add DistributedLock with Acquire/Release/Refresh using Lua scripts - Add AcquireWithRetry with configurable retry interval and max retries - Add integration tests for all lock operations - Add .env for Redis connection (gitignored)
2026-05-29 19:19:01 +08:00
parent 2b568c94bb
commit 6b1f071797
6 changed files with 367 additions and 0 deletions
@@ -0,0 +1,4 @@
+node_modules/
+dist/
+.env
+*.log
@@ -0,0 +1,10 @@
+module redis-rate-limiter
+
+go 1.26.3
+
+require (
+	github.com/cespare/xxhash/v2 v2.3.0 // indirect
+	github.com/joho/godotenv v1.5.1 // indirect
+	github.com/redis/go-redis/v9 v9.20.0 // indirect
+	go.uber.org/atomic v1.11.0 // indirect
+)
@@ -0,0 +1,8 @@
+github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
+github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0=
+github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
+github.com/redis/go-redis/v9 v9.20.0 h1:WnQYxLkgO2xiXTCJY0ldIiI8dNqCDlQAG+AtaH7a2a0=
+github.com/redis/go-redis/v9 v9.20.0/go.mod h1:v/M13XI1PVCDcm01VtPFOADfZtHf8YW3baQf57KlIkA=
+go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE=
+go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0=
@@ -0,0 +1,113 @@
+package lock
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+)
+
+var (
+	ErrLockNotAcquired = errors.New("lock not acquired")
+	ErrLockNotHeld     = errors.New("lock not held by this owner")
+)
+
+// DistributedLock 基于 Redis 的分布式锁
+type DistributedLock struct {
+	client   *redis.Client
+	key      string
+	owner    string // 唯一标识，用于安全释放锁
+	ttl      time.Duration
+}
+
+// New 创建一个分布式锁实例
+//   - client: Redis 客户端
+//   - key: 锁的 key
+//   - owner: 持有者唯一标识（如 UUID）
+//   - ttl: 锁的过期时间，防止死锁
+func New(client *redis.Client, key string, owner string, ttl time.Duration) *DistributedLock {
+	return &DistributedLock{
+		client: client,
+		key:    key,
+		owner:  owner,
+		ttl:    ttl,
+	}
+}
+
+// Acquire 尝试获取锁（非阻塞）
+func (l *DistributedLock) Acquire(ctx context.Context) (bool, error) {
+	ok, err := l.client.SetNX(ctx, l.key, l.owner, l.ttl).Result()
+	if err != nil {
+		return false, err
+	}
+	return ok, nil
+}
+
+// AcquireWithRetry 带重试的获取锁
+//   - retryInterval: 重试间隔
+//   - maxRetries: 最大重试次数，0 表示无限重试
+func (l *DistributedLock) AcquireWithRetry(ctx context.Context, retryInterval time.Duration, maxRetries int) (bool, error) {
+	for i := 0; maxRetries == 0 || i < maxRetries; i++ {
+		acquired, err := l.Acquire(ctx)
+		if err != nil {
+			return false, err
+		}
+		if acquired {
+			return true, nil
+		}
+
+		select {
+		case <-ctx.Done():
+			return false, ctx.Err()
+		case <-time.After(retryInterval):
+		}
+	}
+	return false, ErrLockNotAcquired
+}
+
+// Release 释放锁（仅当自己持有时才释放，使用 Lua 脚本保证原子性）
+func (l *DistributedLock) Release(ctx context.Context) error {
+	script := redis.NewScript(`
+		if redis.call("GET", KEYS[1]) == ARGV[1] then
+			return redis.call("DEL", KEYS[1])
+		else
+			return 0
+		end
+	`)
+
+	result, err := script.Run(ctx, l.client, []string{l.key}, l.owner).Int64()
+	if err != nil {
+		return err
+	}
+	if result == 0 {
+		return ErrLockNotHeld
+	}
+	return nil
+}
+
+// Refresh 续期锁（延长 TTL，仅当自己持有时才续期）
+func (l *DistributedLock) Refresh(ctx context.Context) (bool, error) {
+	script := redis.NewScript(`
+		if redis.call("GET", KEYS[1]) == ARGV[1] then
+			return redis.call("PEXPIRE", KEYS[1], ARGV[2])
+		else
+			return 0
+		end
+	`)
+
+	result, err := script.Run(ctx, l.client, []string{l.key}, l.owner, l.ttl.Milliseconds()).Int64()
+	if err != nil {
+		return false, err
+	}
+	return result == 1, nil
+}
+
+// IsHeld 检查锁是否被持有（不一定是自己）
+func (l *DistributedLock) IsHeld(ctx context.Context) (bool, error) {
+	val, err := l.client.Exists(ctx, l.key).Result()
+	if err != nil {
+		return false, err
+	}
+	return val > 0, nil
+}
@@ -0,0 +1,172 @@
+package lock
+
+import (
+	"context"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/joho/godotenv"
+	"github.com/redis/go-redis/v9"
+)
+
+func setupClient(t *testing.T) *redis.Client {
+	t.Helper()
+
+	// 尝试加载 .env（兼容从项目根目录或子目录运行）
+	_ = godotenv.Load("../../.env")
+	_ = godotenv.Load("../.env")
+	_ = godotenv.Load(".env")
+
+	url := os.Getenv("REDIS_URL")
+	if url == "" {
+		t.Skip("REDIS_URL not set, skipping integration test")
+	}
+
+	opt, err := redis.ParseURL(url)
+	if err != nil {
+		t.Fatalf("Invalid REDIS_URL: %v", err)
+	}
+
+	client := redis.NewClient(opt)
+	if err := client.Ping(context.Background()).Err(); err != nil {
+		t.Fatalf("Redis connection failed: %v", err)
+	}
+
+	t.Cleanup(func() { client.Close() })
+	return client
+}
+
+func TestAcquireAndRelease(t *testing.T) {
+	client := setupClient(t)
+	ctx := context.Background()
+
+	l := New(client, "test-lock-acquire", "owner-1", 10*time.Second)
+
+	// 清理可能残留的 key
+	_ = l.Release(ctx)
+
+	// 第一次获取应该成功
+	ok, err := l.Acquire(ctx)
+	if err != nil {
+		t.Fatalf("Acquire error: %v", err)
+	}
+	if !ok {
+		t.Fatal("Expected to acquire lock")
+	}
+
+	// 释放
+	if err := l.Release(ctx); err != nil {
+		t.Fatalf("Release error: %v", err)
+	}
+
+	// 释放后再次获取应该成功
+	ok, err = l.Acquire(ctx)
+	if err != nil {
+		t.Fatalf("Acquire after release error: %v", err)
+	}
+	if !ok {
+		t.Fatal("Expected to acquire lock after release")
+	}
+
+	_ = l.Release(ctx)
+}
+
+func TestAcquireConflict(t *testing.T) {
+	client := setupClient(t)
+	ctx := context.Background()
+
+	l1 := New(client, "test-lock-conflict", "owner-1", 10*time.Second)
+	l2 := New(client, "test-lock-conflict", "owner-2", 10*time.Second)
+
+	_ = l1.Release(ctx)
+	_ = l2.Release(ctx)
+
+	// owner-1 获取锁
+	ok, err := l1.Acquire(ctx)
+	if err != nil {
+		t.Fatalf("Acquire error: %v", err)
+	}
+	if !ok {
+		t.Fatal("Expected owner-1 to acquire lock")
+	}
+
+	// owner-2 获取同一把锁应该失败
+	ok, err = l2.Acquire(ctx)
+	if err != nil {
+		t.Fatalf("Acquire error: %v", err)
+	}
+	if ok {
+		t.Fatal("Expected owner-2 to fail acquiring lock")
+	}
+
+	// owner-2 释放别人的锁应该失败
+	err = l2.Release(ctx)
+	if err != ErrLockNotHeld {
+		t.Fatalf("Expected ErrLockNotHeld, got: %v", err)
+	}
+
+	// owner-1 正常释放
+	if err := l1.Release(ctx); err != nil {
+		t.Fatalf("Release error: %v", err)
+	}
+}
+
+func TestRefresh(t *testing.T) {
+	client := setupClient(t)
+	ctx := context.Background()
+
+	l := New(client, "test-lock-refresh", "owner-1", 5*time.Second)
+	_ = l.Release(ctx)
+
+	ok, err := l.Acquire(ctx)
+	if err != nil {
+		t.Fatalf("Acquire error: %v", err)
+	}
+	if !ok {
+		t.Fatal("Expected to acquire lock")
+	}
+
+	// 续期
+	ok, err = l.Refresh(ctx)
+	if err != nil {
+		t.Fatalf("Refresh error: %v", err)
+	}
+	if !ok {
+		t.Fatal("Expected refresh to succeed")
+	}
+
+	_ = l.Release(ctx)
+}
+
+func TestAcquireWithRetry(t *testing.T) {
+	client := setupClient(t)
+	ctx := context.Background()
+
+	l1 := New(client, "test-lock-retry", "owner-1", 10*time.Second)
+	l2 := New(client, "test-lock-retry", "owner-2", 10*time.Second)
+
+	_ = l1.Release(ctx)
+
+	// owner-1 先获取锁
+	ok, _ := l1.Acquire(ctx)
+	if !ok {
+		t.Fatal("Expected owner-1 to acquire lock")
+	}
+
+	// owner-2 带重试，100ms 后 owner-1 释放，owner-2 应该能拿到
+	go func() {
+		time.Sleep(100 * time.Millisecond)
+		_ = l1.Release(ctx)
+	}()
+
+	ok, err := l2.AcquireWithRetry(ctx, 50*time.Millisecond, 10)
+	if err != nil {
+		t.Fatalf("AcquireWithRetry error: %v", err)
+	}
+	if !ok {
+		t.Fatal("Expected owner-2 to acquire lock after retry")
+	}
+
+	_ = l2.Release(ctx)
+}
@@ -0,0 +1,60 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"os"
+	"time"
+
+	"github.com/joho/godotenv"
+	"github.com/redis/go-redis/v9"
+	"redis-rate-limiter/lock"
+)
+
+func main() {
+	// 加载 .env
+	if err := godotenv.Load(); err != nil {
+		log.Fatal("Error loading .env file")
+	}
+
+	// 连接 Redis
+	opt, err := redis.ParseURL(os.Getenv("REDIS_URL"))
+	if err != nil {
+		log.Fatalf("Invalid REDIS_URL: %v", err)
+	}
+
+	client := redis.NewClient(opt)
+	defer client.Close()
+
+	ctx := context.Background()
+
+	// 测试连接
+	if err := client.Ping(ctx).Err(); err != nil {
+		log.Fatalf("Redis connection failed: %v", err)
+	}
+	fmt.Println("✅ Redis connected")
+
+	// 示例：使用分布式锁
+	owner := fmt.Sprintf("instance-%d", time.Now().UnixNano())
+	l := lock.New(client, "my-resource-lock", owner, 10*time.Second)
+
+	// 尝试获取锁
+	acquired, err := l.Acquire(ctx)
+	if err != nil {
+		log.Fatalf("Acquire error: %v", err)
+	}
+	fmt.Printf("Lock acquired: %v\n", acquired)
+
+	if acquired {
+		// 模拟业务处理
+		fmt.Println("🔒 Holding lock, doing work...")
+		time.Sleep(2 * time.Second)
+
+		// 释放锁
+		if err := l.Release(ctx); err != nil {
+			log.Fatalf("Release error: %v", err)
+		}
+		fmt.Println("🔓 Lock released")
+	}
+}