quick_test.go

package statedb

import (
	"bytes"
	"cmp"
	"fmt"
	"iter"
	"math/rand"
	"reflect"
	"strings"
	"testing"
	"testing/quick"

	"github.com/cilium/statedb/index"
	"github.com/stretchr/testify/require"
)

// Use an object with strings for both primary and secondary
// indexing. With non-unique indexes we'll create a composite
// key out of them by concatenation and thus we need test that
// this does not break iteration order etc.
type quickObj struct {
	A, B string
}

// TableHeader implements TableWritable.
func (q quickObj) TableHeader() []string {
	return []string{"A", "B"}
}

// TableRow implements TableWritable.
func (q quickObj) TableRow() []string {
	return []string{q.A, q.B}
}

func (q quickObj) String() string {
	return fmt.Sprintf("%x %x", []byte(q.A), []byte(q.B))
}

var _ TableWritable = quickObj{}

var (
	aIndex = Index[quickObj, string]{
		Name: "a",
		FromObject: func(t quickObj) index.KeySet {
			return index.NewKeySet(index.String(t.A))
		},
		FromKey:    index.String,
		FromString: index.FromString,
		Unique:     true,
	}

	bIndex = Index[quickObj, string]{
		Name: "b",
		FromObject: func(t quickObj) index.KeySet {
			return index.NewKeySet(index.String(t.B))
		},
		FromKey:    index.String,
		FromString: index.FromString,
		Unique:     false,
	}
)

func isOrdered(t *testing.T, aFirst bool, it iter.Seq2[quickObj, Revision]) bool {
	var prev quickObj
	for a := range it {
		if aFirst {
			if ret := cmp.Compare(a.A, prev.A); ret < 0 {
				t.Logf("isOrdered(A): %s < %s!", a, prev)
				return false
			} else if ret == 0 && cmp.Compare(a.B, prev.B) < 0 {
				t.Logf("isOrdered(B): %s < %s!", a, prev)
				return false
			}
		} else {
			if ret := cmp.Compare(a.B, prev.B); ret < 0 {
				t.Logf("isOrdered(B): %s < %s!", a, prev)
				return false
			} else if ret == 0 && cmp.Compare(a.A, prev.A) < 0 {
				t.Logf("isOrdered(A): %s < %s!", a, prev)
				return false
			}
		}

		prev = a
	}
	return true
}

func seqLen[A, B any](it iter.Seq2[A, B]) int {
	n := 0
	for range it {
		n++
	}
	return n
}

