feat(distlock): implement distributed Locker based on Redlock algorithm

2026-08-28 20:13:33 +00:00 · 2026-08-28 20:34:06 +10:00
parent 613cd6e418
commit 43bf2baaf7
4 changed files with 872 additions and 0 deletions
@@ -0,0 +1,445 @@
+package distlock_test
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/cenkalti/backoff/v4"
+	"github.com/psviderski/uncloud/pkg/distlock"
+	"github.com/stretchr/testify/require"
+)
+
+var errMemoryNodeUnavailable = errors.New("memory node unavailable")
+
+type memoryCluster struct {
+	nodes []*memoryNode
+}
+
+func newMemoryCluster(size int) *memoryCluster {
+	cluster := &memoryCluster{nodes: make([]*memoryNode, size)}
+	for i := range cluster.nodes {
+		cluster.nodes[i] = &memoryNode{store: distlock.NewMemoryStore(), available: true}
+	}
+	return cluster
+}
+
+func (c *memoryCluster) Nodes(ctx context.Context) ([]distlock.Node, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+
+	nodes := make([]distlock.Node, len(c.nodes))
+	for i, node := range c.nodes {
+		nodes[i] = node
+	}
+	return nodes, nil
+}
+
+type memoryNode struct {
+	mu                  sync.Mutex
+	store               distlock.Store
+	available           bool
+	unavailableObserved chan struct{}
+}
+
+func (n *memoryNode) setAvailable(available bool) {
+	if !available {
+		n.makeUnavailable()
+		return
+	}
+
+	n.mu.Lock()
+	defer n.mu.Unlock()
+	n.available = true
+}
+
+func (n *memoryNode) makeUnavailable() <-chan struct{} {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+	n.available = false
+	n.unavailableObserved = make(chan struct{})
+	return n.unavailableObserved
+}
+
+func (n *memoryNode) restart() {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+	n.store = distlock.NewMemoryStore()
+}
+
+func (n *memoryNode) currentStore(ctx context.Context) (distlock.Store, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+
+	n.mu.Lock()
+	defer n.mu.Unlock()
+	if !n.available {
+		if n.unavailableObserved != nil {
+			close(n.unavailableObserved)
+			n.unavailableObserved = nil
+		}
+		return nil, errMemoryNodeUnavailable
+	}
+	return n.store, nil
+}
+
+func (n *memoryNode) Acquire(
+	ctx context.Context, resource string, token []byte, ttl time.Duration,
+) (bool, error) {
+	store, err := n.currentStore(ctx)
+	if err != nil {
+		return false, err
+	}
+	return store.Acquire(ctx, resource, token, ttl)
+}
+
+func (n *memoryNode) Renew(
+	ctx context.Context, resource string, token []byte, ttl time.Duration,
+) (bool, error) {
+	store, err := n.currentStore(ctx)
+	if err != nil {
+		return false, err
+	}
+	return store.Renew(ctx, resource, token, ttl)
+}
+
+func (n *memoryNode) Release(ctx context.Context, resource string, token []byte) (bool, error) {
+	store, err := n.currentStore(ctx)
+	if err != nil {
+		return false, err
+	}
+	return store.Release(ctx, resource, token)
+}
+
+func retryBackOff() backoff.BackOff {
+	return backoff.NewConstantBackOff(5 * time.Millisecond)
+}
+
+func oneAttemptBackOff() backoff.BackOff {
+	return &backoff.StopBackOff{}
+}
+
+func newTestLocker(
+	t *testing.T, cluster distlock.Cluster, leaseDuration time.Duration, newBackOff func() backoff.BackOff,
+) *distlock.Locker {
+	t.Helper()
+	locker, err := distlock.New(cluster, distlock.Config{
+		LeaseDuration:   leaseDuration,
+		NodeCallTimeout: leaseDuration / 3,
+		NewBackOff:      newBackOff,
+	})
+	require.NoError(t, err)
+	return locker
+}
+
+func acquireLease(t *testing.T, locker *distlock.Locker, resource string) *distlock.Lease {
+	t.Helper()
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	lease, err := locker.Acquire(ctx, resource)
+	require.NoError(t, err)
+	return lease
+}
+
+func releaseLease(t *testing.T, lease *distlock.Lease) {
+	t.Helper()
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	require.NoError(t, lease.Release(ctx))
+}
+
+func requireNoReceive[T any](t *testing.T, ch <-chan T, timeout time.Duration, message string) {
+	t.Helper()
+	select {
+	case <-ch:
+		require.FailNow(t, message)
+	case <-time.After(timeout):
+	}
+}
+
+func requireReceive[T any](t *testing.T, ch <-chan T, timeout time.Duration, message string) T {
+	t.Helper()
+	select {
+	case value := <-ch:
+		return value
+	case <-time.After(timeout):
+		require.FailNow(t, message)
+		var zero T
+		return zero
+	}
+}
+
+func requireContextActive(t *testing.T, ctx context.Context, message string) {
+	t.Helper()
+	select {
+	case <-ctx.Done():
+		require.FailNow(t, message, "cause: %v", context.Cause(ctx))
+	default:
+	}
+}
+
+func requireSignal(t *testing.T, ch <-chan struct{}, timeout time.Duration, message string) {
+	t.Helper()
+	select {
+	case <-ch:
+	case <-time.After(timeout):
+		require.FailNow(t, message)
+	}
+}
+
+func TestLockerAcquireQuorumBoundaries(t *testing.T) {
+	tests := []struct {
+		name        string
+		nodes       int
+		unavailable int
+		wantAcquire bool
+	}{
+		{name: "two nodes at quorum", nodes: 2, wantAcquire: true},
+		{name: "two nodes below quorum", nodes: 2, unavailable: 1, wantAcquire: false},
+		{name: "three nodes at quorum", nodes: 3, unavailable: 1, wantAcquire: true},
+		{name: "three nodes below quorum", nodes: 3, unavailable: 2, wantAcquire: false},
+		{name: "four nodes at quorum", nodes: 4, unavailable: 1, wantAcquire: true},
+		{name: "four nodes below quorum", nodes: 4, unavailable: 2, wantAcquire: false},
+		{name: "five nodes at quorum", nodes: 5, unavailable: 2, wantAcquire: true},
+		{name: "five nodes below quorum", nodes: 5, unavailable: 3, wantAcquire: false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cluster := newMemoryCluster(tt.nodes)
+			unavailableObserved := make([]<-chan struct{}, 0, tt.unavailable)
+			for i := range tt.unavailable {
+				unavailableObserved = append(unavailableObserved, cluster.nodes[tt.nodes-1-i].makeUnavailable())
+			}
+			locker := newTestLocker(t, cluster, 300*time.Millisecond, oneAttemptBackOff)
+
+			acquireCtx, cancelAcquire := context.WithTimeout(context.Background(), 2*time.Second)
+			lease, err := locker.Acquire(acquireCtx, "resource")
+			cancelAcquire()
+			for _, observed := range unavailableObserved {
+				requireSignal(t, observed, time.Second, "unavailable node did not receive acquisition")
+			}
+			for i := range tt.unavailable {
+				cluster.nodes[tt.nodes-1-i].setAvailable(true)
+			}
+
+			if tt.wantAcquire {
+				require.NoError(t, err)
+				require.NotNil(t, lease)
+				releaseLease(t, lease)
+			} else {
+				if lease != nil {
+					releaseLease(t, lease)
+				}
+				require.Error(t, err)
+				require.Nil(t, lease)
+			}
+		})
+	}
+}
+
+func TestLockerSingleNodeLifecycle(t *testing.T) {
+	const leaseDuration = 300 * time.Millisecond
+	cluster := newMemoryCluster(1)
+	locker := newTestLocker(t, cluster, leaseDuration, oneAttemptBackOff)
+
+	lease := acquireLease(t, locker, "resource")
+	defer func() {
+		if lease.Context().Err() != nil {
+			return
+		}
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+		_ = lease.Release(ctx)
+	}()
+
+	time.Sleep(leaseDuration + 100*time.Millisecond)
+	requireContextActive(t, lease.Context(), "single-node lease was not renewed")
+	releaseLease(t, lease)
+
+	secondLease := acquireLease(t, locker, "resource")
+	releaseLease(t, secondLease)
+}
+
+func TestLockerFailedAcquireCleansUpPartialLease(t *testing.T) {
+	cluster := newMemoryCluster(3)
+	cluster.nodes[1].setAvailable(false)
+	cluster.nodes[2].setAvailable(false)
+	locker := newTestLocker(t, cluster, 300*time.Millisecond, oneAttemptBackOff)
+
+	lease, err := locker.Acquire(context.Background(), "resource")
+	require.Error(t, err)
+	require.Nil(t, lease)
+
+	cluster.nodes[1].setAvailable(true)
+	cluster.nodes[2].setAvailable(true)
+	lease = acquireLease(t, locker, "resource")
+	releaseLease(t, lease)
+}
+
+func TestLockerAcquiresIndependentResourcesConcurrently(t *testing.T) {
+	cluster := newMemoryCluster(3)
+	locker := newTestLocker(t, cluster, 300*time.Millisecond, oneAttemptBackOff)
+	resources := []string{"database", "deployment", "network", "volume"}
+
+	type acquireResult struct {
+		resource string
+		lease    *distlock.Lease
+		err      error
+	}
+	resultCh := make(chan acquireResult, len(resources))
+	start := make(chan struct{})
+	acquireCtx, cancelAcquire := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancelAcquire()
+	for _, resource := range resources {
+		go func() {
+			<-start
+			lease, err := locker.Acquire(acquireCtx, resource)
+			resultCh <- acquireResult{resource: resource, lease: lease, err: err}
+		}()
+	}
+	close(start)
+
+	results := make([]acquireResult, 0, len(resources))
+	leases := make([]*distlock.Lease, 0, len(resources))
+	for range resources {
+		result := requireReceive(t, resultCh, 2*time.Second, "concurrent acquisition did not finish")
+		results = append(results, result)
+		if result.lease != nil {
+			leases = append(leases, result.lease)
+		}
+	}
+	defer func() {
+		for _, lease := range leases {
+			releaseLease(t, lease)
+		}
+	}()
+
+	for _, result := range results {
+		require.NoErrorf(t, result.err, "acquire %q", result.resource)
+		require.NotNilf(t, result.lease, "acquire %q", result.resource)
+		requireContextActive(t, result.lease.Context(), "independent lease was lost")
+	}
+}
+
+func TestLockerExcludesCompetingLeaseUntilRelease(t *testing.T) {
+	cluster := newMemoryCluster(3)
+	firstLocker := newTestLocker(t, cluster, 300*time.Millisecond, retryBackOff)
+	secondLocker := newTestLocker(t, cluster, 300*time.Millisecond, retryBackOff)
+
+	firstLease := acquireLease(t, firstLocker, "resource")
+
+	type acquireResult struct {
+		lease *distlock.Lease
+		err   error
+	}
+	resultCh := make(chan acquireResult, 1)
+	acquireCtx, cancelAcquire := context.WithTimeout(context.Background(), time.Second)
+	defer cancelAcquire()
+	go func() {
+		lease, acquireErr := secondLocker.Acquire(acquireCtx, "resource")
+		resultCh <- acquireResult{lease: lease, err: acquireErr}
+	}()
+
+	requireNoReceive(t, resultCh, 50*time.Millisecond, "competing acquisition returned before release")
+	releaseLease(t, firstLease)
+	require.ErrorIs(t, context.Cause(firstLease.Context()), distlock.ErrLeaseReleased)
+
+	result := requireReceive(t, resultCh, time.Second, "competing acquisition did not finish after release")
+	require.NoError(t, result.err)
+	require.NotNil(t, result.lease)
+	releaseLease(t, result.lease)
+}
+
+func TestLockerContendingAcquireRespectsContext(t *testing.T) {
+	cluster := newMemoryCluster(3)
+	firstLocker := newTestLocker(t, cluster, 300*time.Millisecond, retryBackOff)
+	secondLocker := newTestLocker(t, cluster, 300*time.Millisecond, retryBackOff)
+
+	firstLease := acquireLease(t, firstLocker, "resource")
+	t.Cleanup(func() {
+		if firstLease.Context().Err() != nil {
+			return
+		}
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+		_ = firstLease.Release(ctx)
+	})
+
+	acquireCtx, cancelAcquire := context.WithTimeout(context.Background(), 75*time.Millisecond)
+	defer cancelAcquire()
+	competingLease, err := secondLocker.Acquire(acquireCtx, "resource")
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	require.Nil(t, competingLease)
+	requireContextActive(t, firstLease.Context(), "holding lease was affected by a competing acquisition")
+
+	releaseLease(t, firstLease)
+	secondLease := acquireLease(t, secondLocker, "resource")
+	releaseLease(t, secondLease)
+}
+
+func TestLockerAutomaticallyRenewsLease(t *testing.T) {
+	const leaseDuration = 300 * time.Millisecond
+	cluster := newMemoryCluster(3)
+	locker := newTestLocker(t, cluster, leaseDuration, retryBackOff)
+	competingLocker := newTestLocker(t, cluster, leaseDuration, oneAttemptBackOff)
+
+	lease := acquireLease(t, locker, "resource")
+	defer releaseLease(t, lease)
+
+	time.Sleep(leaseDuration + 100*time.Millisecond)
+	requireContextActive(t, lease.Context(), "lease was lost instead of renewed")
+
+	competingLease, err := competingLocker.Acquire(context.Background(), "resource")
+	require.Error(t, err)
+	require.Nil(t, competingLease)
+}
+
+func TestLockerRenewsLeaseWithMinorityUnavailable(t *testing.T) {
+	const leaseDuration = 300 * time.Millisecond
+	cluster := newMemoryCluster(3)
+	locker := newTestLocker(t, cluster, leaseDuration, oneAttemptBackOff)
+	competingLocker := newTestLocker(t, cluster, leaseDuration, oneAttemptBackOff)
+
+	lease := acquireLease(t, locker, "resource")
+	unavailableObserved := cluster.nodes[2].makeUnavailable()
+	originalValidityElapsed := time.NewTimer(leaseDuration + 100*time.Millisecond)
+	defer originalValidityElapsed.Stop()
+	defer func() {
+		cluster.nodes[2].setAvailable(true)
+		releaseLease(t, lease)
+	}()
+
+	requireSignal(t, unavailableObserved, time.Second, "unavailable node did not receive renewal")
+	<-originalValidityElapsed.C
+	requireContextActive(t, lease.Context(), "lease was lost after a minority node became unavailable")
+
+	competingLease, err := competingLocker.Acquire(context.Background(), "resource")
+	require.Error(t, err)
+	require.Nil(t, competingLease)
+}
+
+func TestLockerLosesLeaseWhenNodesRestart(t *testing.T) {
+	const leaseDuration = 300 * time.Millisecond
+	cluster := newMemoryCluster(3)
+	locker := newTestLocker(t, cluster, leaseDuration, oneAttemptBackOff)
+
+	lease := acquireLease(t, locker, "resource")
+	for _, node := range cluster.nodes {
+		node.restart()
+	}
+
+	select {
+	case <-lease.Context().Done():
+		require.ErrorIs(t, context.Cause(lease.Context()), distlock.ErrLeaseLost)
+	case <-time.After(time.Second):
+		require.FailNow(t, "lease was not lost after its node state disappeared")
+	}
+
+	releaseLease(t, lease)
+	require.ErrorIs(t, context.Cause(lease.Context()), distlock.ErrLeaseLost)
+}