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,393 @@
+package distlock
+
+import (
+	"context"
+	"crypto/rand"
+	"errors"
+	"fmt"
+	"math"
+	"slices"
+	"sync"
+	"time"
+
+	"github.com/cenkalti/backoff/v4"
+)
+
+const (
+	DefaultLeaseDuration      = 10 * time.Second
+	DefaultClockDriftFactor   = 0.01
+	DefaultMaxNodeCallTimeout = 5 * time.Second
+)
+
+var (
+	// ErrLeaseLost is the cancellation cause when automatic renewal can no longer maintain a lease.
+	ErrLeaseLost = errors.New("distributed lease lost")
+	// ErrLeaseReleased is the cancellation cause of an explicitly released lease.
+	ErrLeaseReleased = errors.New("distributed lease released")
+)
+
+// Config configures a Locker.
+type Config struct {
+	// LeaseDuration is the TTL used for acquisitions and renewals. The default is DefaultLeaseDuration.
+	LeaseDuration time.Duration
+	// ClockDriftFactor is the fraction of LeaseDuration reserved for differences in clock rates between the Locker and
+	// nodes. The default is DefaultClockDriftFactor.
+	ClockDriftFactor float64
+	// NodeCallTimeout sets the context timeout for an Acquire, Renew, or Release call to one node.
+	// The default is the smaller of DefaultMaxNodeCallTimeout and one third of the lease duration.
+	NodeCallTimeout time.Duration
+	// NewBackOff creates independent retry policies for acquisitions and renewal cycles. The default is an exponential
+	// backoff starting at 100ms and capped at 1s. A policy should not impose its own elapsed-time limit because the
+	// acquisition context and current lease validity already bound retries.
+	NewBackOff func() backoff.BackOff
+}
+
+func (c Config) clockDrift() time.Duration {
+	return time.Duration(math.Ceil(float64(c.LeaseDuration) * c.ClockDriftFactor))
+}
+
+// Locker acquires automatically renewed distributed leases over a Cluster.
+type Locker struct {
+	config  Config
+	cluster Cluster
+}
+
+// New creates a Locker over cluster.
+func New(cluster Cluster, config Config) (*Locker, error) {
+	if cluster == nil {
+		return nil, fmt.Errorf("cluster is nil")
+	}
+
+	if config.LeaseDuration == 0 {
+		config.LeaseDuration = DefaultLeaseDuration
+	}
+	if config.LeaseDuration < 0 {
+		return nil, fmt.Errorf("lease duration must be positive")
+	}
+	if config.ClockDriftFactor == 0 {
+		config.ClockDriftFactor = DefaultClockDriftFactor
+	}
+	if config.ClockDriftFactor <= 0 || config.ClockDriftFactor >= 1 {
+		return nil, fmt.Errorf("clock drift factor must be greater than 0 and less than 1")
+	}
+	if config.NodeCallTimeout == 0 {
+		config.NodeCallTimeout = min(DefaultMaxNodeCallTimeout, config.LeaseDuration/3)
+	}
+	if config.NodeCallTimeout < 0 {
+		return nil, fmt.Errorf("node call timeout must be positive")
+	}
+	if config.NewBackOff == nil {
+		config.NewBackOff = defaultBackOff
+	}
+
+	return &Locker{config: config, cluster: cluster}, nil
+}
+
+func defaultBackOff() backoff.BackOff {
+	return backoff.NewExponentialBackOff(
+		backoff.WithInitialInterval(100*time.Millisecond),
+		backoff.WithMaxInterval(time.Second),
+		backoff.WithMaxElapsedTime(0),
+	)
+}
+
+// Acquire waits until it acquires a lease for resource or ctx ends.
+func (l *Locker) Acquire(ctx context.Context, resource string) (*Lease, error) {
+	if resource == "" {
+		return nil, fmt.Errorf("resource is empty")
+	}
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+
+	nodes, err := l.cluster.Nodes(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("get cluster nodes: %w", err)
+	}
+	if len(nodes) == 0 {
+		return nil, fmt.Errorf("cluster has no nodes")
+	}
+	nodes = slices.Clone(nodes)
+
+	boff := backoff.WithContext(l.config.NewBackOff(), ctx)
+	var lease *Lease
+	err = backoff.Retry(func() error {
+		token, tokenErr := newOwnershipToken()
+		if tokenErr != nil {
+			return backoff.Permanent(fmt.Errorf("generate lease token: %w", tokenErr))
+		}
+
+		candidate := newLease(l, nodes, resource, token)
+		if err := candidate.acquire(ctx); err != nil {
+			return err
+		}
+
+		lease = candidate
+		return nil
+	}, boff)
+	if err != nil {
+		return nil, fmt.Errorf("acquire distributed lease for %q: %w", resource, err)
+	}
+
+	return lease, nil
+}
+
+// newOwnershipToken generates a unique 128-bit random ownership token for a lease.
+func newOwnershipToken() ([]byte, error) {
+	token := make([]byte, 16)
+	if _, err := rand.Read(token); err != nil {
+		return nil, err
+	}
+	return token, nil
+}
+
+// Lease is an automatically renewed distributed lease.
+type Lease struct {
+	config   Config
+	nodes    []Node
+	resource string
+	token    []byte
+
+	ctx    context.Context
+	cancel context.CancelCauseFunc
+	// done is closed when the renewal goroutine exits.
+	done chan struct{}
+
+	// operationMu prevents acquisition, renewal, and release operations for the lease from overlapping.
+	operationMu sync.Mutex
+	quorum      int
+}
+
+func newLease(locker *Locker, nodes []Node, resource string, token []byte) *Lease {
+	ctx, cancel := context.WithCancelCause(context.Background())
+	return &Lease{
+		config:   locker.config,
+		nodes:    nodes,
+		resource: resource,
+		token:    token,
+		ctx:      ctx,
+		cancel:   cancel,
+		done:     make(chan struct{}),
+		quorum:   len(nodes)/2 + 1,
+	}
+}
+
+// Context returns a context that is cancelled when the lease is lost or explicitly released.
+func (l *Lease) Context() context.Context {
+	return l.ctx
+}
+
+// Release stops automatic renewal and attempts to remove the lease from every node in its acquisition snapshot.
+func (l *Lease) Release(ctx context.Context) error {
+	l.cancel(ErrLeaseReleased)
+	select {
+	case <-l.done:
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+
+	if err := l.release(ctx); err != nil {
+		return fmt.Errorf("release distributed lease for %q: %w", l.resource, err)
+	}
+	return nil
+}
+
+type nodeResult struct {
+	success bool
+	err     error
+}
+
+func collectNodeResults(results <-chan nodeResult) (successes int, err error) {
+	var errs []error
+	for result := range results {
+		if result.err != nil {
+			errs = append(errs, result.err)
+		} else if result.success {
+			successes++
+		}
+	}
+	return successes, errors.Join(errs...)
+}
+
+func (l *Lease) executeNodes(ctx context.Context, fn func(context.Context, Node) (bool, error)) <-chan nodeResult {
+	results := make(chan nodeResult, len(l.nodes))
+	var wg sync.WaitGroup
+	for _, node := range l.nodes {
+		wg.Go(func() {
+			callCtx, cancel := context.WithTimeout(ctx, l.config.NodeCallTimeout)
+			defer cancel()
+
+			success, err := fn(callCtx, node)
+			results <- nodeResult{success: success, err: err}
+		})
+	}
+	go func() {
+		wg.Wait()
+		close(results)
+	}()
+	return results
+}
+
+// acquire makes one attempt to obtain the lease from a quorum of nodes and starts renewal on success.
+func (l *Lease) acquire(ctx context.Context) error {
+	startedAt := time.Now()
+	validUntil := startedAt.Add(l.config.LeaseDuration - l.config.clockDrift())
+	resultCh := make(chan error, 1)
+	// Aggregate asynchronously so acquire can return at quorum while this goroutine drains the remaining results and
+	// holds operationMu until every node call has finished.
+	go func() {
+		l.operationMu.Lock()
+		defer l.operationMu.Unlock()
+
+		// Use the lease context so cancelling the context passed to Acquire after it succeeds does not stop node calls
+		// still pending after quorum. On failure, acquire cancels the lease context below. The validity deadline and
+		// per-node NodeCallTimeout bound these calls.
+		operationCtx, cancel := context.WithDeadline(l.ctx, validUntil)
+		defer cancel()
+		results := l.executeNodes(operationCtx, func(ctx context.Context, node Node) (bool, error) {
+			return node.Acquire(ctx, l.resource, l.token, l.config.LeaseDuration)
+		})
+
+		successes := 0
+		errs := make([]error, 0, len(l.nodes))
+		reported := false
+		for result := range results {
+			if result.err != nil {
+				errs = append(errs, result.err)
+			} else if result.success {
+				successes++
+			}
+			if !reported && successes >= l.quorum {
+				resultCh <- nil
+				reported = true
+			}
+		}
+		if !reported {
+			quorumErr := fmt.Errorf("lease acquired on %d of %d nodes, need at least %d",
+				successes, len(l.nodes), l.quorum)
+			resultCh <- errors.Join(quorumErr, errors.Join(errs...))
+		}
+	}()
+
+	var acquireErr error
+	select {
+	case acquireErr = <-resultCh:
+	case <-ctx.Done():
+		acquireErr = ctx.Err()
+	}
+
+	if acquireErr == nil {
+		if !time.Now().Before(validUntil) {
+			acquireErr = fmt.Errorf("lease validity expired during acquisition")
+		} else {
+			go l.runRenew(validUntil)
+			return nil
+		}
+	}
+
+	// Cancel any node Acquire calls still in progress. release waits for them to finish before removing partial leases,
+	// so no node can create this lease after cleanup has run.
+	l.cancel(acquireErr)
+	_ = l.release(context.WithoutCancel(ctx))
+	return acquireErr
+}
+
+// release waits for any in-progress lease operation, removes the lease from every node, and returns any errors.
+func (l *Lease) release(ctx context.Context) error {
+	l.operationMu.Lock()
+	defer l.operationMu.Unlock()
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+
+	results := l.executeNodes(ctx, func(ctx context.Context, node Node) (bool, error) {
+		return node.Release(ctx, l.resource, l.token)
+	})
+	_, err := collectNodeResults(results)
+	return err
+}
+
+// runRenew periodically renews the lease until it is released or lost.
+func (l *Lease) runRenew(validUntil time.Time) {
+	defer close(l.done)
+
+	for {
+		remaining := time.Until(validUntil)
+		if remaining <= 0 {
+			break
+		}
+
+		// Start renewal with two thirds of the current validity remaining
+		// to leave time for slow node calls and retries.
+		timer := time.NewTimer(remaining / 3)
+		select {
+		case <-l.ctx.Done():
+			timer.Stop()
+			return
+		case <-timer.C:
+		}
+
+		renewedUntil, err := l.renew(validUntil)
+		if err != nil {
+			if l.ctx.Err() != nil {
+				return
+			}
+			break
+		}
+		validUntil = renewedUntil
+	}
+
+	l.cancel(ErrLeaseLost)
+	_ = l.release(context.Background())
+}
+
+// renew retries node renewals until a quorum succeeds, the configured backoff stops, or the current lease validity
+// ends. It returns the new validity deadline after reaching quorum.
+func (l *Lease) renew(currentValidUntil time.Time) (time.Time, error) {
+	ctx, cancel := context.WithDeadline(l.ctx, currentValidUntil)
+	defer cancel()
+
+	var validUntil time.Time
+	resultCh := make(chan error, 1)
+	// Coordinate in the background so renew can report expiry even if a node call does not return after cancellation.
+	// Keep operationMu held until every call finishes so cleanup cannot race a pending renewal.
+	go func() {
+		l.operationMu.Lock()
+		defer l.operationMu.Unlock()
+
+		if err := ctx.Err(); err != nil {
+			resultCh <- err
+			return
+		}
+
+		boff := backoff.WithContext(l.config.NewBackOff(), ctx)
+		resultCh <- backoff.Retry(func() error {
+			startedAt := time.Now()
+			results := l.executeNodes(ctx, func(ctx context.Context, node Node) (bool, error) {
+				return node.Renew(ctx, l.resource, l.token, l.config.LeaseDuration)
+			})
+			successes, nodeErr := collectNodeResults(results)
+			if ctx.Err() != nil {
+				return ctx.Err()
+			}
+			if !time.Now().Before(currentValidUntil) {
+				return backoff.Permanent(fmt.Errorf("renewal attempt took longer than its validity window"))
+			}
+
+			if successes >= l.quorum {
+				validUntil = startedAt.Add(l.config.LeaseDuration - l.config.clockDrift())
+				return nil
+			}
+			quorumErr := fmt.Errorf("lease renewed on %d of %d nodes, need at least %d",
+				successes, len(l.nodes), l.quorum)
+			return errors.Join(quorumErr, nodeErr)
+		}, boff)
+	}()
+
+	select {
+	case err := <-resultCh:
+		return validUntil, err
+	case <-ctx.Done():
+		return time.Time{}, ctx.Err()
+	}
+}