初始提交: Gitea 项目代码
This commit is contained in:
root
@@ -0,0 +1,82 @@
|
||||
// Copyright 2024 The Gitea Authors. All rights reserved.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package globallock
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
|
||||
"gitea.dev/modules/setting"
|
||||
)
|
||||
|
||||
var (
|
||||
defaultLocker atomic.Pointer[Locker]
|
||||
defaultMutex sync.Mutex
|
||||
)
|
||||
|
||||
func initDefaultLocker() Locker {
|
||||
switch setting.GlobalLock.ServiceType {
|
||||
case "redis":
|
||||
return NewRedisLocker(setting.GlobalLock.ServiceConnStr)
|
||||
default: // "memory"
|
||||
return NewMemoryLocker()
|
||||
}
|
||||
}
|
||||
|
||||
// DefaultLocker returns the default locker.
|
||||
func DefaultLocker() Locker {
|
||||
ptr := defaultLocker.Load()
|
||||
if ptr == nil {
|
||||
defaultMutex.Lock()
|
||||
ptr = defaultLocker.Load()
|
||||
if ptr == nil {
|
||||
ptr = new(initDefaultLocker())
|
||||
defaultLocker.Store(ptr)
|
||||
}
|
||||
defaultMutex.Unlock()
|
||||
ptr = defaultLocker.Load()
|
||||
}
|
||||
return *ptr
|
||||
}
|
||||
|
||||
// Lock tries to acquire a lock for the given key, it uses the default locker.
|
||||
// Read the documentation of Locker.Lock for more information about the behavior.
|
||||
func Lock(ctx context.Context, key string) (ReleaseFunc, error) {
|
||||
return DefaultLocker().Lock(ctx, key)
|
||||
}
|
||||
|
||||
// TryLock tries to acquire a lock for the given key, it uses the default locker.
|
||||
// Read the documentation of Locker.TryLock for more information about the behavior.
|
||||
func TryLock(ctx context.Context, key string) (bool, ReleaseFunc, error) {
|
||||
return DefaultLocker().TryLock(ctx, key)
|
||||
}
|
||||
|
||||
// LockAndDo tries to acquire a lock for the given key and then calls the given function.
|
||||
// It uses the default locker, and it will return an error if failed to acquire the lock.
|
||||
func LockAndDo(ctx context.Context, key string, f func(context.Context) error) error {
|
||||
release, err := Lock(ctx, key)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer release()
|
||||
|
||||
return f(ctx)
|
||||
}
|
||||
|
||||
// TryLockAndDo tries to acquire a lock for the given key and then calls the given function.
|
||||
// It uses the default locker, and it will return false if failed to acquire the lock.
|
||||
func TryLockAndDo(ctx context.Context, key string, f func(context.Context) error) (bool, error) {
|
||||
ok, release, err := TryLock(ctx, key)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
defer release()
|
||||
|
||||
if !ok {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
return true, f(ctx)
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
// Copyright 2024 The Gitea Authors. All rights reserved.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package globallock
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestLockAndDo(t *testing.T) {
|
||||
t.Run("redis", func(t *testing.T) {
|
||||
locker := newTestRedisLocker(t)
|
||||
defaultLocker.Store(new(locker))
|
||||
testLockAndDo(t)
|
||||
require.NoError(t, locker.(*redisLocker).Close())
|
||||
})
|
||||
t.Run("memory", func(t *testing.T) {
|
||||
defaultLocker.Store(new(NewMemoryLocker()))
|
||||
testLockAndDo(t)
|
||||
})
|
||||
}
|
||||
|
||||
func testLockAndDo(t *testing.T) {
|
||||
const concurrency = 50
|
||||
|
||||
ctx := t.Context()
|
||||
count := 0
|
||||
wg := sync.WaitGroup{}
|
||||
for range concurrency {
|
||||
wg.Go(func() {
|
||||
err := LockAndDo(ctx, "test", func(ctx context.Context) error {
|
||||
count++
|
||||
// It's impossible to acquire the lock inner the function
|
||||
ok, err := TryLockAndDo(ctx, "test", func(ctx context.Context) error {
|
||||
assert.Fail(t, "should not acquire the lock")
|
||||
return nil
|
||||
})
|
||||
assert.False(t, ok)
|
||||
assert.NoError(t, err)
|
||||
return nil
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
assert.Equal(t, concurrency, count)
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
// Copyright 2024 The Gitea Authors. All rights reserved.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package globallock
|
||||
|
||||
import (
|
||||
"context"
|
||||
)
|
||||
|
||||
type Locker interface {
|
||||
// Lock tries to acquire a lock for the given key, it blocks until the lock is acquired or the context is canceled.
|
||||
//
|
||||
// Lock returns a ReleaseFunc to release the lock, it cannot be nil.
|
||||
// It's always safe to call this function even if it fails to acquire the lock, and it will do nothing in that case.
|
||||
// And it's also safe to call it multiple times, but it will only release the lock once.
|
||||
// That's why it's called ReleaseFunc, not UnlockFunc.
|
||||
// But be aware that it's not safe to not call it at all; it could lead to a memory leak.
|
||||
// So a recommended pattern is to use defer to call it:
|
||||
// release, err := locker.Lock(ctx, "key")
|
||||
// if err != nil {
|
||||
// return err
|
||||
// }
|
||||
// defer release()
|
||||
//
|
||||
// Lock returns an error if failed to acquire the lock.
|
||||
// Be aware that even the context is not canceled, it's still possible to fail to acquire the lock.
|
||||
// For example, redis is down, or it reached the maximum number of tries.
|
||||
Lock(ctx context.Context, key string) (ReleaseFunc, error)
|
||||
|
||||
// TryLock tries to acquire a lock for the given key, it returns immediately.
|
||||
// It follows the same pattern as Lock, but it doesn't block.
|
||||
// And if it fails to acquire the lock because it's already locked, not other reasons like redis is down,
|
||||
// it will return false without any error.
|
||||
TryLock(ctx context.Context, key string) (bool, ReleaseFunc, error)
|
||||
}
|
||||
|
||||
// ReleaseFunc is a function that releases a lock.
|
||||
type ReleaseFunc func()
|
||||
@@ -0,0 +1,180 @@
|
||||
// Copyright 2024 The Gitea Authors. All rights reserved.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package globallock
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.dev/modules/test"
|
||||
"gitea.dev/modules/util"
|
||||
|
||||
"github.com/go-redsync/redsync/v4"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newTestRedisLocker(t *testing.T) Locker {
|
||||
t.Helper()
|
||||
redisURL := util.IfZero(os.Getenv("TEST_REDIS_URL"), "redis://127.0.0.1:6379/0")
|
||||
rl := NewRedisLocker(redisURL).(*redisLocker)
|
||||
err := rl.conn.Ping(t.Context()).Err()
|
||||
if err != nil && test.AllowSkipExternalService() {
|
||||
t.Skip("no redis server for testing, skipped")
|
||||
}
|
||||
require.NoError(t, err, "redis error for testing: %v", err)
|
||||
return rl
|
||||
}
|
||||
|
||||
func TestLocker(t *testing.T) {
|
||||
t.Run("redis", func(t *testing.T) {
|
||||
defer test.MockVariableValue(&redisLockExpiry, 5*time.Second)() // make it shorter for testing
|
||||
locker := newTestRedisLocker(t)
|
||||
testLocker(t, locker)
|
||||
testRedisLocker(t, locker.(*redisLocker))
|
||||
require.NoError(t, locker.(*redisLocker).Close())
|
||||
})
|
||||
t.Run("memory", func(t *testing.T) {
|
||||
locker := NewMemoryLocker()
|
||||
testLocker(t, locker)
|
||||
testMemoryLocker(t, locker.(*memoryLocker))
|
||||
})
|
||||
}
|
||||
|
||||
func testLocker(t *testing.T, locker Locker) {
|
||||
t.Run("lock", func(t *testing.T) {
|
||||
parentCtx := t.Context()
|
||||
release, err := locker.Lock(parentCtx, "test")
|
||||
defer release()
|
||||
|
||||
assert.NoError(t, err)
|
||||
|
||||
func() {
|
||||
ctx, cancel := context.WithTimeout(t.Context(), time.Second)
|
||||
defer cancel()
|
||||
release, err := locker.Lock(ctx, "test")
|
||||
defer release()
|
||||
|
||||
assert.Error(t, err)
|
||||
}()
|
||||
|
||||
release()
|
||||
|
||||
func() {
|
||||
release, err := locker.Lock(t.Context(), "test")
|
||||
defer release()
|
||||
|
||||
assert.NoError(t, err)
|
||||
}()
|
||||
})
|
||||
|
||||
t.Run("try lock", func(t *testing.T) {
|
||||
parentCtx := t.Context()
|
||||
ok, release, err := locker.TryLock(parentCtx, "test")
|
||||
defer release()
|
||||
|
||||
assert.True(t, ok)
|
||||
assert.NoError(t, err)
|
||||
|
||||
func() {
|
||||
ctx, cancel := context.WithTimeout(t.Context(), time.Second)
|
||||
defer cancel()
|
||||
ok, release, err := locker.TryLock(ctx, "test")
|
||||
defer release()
|
||||
|
||||
assert.False(t, ok)
|
||||
assert.NoError(t, err)
|
||||
}()
|
||||
|
||||
release()
|
||||
|
||||
func() {
|
||||
ok, release, _ := locker.TryLock(t.Context(), "test")
|
||||
defer release()
|
||||
|
||||
assert.True(t, ok)
|
||||
}()
|
||||
})
|
||||
|
||||
t.Run("wait and acquired", func(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
release, err := locker.Lock(ctx, "test")
|
||||
require.NoError(t, err)
|
||||
|
||||
wg := &sync.WaitGroup{}
|
||||
wg.Go(func() {
|
||||
started := time.Now()
|
||||
release, err := locker.Lock(t.Context(), "test") // should be blocked for seconds
|
||||
defer release()
|
||||
assert.Greater(t, time.Since(started), time.Second)
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
|
||||
time.Sleep(2 * time.Second)
|
||||
release()
|
||||
|
||||
wg.Wait()
|
||||
})
|
||||
|
||||
t.Run("multiple release", func(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
|
||||
release1, err := locker.Lock(ctx, "test")
|
||||
require.NoError(t, err)
|
||||
|
||||
release1()
|
||||
|
||||
release2, err := locker.Lock(ctx, "test")
|
||||
defer release2()
|
||||
require.NoError(t, err)
|
||||
|
||||
// Call release1 again,
|
||||
// it should not panic or block,
|
||||
// and it shouldn't affect the other lock
|
||||
release1()
|
||||
|
||||
ok, release3, err := locker.TryLock(ctx, "test")
|
||||
defer release3()
|
||||
require.NoError(t, err)
|
||||
// It should be able to acquire the lock;
|
||||
// otherwise, it means the lock has been released by release1
|
||||
assert.False(t, ok)
|
||||
})
|
||||
}
|
||||
|
||||
// testMemoryLocker does specific tests for memoryLocker
|
||||
func testMemoryLocker(t *testing.T, locker *memoryLocker) {
|
||||
// nothing to do
|
||||
}
|
||||
|
||||
// testRedisLocker does specific tests for redisLocker
|
||||
func testRedisLocker(t *testing.T, locker *redisLocker) {
|
||||
defer func() {
|
||||
// This case should be tested at the end.
|
||||
// Otherwise, it will affect other tests.
|
||||
t.Run("close", func(t *testing.T) {
|
||||
assert.NoError(t, locker.Close())
|
||||
_, err := locker.Lock(t.Context(), "test")
|
||||
assert.Error(t, err)
|
||||
})
|
||||
}()
|
||||
|
||||
t.Run("failed extend", func(t *testing.T) {
|
||||
release, err := locker.Lock(t.Context(), "test")
|
||||
defer release()
|
||||
require.NoError(t, err)
|
||||
|
||||
// It simulates that there are some problems with extending like network issues or redis server down.
|
||||
v, ok := locker.mutexM.Load("test")
|
||||
require.True(t, ok)
|
||||
m := v.(*redsync.Mutex)
|
||||
_, _ = m.Unlock() // release it to make it impossible to extend
|
||||
|
||||
// In current design, callers can't know the lock can't be extended.
|
||||
// Just keep this case to improve the test coverage.
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
// Copyright 2024 The Gitea Authors. All rights reserved.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package globallock
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type memoryLocker struct {
|
||||
locks sync.Map
|
||||
}
|
||||
|
||||
var _ Locker = &memoryLocker{}
|
||||
|
||||
func NewMemoryLocker() Locker {
|
||||
return &memoryLocker{}
|
||||
}
|
||||
|
||||
func (l *memoryLocker) Lock(ctx context.Context, key string) (ReleaseFunc, error) {
|
||||
if l.tryLock(key) {
|
||||
releaseOnce := sync.Once{}
|
||||
return func() {
|
||||
releaseOnce.Do(func() {
|
||||
l.locks.Delete(key)
|
||||
})
|
||||
}, nil
|
||||
}
|
||||
|
||||
ticker := time.NewTicker(time.Millisecond * 100)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return func() {}, ctx.Err()
|
||||
case <-ticker.C:
|
||||
if l.tryLock(key) {
|
||||
releaseOnce := sync.Once{}
|
||||
return func() {
|
||||
releaseOnce.Do(func() {
|
||||
l.locks.Delete(key)
|
||||
})
|
||||
}, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (l *memoryLocker) TryLock(_ context.Context, key string) (bool, ReleaseFunc, error) {
|
||||
if l.tryLock(key) {
|
||||
releaseOnce := sync.Once{}
|
||||
return true, func() {
|
||||
releaseOnce.Do(func() {
|
||||
l.locks.Delete(key)
|
||||
})
|
||||
}, nil
|
||||
}
|
||||
|
||||
return false, func() {}, nil
|
||||
}
|
||||
|
||||
func (l *memoryLocker) tryLock(key string) bool {
|
||||
_, loaded := l.locks.LoadOrStore(key, struct{}{})
|
||||
return !loaded
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
// Copyright 2024 The Gitea Authors. All rights reserved.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package globallock
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"gitea.dev/modules/nosql"
|
||||
|
||||
"github.com/go-redsync/redsync/v4"
|
||||
"github.com/go-redsync/redsync/v4/redis/goredis/v9"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
const redisLockKeyPrefix = "gitea:globallock:"
|
||||
|
||||
// redisLockExpiry is the default expiry time for a lock.
|
||||
// Define it as a variable to make it possible to change it in tests.
|
||||
var redisLockExpiry = 30 * time.Second
|
||||
|
||||
type redisLocker struct {
|
||||
conn redis.UniversalClient
|
||||
rs *redsync.Redsync
|
||||
|
||||
mutexM sync.Map
|
||||
closed atomic.Bool
|
||||
extendWg sync.WaitGroup
|
||||
}
|
||||
|
||||
var _ Locker = &redisLocker{}
|
||||
|
||||
func NewRedisLocker(connection string) Locker {
|
||||
conn := nosql.GetManager().GetRedisClient(connection)
|
||||
l := &redisLocker{
|
||||
conn: conn,
|
||||
rs: redsync.New(goredis.NewPool(conn)),
|
||||
}
|
||||
l.extendWg.Add(1)
|
||||
l.startExtend()
|
||||
return l
|
||||
}
|
||||
|
||||
func (l *redisLocker) Lock(ctx context.Context, key string) (ReleaseFunc, error) {
|
||||
return l.lock(ctx, key, 0)
|
||||
}
|
||||
|
||||
func (l *redisLocker) TryLock(ctx context.Context, key string) (bool, ReleaseFunc, error) {
|
||||
f, err := l.lock(ctx, key, 1)
|
||||
|
||||
var (
|
||||
errTaken *redsync.ErrTaken
|
||||
errNodeTaken *redsync.ErrNodeTaken
|
||||
)
|
||||
if errors.As(err, &errTaken) || errors.As(err, &errNodeTaken) {
|
||||
return false, f, nil
|
||||
}
|
||||
return err == nil, f, err
|
||||
}
|
||||
|
||||
// Close closes the locker.
|
||||
// It will stop extending the locks and refuse to acquire new locks.
|
||||
// In actual use, it is not necessary to call this function.
|
||||
// But it's useful in tests to release resources.
|
||||
// It could take some time since it waits for the extending goroutine to finish.
|
||||
func (l *redisLocker) Close() error {
|
||||
l.closed.Store(true)
|
||||
l.extendWg.Wait()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (l *redisLocker) lock(ctx context.Context, key string, tries int) (ReleaseFunc, error) {
|
||||
if l.closed.Load() {
|
||||
return func() {}, errors.New("locker is closed")
|
||||
}
|
||||
|
||||
options := []redsync.Option{
|
||||
redsync.WithExpiry(redisLockExpiry),
|
||||
}
|
||||
if tries > 0 {
|
||||
options = append(options, redsync.WithTries(tries))
|
||||
}
|
||||
mutex := l.rs.NewMutex(redisLockKeyPrefix+key, options...)
|
||||
if err := mutex.LockContext(ctx); err != nil {
|
||||
return func() {}, err
|
||||
}
|
||||
|
||||
l.mutexM.Store(key, mutex)
|
||||
|
||||
releaseOnce := sync.Once{}
|
||||
return func() {
|
||||
releaseOnce.Do(func() {
|
||||
l.mutexM.Delete(key)
|
||||
|
||||
// It's safe to ignore the error here,
|
||||
// if it failed to unlock, it will be released automatically after the lock expires.
|
||||
// Do not call mutex.UnlockContext(ctx) here, or it will fail to release when ctx has timed out.
|
||||
_, _ = mutex.Unlock()
|
||||
})
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (l *redisLocker) startExtend() {
|
||||
if l.closed.Load() {
|
||||
l.extendWg.Done()
|
||||
return
|
||||
}
|
||||
|
||||
toExtend := make([]*redsync.Mutex, 0)
|
||||
l.mutexM.Range(func(_, value any) bool {
|
||||
m := value.(*redsync.Mutex)
|
||||
|
||||
// Extend the lock if it is not expired.
|
||||
// Although the mutex will be removed from the map before it is released,
|
||||
// it still can be expired because of a failed extension.
|
||||
// If it happens, it does not need to be extended anymore.
|
||||
if time.Now().After(m.Until()) {
|
||||
return true
|
||||
}
|
||||
|
||||
toExtend = append(toExtend, m)
|
||||
return true
|
||||
})
|
||||
for _, v := range toExtend {
|
||||
// If it failed to extend, it will be released automatically after the lock expires.
|
||||
_, _ = v.Extend()
|
||||
}
|
||||
|
||||
time.AfterFunc(redisLockExpiry/2, l.startExtend)
|
||||
}
|
||||
Reference in New Issue
Block a user