package lease
import (
"sync/atomic"
"testing"
"time"
"github.com/nats-io/nats.go"
"craftdesign.group/quarry/internal/events"
)
// countingKV counts every round trip the manager makes to the bucket.
type countingKV struct {
nats.KeyValue
calls atomic.Int64
}
func (c *countingKV) Get(key string) (nats.KeyValueEntry, error) {
c.calls.Add(1)
return c.KeyValue.Get(key)
}
func (c *countingKV) Create(key string, value []byte) (uint64, error) {
c.calls.Add(1)
return c.KeyValue.Create(key, value)
}
func (c *countingKV) Update(key string, value []byte, last uint64) (uint64, error) {
c.calls.Add(1)
return c.KeyValue.Update(key, value, last)
}
// Background jobs ask Held on every tick for every store; the answer
// must come from state the renewal loop keeps, not a bucket round trip.
func TestHeldAnswersWithoutARoundTrip(t *testing.T) {
ev, err := events.Start(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(ev.Close)
kv, err := ev.JS().CreateKeyValue(&nats.KeyValueConfig{Bucket: Bucket, TTL: time.Minute})
if err != nil {
t.Fatal(err)
}
counted := &countingKV{KeyValue: kv}
m := start(counted, "node-a", time.Minute)
t.Cleanup(m.Close)
if !m.Held("repo") || !m.Held("other") {
t.Fatal("FAIL: a lone node does not hold its stores")
}
before := counted.calls.Load()
for i := 0; i < 1000; i++ {
if !m.Held("repo") || !m.Held("other") {
t.Fatal("FAIL: a held lease read as not held")
}
}
if n := counted.calls.Load() - before; n != 0 {
t.Fatalf("FAIL: 2000 Held reads made %d bucket round trips, want 0", n)
}
t.Log("ok: repeated Held reads answer from state")
}
// A closed manager runs no more jobs: Held is false at once, not a TTL
// later, and the lease moves to a node that only ever asked Held.
func TestHeldMovesWhenTheHolderCloses(t *testing.T) {
ev, err := events.Start(t.TempDir())
if err != nil {
t.Fatal(err)
}
t.Cleanup(ev.Close)
ttl := time.Second
a := manager(t, ev, "node-a", ttl)
b := manager(t, ev, "node-b", ttl)
if !a.Held("fleet") {
t.Fatal("first asker does not hold")
}
if b.Held("fleet") {
t.Fatal("FAIL: two nodes hold one lease")
}
a.Close()
if a.Held("fleet") {
t.Fatal("FAIL: a closed manager still reports holding")
}
deadline := time.Now().Add(5 * ttl)
for !b.Held("fleet") {
if time.Now().After(deadline) {
t.Fatal("FAIL: the lease never moved to the other node")
}
time.Sleep(50 * time.Millisecond)
}
t.Log("ok: the lease moved once the holder closed")
}
// A holder whose renewals stop succeeding must stop answering true
// once its last renewal is a TTL old, whatever the cache last said.
func TestHeldExpiresWithoutRenewal(t *testing.T) {
ev, err := events.Start(t.TempDir())
if err != nil {
t.Fatal(err)
}
ttl := time.Second
a := manager(t, ev, "node-a", ttl)
if !a.Held("repo") {
t.Fatal("lone node does not hold")
}
ev.Close() // coordinator gone: every renewal fails from here
deadline := time.Now().Add(3 * ttl)
for a.Held("repo") {
if time.Now().After(deadline) {
t.Fatal("FAIL: Held stayed true with no coordinator past the TTL")
}
time.Sleep(50 * time.Millisecond)
}
t.Log("ok: Held went false once renewals stopped")
}