From 6b1f071797f66ae146951edc31e6bb85a2428d4f Mon Sep 17 00:00:00 2001 From: wonder Date: Fri, 29 May 2026 19:19:01 +0800 Subject: [PATCH] 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) --- .gitignore | 4 ++ go.mod | 10 +++ go.sum | 8 +++ lock/lock.go | 113 ++++++++++++++++++++++++++++++ lock/lock_test.go | 172 ++++++++++++++++++++++++++++++++++++++++++++++ main.go | 60 ++++++++++++++++ 6 files changed, 367 insertions(+) create mode 100644 .gitignore create mode 100644 go.mod create mode 100644 go.sum create mode 100644 lock/lock.go create mode 100644 lock/lock_test.go create mode 100644 main.go diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..aa0926a --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +node_modules/ +dist/ +.env +*.log diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..ea4cc79 --- /dev/null +++ b/go.mod @@ -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 +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..d635c3e --- /dev/null +++ b/go.sum @@ -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= diff --git a/lock/lock.go b/lock/lock.go new file mode 100644 index 0000000..fbdbc92 --- /dev/null +++ b/lock/lock.go @@ -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 +} diff --git a/lock/lock_test.go b/lock/lock_test.go new file mode 100644 index 0000000..16c975e --- /dev/null +++ b/lock/lock_test.go @@ -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) +} diff --git a/main.go b/main.go new file mode 100644 index 0000000..c3806c2 --- /dev/null +++ b/main.go @@ -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") + } +}