feat(distlock): implement in-memory lease store

2026-08-28 12:03:33 +00:00 · 2026-08-26 16:53:21 +10:00
parent 5a956e72de
commit a5ab6e3788
4 changed files with 403 additions and 0 deletions
@@ -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
@@ -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
+}
@@ -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
+}
@@ -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)
+}