func TestDB_Quick(t *testing.T) {
	t.Parallel()
	db := New()
	table, err := NewTable(db, "test", aIndex, bIndex)
	require.NoError(t, err, "NewTable")

	anyTable := AnyTable{table}

	values := map[string]string{}
	pickRandom := func() string {
		if len(values) == 0 || rand.Intn(10) == 0 {
			return "nonexisting"
		}
		n := rand.Intn(len(values))
		for a := range values {
			if n <= 0 {
				return a
			}
			n--
		}
		return ""
	}

	numInserted, numRemoved := 0, 0

	check := func(a, b string, remove bool, useModify bool) bool {
		txn := db.WriteTxn(table)
		if remove {
			key := pickRandom()
			expected, found := values[key]
			actual, hadOld, err := table.Delete(txn, quickObj{A: key})
			if err != nil {
				t.Logf("delete error: %s", err)
				return false
			}
			if found {
				if !hadOld {
					t.Logf("object was inserted but not found when removing")
					return false
				}
				if actual.B != expected {
					t.Logf("removed object value mismatch %q vs %q", expected, actual.B)
				}
				delete(values, key)
				numRemoved++
			} else {
				if hadOld {
					t.Logf("object was removed when it should have not existed")
					return false
				}
			}
			txn.Commit()
			return true
		}

		getObj, _, getWatch, getFound := table.GetWatch(txn, aIndex.Query(a))

		expected, found := values[a]
		expectedB := b

		var (
			old    quickObj
			hadOld bool
			err    error
		)
		if useModify {
			old, hadOld, err = table.Modify(txn, quickObj{a, b}, func(old, new quickObj) quickObj {
				new.B = old.B + new.B
				return new
			})
			require.NoError(t, err, "Modify")
			if found {
				expectedB = expected + b
			}
		} else {
			old, hadOld, err = table.Insert(txn, quickObj{a, b})
			require.NoError(t, err, "Insert")
		}
		numInserted++

		require.Equal(t, getFound, hadOld)

		if found {
			if !hadOld {
				t.Logf("object was updated but old value not returned")
				return false
			}
			if old.B != expected {
				t.Logf("insert returned wrong old object, %q vs %q", expected, old.B)
				return false
			}
		} else {
			if hadOld {
				t.Logf("insert returned old value when it should not have existed")
				return false
			}
		}
		values[a] = expectedB

		if len(values) != table.NumObjects(txn) {
			t.Logf("wrong object count")
			return false
		}

		txn.Commit()

		if hadOld {
			select {
			case <-getWatch:
			default:
				t.Logf("get watch channel not closed after insert")
				return false
			}
			require.Equal(t, old, getObj)
		}

		rtxn := db.ReadTxn()

		if len(values) != table.NumObjects(rtxn) {
			t.Logf("wrong object count")
			return false
		}

		//
		// Check queries against the primary index
		//

		if len(values) != seqLen(table.All(rtxn)) {
			t.Logf("All() via aIndex wrong length")
			return false
		}
		if len(values) != seqLen(table.Prefix(rtxn, aIndex.Query(""))) {
			t.Logf("Prefix() via aIndex wrong length")
			return false
		}
		if len(values) != seqLen(table.LowerBound(rtxn, aIndex.Query(""))) {
			t.Logf("LowerBound() via aIndex wrong length")
			return false
		}

		obj, _, found := table.Get(rtxn, aIndex.Query(a))
		if !found || obj.A != a {
			t.Logf("Get() via aIndex")
			return false
		}
		for obj := range table.Prefix(rtxn, aIndex.Query(a)) {
			if !strings.HasPrefix(obj.A, a) {
				t.Logf("Prefix() returned object with wrong prefix via aIndex")
				return false
			}
		}
		anyObjs, err := anyTable.Prefix(rtxn, "a", a)
		require.NoError(t, err, "AnyTable.Prefix")
		for anyObj := range anyObjs {
			obj := anyObj.(quickObj)
			if !strings.HasPrefix(obj.A, a) {
				t.Logf("AnyTable.Prefix() returned object with wrong prefix via aIndex")
				return false
			}
		}

		for obj := range table.LowerBound(rtxn, aIndex.Query(a)) {
			if cmp.Compare(obj.A, a) < 0 {
				t.Logf("LowerBound() order wrong")
				return false
			}
		}
		anyObjs, err = anyTable.LowerBound(rtxn, "a", a)
		require.NoError(t, err, "AnyTable.LowerBound")
		for anyObj := range anyObjs {
			obj := anyObj.(quickObj)
			if cmp.Compare(obj.A, a) < 0 {
				t.Logf("AnyTable.LowerBound(%x) order wrong: %x < %x", []byte(a), []byte(obj.A), []byte(a))
				return false
			}
		}

		if !isOrdered(t, true, table.All(rtxn)) {
			t.Logf("All() wrong order")
			return false
		}
		if !isOrdered(t, true, table.Prefix(rtxn, aIndex.Query(""))) {
			t.Logf("Prefix() via aIndex wrong order")
			return false
		}
		if !isOrdered(t, true, table.LowerBound(rtxn, aIndex.Query(""))) {
			t.Logf("LowerBound() via aIndex wrong order")
			return false
		}

		//
		// Check against the secondary (non-unique index)
		//
		queryB := expectedB

		// Non-unique indexes return the same number of objects as we've inserted.
		if len(values) != seqLen(table.Prefix(rtxn, bIndex.Query(""))) {
			t.Logf("Prefix() via bIndex wrong length")
			return false
		}

		if len(values) != seqLen(table.LowerBound(rtxn, bIndex.Query(""))) {
			t.Logf("LowerBound() via bIndex wrong length")
			return false
		}

		// Get returns the first match, but since the index is non-unique, this might
		// not be the one that we just inserted.
		obj, _, found = table.Get(rtxn, bIndex.Query(queryB))
		if !found || obj.B != queryB {
			t.Logf("Get(%q) via bIndex not found (%v) or wrong B (%q vs %q)", queryB, found, obj.B, queryB)
			return false
		}

		found = false
		for obj := range table.List(rtxn, bIndex.Query(queryB)) {
			if obj.B != queryB {
				t.Logf("List() via bIndex wrong B")
				return false
			}
			if obj.A == a {
				found = true
			}
		}
		if !found {
			t.Logf("List() via bIndex, object not found")
			return false
		}

		visited := map[string]struct{}{}
		for obj := range table.Prefix(rtxn, bIndex.Query(queryB)) {
			if !strings.HasPrefix(obj.B, queryB) {
				t.Logf("Prefix() via bIndex has wrong prefix")
				return false
			}
			if _, found := visited[obj.A]; found {
				t.Logf("Prefix() visited object %q twice", obj.A)
				return false
			}
			visited[obj.A] = struct{}{}
		}

		anyObjs, err = anyTable.Prefix(rtxn, "b", queryB)
		require.NoError(t, err, "AnyTable.Prefix")
		for anyObj := range anyObjs {
			obj := anyObj.(quickObj)
			if !strings.HasPrefix(obj.B, queryB) {
				t.Logf("AnyTable.Prefix() via bIndex has wrong prefix: %q vs %q", obj.B, queryB)
				return false
			}
		}

		visited = map[string]struct{}{}
		for obj := range table.LowerBound(rtxn, bIndex.Query(queryB)) {
			if cmp.Compare(obj.B, queryB) < 0 {
				t.Logf("LowerBound() via bIndex has wrong objects, expected %v >= %v", []byte(obj.B), []byte(queryB))
				return false
			}
			if _, found := visited[obj.A]; found {
				t.Logf("Prefix() visited object %q twice", obj.A)
				return false
			}
			visited[obj.A] = struct{}{}
		}

		anyObjs, err = anyTable.LowerBound(rtxn, "b", queryB)
		require.NoError(t, err, "AnyTable.LowerBound")
		for anyObj := range anyObjs {
			obj := anyObj.(quickObj)
			if cmp.Compare(obj.B, queryB) < 0 {
				t.Logf("AnyTable.LowerBound() via bIndex has wrong objects, expected %v >= %v", []byte(obj.B), []byte(queryB))
				return false
			}
		}

		// Iterating over the secondary index returns the objects in order
		// defined by the "B" key first and then by the "A" key.
		if !isOrdered(t, false, table.Prefix(rtxn, bIndex.Query(""))) {
			t.Logf("Prefix() via bIndex wrong order")
			return false
		}
		if !isOrdered(t, false, table.LowerBound(rtxn, bIndex.Query(""))) {
			t.Logf("Prefix() via bIndex wrong order")
			return false
		}
		return true
	}

	require.NoError(t, quick.Check(check, &quick.Config{
		MaxCount: 5000,

		// Make 1-8 byte long strings as input data. Keep the strings shorter
		// than the default quick value generation to hit the more interesting cases
		// often.
		Values: func(args []reflect.Value, rand *rand.Rand) {
			if len(args) != 4 {
				panic("unexpected args count")
			}
			for i := range args[:2] {
				numBytes := 1 + rand.Intn(8)
				bs := make([]byte, numBytes)
				rand.Read(bs)
				args[i] = reflect.ValueOf(string(bs))
			}
			// Remove 33% of the time
			args[2] = reflect.ValueOf(rand.Intn(3) == 1)
			// Use Modify 50% of the time
			args[3] = reflect.ValueOf(rand.Intn(2) == 0)
		},
	}))

	t.Logf("%d objects after check, %d inserts, %d removals", len(values), numInserted, numRemoved)
}

