zk/zk_test.go

package zk

import (
	"errors"
	"reflect"
	"sync"
	"testing"
)

func must[T any](t *testing.T) func(T, error) T {
	return func(v T, err error) T {
		t.Helper()
		if err != nil {
			t.Fatal(err)
		}
		return v
	}
}

func expectErr(t *testing.T, got, want error) {
	t.Helper()
	if !errors.Is(got, want) {
		t.Fatalf("got %v, want %v", got, want)
	}
}

func recv(t *testing.T, ch <-chan Event, want Event) {
	t.Helper()
	select {
	case ev := <-ch:
		if ev != want {
			t.Fatalf("event %+v, want %+v", ev, want)
		}
	default:
		t.Fatalf("no event, want %+v", want)
	}
}

func quiet(t *testing.T, ch <-chan Event) {
	t.Helper()
	select {
	case ev := <-ch:
		t.Fatalf("unexpected event %+v", ev)
	default:
	}
}

func TestCRUDAndVersions(t *testing.T) {
	c := New().NewSession()
	must[string](t)(c.Create("/a", []byte("x"), Persistent))
	must[string](t)(c.Create("/a/b", nil, Persistent))

	_, err := c.Create("/a", nil, Persistent)
	expectErr(t, err, ErrNodeExists)
	_, err = c.Create("/missing/b", nil, Persistent)
	expectErr(t, err, ErrNoNode)
	_, err = c.Create("a", nil, Persistent)
	expectErr(t, err, ErrBadPath)

	data, st, _, err := c.Get("/a", false)
	if err != nil || string(data) != "x" || st.Version != 0 || st.NumChildren != 1 {
		t.Fatalf("get: %q %+v %v", data, st, err)
	}

	_, err = c.Set("/a", []byte("y"), 5)
	expectErr(t, err, ErrBadVersion)
	st = must[Stat](t)(c.Set("/a", []byte("y"), 0))
	if st.Version != 1 {
		t.Fatalf("version %d", st.Version)
	}

	expectErr(t, c.Delete("/a", AnyVersion), ErrNotEmpty)
	expectErr(t, c.Delete("/a/b", 3), ErrBadVersion)
	if err := c.Delete("/a/b", 0); err != nil {
		t.Fatal(err)
	}
	if err := c.Delete("/a", 1); err != nil {
		t.Fatal(err)
	}
	if st, _, _ := c.Exists("/a", false); st != nil {
		t.Fatal("still exists")
	}
}

func TestGetReturnsCopy(t *testing.T) {
	c := New().NewSession()
	in := []byte("abc")
	must[string](t)(c.Create("/n", in, Persistent))
	in[0] = 'z'
	out, _, _, _ := c.Get("/n", false)
	out[1] = 'z'
	again, _, _, _ := c.Get("/n", false)
	if string(again) != "abc" {
		t.Fatalf("aliasing: %q", again)
	}
}

func TestSequential(t *testing.T) {
	c := New().NewSession()
	must[string](t)(c.Create("/q", nil, Persistent))
	p0 := must[string](t)(c.Create("/q/item-", nil, PersistentSequential))
	p1 := must[string](t)(c.Create("/q/item-", nil, EphemeralSequential))
	if p0 != "/q/item-0000000000" || p1 != "/q/item-0000000001" {
		t.Fatalf("%s %s", p0, p1)
	}
	// Counter never reuses after delete.
	c.Delete(p1, AnyVersion)
	p2 := must[string](t)(c.Create("/q/item-", nil, PersistentSequential))
	if p2 != "/q/item-0000000002" {
		t.Fatal(p2)
	}
	// Sequential directly under root.
	r := must[string](t)(c.Create("/r-", nil, PersistentSequential))
	if r[:3] != "/r-" || len(r) != 13 {
		t.Fatal(r)
	}
}

func TestEphemeralLifecycle(t *testing.T) {
	s := New()
	a, b := s.NewSession(), s.NewSession()
	must[string](t)(a.Create("/p", nil, Persistent))
	must[string](t)(a.Create("/p/e1", nil, Ephemeral))
	must[string](t)(a.Create("/p/e2", nil, Ephemeral))

	_, err := a.Create("/p/e1/child", nil, Persistent)
	expectErr(t, err, ErrNoChildrenForEphemerals)

	st, _, _ := b.Exists("/p/e1", false)
	if st == nil || st.EphemeralOwner != a.ID() {
		t.Fatalf("owner: %+v", st)
	}

	// Another session deleting it removes it from a's ownership set.
	b.Delete("/p/e2", AnyVersion)

	_, w, _ := b.Children("/p", true)
	a.Close()
	recv(t, w, Event{NodeChildrenChanged, "/p"})

	kids, _, _ := b.Children("/p", false)
	if len(kids) != 0 {
		t.Fatalf("leftover: %v", kids)
	}
	_, err = a.Create("/x", nil, Persistent)
	expectErr(t, err, ErrSessionClosed)
	a.Close() // idempotent
}

