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)
This commit is contained in:
@@ -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=
|
||||||
+113
@@ -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")
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user