diff --git a/internal/distlock/doc.go b/internal/distlock/doc.go new file mode 100644 index 00000000..7282f683 --- /dev/null +++ b/internal/distlock/doc.go @@ -0,0 +1,5 @@ +// Package distlock provides distributed, automatically renewed leases across Uncloud machines. +// +// The initial in-memory store does not preserve leases across machine daemon restarts. Cluster membership must remain +// stable while leases are active. Callers must stop protected work when the context returned by Lease.Context is done. +package distlock diff --git a/internal/distlock/memory.go b/internal/distlock/memory.go new file mode 100644 index 00000000..3ee51b4b --- /dev/null +++ b/internal/distlock/memory.go @@ -0,0 +1,139 @@ +package distlock + +import ( + "bytes" + "context" + "fmt" + "sync" + "time" +) + +type memoryLease struct { + token []byte + expiresAt time.Time +} + +// MemoryStore stores leases in memory. Its contents are lost when the process exits. +type MemoryStore struct { + mu sync.Mutex + leases map[string]memoryLease + now func() time.Time +} + +// NewMemoryStore creates an empty in-memory lease store. +func NewMemoryStore() *MemoryStore { + return &MemoryStore{ + leases: make(map[string]memoryLease), + now: time.Now, + } +} + +// Acquire creates a lease when the resource does not have an unexpired lease. +func (s *MemoryStore) Acquire( + ctx context.Context, resource string, token []byte, ttl time.Duration, +) (bool, error) { + if err := validateStoreInput(ctx, resource, token, ttl); err != nil { + return false, err + } + + s.mu.Lock() + defer s.mu.Unlock() + if err := ctx.Err(); err != nil { + return false, err + } + now := s.now() + + if lease, ok := s.leases[resource]; ok && now.Before(lease.expiresAt) { + return false, nil + } + + s.leases[resource] = memoryLease{ + token: bytes.Clone(token), + expiresAt: now.Add(ttl), + } + return true, nil +} + +// Renew extends an unexpired lease when its ownership token matches. +func (s *MemoryStore) Renew( + ctx context.Context, resource string, token []byte, ttl time.Duration, +) (bool, error) { + if err := validateStoreInput(ctx, resource, token, ttl); err != nil { + return false, err + } + + s.mu.Lock() + defer s.mu.Unlock() + if err := ctx.Err(); err != nil { + return false, err + } + now := s.now() + + lease, ok := s.leases[resource] + if !ok { + return false, nil + } + if !now.Before(lease.expiresAt) { + delete(s.leases, resource) + return false, nil + } + if !bytes.Equal(lease.token, token) { + return false, nil + } + + lease.expiresAt = now.Add(ttl) + s.leases[resource] = lease + return true, nil +} + +// Release removes a lease when its ownership token matches. +func (s *MemoryStore) Release(ctx context.Context, resource string, token []byte) (bool, error) { + if err := validateStoreResourceToken(ctx, resource, token); err != nil { + return false, err + } + + s.mu.Lock() + defer s.mu.Unlock() + if err := ctx.Err(); err != nil { + return false, err + } + now := s.now() + + lease, ok := s.leases[resource] + if !ok { + return false, nil + } + if !now.Before(lease.expiresAt) { + delete(s.leases, resource) + return false, nil + } + if !bytes.Equal(lease.token, token) { + return false, nil + } + + delete(s.leases, resource) + return true, nil +} + +func validateStoreInput(ctx context.Context, resource string, token []byte, ttl time.Duration) error { + if err := validateStoreResourceToken(ctx, resource, token); err != nil { + return err + } + if ttl <= 0 { + return fmt.Errorf("TTL must be positive") + } + return nil +} + +func validateStoreResourceToken(ctx context.Context, resource string, token []byte) error { + if err := ctx.Err(); err != nil { + return err + } + if resource == "" { + return fmt.Errorf("resource is empty") + } + if len(token) == 0 { + return fmt.Errorf("token is empty") + } + return nil +} diff --git a/internal/distlock/memory_test.go b/internal/distlock/memory_test.go new file mode 100644 index 00000000..ab5accfe --- /dev/null +++ b/internal/distlock/memory_test.go @@ -0,0 +1,243 @@ +package distlock + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestMemoryStoreAcquire(t *testing.T) { + store := newTestMemoryStore() + ctx := context.Background() + + acquired, err := store.Acquire(ctx, "resource-a", []byte("owner-a"), time.Minute) + require.NoError(t, err) + require.True(t, acquired) + + acquired, err = store.Acquire(ctx, "resource-a", []byte("owner-b"), time.Minute) + require.NoError(t, err) + require.False(t, acquired) + + acquired, err = store.Acquire(ctx, "resource-b", []byte("owner-b"), time.Minute) + require.NoError(t, err) + require.True(t, acquired) +} + +func TestMemoryStoreExpirationAndStaleOwner(t *testing.T) { + now := time.Date(2026, time.August, 26, 12, 0, 0, 0, time.UTC) + store := newTestMemoryStore() + store.now = func() time.Time { return now } + ctx := context.Background() + oldToken := []byte("old-owner") + newToken := []byte("new-owner") + + acquired, err := store.Acquire(ctx, "resource", oldToken, time.Minute) + require.NoError(t, err) + require.True(t, acquired) + + // A lease is expired at its expiration time, not only after it. + now = now.Add(time.Minute) + acquired, err = store.Acquire(ctx, "resource", newToken, time.Minute) + require.NoError(t, err) + require.True(t, acquired) + + released, err := store.Release(ctx, "resource", oldToken) + require.NoError(t, err) + require.False(t, released) + + renewed, err := store.Renew(ctx, "resource", oldToken, time.Minute) + require.NoError(t, err) + require.False(t, renewed) + + released, err = store.Release(ctx, "resource", newToken) + require.NoError(t, err) + require.True(t, released) +} + +func TestMemoryStoreRenew(t *testing.T) { + now := time.Date(2026, time.August, 26, 12, 0, 0, 0, time.UTC) + store := newTestMemoryStore() + store.now = func() time.Time { return now } + ctx := context.Background() + token := []byte("owner") + + acquired, err := store.Acquire(ctx, "resource", token, time.Minute) + require.NoError(t, err) + require.True(t, acquired) + + now = now.Add(30 * time.Second) + renewed, err := store.Renew(ctx, "resource", []byte("another-owner"), time.Minute) + require.NoError(t, err) + require.False(t, renewed) + + renewed, err = store.Renew(ctx, "resource", token, time.Minute) + require.NoError(t, err) + require.True(t, renewed) + + // The renewal extends the lease from the renewal time. + now = now.Add(30 * time.Second) + acquired, err = store.Acquire(ctx, "resource", []byte("another-owner"), time.Minute) + require.NoError(t, err) + require.False(t, acquired) + + now = now.Add(30 * time.Second) + renewed, err = store.Renew(ctx, "resource", token, time.Minute) + require.NoError(t, err) + require.False(t, renewed) +} + +func TestMemoryStoreRelease(t *testing.T) { + store := newTestMemoryStore() + ctx := context.Background() + token := []byte("owner") + + released, err := store.Release(ctx, "resource", token) + require.NoError(t, err) + require.False(t, released) + + acquired, err := store.Acquire(ctx, "resource", token, time.Minute) + require.NoError(t, err) + require.True(t, acquired) + + released, err = store.Release(ctx, "resource", []byte("another-owner")) + require.NoError(t, err) + require.False(t, released) + + released, err = store.Release(ctx, "resource", token) + require.NoError(t, err) + require.True(t, released) + + released, err = store.Release(ctx, "resource", token) + require.NoError(t, err) + require.False(t, released) +} + +func TestMemoryStoreCopiesToken(t *testing.T) { + store := newTestMemoryStore() + ctx := context.Background() + token := []byte("owner") + originalToken := append([]byte(nil), token...) + + acquired, err := store.Acquire(ctx, "resource", token, time.Minute) + require.NoError(t, err) + require.True(t, acquired) + + token[0] = 'x' + released, err := store.Release(ctx, "resource", originalToken) + require.NoError(t, err) + require.True(t, released) +} + +func TestMemoryStoreValidation(t *testing.T) { + tests := []struct { + name string + run func(*MemoryStore) error + }{ + { + name: "acquire with empty resource", + run: func(store *MemoryStore) error { + _, err := store.Acquire(context.Background(), "", []byte("owner"), time.Minute) + return err + }, + }, + { + name: "acquire with empty token", + run: func(store *MemoryStore) error { + _, err := store.Acquire(context.Background(), "resource", nil, time.Minute) + return err + }, + }, + { + name: "acquire with zero TTL", + run: func(store *MemoryStore) error { + _, err := store.Acquire(context.Background(), "resource", []byte("owner"), 0) + return err + }, + }, + { + name: "renew with negative TTL", + run: func(store *MemoryStore) error { + _, err := store.Renew(context.Background(), "resource", []byte("owner"), -time.Second) + return err + }, + }, + { + name: "release with empty token", + run: func(store *MemoryStore) error { + _, err := store.Release(context.Background(), "resource", nil) + return err + }, + }, + { + name: "cancelled context", + run: func(store *MemoryStore) error { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := store.Acquire(ctx, "resource", []byte("owner"), time.Minute) + return err + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + store := newTestMemoryStore() + require.Error(t, test.run(store)) + + // Invalid operations must not create or replace a lease. + acquired, err := store.Acquire(context.Background(), "resource", []byte("valid-owner"), time.Minute) + require.NoError(t, err) + require.True(t, acquired) + }) + } +} + +func TestMemoryStoreConcurrentAcquire(t *testing.T) { + const attempts = 100 + + store := newTestMemoryStore() + ctx := context.Background() + start := make(chan struct{}) + errCh := make(chan error, attempts) + var acquired atomic.Int64 + var wg sync.WaitGroup + + for i := range attempts { + wg.Add(1) + go func() { + defer wg.Done() + <-start + + ok, err := store.Acquire(ctx, "resource", []byte(fmt.Sprintf("owner-%d", i)), time.Minute) + if err != nil { + errCh <- err + return + } + if ok { + acquired.Add(1) + } + }() + } + + close(start) + wg.Wait() + close(errCh) + + for err := range errCh { + require.NoError(t, err) + } + require.EqualValues(t, 1, acquired.Load()) +} + +func newTestMemoryStore() *MemoryStore { + store := NewMemoryStore() + store.now = func() time.Time { + return time.Date(2026, time.August, 26, 12, 0, 0, 0, time.UTC) + } + return store +} diff --git a/internal/distlock/store.go b/internal/distlock/store.go new file mode 100644 index 00000000..55cb063d --- /dev/null +++ b/internal/distlock/store.go @@ -0,0 +1,16 @@ +package distlock + +import ( + "context" + "time" +) + +// Store holds machine-local lease state and provides atomic operations over it. +type Store interface { + // Acquire creates a lease when the resource does not have an unexpired lease. + Acquire(ctx context.Context, resource string, token []byte, ttl time.Duration) (bool, error) + // Renew extends an unexpired lease when its ownership token matches. + Renew(ctx context.Context, resource string, token []byte, ttl time.Duration) (bool, error) + // Release removes a lease when its ownership token matches. + Release(ctx context.Context, resource string, token []byte) (bool, error) +}