Skip to content
quarry Artifact [233d73850c]
Login

Artifact [233d73850c]

Check-ins Using Download Hex

Artifact 233d73850c97f89d6375c76f2d731d4d48e8bb69cd79b77d7eb11f5c4ea2bef1:

// Package lease is the 0001 §3 single-writer coordinator: per store, a
// NATS KV entry held by the node currently allowed to invoke writes.
// Writes route to the holder; lease loss or an unreachable coordinator
// refuses writes (correctness over availability — the lease is
// coordination, not truth, and zit's on-disk flock is the second line).
//
// The KV revision is the fencing epoch. A TTL lease cannot truly fence a
// POSIX filesystem, so the window is bounded, not closed (RFC §3): a
// commit-point re-check just before each write shrinks the danger to one
// operation, and a failover grace makes a displaced holder finish first.
// The append-only log with prefix-validated reads is the backstop that
// makes residual overlap an ordering wart, not corruption.
package lease

import (
	"errors"
	"fmt"
	"sync"
	"time"

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

const (
	// Bucket holds one entry per store, keyed by project code.
	Bucket = "leases"
	// DefaultTTL is the lease lifetime; a holder renews within it.
	DefaultTTL = 15 * time.Second
)

// ErrUnavailable means the coordinator could not be reached — writes
// refuse rather than proceed uncoordinated.
var ErrUnavailable = errors.New("lease coordinator unavailable")

// ErrNotHolder means another node holds the store's lease.
var ErrNotHolder = errors.New("store lease held by another node")

// Manager holds this node's leases over one KV bucket.
type Manager struct {
	kv   nats.KeyValue
	node string
	ttl  time.Duration

	mu     sync.Mutex
	state  map[string]*hold
	stop   chan struct{}
	closed sync.Once
}

type hold struct {
	rev        uint64    // KV revision we hold at (fencing epoch)
	notBefore  time.Time // first write allowed after this (failover grace)
	sawForeign bool      // we observed a different node holding this code
	ok         bool      // we currently hold
}

// New opens (creating if needed) the lease bucket and starts the
// renewal loop. node identifies this process across the cluster.
func New(js nats.JetStreamContext, node string, ttl time.Duration) (*Manager, error) {
	if ttl <= 0 {
		ttl = DefaultTTL
	}
	kv, err := js.CreateKeyValue(&nats.KeyValueConfig{Bucket: Bucket, TTL: ttl})
	if err != nil {
		if kv, err = js.KeyValue(Bucket); err != nil {
			return nil, fmt.Errorf("lease: bucket: %w", err)
		}
	}
	m := &Manager{
		kv: kv, node: node, ttl: ttl,
		state: map[string]*hold{}, stop: make(chan struct{}),
	}
	go m.renewLoop()
	return m, nil
}

// Close stops renewal. Held leases then age out of the bucket after the
// TTL, releasing them to other nodes.
func (m *Manager) Close() { m.closed.Do(func() { close(m.stop) }) }

func (m *Manager) renewLoop() {
	t := time.NewTicker(m.ttl / 3)
	defer t.Stop()
	for {
		select {
		case <-m.stop:
			return
		case <-t.C:
			m.mu.Lock()
			for code, h := range m.state {
				if h.ok {
					m.ensure(code) // renew; drops h.ok on loss
				}
			}
			m.mu.Unlock()
		}
	}
}

// ensure brings this node's lease on code up to date, returning the
// current hold. Caller holds m.mu.
func (m *Manager) ensure(code string) *hold {
	h := m.state[code]
	if h == nil {
		h = &hold{}
		m.state[code] = h
	}
	// Renew if we believe we hold.
	if h.ok {
		if rev, err := m.kv.Update(code, []byte(m.node), h.rev); err == nil {
			h.rev = rev
			return h
		}
		h.ok = false // lost the CAS: someone took over, or coordinator gone
	}
	// Try a clean acquire (Create fails if the key exists).
	if rev, err := m.kv.Create(code, []byte(m.node)); err == nil {
		h.rev, h.ok = rev, true
		// Grace only after a failover: if we never saw a foreign holder,
		// there is nobody to wait out (clean first start).
		if h.sawForeign {
			h.notBefore = time.Now().Add(m.ttl)
		}
		return h
	}
	// Key exists — read it. Ours (adopt) or foreign (not holder)?
	entry, err := m.kv.Get(code)
	if err != nil {
		h.ok = false
		return h // coordinator trouble: caller treats as unavailable
	}
	if string(entry.Value()) == m.node {
		if rev, uerr := m.kv.Update(code, []byte(m.node), entry.Revision()); uerr == nil {
			h.rev, h.ok = rev, true
			return h
		}
	} else {
		h.sawForeign = true
	}
	h.ok = false
	return h
}

// Held reports whether this node currently holds code's lease (best
// effort; the authoritative check is inside Guard).
func (m *Manager) Held(code string) bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.ensure(code).ok
}

// Guard runs write under the store's lease. It acquires or renews,
// enforces the failover grace, re-reads the lease at its commit point,
// and only then invokes write — so a lease lost between acquisition and
// the write aborts instead of racing another holder.
func (m *Manager) Guard(code string, write func() error) error {
	m.mu.Lock()
	h := m.ensure(code)
	if !h.ok {
		m.mu.Unlock()
		// Distinguish "someone else holds" from "can't reach coordinator":
		// a foreign holder is ErrNotHolder, everything else unavailable.
		if _, err := m.kv.Get(code); err == nil {
			return ErrNotHolder
		}
		return ErrUnavailable
	}
	if wait := time.Until(h.notBefore); wait > 0 {
		m.mu.Unlock()
		return fmt.Errorf("lease: acquisition grace, retry in %s: %w", wait.Round(time.Millisecond), ErrNotHolder)
	}
	// Commit-point re-check: the epoch must be unchanged and still ours.
	entry, err := m.kv.Get(code)
	if err != nil {
		h.ok = false
		m.mu.Unlock()
		return ErrUnavailable
	}
	if entry.Revision() != h.rev || string(entry.Value()) != m.node {
		h.ok = false
		m.mu.Unlock()
		return fmt.Errorf("lease: lost at commit point: %w", ErrNotHolder)
	}
	m.mu.Unlock()
	return write()
}