309 lines
11 KiB
Go
309 lines
11 KiB
Go
package lease_test
|
|
|
|
import (
|
|
"errors"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.wntrmute.dev/kyle/crossbar/internal/lease"
|
|
"git.wntrmute.dev/kyle/crossbar/internal/store"
|
|
)
|
|
|
|
// memPersister is an in-memory Persister that also counts writes.
|
|
type memPersister struct {
|
|
mu sync.Mutex
|
|
leases map[[3]string]store.Lease
|
|
events []store.LeaseEvent
|
|
saves int
|
|
}
|
|
|
|
func newPersister() *memPersister { return &memPersister{leases: map[[3]string]store.Lease{}} }
|
|
|
|
func (m *memPersister) SaveLease(l store.Lease) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
m.saves++
|
|
m.leases[[3]string{l.Route, l.FP, l.Model}] = l
|
|
return nil
|
|
}
|
|
func (m *memPersister) DeleteLease(route, fp, model string) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
delete(m.leases, [3]string{route, fp, model})
|
|
return nil
|
|
}
|
|
func (m *memPersister) ListLeases() ([]store.Lease, error) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
out := []store.Lease{}
|
|
for _, l := range m.leases {
|
|
out = append(out, l)
|
|
}
|
|
return out, nil
|
|
}
|
|
func (m *memPersister) RecordEvent(e store.LeaseEvent) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
m.events = append(m.events, e)
|
|
return nil
|
|
}
|
|
func (m *memPersister) reasons() []string {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
var r []string
|
|
for _, e := range m.events {
|
|
r = append(r, e.Reason)
|
|
}
|
|
return r
|
|
}
|
|
|
|
// world is a hand-set view of hosts plus a chooser that returns a fixed answer.
|
|
type world struct {
|
|
mu sync.Mutex
|
|
healthy map[string]bool
|
|
draining map[string]bool
|
|
pick string
|
|
picks []string // candidates seen by Choose, for assertions
|
|
}
|
|
|
|
func (w *world) Healthy(name string) bool { w.mu.Lock(); defer w.mu.Unlock(); return w.healthy[name] }
|
|
func (w *world) Draining(name string) bool { w.mu.Lock(); defer w.mu.Unlock(); return w.draining[name] }
|
|
func (w *world) Choose(candidates []string, model string) (string, bool) {
|
|
w.mu.Lock()
|
|
defer w.mu.Unlock()
|
|
w.picks = append([]string{}, candidates...)
|
|
for _, c := range candidates {
|
|
if c == w.pick {
|
|
return c, true
|
|
}
|
|
}
|
|
if len(candidates) > 0 {
|
|
return candidates[0], true
|
|
}
|
|
return "", false
|
|
}
|
|
|
|
var t0 = time.Date(2026, 9, 25, 10, 0, 0, 0, time.UTC)
|
|
|
|
func newTable(t *testing.T, p *memPersister, w *world) *lease.Table {
|
|
tbl, err := lease.New(p, w, w, 30*time.Minute)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return tbl
|
|
}
|
|
|
|
func TestNewLeaseThenSticky(t *testing.T) {
|
|
p, w := newPersister(), &world{healthy: map[string]bool{"alpha": true, "beta": true}, pick: "beta"}
|
|
tbl := newTable(t, p, w)
|
|
k := lease.Key{Route: "r", FP: "conv1", Model: "m"}
|
|
host, reused, err := tbl.Acquire(k, []string{"alpha", "beta"}, t0)
|
|
if err != nil || host != "beta" || reused {
|
|
t.Fatalf("first: %q %v %v", host, reused, err)
|
|
}
|
|
w.pick = "alpha" // the chooser would now prefer alpha; the lease must hold
|
|
for i := 1; i <= 5; i++ {
|
|
host, reused, err = tbl.Acquire(k, []string{"alpha", "beta"}, t0.Add(time.Duration(i)*time.Minute))
|
|
if err != nil || host != "beta" || !reused {
|
|
t.Fatalf("turn %d: %q reused=%v %v, want beta reused", i, host, reused, err)
|
|
}
|
|
}
|
|
if got := p.reasons(); len(got) != 1 || got[0] != store.ReasonNew {
|
|
t.Errorf("events = %v, want one 'new'", got)
|
|
}
|
|
snap := tbl.Snapshot()
|
|
if len(snap) != 1 || snap[0].Host != "beta" || !snap[0].LastUsed.Equal(t0.Add(5*time.Minute)) {
|
|
t.Errorf("snapshot = %+v", snap)
|
|
}
|
|
if p.saves < 2 {
|
|
t.Errorf("LastUsed must be written through (saves=%d)", p.saves)
|
|
}
|
|
}
|
|
|
|
func TestUnhealthyHostMovesTheLease(t *testing.T) {
|
|
p, w := newPersister(), &world{healthy: map[string]bool{"alpha": true, "beta": true}, pick: "alpha"}
|
|
tbl := newTable(t, p, w)
|
|
k := lease.Key{Route: "r", FP: "c", Model: "m"}
|
|
if host, _, _ := tbl.Acquire(k, []string{"alpha", "beta"}, t0); host != "alpha" {
|
|
t.Fatalf("first: %q", host)
|
|
}
|
|
w.mu.Lock()
|
|
w.healthy["alpha"] = false
|
|
w.pick = "beta"
|
|
w.mu.Unlock()
|
|
host, reused, err := tbl.Acquire(k, []string{"alpha", "beta"}, t0.Add(time.Minute))
|
|
if err != nil || host != "beta" || reused {
|
|
t.Fatalf("after alpha down: %q reused=%v %v", host, reused, err)
|
|
}
|
|
if got := p.reasons(); len(got) != 2 || got[1] != store.ReasonUnhealthy {
|
|
t.Errorf("events = %v, want [new unhealthy]", got)
|
|
}
|
|
if len(w.picks) != 1 || w.picks[0] != "beta" {
|
|
t.Errorf("Choose must not see the unhealthy host: %v", w.picks)
|
|
}
|
|
}
|
|
|
|
func TestIdleExpiry(t *testing.T) {
|
|
p, w := newPersister(), &world{healthy: map[string]bool{"alpha": true, "beta": true}, pick: "alpha"}
|
|
tbl := newTable(t, p, w)
|
|
k := lease.Key{Route: "r", FP: "c", Model: "m"}
|
|
tbl.Acquire(k, []string{"alpha", "beta"}, t0)
|
|
if n := tbl.ExpireIdle(t0.Add(29 * time.Minute)); n != 0 {
|
|
t.Errorf("expired %d before lease_idle", n)
|
|
}
|
|
if n := tbl.ExpireIdle(t0.Add(31 * time.Minute)); n != 1 {
|
|
t.Errorf("expired %d after lease_idle, want 1", n)
|
|
}
|
|
if got := p.reasons(); got[len(got)-1] != store.ReasonIdle {
|
|
t.Errorf("events = %v, want idle last", got)
|
|
}
|
|
w.pick = "beta"
|
|
if host, reused, _ := tbl.Acquire(k, []string{"alpha", "beta"}, t0.Add(32*time.Minute)); host != "beta" || reused {
|
|
t.Errorf("after expiry a new lease is chosen: %q reused=%v", host, reused)
|
|
}
|
|
if l, _ := p.ListLeases(); len(l) != 1 {
|
|
t.Errorf("persister holds %d leases, want 1", len(l))
|
|
}
|
|
}
|
|
|
|
func TestFingerprintInheritsRouteLease(t *testing.T) {
|
|
p, w := newPersister(), &world{healthy: map[string]bool{"alpha": true, "beta": true}, pick: "beta"}
|
|
tbl := newTable(t, p, w)
|
|
// A request without a fingerprint (no user message) leases the route itself…
|
|
if host, _, _ := tbl.Acquire(lease.Key{Route: "r", FP: "", Model: "m"}, []string{"alpha", "beta"}, t0); host != "beta" {
|
|
t.Fatalf("route lease: %q", host)
|
|
}
|
|
w.pick = "alpha"
|
|
// …and a new conversation on that route starts where the route already is.
|
|
host, reused, err := tbl.Acquire(lease.Key{Route: "r", FP: "conv", Model: "m"}, []string{"alpha", "beta"}, t0.Add(time.Second))
|
|
if err != nil || host != "beta" || !reused {
|
|
t.Errorf("fingerprint lease must inherit the route's host: %q reused=%v %v", host, reused, err)
|
|
}
|
|
if len(tbl.Snapshot()) != 2 {
|
|
t.Errorf("both the route lease and the conversation lease exist: %+v", tbl.Snapshot())
|
|
}
|
|
}
|
|
|
|
func TestPinAndUnpin(t *testing.T) {
|
|
p, w := newPersister(), &world{healthy: map[string]bool{"alpha": true, "beta": true}, pick: "alpha"}
|
|
tbl := newTable(t, p, w)
|
|
k := lease.Key{Route: "r", FP: "c", Model: "m"}
|
|
tbl.Acquire(k, []string{"alpha", "beta"}, t0)
|
|
if err := tbl.Pin("r", "beta", t0.Add(time.Minute)); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
host, _, err := tbl.Acquire(k, []string{"alpha", "beta"}, t0.Add(2*time.Minute))
|
|
if err != nil || host != "beta" {
|
|
t.Fatalf("pinned route must go to beta: %q %v", host, err)
|
|
}
|
|
host, _, err = tbl.Acquire(lease.Key{Route: "r", FP: "other", Model: "m"}, []string{"alpha", "beta"}, t0.Add(2*time.Minute))
|
|
if err != nil || host != "beta" {
|
|
t.Fatalf("new conversations on a pinned route go to the pin too: %q %v", host, err)
|
|
}
|
|
w.mu.Lock()
|
|
w.healthy["beta"] = false
|
|
w.mu.Unlock()
|
|
if _, _, err := tbl.Acquire(k, []string{"alpha", "beta"}, t0.Add(3*time.Minute)); !errors.Is(err, lease.ErrPinnedDown) {
|
|
t.Errorf("a pinned host that is down is ErrPinnedDown, never a silent move: %v", err)
|
|
}
|
|
if err := tbl.Pin("r", "nobody", t0); !errors.Is(err, lease.ErrUnknownHost) {
|
|
t.Errorf("pinning to a host not in the candidates of any lease: %v, want ErrUnknownHost", err)
|
|
}
|
|
tbl.Unpin("r")
|
|
w.mu.Lock()
|
|
w.healthy["beta"] = true
|
|
w.mu.Unlock()
|
|
if host, _, _ := tbl.Acquire(k, []string{"alpha", "beta"}, t0.Add(4*time.Minute)); host != "beta" {
|
|
t.Errorf("after unpin the existing lease (on beta) simply continues: %q", host)
|
|
}
|
|
// Events: a pin event naming beta must exist, and the unpin's release event must come after it.
|
|
// Acquires under the pin may record their own events in between; their number is not fixed here.
|
|
got := p.reasons()
|
|
pinAt, releaseAt := -1, -1
|
|
for i, r := range got {
|
|
if r == store.ReasonPin && pinAt < 0 {
|
|
pinAt = i
|
|
}
|
|
if r == store.ReasonRelease {
|
|
releaseAt = i
|
|
}
|
|
}
|
|
if pinAt < 0 || releaseAt < pinAt {
|
|
t.Errorf("events = %v, want a pin event followed later by a release event", got)
|
|
}
|
|
p.mu.Lock()
|
|
if pinAt >= 0 && p.events[pinAt].ToHost != "beta" {
|
|
t.Errorf("pin event = %+v, want ToHost beta", p.events[pinAt])
|
|
}
|
|
p.mu.Unlock()
|
|
}
|
|
|
|
func TestDrainKeepsExistingRefusesNew(t *testing.T) {
|
|
p, w := newPersister(), &world{healthy: map[string]bool{"alpha": true, "beta": true}, draining: map[string]bool{}, pick: "alpha"}
|
|
tbl := newTable(t, p, w)
|
|
k := lease.Key{Route: "r", FP: "c", Model: "m"}
|
|
tbl.Acquire(k, []string{"alpha", "beta"}, t0)
|
|
w.mu.Lock()
|
|
w.draining["alpha"] = true
|
|
w.mu.Unlock()
|
|
if host, reused, _ := tbl.Acquire(k, []string{"alpha", "beta"}, t0.Add(time.Minute)); host != "alpha" || !reused {
|
|
t.Errorf("an existing lease on a draining host continues: %q reused=%v", host, reused)
|
|
}
|
|
host, _, err := tbl.Acquire(lease.Key{Route: "r2", FP: "x", Model: "m"}, []string{"alpha", "beta"}, t0.Add(time.Minute))
|
|
if err != nil || host != "beta" {
|
|
t.Errorf("a new lease avoids the draining host: %q %v", host, err)
|
|
}
|
|
if len(w.picks) != 1 || w.picks[0] != "beta" {
|
|
t.Errorf("Choose must not see the draining host: %v", w.picks)
|
|
}
|
|
if _, _, err := tbl.Acquire(lease.Key{Route: "r3", FP: "y", Model: "m"}, []string{"alpha"}, t0); !errors.Is(err, lease.ErrNoHost) {
|
|
t.Errorf("only draining candidates: %v, want ErrNoHost", err)
|
|
}
|
|
}
|
|
|
|
func TestReleaseRoute(t *testing.T) {
|
|
p, w := newPersister(), &world{healthy: map[string]bool{"alpha": true, "beta": true}, pick: "alpha"}
|
|
tbl := newTable(t, p, w)
|
|
tbl.Acquire(lease.Key{Route: "r", FP: "a", Model: "m"}, []string{"alpha", "beta"}, t0)
|
|
tbl.Acquire(lease.Key{Route: "r", FP: "b", Model: "m"}, []string{"alpha", "beta"}, t0)
|
|
tbl.Acquire(lease.Key{Route: "other", FP: "c", Model: "m"}, []string{"alpha", "beta"}, t0)
|
|
if n := tbl.Release("r"); n != 2 {
|
|
t.Errorf("Release removed %d, want 2", n)
|
|
}
|
|
if n := tbl.Release("r"); n != 0 {
|
|
t.Errorf("second Release removed %d", n)
|
|
}
|
|
if l, _ := p.ListLeases(); len(l) != 1 || l[0].Route != "other" {
|
|
t.Errorf("persister after release: %+v", l)
|
|
}
|
|
w.pick = "beta"
|
|
if host, reused, _ := tbl.Acquire(lease.Key{Route: "r", FP: "a", Model: "m"}, []string{"alpha", "beta"}, t0); host != "beta" || reused {
|
|
t.Errorf("after release the route is re-chosen: %q reused=%v", host, reused)
|
|
}
|
|
}
|
|
|
|
func TestLoadsFromPersister(t *testing.T) {
|
|
p, w := newPersister(), &world{healthy: map[string]bool{"alpha": true, "beta": true}, pick: "alpha"}
|
|
_ = p.SaveLease(store.Lease{Route: "r", FP: "c", Model: "m", Host: "beta", State: store.Active, Created: t0, LastUsed: t0})
|
|
_ = p.SaveLease(store.Lease{Route: "pinned", FP: "", Model: "", Host: "beta", State: store.Pinned, Created: t0, LastUsed: t0})
|
|
tbl := newTable(t, p, w)
|
|
if host, reused, _ := tbl.Acquire(lease.Key{Route: "r", FP: "c", Model: "m"}, []string{"alpha", "beta"}, t0.Add(time.Second)); host != "beta" || !reused {
|
|
t.Errorf("a restart must not reshuffle: %q reused=%v", host, reused)
|
|
}
|
|
if host, _, _ := tbl.Acquire(lease.Key{Route: "pinned", FP: "new", Model: "m"}, []string{"alpha", "beta"}, t0.Add(time.Second)); host != "beta" {
|
|
t.Errorf("a pin survives a restart: %q", host)
|
|
}
|
|
}
|
|
|
|
func TestNoCandidates(t *testing.T) {
|
|
p, w := newPersister(), &world{healthy: map[string]bool{}, pick: ""}
|
|
tbl := newTable(t, p, w)
|
|
if _, _, err := tbl.Acquire(lease.Key{Route: "r", FP: "c", Model: "m"}, []string{"alpha"}, t0); !errors.Is(err, lease.ErrNoHost) {
|
|
t.Errorf("no healthy host: %v, want ErrNoHost", err)
|
|
}
|
|
if len(tbl.Snapshot()) != 0 {
|
|
t.Errorf("a failed acquire must not create a lease")
|
|
}
|
|
}
|