anchorage/internal/pkg/leader/leader.go

// Package leader implements TTL-based leader election via NATS JetStream KV.
//
// The pattern, lifted from the sibling kanrisha project:
//
//  1. A JetStream KV bucket "ANCHORAGE_LEADER" holds a single key "leader"
//     whose value is the current leader's node ID. The bucket TTL is set to
//     a few seconds (5s default).
//  2. Every candidate tries `kv.Create("leader", myNodeID)` on a loop. Exactly
//     one CAS succeeds — that node is the leader.
//  3. The leader renews the key (via kv.Update with the known revision)
//     roughly once per second. Missing three renews lets the TTL expire; a
//     candidate then wins the next Create.
//  4. Losing leadership fires OnDemote; winning it fires OnPromote. Both
//     callbacks run synchronously in the Elector goroutine.
//
// The design accepts brief periods of "two leaders" during a partition as
// an acceptable cost: every leader-only job in anchorage (rebalancer,
// sweeper) is idempotent at the Postgres level.
package leader

import (
	"context"
	"errors"
	"log/slog"
	"sync/atomic"
	"time"

	"github.com/nats-io/nats.go/jetstream"

	"anchorage/internal/pkg/metrics"
)

// DefaultBucket is the JetStream KV bucket name.
const DefaultBucket = "ANCHORAGE_LEADER"

// DefaultKey is the sole key in the bucket.
const DefaultKey = "leader"

// Options tunes election behavior.
type Options struct {
	// NodeID is the value written to the KV. Must be unique per candidate.
	NodeID string
	// Bucket overrides DefaultBucket.
	Bucket string
	// Key overrides DefaultKey.
	Key string
	// TTL is the key expiry. If renewals lapse for >TTL, the key vanishes
	// and another candidate can claim it. Default 5s.
	TTL time.Duration
	// RenewInterval is how often the leader refreshes the key. Should be
	// comfortably under TTL — default TTL/5.
	RenewInterval time.Duration
	// ClaimInterval is how often candidates try to claim when there's no
	// leader. Default equal to RenewInterval.
	ClaimInterval time.Duration
	// OnPromote fires when this node becomes leader.
	OnPromote func(ctx context.Context)
	// OnDemote fires when this node stops being leader (voluntarily or
	// because a renewal failed).
	OnDemote func(ctx context.Context)
}

// Elector runs a leader-election loop in its own goroutine.
type Elector struct {
	opts    Options
	kv      jetstream.KeyValue
	leading atomic.Bool
}

// New creates (or opens) the KV bucket and returns an Elector. It does
// not start the loop — call Run to begin claiming.
func New(ctx context.Context, js jetstream.JetStream, opts Options) (*Elector, error) {
	if opts.NodeID == "" {
		return nil, errors.New("leader: NodeID is required")
	}
	if opts.Bucket == "" {
		opts.Bucket = DefaultBucket
	}
	if opts.Key == "" {
		opts.Key = DefaultKey
	}
	if opts.TTL <= 0 {
		opts.TTL = 5 * time.Second
	}
	if opts.RenewInterval <= 0 {
		opts.RenewInterval = opts.TTL / 5
		if opts.RenewInterval < 250*time.Millisecond {
			opts.RenewInterval = 250 * time.Millisecond
		}
	}
	if opts.ClaimInterval <= 0 {
		opts.ClaimInterval = opts.RenewInterval
	}

	kv, err := js.CreateOrUpdateKeyValue(ctx, jetstream.KeyValueConfig{
		Bucket: opts.Bucket,
		TTL:    opts.TTL,
	})
	if err != nil {
		return nil, err
	}
	return &Elector{opts: opts, kv: kv}, nil
}

// IsLeader reports whether this node currently holds leadership.
func (e *Elector) IsLeader() bool { return e.leading.Load() }

// Run blocks until ctx is cancelled. It acquires, renews, or fails over
// leadership, invoking OnPromote / OnDemote on transitions.
//
// Return value: ctx.Err() when cancelled, or a persistent infrastructure
// error (e.g., bucket deletion) — transient KV errors are logged and
// retried.
func (e *Elector) Run(ctx context.Context) error {
	var revision uint64

	// Loop pacing: use separate tickers for claim and renew so the code
	// doesn't need to track "leading vs not" state twice.
	tryTicker := time.NewTicker(e.opts.ClaimInterval)
	defer tryTicker.Stop()
	renewTicker := time.NewTicker(e.opts.RenewInterval)
	defer renewTicker.Stop()

	for {
		if e.leading.Load() {
			// As leader, renew on every renewTicker tick.
			select {
			case <-ctx.Done():
				e.stepDown(ctx, revision)
				return ctx.Err()
			case <-renewTicker.C:
				rev, err := e.kv.Update(ctx, e.opts.Key, []byte(e.opts.NodeID), revision)
				if err != nil {
					slog.Warn("leader: renew failed — stepping down", "err", err)
					e.demote(ctx)
					revision = 0
					continue
				}
				revision = rev
			}
		} else {
			// As follower, try to claim.
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-tryTicker.C:
				rev, err := e.kv.Create(ctx, e.opts.Key, []byte(e.opts.NodeID))
				if err != nil {
					// Key already held by someone else — normal follower state.
					continue
				}
				revision = rev
				e.promote(ctx)
			}
		}
	}
}

// StepDown relinquishes leadership if held. Safe to call from any goroutine.
func (e *Elector) StepDown(ctx context.Context) {
	if !e.leading.CompareAndSwap(true, false) {
		return
	}
	// Best effort: delete the key so the next follower claims instantly.
	_ = e.kv.Delete(ctx, e.opts.Key)
	if e.opts.OnDemote != nil {
		e.opts.OnDemote(ctx)
	}
}

func (e *Elector) promote(ctx context.Context) {
	if !e.leading.CompareAndSwap(false, true) {
		return
	}
	metrics.LeaderIsElected.Set(1)
	slog.Info("leader: promoted", "node_id", e.opts.NodeID, "bucket", e.opts.Bucket)
	if e.opts.OnPromote != nil {
		e.opts.OnPromote(ctx)
	}
}

func (e *Elector) demote(ctx context.Context) {
	if !e.leading.CompareAndSwap(true, false) {
		return
	}
	metrics.LeaderIsElected.Set(0)
	slog.Info("leader: demoted", "node_id", e.opts.NodeID)
	if e.opts.OnDemote != nil {
		e.opts.OnDemote(ctx)
	}
}

func (e *Elector) stepDown(ctx context.Context, revision uint64) {
	_ = revision
	e.demote(ctx)
	_ = e.kv.Delete(ctx, e.opts.Key)
}