diff --git a/go.mod b/go.mod index e0404a9..8ec3c7c 100644 --- a/go.mod +++ b/go.mod @@ -1,10 +1,18 @@ module github.com/yuninks/lockx -go 1.20 +go 1.24 -require github.com/redis/go-redis/v9 v9.14.0 +require ( + github.com/google/uuid v1.6.0 + github.com/redis/go-redis/v9 v9.14.0 + github.com/stretchr/testify v1.8.4 +) require ( github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/davecgh/go-spew v1.1.1 // indirect github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect + github.com/stretchr/objx v0.5.2 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index 2bfd88e..6f61dd0 100644 --- a/go.sum +++ b/go.sum @@ -1,8 +1,24 @@ github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= +github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= +github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= 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/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/redis/go-redis/v9 v9.14.0 h1:u4tNCjXOyzfgeLN+vAZaW1xUooqWDqVEsZN0U01jfAE= github.com/redis/go-redis/v9 v9.14.0/go.mod h1:huWgSWd8mW6+m0VPhJjSSQ+d6Nh1VICQ6Q5lHuCH/Iw= +github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= +github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk= +github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/locks.go b/locks.go new file mode 100644 index 0000000..21fad47 --- /dev/null +++ b/locks.go @@ -0,0 +1,28 @@ +package lockx + +import ( + "context" + "fmt" + + "github.com/redis/go-redis/v9" +) + +// 简单使用 + +var redisConn redis.UniversalClient + +// 初始化redis连接 +func Init(ctx context.Context, redis redis.UniversalClient, opts ...Option) error { + redisConn = redis + InitOption(opts...) + return nil +} + +// 新起一个锁对象 +// 先Init后New再Lock +func New(ctx context.Context, uniqueKey string) (*GlobalLock, error) { + if redisConn == nil { + return nil, fmt.Errorf("redis client is nil") + } + return NewGlobalLock(ctx, redisConn, uniqueKey, globalOpts...) +} diff --git a/locks_test.go b/locks_test.go new file mode 100644 index 0000000..c9f550d --- /dev/null +++ b/locks_test.go @@ -0,0 +1,34 @@ +package lockx_test + +import ( + "context" + "fmt" + "testing" + + "github.com/redis/go-redis/v9" + "github.com/yuninks/lockx" +) + +func TestSimpleLock(t *testing.T) { + ctx := context.Background() + client := redis.NewClient(&redis.Options{ + Addr: "127.0.0.1" + ":" + "6379", + Password: "123456", // no password set + DB: 0, // use default DB + }) + if client == nil { + fmt.Println("redis init error") + return + } + lockx.Init(ctx, client) + + l, err := lockx.New(ctx, "lockx:test") + if err != nil { + t.Log(err) + return + } + if b, _ := l.Lock(); b { + fmt.Println("lock success") + l.Unlock() + } +} diff --git a/lockx.go b/lockx.go index a6ff3bc..f275b56 100644 --- a/lockx.go +++ b/lockx.go @@ -3,117 +3,237 @@ package lockx import ( "context" "fmt" + "sync" + "sync/atomic" "time" "github.com/redis/go-redis/v9" + "github.com/google/uuid" ) // 全局锁 -type globalLock struct { - redis redis.UniversalClient - ctx context.Context - cancel context.CancelFunc - uniqueKey string - value string +type GlobalLock struct { + redis redis.UniversalClient + ctx context.Context + cancel context.CancelFunc + uniqueKey string + value string + isClosed bool + closeLock sync.RWMutex + options *option + refreshErrCount int64 // 刷新错误次数 } -func NewGlobalLock(ctx context.Context, red redis.UniversalClient, uniqueKey string) *globalLock { - ctx, cancel := context.WithTimeout(ctx, opt.lockTimeout) - return &globalLock{ +func NewGlobalLock(ctx context.Context, red redis.UniversalClient, uniqueKey string, opts ...Option) (*GlobalLock, error) { + if uniqueKey == "" { + return nil, fmt.Errorf("uniqueKey is empty") + } + if red == nil { + return nil, fmt.Errorf("redis is nil") + } + + options := defaultOption() + for _, opt := range opts { + opt(options) + } + + ctx, cancel := context.WithTimeout(ctx, options.lockTimeout) + + u, err := uuid.NewV7() + if err != nil { + cancel() + return nil, fmt.Errorf("failed to generate UUID: %w", err) + } + + return &GlobalLock{ redis: red, ctx: ctx, cancel: cancel, uniqueKey: uniqueKey, - value: fmt.Sprintf("%d", time.Now().UnixNano()), - } + value: u.String(), + options: options, + }, nil +} + +// 获取上下文 +func (l *GlobalLock) GetCtx() context.Context { + return l.ctx } // 获取锁 -func (g *globalLock) Lock() bool { - +func (g *GlobalLock) Lock() (bool, error) { script := ` - local token = redis.call('get',KEYS[1]) - if token == false - then - return redis.call('set',KEYS[1],ARGV[1],'EX',ARGV[2]) + if redis.call('set', KEYS[1], ARGV[1], 'NX', 'EX', ARGV[2]) then + return 'OK' + else + local current_val = redis.call('get', KEYS[1]) + if current_val == ARGV[1] then + redis.call('expire', KEYS[1], ARGV[2]) + return 'OK' + else + return 'ERROR' + end end - return 'ERROR' ` - resp, err := g.redis.Eval(g.ctx, script, []string{g.uniqueKey}, g.value, 5).Result() - if resp != "OK" { - opt.logger.Errorf(g.ctx, "global lock err resp:%+v err:%+v uniKey:%+v value:%+v", resp, err, g.uniqueKey, g.value) + resp, err := g.redis.Eval(g.ctx, script, []string{g.uniqueKey}, g.value, int(g.options.Expiry.Seconds())).Result() + if err != nil { + if g.options.logger != nil { + g.options.logger.Errorf(g.ctx, "global lock failed: %v, key: %s, value: %s", err, g.uniqueKey, g.value) + } + return false, err } + if resp == "OK" { - g.refresh() - return true + g.startRefresh() + return true, nil } - return false + + return false, nil } // 尝试获取锁 -func (g *globalLock) Try(limitTimes int) bool { - for i := 0; i < limitTimes; i++ { - if g.Lock() { - return true +func (g *GlobalLock) Try() (bool, error) { + for i := 0; i < g.options.MaxRetryTimes; i++ { + success, err := g.Lock() + if err != nil { + return false, err + } + if success { + return true, nil + } + + select { + case <-time.After(g.options.RetryInterval): + continue + case <-g.ctx.Done(): + return false, g.ctx.Err() } - time.Sleep(time.Millisecond * 100) } - return false + return false, nil } // 删除锁 -func (g *globalLock) Unlock() bool { +func (g *GlobalLock) Unlock() error { + // 已经关闭就不需要重复关闭 + if g.setOrGetClose() { + // g.options.logger.Infof(g.ctx, "global lock already closed, key: %s, value: %s", g.uniqueKey, g.value) + return nil + } script := ` - local token = redis.call('get',KEYS[1]) - if token == ARGV[1] - then - redis.call('del',KEYS[1]) - return 'OK' + if redis.call('get', KEYS[1]) == ARGV[1] then + return redis.call('del', KEYS[1]) + else + return 0 end - return 'ERROR' ` - resp, err := g.redis.Eval(g.ctx, script, []string{g.uniqueKey}, g.value).Result() - if resp != "OK" { - opt.logger.Errorf(g.ctx, "global Unlock err resp:%+v err:%+v uniKey:%+v value:%+v", resp, err, g.uniqueKey, g.value) + // 避免上游已取消执行不了 + ctx := context.WithoutCancel(g.ctx) + + resp, err := g.redis.Eval(ctx, script, []string{g.uniqueKey}, g.value).Result() + if err != nil { + if g.options.logger != nil { + g.options.logger.Infof(g.ctx, "global unlock may have failed: %v, key: %s, value: %s", err, g.uniqueKey, g.value) + } + // 即使删除失败也继续执行，因为锁可能会自动过期 } - g.cancel() - return true + + if delCount, ok := resp.(int64); ok && delCount == 1 { + return nil + } + g.options.logger.Infof(g.ctx, "global unlock may have failed: %v, key: %s, value: %s", err, g.uniqueKey, g.value) + return fmt.Errorf("lock was already released or owned by another client") } -// 刷新锁 -func (g *globalLock) refresh() { +// 启动刷新goroutine +func (g *GlobalLock) startRefresh() { + go func() { - t := time.NewTicker(time.Second) + + ticker := time.NewTicker(g.options.RefreshPeriod) + defer ticker.Stop() + for { select { - case <-t.C: - g.refreshExec() + case <-ticker.C: + g.refreshExec() case <-g.ctx.Done(): - t.Stop() + // g.options.logger.Infof(g.ctx, "global lock refresh canceled, key: %s, value: %s", g.uniqueKey, g.value) + g.Unlock() return } } }() } -func (g *globalLock) refreshExec() bool { - script := ` - local token = redis.call('get',KEYS[1]) - if token == ARGV[1] - then - redis.call('set',KEYS[1],ARGV[1],'EX',ARGV[2]) - return 'OK' - end - return 'ERROR' - ` +// 返回原来的 +func (l *GlobalLock) setOrGetClose() (readyClosed bool) { + l.closeLock.Lock() + defer l.closeLock.Unlock() - resp, err := g.redis.Eval(g.ctx, script, []string{g.uniqueKey}, g.value, 5).Result() - if resp != "OK" { - opt.logger.Errorf(g.ctx, "global refresh err resp:%+v err:%+v uniKey:%+v value:%+v", resp, err, g.uniqueKey, g.value) + if l.isClosed { + return true + } + + l.isClosed = true + + l.cancel() + + return false +} + +// 执行刷新操作 +func (g *GlobalLock) refreshExec() bool { + if g.IsClosed() { + g.options.logger.Infof(g.ctx, "global lock already closed, key: %s, value: %s", g.uniqueKey, g.value) return false } + + script := ` + if redis.call('get', KEYS[1]) == ARGV[1] then + redis.call('expire', KEYS[1], ARGV[2]) + return 1 + else + return 0 + end + ` + + resp, err := g.redis.Eval(g.ctx, script, []string{g.uniqueKey}, g.value, int(g.options.Expiry.Seconds())).Result() + if err != nil { + if g.options.logger != nil { + g.options.logger.Errorf(g.ctx, "global refresh failed: %v, key: %s, value: %s", err, g.uniqueKey, g.value) + } + + newCount := atomic.AddInt64(&g.refreshErrCount, 1) + if newCount >= 3 { + // 关闭锁 + g.setOrGetClose() + g.options.logger.Errorf(g.ctx, "global refresh failed, lock may be lost 3 times, key: %s, value: %s", g.uniqueKey, g.value) + } + return false + } + + if refreshed, ok := resp.(int64); !ok || refreshed != 1 { + if g.options.logger != nil { + g.options.logger.Errorf(g.ctx, "global refresh failed, lock may be lost, key: %s, value: %s", g.uniqueKey, g.value) + } + // 关闭锁 + g.setOrGetClose() + return false + } + return true } + +// 检查锁是否已关闭 +func (g *GlobalLock) IsClosed() bool { + g.closeLock.RLock() + defer g.closeLock.RUnlock() + return g.isClosed +} + +func (l *GlobalLock) GetValue() string { + return l.value +} diff --git a/lockx_test.go b/lockx_test.go index c08121d..56350af 100644 --- a/lockx_test.go +++ b/lockx_test.go @@ -2,11 +2,17 @@ package lockx_test import ( "context" + "errors" "fmt" + "sync" "testing" + "time" - "github.com/yuninks/lockx" "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + "github.com/yuninks/lockx" ) var Redis *redis.Client @@ -25,28 +31,554 @@ var Redis *redis.Client // Redis = client // } -func TestLockx(t *testing.T) { +// func TestLockx(t *testing.T) { +// client := redis.NewClient(&redis.Options{ +// Addr: "127.0.0.1" + ":" + "6379", +// Password: "123456", // no password set +// DB: 0, // use default DB +// }) +// if client == nil { +// fmt.Println("redis init error") +// return +// } +// fmt.Println("begin") +// ctx := context.Background() +// ctx, cancel := context.WithCancel(ctx) +// defer cancel() + +// wg := sync.WaitGroup{} + +// for i := 0; i < 20; i++ { +// wg.Add(1) +// go func(i int) { +// defer wg.Done() +// lock, _ := lockx.NewGlobalLock(ctx, client, "lockx:test") +// if b, _ := lock.Lock(); !b { +// fmt.Println("lock error", i) +// return +// } +// defer lock.Unlock() + +// fmt.Println("ssss2", i) + +// time.Sleep(time.Second * 2) +// }(i) +// time.Sleep(time.Second) +// } + +// wg.Wait() +// } + +// MockLogger 用于测试的模拟日志器 +type MockLogger struct { + mock.Mock +} + +func (m *MockLogger) Errorf(ctx context.Context, format string, args ...interface{}) { + m.Called(ctx, format, args) +} + +func (m *MockLogger) Infof(ctx context.Context, format string, args ...interface{}) { + m.Called(ctx, format, args) +} + +// 测试工具函数 +func setupTestRedis(t *testing.T) (redis.UniversalClient, func()) { + client := redis.NewClient(&redis.Options{ Addr: "127.0.0.1" + ":" + "6379", - Password: "", // no password set - DB: 0, // use default DB + Password: "123456", // no password set + DB: 0, // use default DB }) - if client == nil { - fmt.Println("redis init error") - return + + return client, func() { + client.Close() } - fmt.Println("begin") +} + +func TestNewGlobalLock(t *testing.T) { + redisClient, teardown := setupTestRedis(t) + defer teardown() + + t.Run("成功创建锁", func(t *testing.T) { + ctx := context.Background() + lock, err := lockx.NewGlobalLock(ctx, redisClient, "test-key") + + require.NoError(t, err) + assert.NotEmpty(t, lock.GetValue()) + assert.False(t, lock.IsClosed()) + }) + + t.Run("创建锁时UUID生成失败", func(t *testing.T) { + // 这个测试需要模拟uuid.NewV7()失败，在实际中较难触发 + // 通常可以跳过或者使用mock来测试 + t.Skip("很难模拟UUID生成失败的情况") + }) +} + +func TestGlobalLock_Lock(t *testing.T) { + redisClient, teardown := setupTestRedis(t) + defer teardown() + ctx := context.Background() - ctx, cancel := context.WithCancel(ctx) - defer cancel() - lock := lockx.NewGlobalLock(ctx, client, "lockx:test") + t.Run("成功获取锁", func(t *testing.T) { + lock, err := lockx.NewGlobalLock(ctx, redisClient, "test-key-success") + require.NoError(t, err) - if !lock.Lock() { - fmt.Println("lock error") - } - defer lock.Unlock() + success, err := lock.Lock() + require.NoError(t, err) + assert.True(t, success) - fmt.Println("ssss") + defer lock.Unlock() + }) + + t.Run("获取已存在的锁失败", func(t *testing.T) { + // 第一个客户端获取锁 + lock1, err := lockx.NewGlobalLock(ctx, redisClient, "test-key-conflict") + require.NoError(t, err) + + success1, err := lock1.Lock() + require.NoError(t, err) + assert.True(t, success1) + + defer lock1.Unlock() + + // 第二个客户端尝试获取同一个锁 + lock2, err := lockx.NewGlobalLock(ctx, redisClient, "test-key-conflict") + require.NoError(t, err) + + success2, err := lock2.Lock() + require.NoError(t, err) + assert.False(t, success2) + + defer lock2.Unlock() + }) + + t.Run("Redis连接失败", func(t *testing.T) { + // 创建无效的Redis客户端 + invalidClient := redis.NewClient(&redis.Options{ + Addr: "invalid:6379", + }) + + _, err := lockx.NewGlobalLock(ctx, invalidClient, "test-key-fail") + require.NoError(t, err) + + // 设置短超时 + ctxShort, cancel := context.WithTimeout(ctx, 100*time.Millisecond) + defer cancel() + + lockWithCtx, err := lockx.NewGlobalLock(ctxShort, invalidClient, "test-key-fail") + require.NoError(t, err) + + success, err := lockWithCtx.Lock() + assert.Error(t, err) + assert.False(t, success) + }) +} + +func TestGlobalLock_Try(t *testing.T) { + redisClient, teardown := setupTestRedis(t) + defer teardown() + + ctx := context.Background() + + t.Run("尝试获取可用锁成功", func(t *testing.T) { + _, err := lockx.NewGlobalLock(ctx, redisClient, "test-try-success") + require.NoError(t, err) + + // 使用自定义选项 + customLock, err := lockx.NewGlobalLock(ctx, redisClient, "test-try-custom", + lockx.WithMaxRetryTimes(2), + lockx.WithRetryInterval(50*time.Millisecond), + ) + require.NoError(t, err) + + success, err := customLock.Try() + require.NoError(t, err) + assert.True(t, success) + + defer customLock.Unlock() + }) + + t.Run("尝试获取被占用的锁失败", func(t *testing.T) { + // 先获取锁 + lock1, err := lockx.NewGlobalLock(ctx, redisClient, "test-try-busy") + require.NoError(t, err) + + success1, err := lock1.Lock() + require.NoError(t, err) + assert.True(t, success1) + + // 另一个客户端尝试获取 + lock2, err := lockx.NewGlobalLock(ctx, redisClient, "test-try-busy", + lockx.WithMaxRetryTimes(2), + lockx.WithRetryInterval(10*time.Millisecond), + ) + require.NoError(t, err) + + success2, err := lock2.Try() + require.NoError(t, err) + assert.False(t, success2) + + lock1.Unlock() + lock2.Unlock() + }) + + t.Run("上下文取消", func(t *testing.T) { + ctxCancel, cancel := context.WithCancel(ctx) + lock, err := lockx.NewGlobalLock(ctxCancel, redisClient, "test-try-cancel") + require.NoError(t, err) + + cancel() // 立即取消 + + success, err := lock.Try() + assert.True(t, errors.Is(err, context.Canceled)) + assert.False(t, success) + + err = lock.Unlock() + assert.Error(t, err) + }) +} + +func TestGlobalLock_Unlock(t *testing.T) { + redisClient, teardown := setupTestRedis(t) + defer teardown() + + ctx := context.Background() + + t.Run("成功释放锁", func(t *testing.T) { + lock, err := lockx.NewGlobalLock(ctx, redisClient, "test-unlock-success") + require.NoError(t, err) + + success, err := lock.Lock() + require.NoError(t, err) + assert.True(t, success) + + err = lock.Unlock() + assert.NoError(t, err) + assert.True(t, lock.IsClosed()) + + // 验证锁确实被释放了 + val, err := redisClient.Get(ctx, "test-unlock-success").Result() + assert.Equal(t, redis.Nil, err) + assert.Empty(t, val) + }) + + t.Run("重复释放锁", func(t *testing.T) { + lock, err := lockx.NewGlobalLock(ctx, redisClient, "test-unlock-twice") + require.NoError(t, err) + + success, err := lock.Lock() + require.NoError(t, err) + assert.True(t, success) + + err = lock.Unlock() + assert.NoError(t, err) + + // 再次释放应该不会报错 + err = lock.Unlock() + assert.NoError(t, err) + }) + + t.Run("释放其他客户端的锁", func(t *testing.T) { + lock1, err := lockx.NewGlobalLock(ctx, redisClient, "test-unlock-other") + require.NoError(t, err) + + success, err := lock1.Lock() + require.NoError(t, err) + assert.True(t, success) + + // 另一个客户端尝试释放（使用不同的value） + lock2, err := lockx.NewGlobalLock(ctx, redisClient, "test-unlock-other") + require.NoError(t, err) + + err = lock2.Unlock() + assert.Error(t, err) + assert.Contains(t, err.Error(), "lock was already released or owned by another client") + + lock1.Unlock() + lock2.Unlock() + }) +} + +func TestGlobalLock_Refresh(t *testing.T) { + redisClient, teardown := setupTestRedis(t) + defer teardown() + + ctx := context.Background() + + t.Run("自动刷新保持锁", func(t *testing.T) { + lock, err := lockx.NewGlobalLock(ctx, redisClient, "test-refresh", + lockx.WithExpiry(2*time.Second), + lockx.WithRefreshPeriod(500*time.Millisecond), + ) + require.NoError(t, err) + + success, err := lock.Lock() + require.NoError(t, err) + assert.True(t, success) + + // 等待一段时间，确保刷新机制工作 + time.Sleep(3 * time.Second) + + // 检查锁仍然存在 + val, err := redisClient.Get(ctx, "test-refresh").Result() + require.NoError(t, err) + assert.Equal(t, lock.GetValue(), val) + + // 检查TTL应该还在 + ttl, err := redisClient.TTL(ctx, "test-refresh").Result() + require.NoError(t, err) + assert.True(t, ttl > 0 && ttl <= 2*time.Second) + + lock.Unlock() + }) + + t.Run("锁丢失时停止刷新", func(t *testing.T) { + mockLogger := new(MockLogger) + mockLogger.On("Errorf", mock.Anything, mock.Anything, mock.Anything).Return() + mockLogger.On("Infof", mock.Anything, mock.Anything, mock.Anything).Return() + + lock, err := lockx.NewGlobalLock(ctx, redisClient, "test-refresh-lost", + lockx.WithLogger(mockLogger), + lockx.WithExpiry(1*time.Second), + ) + require.NoError(t, err) + + success, err := lock.Lock() + require.NoError(t, err) + assert.True(t, success) + + // 手动删除锁来模拟锁丢失 + err = redisClient.Del(ctx, "test-refresh-lost").Err() + require.NoError(t, err) + + // 等待刷新周期 + time.Sleep(1500 * time.Millisecond) + + // 验证日志被调用 + mockLogger.AssertCalled(t, "Errorf", mock.Anything, mock.Anything, mock.Anything) + + // <-lock.GetCtx().Done() + + lock.Unlock() + }) +} + +func TestGlobalLock_Concurrency(t *testing.T) { + redisClient, teardown := setupTestRedis(t) + defer teardown() + + ctx := context.Background() + + t.Run("多 goroutine 竞争锁", func(t *testing.T) { + const numGoroutines = 100 + var successCount int + var wg sync.WaitGroup + var mu sync.Mutex + + wg.Add(numGoroutines) + + for i := 0; i < numGoroutines; i++ { + go func(id int) { + defer wg.Done() + + lock, err := lockx.NewGlobalLock(ctx, redisClient, "concurrent-test") + require.NoError(t, err) + + success, err := lock.Lock() + require.NoError(t, err) + + if success { + mu.Lock() + successCount++ + mu.Unlock() + + // 持有锁一段时间 + time.Sleep(100 * time.Millisecond) + lock.Unlock() + } + }(i) + } + + wg.Wait() + + // 应该只有一个goroutine成功获取锁 + assert.Equal(t, 1, successCount) + }) +} + +func TestGlobalLock_ContextCancellation(t *testing.T) { + redisClient, teardown := setupTestRedis(t) + defer teardown() + + t.Run("上下文超时", func(t *testing.T) { + ctxShort, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + lock, err := lockx.NewGlobalLock(ctxShort, redisClient, "test-context-timeout") + require.NoError(t, err) + + // 等待上下文超时 + time.Sleep(200 * time.Millisecond) + + // 尝试操作应该失败 + success, err := lock.Lock() + assert.Error(t, err) + assert.False(t, success) + }) + + t.Run("上下文取消时自动释放锁", func(t *testing.T) { + ctxCancel, cancel := context.WithCancel(context.Background()) + lock, err := lockx.NewGlobalLock(ctxCancel, redisClient, "test-context-cancel") + require.NoError(t, err) + + success, err := lock.Lock() + require.NoError(t, err) + assert.True(t, success) + + // 取消上下文 + cancel() + + // lock.Unlock() + + // 等待自动清理 + time.Sleep(5 * time.Second) + + // 检查锁应该被释放 + val, err := redisClient.Get(context.Background(), "test-context-cancel").Result() + assert.Equal(t, redis.Nil, err) + assert.Empty(t, val) + }) +} + +func TestGlobalLock_CustomOptions(t *testing.T) { + redisClient, teardown := setupTestRedis(t) + defer teardown() + + ctx := context.Background() + + t.Run("自定义配置选项", func(t *testing.T) { + mockLogger := new(MockLogger) + mockLogger.On("Errorf", mock.Anything, mock.Anything, mock.Anything).Return() + mockLogger.On("Infof", mock.Anything, mock.Anything, mock.Anything).Return() + + lock, err := lockx.NewGlobalLock(ctx, redisClient, "test-custom-options", + lockx.WithLockTimeout(5*time.Second), + lockx.WithExpiry(10*time.Second), + lockx.WithRefreshPeriod(2*time.Second), + lockx.WithMaxRetryTimes(5), + lockx.WithRetryInterval(200*time.Millisecond), + lockx.WithLogger(mockLogger), + ) + require.NoError(t, err) + + success, err := lock.Lock() + require.NoError(t, err) + assert.True(t, success) + + // 检查TTL是否正确设置 + ttl, err := redisClient.TTL(ctx, "test-custom-options").Result() + require.NoError(t, err) + assert.True(t, ttl > 8*time.Second && ttl <= 10*time.Second) + + lock.Unlock() + }) +} + +func TestGlobalLock_EdgeCases(t *testing.T) { + redisClient, teardown := setupTestRedis(t) + defer teardown() + + ctx := context.Background() + + t.Run("空键名", func(t *testing.T) { + _, err := lockx.NewGlobalLock(ctx, redisClient, "") + require.Error(t, err) + + }) + + t.Run("非常长的键名", func(t *testing.T) { + longKey := string(make([]byte, 1000)) // 很长的键名 + lock, err := lockx.NewGlobalLock(ctx, redisClient, longKey) + require.NoError(t, err) + + success, err := lock.Lock() + require.NoError(t, err) + assert.True(t, success) + + lock.Unlock() + }) +} + +// 基准测试 +func BenchmarkGlobalLock(b *testing.B) { + t := testing.T{} + redisClient, teardown := setupTestRedis(&t) + defer teardown() + + ctx := context.Background() + + b.Run("获取释放锁", func(b *testing.B) { + for i := 0; i < b.N; i++ { + lock, err := lockx.NewGlobalLock(ctx, redisClient, fmt.Sprintf("benchmark-%d", i)) + if err != nil { + b.Fatal(err) + } + + success, err := lock.Lock() + if err != nil { + b.Fatal(err) + } + if !success { + b.Fatal("Failed to acquire lock") + } + + err = lock.Unlock() + if err != nil { + b.Fatal(err) + } + } + }) +} + +func TestTimeout(t *testing.T) { + redisClient, teardown := setupTestRedis(t) + defer teardown() + + ctx := context.Background() + + begin := time.Now() + + lock, err := lockx.NewGlobalLock(ctx, redisClient, "test-timeout", lockx.WithLockTimeout(time.Second)) + require.NoError(t, err) + success, err := lock.Lock() + require.NoError(t, err) + assert.True(t, success) + + <-lock.GetCtx().Done() + t.Log("Done", time.Since(begin)) + +} + +func TestRedisClose(t *testing.T) { + redisClient, teardown := setupTestRedis(t) + defer teardown() + + ctx := context.Background() + + begin := time.Now() + + lock, err := lockx.NewGlobalLock(ctx, redisClient, "test-timeout", lockx.WithLockTimeout(time.Hour)) + require.NoError(t, err) + success, err := lock.Lock() + require.NoError(t, err) + assert.True(t, success) + + redisClient.Close() + + <-lock.GetCtx().Done() + t.Log("Done", time.Since(begin)) } diff --git a/options.go b/options.go index 2bf41f9..6d71419 100644 --- a/options.go +++ b/options.go @@ -8,46 +8,87 @@ import ( type option struct { lockTimeout time.Duration // 锁的超时时间 - logger Logger // 日志 + + MaxRetryTimes int // 尝试次数 + RetryInterval time.Duration // 尝试间隔 + + Expiry time.Duration // 单次刷新有效时间 + RefreshPeriod time.Duration // 刷新间隔 + + logger Logger // 日志 } func defaultOption() *option { return &option{ - lockTimeout: time.Minute * 60, - logger: &print{}, + lockTimeout: time.Minute * 60, + Expiry: 5 * time.Second, + RefreshPeriod: 1 * time.Second, + MaxRetryTimes: 3, + RetryInterval: 100 * time.Millisecond, + logger: &print{}, } } -var opt *option - -func init() { - opt = defaultOption() -} +var globalOpts []Option // 设置 func InitOption(opts ...Option) { - for _, app := range opts { - app(opt) - } + globalOpts = opts } type Option func(*option) -func SetTimeout(t time.Duration) Option { +// 最大锁定时间(默认1h) +func WithLockTimeout(t time.Duration) Option { return func(o *option) { o.lockTimeout = t } } -func SetLogger(logger Logger) Option { +// 日志 +func WithLogger(logger Logger) Option { return func(o *option) { o.logger = logger } } +// key有效时间(会自动刷新) +func WithExpiry(expiry time.Duration) Option { + return func(o *option) { + o.Expiry = expiry + if o.Expiry/3 < o.RefreshPeriod { + o.RefreshPeriod = o.Expiry / 3 + } + } +} + +// 刷新间隔 +func WithRefreshPeriod(period time.Duration) Option { + return func(o *option) { + o.RefreshPeriod = period + if o.RefreshPeriod*3 > o.Expiry { + o.Expiry = o.RefreshPeriod * 3 + } + } +} + +// 最大尝试次数 +func WithMaxRetryTimes(times int) Option { + return func(o *option) { + o.MaxRetryTimes = times + } +} + +// 尝试间隔 +func WithRetryInterval(interval time.Duration) Option { + return func(o *option) { + o.RetryInterval = interval + } +} + type Logger interface { Errorf(ctx context.Context, format string, v ...any) - Printf(ctx context.Context, format string, v ...any) + Infof(ctx context.Context, format string, v ...any) } type print struct{} @@ -56,6 +97,6 @@ func (*print) Errorf(ctx context.Context, format string, v ...any) { log.Printf(format, v...) } -func (*print) Printf(ctx context.Context, format string, v ...any) { +func (*print) Infof(ctx context.Context, format string, v ...any) { log.Printf(format, v...) } diff --git a/readme.md b/readme.md new file mode 100644 index 0000000..1d809af --- /dev/null +++ b/readme.md @@ -0,0 +1,5 @@ +1.这个锁是基于redis的全局锁 + +2.在同一个new下边，多次加锁是可以重入的(因为值是一致) + +3.多次加锁只要执行了一次释放，这个锁将会被释放