Skip to content
quarry Artifact [ecd64dcbc1]
Login

Artifact [ecd64dcbc1]

Artifact ecd64dcbc1660035e2618f870dc2006582c2f9af5fee451c55f93ab103e86fa3:


// 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"
	"sync/atomic"
	"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")

// ErrNoToken means a store write was attempted with a Token that neither
// Guard nor Unguarded made — the zero value, which anyone can spell.
var ErrNoToken = errors.New("store write without a lease token")

// Token is what a store write must present: every zitcli call that
// writes takes one first, so a write outside Guard does not compile.
// The field is unexported so only this package can make a usable one;
// the zero value, the one thing another package can construct, is
// refused at the write by Err.
//
// The token proves a write was reached through Guard; it does not bound
// the write. A closure can keep its token and write with it after Guard
// has returned, when the lease may have passed to another node, and a
// token from one store's Guard is accepted by a write to any store:
// neither Err nor the zitcli writes check the store code or whether the
// lease is still held. Using the token only inside the closure it was
// handed to is a convention the type does not enforce.
type Token struct{ grant *grant }

type grant struct {
	code   string // the store whose lease was held, or "" when unguarded
	reason string // why a write runs without the lease; "" when guarded
}

// Err reports whether t may be used to write: nil for a token from
// Guard or Unguarded, ErrNoToken for anything else.
func (t Token) Err() error {
	if t.grant == nil {
		return ErrNoToken
	}
	return nil
}

// Unguarded is a token for a write that genuinely has no lease to hold:
// a path with no coordinator, such as a single-node forge or a test.
// reason says why, and every use is meant to be found by grepping for
// lease.Unguarded — this is the one sanctioned way around the guard.
func Unguarded(reason string) Token {
	return Token{grant: &grant{reason: reason}}
}

// 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
	done   atomic.Bool

	// held is what Held answers from without the mutex or a round trip:
	// code -> *atomic.Int64, the instant (unix nanos) this node's hold
	// on code runs out, or 0 when it holds nothing. Every ensure
	// refreshes it, and the renewal loop runs ensure on every code this
	// node has asked about, held or not, so the answer stays current.
	held sync.Map
}

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)
		}
	}
	return start(kv, node, ttl), nil
}

func start(kv nats.KeyValue, node string, ttl time.Duration) *Manager {
	m := &Manager{
		kv: kv, node: node, ttl: ttl,
		state: map[string]*hold{}, stop: make(chan struct{}),
	}
	go m.renewLoop()
	return m
}

// Close stops renewal and makes Held false at once, so this node starts
// no more background jobs. Held leases then age out of the bucket after
// the TTL, releasing them to other nodes.
func (m *Manager) Close() {
	m.closed.Do(func() {
		m.done.Store(true)
		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:
			// Held codes renew; the rest try to acquire, which is how a
			// lease whose holder went away moves to this node.
			m.mu.Lock()
			for code := range m.state {
				m.ensure(code)
			}
			m.mu.Unlock()
		}
	}
}

// ensure brings this node's lease on code up to date and publishes the
// answer Held reads, returning the current hold. Caller holds m.mu.
func (m *Manager) ensure(code string) *hold {
	// The bucket's TTL runs from the write, which is after this instant,
	// so a hold published as lasting until began+ttl never outlives the
	// entry that backs it.
	began := time.Now()
	h := m.refresh(code)
	m.publish(code, h, began)
	return h
}

// publish records h as what Held answers for code.
func (m *Manager) publish(code string, h *hold, began time.Time) {
	v, _ := m.held.LoadOrStore(code, new(atomic.Int64))
	until := int64(0)
	if h.ok {
		until = began.Add(m.ttl).UnixNano()
	}
	v.(*atomic.Int64).Store(until)
}

// refresh does ensure's round trips. Caller holds m.mu.
func (m *Manager) refresh(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). It is what decides
// which node runs a background job, asked every tick for every store,
// so it answers from the state the renewal loop keeps: lock-free and
// with no round trip, except the first time a code is asked about.
// The answer lags the bucket by at most one renewal interval, and a
// hold whose renewals have stopped reads false once it is a TTL old.
func (m *Manager) Held(code string) bool {
	if m.done.Load() {
		return false
	}
	v, ok := m.held.Load(code)
	if !ok {
		m.mu.Lock()
		m.ensure(code)
		m.mu.Unlock()
		v, _ = m.held.Load(code)
	}
	return time.Now().UnixNano() < v.(*atomic.Int64).Load()
}

// 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. write receives the
// Token its store writes must present.
func (m *Manager) Guard(code string, write func(Token) 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.publish(code, h, time.Time{})
		m.mu.Unlock()
		return ErrUnavailable
	}
	if entry.Revision() != h.rev || string(entry.Value()) != m.node {
		h.ok = false
		m.publish(code, h, time.Time{})
		m.mu.Unlock()
		return fmt.Errorf("lease: lost at commit point: %w", ErrNotHolder)
	}
	m.mu.Unlock()
	return write(Token{grant: &grant{code: code}})
}