1224 lines
33 KiB
Go
1224 lines
33 KiB
Go
package configstore
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/sqlite"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/logger"
|
|
|
|
"github.com/maximhq/bifrost/core/schemas"
|
|
"github.com/maximhq/bifrost/framework/configstore/tables"
|
|
)
|
|
|
|
// mockLogger implements schemas.Logger for testing
|
|
type mockLogger struct {
|
|
debugMessages []string
|
|
infoMessages []string
|
|
warnMessages []string
|
|
errorMessages []string
|
|
mu sync.Mutex
|
|
}
|
|
|
|
func newMockLogger() *mockLogger {
|
|
return &mockLogger{
|
|
debugMessages: make([]string, 0),
|
|
infoMessages: make([]string, 0),
|
|
warnMessages: make([]string, 0),
|
|
errorMessages: make([]string, 0),
|
|
}
|
|
}
|
|
|
|
func (l *mockLogger) Debug(msg string, args ...any) {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
l.debugMessages = append(l.debugMessages, msg)
|
|
}
|
|
|
|
func (l *mockLogger) Info(msg string, args ...any) {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
l.infoMessages = append(l.infoMessages, msg)
|
|
}
|
|
|
|
func (l *mockLogger) Warn(msg string, args ...any) {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
l.warnMessages = append(l.warnMessages, msg)
|
|
}
|
|
|
|
func (l *mockLogger) Error(msg string, args ...any) {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
l.errorMessages = append(l.errorMessages, msg)
|
|
}
|
|
|
|
func (l *mockLogger) Fatal(msg string, args ...any) {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
l.errorMessages = append(l.errorMessages, msg)
|
|
}
|
|
|
|
func (l *mockLogger) SetLevel(level schemas.LogLevel) {}
|
|
|
|
func (l *mockLogger) SetOutputType(outputType schemas.LoggerOutputType) {}
|
|
func (l *mockLogger) LogHTTPRequest(level schemas.LogLevel, msg string) schemas.LogEventBuilder {
|
|
return schemas.NoopLogEvent
|
|
}
|
|
|
|
// setupLockTestStore creates a test RDBConfigStore with SQLite in-memory database
|
|
func setupLockTestStore(t *testing.T) *RDBConfigStore {
|
|
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{
|
|
Logger: logger.Default.LogMode(logger.Silent),
|
|
})
|
|
require.NoError(t, err, "Failed to create test database")
|
|
|
|
// Force single connection to ensure all operations use the same in-memory database.
|
|
// SQLite in-memory with ":memory:" creates separate DBs per connection, so we must
|
|
// limit to one connection to preserve distributed lock semantics in concurrent tests.
|
|
sqlDB, err := db.DB()
|
|
require.NoError(t, err, "Failed to get underlying sql.DB")
|
|
sqlDB.SetMaxOpenConns(1)
|
|
sqlDB.SetMaxIdleConns(1)
|
|
|
|
err = db.AutoMigrate(&tables.TableDistributedLock{})
|
|
require.NoError(t, err, "Failed to migrate test database")
|
|
|
|
s := &RDBConfigStore{logger: newMockLogger()}
|
|
s.db.Store(db)
|
|
s.migrateOnFreshFn = func(ctx context.Context, fn func(context.Context, *gorm.DB) error) error {
|
|
return fn(ctx, s.DB())
|
|
}
|
|
s.refreshPoolFn = func(ctx context.Context) error { return nil }
|
|
return s
|
|
}
|
|
|
|
// =============================================================================
|
|
// LockStore Interface Tests (RDBConfigStore implementation)
|
|
// =============================================================================
|
|
|
|
func TestTryAcquireLock_Success(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
ctx := context.Background()
|
|
|
|
lock := &tables.TableDistributedLock{
|
|
LockKey: "test-lock",
|
|
HolderID: "holder-1",
|
|
ExpiresAt: time.Now().UTC().Add(30 * time.Second),
|
|
}
|
|
|
|
acquired, err := store.TryAcquireLock(ctx, lock)
|
|
require.NoError(t, err)
|
|
assert.True(t, acquired, "Lock should be acquired")
|
|
|
|
// Verify lock exists in database
|
|
storedLock, err := store.GetLock(ctx, "test-lock")
|
|
require.NoError(t, err)
|
|
assert.NotNil(t, storedLock)
|
|
assert.Equal(t, "holder-1", storedLock.HolderID)
|
|
}
|
|
|
|
func TestTryAcquireLock_AlreadyHeld(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
ctx := context.Background()
|
|
|
|
// First holder acquires the lock
|
|
lock1 := &tables.TableDistributedLock{
|
|
LockKey: "test-lock",
|
|
HolderID: "holder-1",
|
|
ExpiresAt: time.Now().UTC().Add(30 * time.Second),
|
|
}
|
|
acquired, err := store.TryAcquireLock(ctx, lock1)
|
|
require.NoError(t, err)
|
|
assert.True(t, acquired)
|
|
|
|
// Second holder tries to acquire the same lock
|
|
lock2 := &tables.TableDistributedLock{
|
|
LockKey: "test-lock",
|
|
HolderID: "holder-2",
|
|
ExpiresAt: time.Now().UTC().Add(30 * time.Second),
|
|
}
|
|
acquired, err = store.TryAcquireLock(ctx, lock2)
|
|
require.NoError(t, err)
|
|
assert.False(t, acquired, "Lock should not be acquired by second holder")
|
|
|
|
// Verify original holder still owns the lock
|
|
storedLock, err := store.GetLock(ctx, "test-lock")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "holder-1", storedLock.HolderID)
|
|
}
|
|
|
|
func TestTryAcquireLock_MultipleLocks(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
ctx := context.Background()
|
|
|
|
locks := []string{"lock-a", "lock-b", "lock-c"}
|
|
|
|
for _, lockKey := range locks {
|
|
lock := &tables.TableDistributedLock{
|
|
LockKey: lockKey,
|
|
HolderID: "holder-1",
|
|
ExpiresAt: time.Now().UTC().Add(30 * time.Second),
|
|
}
|
|
acquired, err := store.TryAcquireLock(ctx, lock)
|
|
require.NoError(t, err)
|
|
assert.True(t, acquired, "Lock %s should be acquired", lockKey)
|
|
}
|
|
|
|
// Verify all locks exist
|
|
for _, lockKey := range locks {
|
|
storedLock, err := store.GetLock(ctx, lockKey)
|
|
require.NoError(t, err)
|
|
assert.NotNil(t, storedLock)
|
|
}
|
|
}
|
|
|
|
func TestGetLock_NotFound(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
ctx := context.Background()
|
|
|
|
lock, err := store.GetLock(ctx, "non-existent-lock")
|
|
require.NoError(t, err)
|
|
assert.Nil(t, lock, "Should return nil for non-existent lock")
|
|
}
|
|
|
|
func TestUpdateLockExpiry_Success(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
ctx := context.Background()
|
|
|
|
// Create initial lock
|
|
lock := &tables.TableDistributedLock{
|
|
LockKey: "test-lock",
|
|
HolderID: "holder-1",
|
|
ExpiresAt: time.Now().UTC().Add(30 * time.Second),
|
|
}
|
|
acquired, err := store.TryAcquireLock(ctx, lock)
|
|
require.NoError(t, err)
|
|
require.True(t, acquired)
|
|
|
|
// Extend expiry
|
|
newExpiry := time.Now().UTC().Add(60 * time.Second)
|
|
err = store.UpdateLockExpiry(ctx, "test-lock", "holder-1", newExpiry)
|
|
require.NoError(t, err)
|
|
|
|
// Verify expiry was updated
|
|
storedLock, err := store.GetLock(ctx, "test-lock")
|
|
require.NoError(t, err)
|
|
assert.WithinDuration(t, newExpiry, storedLock.ExpiresAt, time.Second)
|
|
}
|
|
|
|
func TestUpdateLockExpiry_WrongHolder(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
ctx := context.Background()
|
|
|
|
// Create lock with holder-1
|
|
lock := &tables.TableDistributedLock{
|
|
LockKey: "test-lock",
|
|
HolderID: "holder-1",
|
|
ExpiresAt: time.Now().UTC().Add(30 * time.Second),
|
|
}
|
|
acquired, err := store.TryAcquireLock(ctx, lock)
|
|
require.NoError(t, err)
|
|
require.True(t, acquired)
|
|
|
|
// Try to extend with wrong holder
|
|
newExpiry := time.Now().UTC().Add(60 * time.Second)
|
|
err = store.UpdateLockExpiry(ctx, "test-lock", "holder-2", newExpiry)
|
|
assert.ErrorIs(t, err, ErrLockNotHeld)
|
|
}
|
|
|
|
func TestUpdateLockExpiry_ExpiredLock(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
ctx := context.Background()
|
|
|
|
// Create lock that's already expired
|
|
lock := &tables.TableDistributedLock{
|
|
LockKey: "test-lock",
|
|
HolderID: "holder-1",
|
|
ExpiresAt: time.Now().UTC().Add(-1 * time.Second),
|
|
}
|
|
// Directly insert the expired lock
|
|
err := store.DB().Create(lock).Error
|
|
require.NoError(t, err)
|
|
|
|
// Try to extend expired lock
|
|
newExpiry := time.Now().UTC().Add(60 * time.Second)
|
|
err = store.UpdateLockExpiry(ctx, "test-lock", "holder-1", newExpiry)
|
|
assert.ErrorIs(t, err, ErrLockNotHeld)
|
|
}
|
|
|
|
func TestReleaseLock_Success(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
ctx := context.Background()
|
|
|
|
// Create lock
|
|
lock := &tables.TableDistributedLock{
|
|
LockKey: "test-lock",
|
|
HolderID: "holder-1",
|
|
ExpiresAt: time.Now().UTC().Add(30 * time.Second),
|
|
}
|
|
acquired, err := store.TryAcquireLock(ctx, lock)
|
|
require.NoError(t, err)
|
|
require.True(t, acquired)
|
|
|
|
// Release lock
|
|
released, err := store.ReleaseLock(ctx, "test-lock", "holder-1")
|
|
require.NoError(t, err)
|
|
assert.True(t, released)
|
|
|
|
// Verify lock is gone
|
|
storedLock, err := store.GetLock(ctx, "test-lock")
|
|
require.NoError(t, err)
|
|
assert.Nil(t, storedLock)
|
|
}
|
|
|
|
func TestReleaseLock_WrongHolder(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
ctx := context.Background()
|
|
|
|
// Create lock with holder-1
|
|
lock := &tables.TableDistributedLock{
|
|
LockKey: "test-lock",
|
|
HolderID: "holder-1",
|
|
ExpiresAt: time.Now().UTC().Add(30 * time.Second),
|
|
}
|
|
acquired, err := store.TryAcquireLock(ctx, lock)
|
|
require.NoError(t, err)
|
|
require.True(t, acquired)
|
|
|
|
// Try to release with wrong holder
|
|
released, err := store.ReleaseLock(ctx, "test-lock", "holder-2")
|
|
require.NoError(t, err)
|
|
assert.False(t, released, "Should not release lock with wrong holder")
|
|
|
|
// Verify lock still exists
|
|
storedLock, err := store.GetLock(ctx, "test-lock")
|
|
require.NoError(t, err)
|
|
assert.NotNil(t, storedLock)
|
|
assert.Equal(t, "holder-1", storedLock.HolderID)
|
|
}
|
|
|
|
func TestReleaseLock_NotExists(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
ctx := context.Background()
|
|
|
|
released, err := store.ReleaseLock(ctx, "non-existent-lock", "holder-1")
|
|
require.NoError(t, err)
|
|
assert.False(t, released)
|
|
}
|
|
|
|
func TestCleanupExpiredLocks_Success(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
ctx := context.Background()
|
|
|
|
// Create some expired locks
|
|
expiredLocks := []tables.TableDistributedLock{
|
|
{LockKey: "expired-1", HolderID: "h1", ExpiresAt: time.Now().UTC().Add(-1 * time.Minute)},
|
|
{LockKey: "expired-2", HolderID: "h2", ExpiresAt: time.Now().UTC().Add(-2 * time.Minute)},
|
|
}
|
|
|
|
// Create some valid locks
|
|
validLocks := []tables.TableDistributedLock{
|
|
{LockKey: "valid-1", HolderID: "h3", ExpiresAt: time.Now().UTC().Add(30 * time.Second)},
|
|
{LockKey: "valid-2", HolderID: "h4", ExpiresAt: time.Now().UTC().Add(60 * time.Second)},
|
|
}
|
|
|
|
for _, l := range expiredLocks {
|
|
err := store.DB().Create(&l).Error
|
|
require.NoError(t, err)
|
|
}
|
|
for _, l := range validLocks {
|
|
err := store.DB().Create(&l).Error
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
// Cleanup expired locks
|
|
count, err := store.CleanupExpiredLocks(ctx)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(2), count, "Should cleanup 2 expired locks")
|
|
|
|
// Verify expired locks are gone
|
|
for _, l := range expiredLocks {
|
|
lock, err := store.GetLock(ctx, l.LockKey)
|
|
require.NoError(t, err)
|
|
assert.Nil(t, lock)
|
|
}
|
|
|
|
// Verify valid locks still exist
|
|
for _, l := range validLocks {
|
|
lock, err := store.GetLock(ctx, l.LockKey)
|
|
require.NoError(t, err)
|
|
assert.NotNil(t, lock)
|
|
}
|
|
}
|
|
|
|
func TestCleanupExpiredLocks_NoExpired(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
ctx := context.Background()
|
|
|
|
// Create only valid locks
|
|
lock := &tables.TableDistributedLock{
|
|
LockKey: "valid-lock",
|
|
HolderID: "holder-1",
|
|
ExpiresAt: time.Now().UTC().Add(30 * time.Second),
|
|
}
|
|
_, err := store.TryAcquireLock(ctx, lock)
|
|
require.NoError(t, err)
|
|
|
|
count, err := store.CleanupExpiredLocks(ctx)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(0), count)
|
|
}
|
|
|
|
func TestCleanupExpiredLockByKey_Success(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
ctx := context.Background()
|
|
|
|
// Create expired lock
|
|
lock := tables.TableDistributedLock{
|
|
LockKey: "expired-lock",
|
|
HolderID: "holder-1",
|
|
ExpiresAt: time.Now().UTC().Add(-1 * time.Minute),
|
|
}
|
|
err := store.DB().Create(&lock).Error
|
|
require.NoError(t, err)
|
|
|
|
// Cleanup specific expired lock
|
|
cleaned, err := store.CleanupExpiredLockByKey(ctx, "expired-lock")
|
|
require.NoError(t, err)
|
|
assert.True(t, cleaned)
|
|
|
|
// Verify lock is gone
|
|
storedLock, err := store.GetLock(ctx, "expired-lock")
|
|
require.NoError(t, err)
|
|
assert.Nil(t, storedLock)
|
|
}
|
|
|
|
func TestCleanupExpiredLockByKey_NotExpired(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
ctx := context.Background()
|
|
|
|
// Create valid lock
|
|
lock := &tables.TableDistributedLock{
|
|
LockKey: "valid-lock",
|
|
HolderID: "holder-1",
|
|
ExpiresAt: time.Now().UTC().Add(30 * time.Second),
|
|
}
|
|
_, err := store.TryAcquireLock(ctx, lock)
|
|
require.NoError(t, err)
|
|
|
|
// Try to cleanup non-expired lock
|
|
cleaned, err := store.CleanupExpiredLockByKey(ctx, "valid-lock")
|
|
require.NoError(t, err)
|
|
assert.False(t, cleaned, "Should not cleanup non-expired lock")
|
|
|
|
// Verify lock still exists
|
|
storedLock, err := store.GetLock(ctx, "valid-lock")
|
|
require.NoError(t, err)
|
|
assert.NotNil(t, storedLock)
|
|
}
|
|
|
|
func TestCleanupExpiredLockByKey_NotExists(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
ctx := context.Background()
|
|
|
|
cleaned, err := store.CleanupExpiredLockByKey(ctx, "non-existent")
|
|
require.NoError(t, err)
|
|
assert.False(t, cleaned)
|
|
}
|
|
|
|
// =============================================================================
|
|
// DistributedLockManager Tests
|
|
// =============================================================================
|
|
|
|
func TestNewDistributedLockManager_DefaultOptions(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
|
|
manager := NewDistributedLockManager(store, logger)
|
|
|
|
assert.NotNil(t, manager)
|
|
assert.Equal(t, DefaultLockTTL, manager.defaultTTL)
|
|
assert.Equal(t, DefaultRetryInterval, manager.retryInterval)
|
|
assert.Equal(t, DefaultMaxRetries, manager.maxRetries)
|
|
}
|
|
|
|
func TestNewDistributedLockManager_CustomOptions(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
|
|
customTTL := 60 * time.Second
|
|
customRetryInterval := 200 * time.Millisecond
|
|
customMaxRetries := 50
|
|
|
|
manager := NewDistributedLockManager(
|
|
store,
|
|
logger,
|
|
WithDefaultTTL(customTTL),
|
|
WithRetryInterval(customRetryInterval),
|
|
WithMaxRetries(customMaxRetries),
|
|
)
|
|
|
|
assert.Equal(t, customTTL, manager.defaultTTL)
|
|
assert.Equal(t, customRetryInterval, manager.retryInterval)
|
|
assert.Equal(t, customMaxRetries, manager.maxRetries)
|
|
}
|
|
|
|
func TestDistributedLockManager_NewLock(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
|
|
lock, err := manager.NewLock("my-lock")
|
|
require.NoError(t, err)
|
|
|
|
assert.NotNil(t, lock)
|
|
assert.Equal(t, "my-lock", lock.Key())
|
|
assert.NotEmpty(t, lock.HolderID())
|
|
assert.Equal(t, DefaultLockTTL, lock.ttl)
|
|
}
|
|
|
|
func TestDistributedLockManager_NewLockWithTTL(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
|
|
customTTL := 5 * time.Minute
|
|
lock, err := manager.NewLockWithTTL("my-lock", customTTL)
|
|
require.NoError(t, err)
|
|
|
|
assert.Equal(t, customTTL, lock.ttl)
|
|
}
|
|
|
|
func TestDistributedLockManager_CleanupExpiredLocks(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
// Create expired lock directly
|
|
lock := tables.TableDistributedLock{
|
|
LockKey: "expired-lock",
|
|
HolderID: "holder-1",
|
|
ExpiresAt: time.Now().UTC().Add(-1 * time.Minute),
|
|
}
|
|
err := store.DB().Create(&lock).Error
|
|
require.NoError(t, err)
|
|
|
|
count, err := manager.CleanupExpiredLocks(ctx)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(1), count)
|
|
}
|
|
|
|
// =============================================================================
|
|
// DistributedLock Tests
|
|
// =============================================================================
|
|
|
|
func TestDistributedLock_TryLock_Success(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
lock, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
|
|
acquired, err := lock.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
assert.True(t, acquired)
|
|
}
|
|
|
|
func TestDistributedLock_TryLock_AlreadyHeld(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
lock1, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
lock2, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
|
|
// First lock succeeds
|
|
acquired, err := lock1.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
assert.True(t, acquired)
|
|
|
|
// Second lock fails
|
|
acquired, err = lock2.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
assert.False(t, acquired)
|
|
}
|
|
|
|
func TestDistributedLock_TryLock_CleansUpExpired(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
// Create expired lock directly in database
|
|
expiredLock := tables.TableDistributedLock{
|
|
LockKey: "test-lock",
|
|
HolderID: "old-holder",
|
|
ExpiresAt: time.Now().UTC().Add(-1 * time.Minute),
|
|
}
|
|
err := store.DB().Create(&expiredLock).Error
|
|
require.NoError(t, err)
|
|
|
|
// New lock should be able to acquire after cleanup
|
|
lock, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
acquired, err := lock.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
assert.True(t, acquired, "Should acquire lock after expired cleanup")
|
|
}
|
|
|
|
func TestDistributedLock_Lock_Success(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger,
|
|
WithRetryInterval(10*time.Millisecond),
|
|
WithMaxRetries(3),
|
|
)
|
|
ctx := context.Background()
|
|
|
|
lock, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
|
|
err = lock.Lock(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// Verify lock is held
|
|
held, err := lock.IsHeld(ctx)
|
|
require.NoError(t, err)
|
|
assert.True(t, held)
|
|
}
|
|
|
|
func TestDistributedLock_Lock_ContextCancelled(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
|
|
// First acquire the lock
|
|
lock1, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
ctx := context.Background()
|
|
_, err = lock1.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// Try to acquire with cancelled context
|
|
cancelCtx, cancel := context.WithCancel(context.Background())
|
|
cancel() // Cancel immediately
|
|
|
|
lock2, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
err = lock2.Lock(cancelCtx)
|
|
assert.ErrorIs(t, err, context.Canceled)
|
|
}
|
|
|
|
func TestDistributedLock_Lock_Timeout(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger,
|
|
WithRetryInterval(10*time.Millisecond),
|
|
WithMaxRetries(3),
|
|
)
|
|
ctx := context.Background()
|
|
|
|
// First lock holds
|
|
lock1, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
_, err = lock1.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// Second lock should fail after retries
|
|
lock2, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
err = lock2.Lock(ctx)
|
|
assert.ErrorIs(t, err, ErrLockNotAcquired)
|
|
}
|
|
|
|
func TestDistributedLock_LockWithRetry_ExponentialBackoff(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
// First lock holds
|
|
lock1, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
_, err = lock1.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// Second lock with limited retries
|
|
lock2, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
start := time.Now()
|
|
err = lock2.LockWithRetry(ctx, 2)
|
|
elapsed := time.Since(start)
|
|
|
|
assert.ErrorIs(t, err, ErrLockNotAcquired)
|
|
// With exponential backoff: 1s + 2s = 3s minimum
|
|
assert.True(t, elapsed >= 3*time.Second, "Expected at least 3s delay, got %v", elapsed)
|
|
}
|
|
|
|
func TestDistributedLock_Unlock_Success(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
lock, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
_, err = lock.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
|
|
err = lock.Unlock(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// Verify lock is released
|
|
storedLock, err := store.GetLock(ctx, "test-lock")
|
|
require.NoError(t, err)
|
|
assert.Nil(t, storedLock)
|
|
}
|
|
|
|
func TestDistributedLock_Unlock_NotHeld(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
lock, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
// Never acquired
|
|
|
|
err = lock.Unlock(ctx)
|
|
assert.ErrorIs(t, err, ErrLockNotHeld)
|
|
}
|
|
|
|
func TestDistributedLock_Unlock_AlreadyReleased(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
lock, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
_, err = lock.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// First unlock succeeds
|
|
err = lock.Unlock(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// Second unlock fails
|
|
err = lock.Unlock(ctx)
|
|
assert.ErrorIs(t, err, ErrLockNotHeld)
|
|
}
|
|
|
|
func TestDistributedLock_Extend_Success(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger, WithDefaultTTL(10*time.Second))
|
|
ctx := context.Background()
|
|
|
|
lock, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
_, err = lock.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// Get original expiry
|
|
storedLock, err := store.GetLock(ctx, "test-lock")
|
|
require.NoError(t, err)
|
|
originalExpiry := storedLock.ExpiresAt
|
|
|
|
// Wait a bit and extend
|
|
time.Sleep(100 * time.Millisecond)
|
|
err = lock.Extend(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// Verify expiry was extended
|
|
storedLock, err = store.GetLock(ctx, "test-lock")
|
|
require.NoError(t, err)
|
|
assert.True(t, storedLock.ExpiresAt.After(originalExpiry))
|
|
}
|
|
|
|
func TestDistributedLock_Extend_NotHeld(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
lock, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
// Never acquired
|
|
|
|
err = lock.Extend(ctx)
|
|
assert.ErrorIs(t, err, ErrLockNotHeld)
|
|
}
|
|
|
|
func TestDistributedLock_Extend_StolenLock(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
lock, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
_, err = lock.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// Simulate lock being stolen by another process
|
|
err = store.DB().Model(&tables.TableDistributedLock{}).
|
|
Where("lock_key = ?", "test-lock").
|
|
Update("holder_id", "another-holder").Error
|
|
require.NoError(t, err)
|
|
|
|
// Extend should fail
|
|
err = lock.Extend(ctx)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestDistributedLock_IsHeld_True(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
lock, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
_, err = lock.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
|
|
held, err := lock.IsHeld(ctx)
|
|
require.NoError(t, err)
|
|
assert.True(t, held)
|
|
}
|
|
|
|
func TestDistributedLock_IsHeld_NotAcquired(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
lock, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
// Never acquired
|
|
|
|
held, err := lock.IsHeld(ctx)
|
|
require.NoError(t, err)
|
|
assert.False(t, held)
|
|
}
|
|
|
|
func TestDistributedLock_IsHeld_Expired(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger, WithDefaultTTL(50*time.Millisecond))
|
|
ctx := context.Background()
|
|
|
|
lock, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
_, err = lock.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// Wait for lock to expire
|
|
time.Sleep(100 * time.Millisecond)
|
|
|
|
held, err := lock.IsHeld(ctx)
|
|
require.NoError(t, err)
|
|
assert.False(t, held, "Lock should not be held after expiry")
|
|
}
|
|
|
|
func TestDistributedLock_IsHeld_StolenByAnotherHolder(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
lock, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
_, err = lock.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// Simulate lock being stolen by another process
|
|
err = store.DB().Model(&tables.TableDistributedLock{}).
|
|
Where("lock_key = ?", "test-lock").
|
|
Update("holder_id", "another-holder").Error
|
|
require.NoError(t, err)
|
|
|
|
held, err := lock.IsHeld(ctx)
|
|
require.NoError(t, err)
|
|
assert.False(t, held)
|
|
}
|
|
|
|
func TestDistributedLock_IsHeld_DeletedFromDB(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
lock, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
_, err = lock.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// Delete lock directly from database
|
|
err = store.DB().Where("lock_key = ?", "test-lock").Delete(&tables.TableDistributedLock{}).Error
|
|
require.NoError(t, err)
|
|
|
|
held, err := lock.IsHeld(ctx)
|
|
require.NoError(t, err)
|
|
assert.False(t, held)
|
|
}
|
|
|
|
// =============================================================================
|
|
// Concurrent Access Tests
|
|
// =============================================================================
|
|
|
|
func TestDistributedLock_ConcurrentAcquire(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger,
|
|
WithRetryInterval(10*time.Millisecond),
|
|
WithMaxRetries(5),
|
|
)
|
|
ctx := context.Background()
|
|
|
|
const numGoroutines = 10
|
|
successCount := 0
|
|
var mu sync.Mutex
|
|
var wg sync.WaitGroup
|
|
|
|
for i := 0; i < numGoroutines; i++ {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
lock, err := manager.NewLock("contended-lock")
|
|
if err != nil {
|
|
return
|
|
}
|
|
acquired, err := lock.TryLock(ctx)
|
|
if err == nil && acquired {
|
|
mu.Lock()
|
|
successCount++
|
|
mu.Unlock()
|
|
}
|
|
}()
|
|
}
|
|
|
|
wg.Wait()
|
|
|
|
// Only one goroutine should have acquired the lock
|
|
assert.Equal(t, 1, successCount, "Exactly one goroutine should acquire the lock")
|
|
}
|
|
|
|
func TestDistributedLock_ConcurrentLockUnlock(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger,
|
|
WithRetryInterval(50*time.Millisecond),
|
|
WithMaxRetries(100),
|
|
WithDefaultTTL(5*time.Second),
|
|
)
|
|
ctx := context.Background()
|
|
|
|
const numGoroutines = 5
|
|
counter := 0
|
|
var mu sync.Mutex
|
|
var wg sync.WaitGroup
|
|
|
|
for i := 0; i < numGoroutines; i++ {
|
|
wg.Add(1)
|
|
go func(id int) {
|
|
defer wg.Done()
|
|
lock, err := manager.NewLock("counter-lock")
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
// Each goroutine tries to increment the counter with lock protection
|
|
err = lock.Lock(ctx)
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
mu.Lock()
|
|
counter++
|
|
mu.Unlock()
|
|
|
|
// Simulate some work
|
|
time.Sleep(10 * time.Millisecond)
|
|
|
|
_ = lock.Unlock(ctx)
|
|
}(i)
|
|
}
|
|
|
|
wg.Wait()
|
|
|
|
// All goroutines should have incremented the counter
|
|
assert.Equal(t, numGoroutines, counter, "All goroutines should complete")
|
|
}
|
|
|
|
func TestDistributedLock_MultipleLocksPerManager(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
const numLocks = 10
|
|
var wg sync.WaitGroup
|
|
errCh := make(chan error, numLocks*2) // Buffer for potential TryLock and Unlock errors
|
|
|
|
for i := 0; i < numLocks; i++ {
|
|
wg.Add(1)
|
|
go func(id int) {
|
|
defer wg.Done()
|
|
lockKey := fmt.Sprintf("lock-%d", id)
|
|
lock, err := manager.NewLock(lockKey)
|
|
if err != nil {
|
|
errCh <- fmt.Errorf("lock %s NewLock error: %w", lockKey, err)
|
|
return
|
|
}
|
|
|
|
acquired, err := lock.TryLock(ctx)
|
|
if err != nil {
|
|
errCh <- fmt.Errorf("lock %s TryLock error: %w", lockKey, err)
|
|
return
|
|
}
|
|
if !acquired {
|
|
errCh <- fmt.Errorf("lock %s should be acquired", lockKey)
|
|
return
|
|
}
|
|
|
|
time.Sleep(10 * time.Millisecond)
|
|
|
|
if err := lock.Unlock(ctx); err != nil {
|
|
errCh <- fmt.Errorf("lock %s Unlock error: %w", lockKey, err)
|
|
return
|
|
}
|
|
}(i)
|
|
}
|
|
|
|
wg.Wait()
|
|
close(errCh)
|
|
|
|
// Collect and report any errors from goroutines
|
|
var errs []error
|
|
for err := range errCh {
|
|
errs = append(errs, err)
|
|
}
|
|
require.Empty(t, errs, "goroutines reported errors: %v", errs)
|
|
}
|
|
|
|
// =============================================================================
|
|
// Edge Case Tests
|
|
// =============================================================================
|
|
|
|
func TestDistributedLock_EmptyLockKey(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
|
|
lock, err := manager.NewLock("")
|
|
|
|
assert.Nil(t, lock, "Empty lock key should return nil lock")
|
|
assert.ErrorIs(t, err, ErrEmptyLockKey, "Empty lock key should return ErrEmptyLockKey")
|
|
}
|
|
|
|
func TestDistributedLock_LongLockKey(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
// Create a 255 character lock key (max size per schema)
|
|
longKey := ""
|
|
for i := 0; i < 255; i++ {
|
|
longKey += "a"
|
|
}
|
|
|
|
lock, err := manager.NewLock(longKey)
|
|
require.NoError(t, err)
|
|
|
|
acquired, err := lock.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
assert.True(t, acquired)
|
|
}
|
|
|
|
func TestDistributedLock_SpecialCharactersInKey(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
specialKeys := []string{
|
|
"lock:with:colons",
|
|
"lock/with/slashes",
|
|
"lock-with-dashes",
|
|
"lock_with_underscores",
|
|
"lock.with.dots",
|
|
"lock with spaces",
|
|
"lock\twith\ttabs",
|
|
}
|
|
|
|
for _, key := range specialKeys {
|
|
lock, err := manager.NewLock(key)
|
|
require.NoError(t, err, "Key: %s", key)
|
|
acquired, err := lock.TryLock(ctx)
|
|
require.NoError(t, err, "Key: %s", key)
|
|
assert.True(t, acquired, "Lock with key %s should be acquired", key)
|
|
|
|
err = lock.Unlock(ctx)
|
|
require.NoError(t, err)
|
|
}
|
|
}
|
|
|
|
func TestDistributedLock_ZeroTTL(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger, WithDefaultTTL(0))
|
|
ctx := context.Background()
|
|
|
|
lock, err := manager.NewLock("zero-ttl-lock")
|
|
require.NoError(t, err)
|
|
|
|
// Zero TTL should still work but lock will be immediately expired
|
|
acquired, err := lock.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
assert.True(t, acquired)
|
|
|
|
// Lock should immediately appear not held due to zero TTL
|
|
held, err := lock.IsHeld(ctx)
|
|
require.NoError(t, err)
|
|
assert.False(t, held, "Zero TTL lock should not be held")
|
|
}
|
|
|
|
func TestDistributedLock_VeryShortTTL(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger, WithDefaultTTL(1*time.Millisecond))
|
|
ctx := context.Background()
|
|
|
|
lock, err := manager.NewLock("short-ttl-lock")
|
|
require.NoError(t, err)
|
|
|
|
acquired, err := lock.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
assert.True(t, acquired)
|
|
|
|
// Wait for TTL to expire
|
|
time.Sleep(10 * time.Millisecond)
|
|
|
|
// Another lock should be able to acquire
|
|
lock2, err := manager.NewLock("short-ttl-lock")
|
|
require.NoError(t, err)
|
|
acquired, err = lock2.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
assert.True(t, acquired, "Should acquire lock after TTL expires")
|
|
}
|
|
|
|
func TestDistributedLock_ReacquireAfterUnlock(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
ctx := context.Background()
|
|
|
|
lock, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
|
|
// First acquire
|
|
acquired, err := lock.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
assert.True(t, acquired)
|
|
|
|
// Release
|
|
err = lock.Unlock(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// Same lock instance should NOT be able to reacquire (new holder ID needed)
|
|
// But a new lock should work
|
|
lock2, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
acquired, err = lock2.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
assert.True(t, acquired, "New lock instance should acquire after release")
|
|
}
|
|
|
|
func TestDistributedLock_ExtendMultipleTimes(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger, WithDefaultTTL(100*time.Millisecond))
|
|
ctx := context.Background()
|
|
|
|
lock, err := manager.NewLock("test-lock")
|
|
require.NoError(t, err)
|
|
_, err = lock.TryLock(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// Extend multiple times
|
|
for i := 0; i < 5; i++ {
|
|
time.Sleep(20 * time.Millisecond)
|
|
err = lock.Extend(ctx)
|
|
require.NoError(t, err, "Extension %d failed", i+1)
|
|
}
|
|
|
|
// Lock should still be held
|
|
held, err := lock.IsHeld(ctx)
|
|
require.NoError(t, err)
|
|
assert.True(t, held)
|
|
}
|
|
|
|
// =============================================================================
|
|
// Constants and Defaults Tests
|
|
// =============================================================================
|
|
|
|
func TestConstants(t *testing.T) {
|
|
assert.Equal(t, 30*time.Second, DefaultLockTTL)
|
|
assert.Equal(t, 100*time.Millisecond, DefaultRetryInterval)
|
|
assert.Equal(t, 100, DefaultMaxRetries)
|
|
assert.Equal(t, 5*time.Minute, DefaultCleanupInterval)
|
|
}
|
|
|
|
func TestErrors(t *testing.T) {
|
|
assert.Equal(t, "failed to acquire lock", ErrLockNotAcquired.Error())
|
|
assert.Equal(t, "lock not held by this holder", ErrLockNotHeld.Error())
|
|
assert.Equal(t, "lock has expired", ErrLockExpired.Error())
|
|
assert.Equal(t, "empty lock key", ErrEmptyLockKey.Error())
|
|
}
|
|
|
|
// =============================================================================
|
|
// Key and HolderID Accessor Tests
|
|
// =============================================================================
|
|
|
|
func TestDistributedLock_Key(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
|
|
lock, err := manager.NewLock("my-unique-lock")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "my-unique-lock", lock.Key())
|
|
}
|
|
|
|
func TestDistributedLock_HolderID(t *testing.T) {
|
|
store := setupLockTestStore(t)
|
|
logger := newMockLogger()
|
|
manager := NewDistributedLockManager(store, logger)
|
|
|
|
lock1, err := manager.NewLock("lock-1")
|
|
require.NoError(t, err)
|
|
lock2, err := manager.NewLock("lock-2")
|
|
require.NoError(t, err)
|
|
|
|
// Each lock should have a unique holder ID
|
|
assert.NotEmpty(t, lock1.HolderID())
|
|
assert.NotEmpty(t, lock2.HolderID())
|
|
assert.NotEqual(t, lock1.HolderID(), lock2.HolderID())
|
|
}
|