func TestWatches(t *testing.T) {
	s := New()
	c, other := s.NewSession(), s.NewSession()

	// Exists watch on missing node fires on create.
	st, w, _ := c.Exists("/w", true)
	if st != nil {
		t.Fatal("should not exist")
	}
	must[string](t)(other.Create("/w", nil, Persistent))
	recv(t, w, Event{NodeCreated, "/w"})

	// Data watch fires once.
	_, _, w, _ = c.Get("/w", true)
	other.Set("/w", []byte("1"), AnyVersion)
	recv(t, w, Event{NodeDataChanged, "/w"})
	other.Set("/w", []byte("2"), AnyVersion)
	quiet(t, w)

	// Child watch fires on create; not on child's data change.
	_, cw, _ := c.Children("/w", true)
	must[string](t)(other.Create("/w/k", nil, Persistent))
	recv(t, cw, Event{NodeChildrenChanged, "/w"})

	_, cw, _ = c.Children("/w", true)
	other.Set("/w/k", []byte("z"), AnyVersion)
	quiet(t, cw)
	other.Delete("/w/k", AnyVersion)
	recv(t, cw, Event{NodeChildrenChanged, "/w"})

	// Deleting a node fires its data and child watches with NodeDeleted.
	_, _, dw, _ := c.Get("/w", true)
	_, cw, _ = c.Children("/w", true)
	other.Delete("/w", AnyVersion)
	recv(t, dw, Event{NodeDeleted, "/w"})
	recv(t, cw, Event{NodeDeleted, "/w"})
}

func TestStatZxids(t *testing.T) {
	c := New().NewSession()
	must[string](t)(c.Create("/a", nil, Persistent))
	must[string](t)(c.Create("/a/b", nil, Persistent))
	c.Set("/a", []byte("v"), AnyVersion)
	_, st, _, _ := c.Get("/a", false)
	want := Stat{Czxid: 1, Mzxid: 3, Pzxid: 2, Version: 1, Cversion: 1, NumChildren: 1, DataLength: 1}
	if !reflect.DeepEqual(st, want) {
		t.Fatalf("%+v\nwant %+v", st, want)
	}
}

// Leader election: lowest sequential ephemeral wins; when it goes away the
// next one sees its predecessor's delete.
func TestLeaderElectionRecipe(t *testing.T) {
	s := New()
	setup := s.NewSession()
	must[string](t)(setup.Create("/election", nil, Persistent))

	var sessions []*Session
	var paths []string
	for i := 0; i < 3; i++ {
		c := s.NewSession()
		sessions = append(sessions, c)
		paths = append(paths, must[string](t)(c.Create("/election/n-", nil, EphemeralSequential)))
	}
	kids, _, _ := setup.Children("/election", false)
	if "/election/"+kids[0] != paths[0] {
		t.Fatalf("leader %s", kids[0])
	}
	// Candidate 1 watches its predecessor (candidate 0).
	_, w, _ := sessions[1].Exists(paths[0], true)
	sessions[0].Close()
	recv(t, w, Event{NodeDeleted, paths[0]})
	kids, _, _ = setup.Children("/election", false)
	if "/election/"+kids[0] != paths[1] {
		t.Fatalf("new leader %s", kids[0])
	}
}

func TestConcurrentSequential(t *testing.T) {
	s := New()
	setup := s.NewSession()
	must[string](t)(setup.Create("/c", nil, Persistent))
	var wg sync.WaitGroup
	seen := sync.Map{}
	for i := 0; i < 50; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			c := s.NewSession()
			for j := 0; j < 20; j++ {
				p, err := c.Create("/c/x-", nil, PersistentSequential)
				if err != nil {
					t.Error(err)
					return
				}
				if _, dup := seen.LoadOrStore(p, true); dup {
					t.Errorf("duplicate %s", p)
				}
			}
		}()
	}
	wg.Wait()
	kids, _, _ := setup.Children("/c", false)
	if len(kids) != 1000 {
		t.Fatalf("%d children", len(kids))
	}
}