func Test_Quick_nonUniqueKey(t *testing.T) {
	t.Parallel()
	check := func(p1, s1, p2, s2 []byte) bool {
		key1 := encodeNonUniqueKey(p1, s1)
		expectedLen := encodedLength(p1) + 1 + encodedLength(s1) + 2
		minLen := len(p1) + 1 + len(s1) + 2
		if expectedLen < minLen {
			t.Logf("expected length too short (%d), must be >= %d", expectedLen, minLen)
			return false
		}
		if len(key1) != expectedLen {
			t.Logf("length mismatch, expected %d, got %d", expectedLen, len(key1))
			return false
		}

		nuk1 := nonUniqueKey(key1)
		if len(nuk1.encodedPrimary()) < len(p1) {
			t.Logf("encodedPrimary() length (%d) shorter than original (%d)", len(nuk1.encodedPrimary()), len(p1))
			return false
		}
		if len(nuk1.encodedPrimary()) != encodedLength(p1) {
			t.Logf("encodedPrimary() length (%d) does not match encodedLength() (%d)", len(nuk1.encodedPrimary()), len(p1))
			return false
		}
		if len(nuk1.encodedSecondary()) < len(s1) {
			t.Logf("encodedSecondary() length (%d) shorter than original (%d)", len(nuk1.encodedSecondary()), len(s1))
		}
		if len(nuk1.encodedSecondary()) != encodedLength(s1) {
			t.Logf("encodedSecondary() length (%d) does not match encodedLength() (%d)", len(nuk1.encodedSecondary()), len(s1))
			return false
		}

		// Do another key and check that ordering is preserved.
		key2 := encodeNonUniqueKey(p2, s2)
		scmp := bytes.Compare(s1, s2)
		pcmp := bytes.Compare(p1, p2)
		kcmp := bytes.Compare(key1, key2)

		if (scmp == 0 && pcmp != kcmp) /* secondary key matches, primary key determines order */ ||
			(scmp != 0 && scmp != kcmp) /* secondary key determines order */ {
			t.Logf("ordering not preserved: key1=%v key2=%v, p1=%v, s1=%v, p2=%v, s2=%v", key1, key2, p1, s1, p2, s2)
			return false
		}
		return true
	}
	require.NoError(t, quick.Check(check, &quick.Config{
		MaxCount: 50000,
		Values: func(args []reflect.Value, rand *rand.Rand) {
			for i := range args {
				numBytes := 1 + rand.Intn(8)
				bs := make([]byte, numBytes)
				rand.Read(bs)
				args[i] = reflect.ValueOf(bs)
			}
		},
	}))
}